feat(session): support local logs and command cancellation

Author: RealWhyKnot

cmd/new.go

@@ -3,18 +3,21 @@ package cmd
 
 import (
 	"context"
+	"encoding/json"
 	"errors"
 	"fmt"
 	"io"
 	"os"
 	"os/signal"
+	"sync"
 	"syscall"
 	"time"
 
 	"github.com/RealWhyKnot/Handoff/internal/audit"
 	"github.com/RealWhyKnot/Handoff/internal/capabilities"
 	"github.com/RealWhyKnot/Handoff/internal/dispatch"
 	"github.com/RealWhyKnot/Handoff/internal/relay"
+	"github.com/RealWhyKnot/Handoff/internal/supportlog"
 )
 
 // Version is stamped from main.go.
@@ -24,6 +27,7 @@ var Version = "0.1.0"
 // and runs the host agent loop until Ctrl+C or relay disconnect.
 func New(args []string) {
 	relayBase := defaultRelay()
+	supportlog.Printf("session start relay=%s version=%s", relayBase, Version)
 	fmt.Println("relay:", relayBase)
 
 	ctx, cancel := context.WithCancel(context.Background())
@@ -34,10 +38,12 @@ func New(args []string) {
 	mint, err := relay.Mint(mintCtx, relayBase)
 	mintCancel()
 	if err != nil {
+		supportlog.Printf("mint failed: %v", err)
 		fmt.Fprintln(os.Stderr, "could not mint session:", err)
 		os.Exit(1)
 	}
 	sid := shortSid(mint.ViewToken)
+	supportlog.Printf("mint ok sid=%s view_url=%s", sid, mint.ViewURL)
 	fmt.Println()
 	fmt.Println("session live -- share the view URL with your helper:")
 	fmt.Println("  ", mint.ViewURL)
@@ -48,6 +54,7 @@ func New(args []string) {
 	// Audit log.
 	log, err := audit.New()
 	if err != nil {
+		supportlog.Printf("audit log unavailable: %v", err)
 		fmt.Fprintln(os.Stderr, "warning: audit log unavailable:", err)
 	}
 	defer func() {
@@ -59,13 +66,15 @@ func New(args []string) {
 	// Dispatcher with all capabilities registered.
 	router := dispatch.New()
 	capabilities.RegisterAll(router)
+	supportlog.Printf("capabilities registered count=%d", len(router.Kinds()))
 	fmt.Printf("ready -- %d capabilities registered\n\n", len(router.Kinds()))
 
 	// Handle Ctrl+C.
 	sig := make(chan os.Signal, 1)
 	signal.Notify(sig, os.Interrupt, syscall.SIGTERM)
 	go func() {
 		<-sig
+		supportlog.Printf("shutdown requested by signal")
 		fmt.Println("\nshutting down...")
 		cancel()
 	}()
@@ -75,55 +84,174 @@ func New(args []string) {
 	bridge, err := relay.Dial(dialCtx, relayBase, mint.WriteToken)
 	dialCancel()
 	if err != nil {
+		supportlog.Printf("bridge dial failed sid=%s: %v", sid, err)
 		fmt.Fprintln(os.Stderr, "could not open bridge:", err)
 		os.Exit(1)
 	}
 	defer bridge.Close()
+	supportlog.Printf("bridge connected sid=%s", sid)
 
 	hostname, _ := os.Hostname()
 	if err := bridge.SendHello(ctx, hostname, Version, router.Kinds()); err != nil {
+		supportlog.Printf("hello failed sid=%s: %v", sid, err)
 		fmt.Fprintln(os.Stderr, "hello failed:", err)
 		os.Exit(1)
 	}
+	supportlog.Printf("hello sent sid=%s hostname=%s", sid, hostname)
 
-	// Main loop: receive commands, dispatch, send results.
+	jobs := newJobRunner(ctx, sid, router, bridge, log)
+	defer jobs.CancelAll()
+
+	// Main loop: receive commands and hand them to cancellable workers.
 	for {
 		cmd, err := bridge.Recv(ctx)
 		if err != nil {
 			if errors.Is(err, context.Canceled) || errors.Is(err, io.EOF) {
+				supportlog.Printf("recv ended sid=%s err=%v", sid, err)
 				return
 			}
+			supportlog.Printf("recv error sid=%s: %v", sid, err)
 			fmt.Fprintln(os.Stderr, "recv error:", err)
 			return
 		}
+		supportlog.Printf("command received sid=%s id=%s kind=%s", sid, cmd.ID, cmd.Kind)
+		if cmd.Kind == "control.cancel" {
+			targetID := readStringExtra(cmd.Extras, "target_id")
+			if targetID == "" {
+				supportlog.Printf("cancel control missing target sid=%s id=%s", sid, cmd.ID)
+				continue
+			}
+			if jobs.Cancel(targetID) {
+				fmt.Printf("[cancel] %s\n", targetID)
+				supportlog.Printf("command cancel requested sid=%s id=%s", sid, targetID)
+			} else {
+				supportlog.Printf("cancel target not active sid=%s id=%s", sid, targetID)
+			}
+			continue
+		}
 		fmt.Printf("[cmd] %s  kind=%s\n", cmd.ID, cmd.Kind)
+		jobs.Start(cmd)
+	}
+}
 
-		out := router.Dispatch(ctx, cmd.Kind, cmd.Extras)
+type jobRunner struct {
+	rootCtx context.Context
+	sid     string
+	router  *dispatch.Router
+	bridge  *relay.Bridge
+	audit   *audit.Logger
 
-		// Audit.
-		if log != nil {
-			res := "ok"
-			if !out.OK {
-				res = "err"
+	mu     sync.Mutex
+	active map[string]context.CancelFunc
+}
+
+func newJobRunner(rootCtx context.Context, sid string, router *dispatch.Router, bridge *relay.Bridge, auditLog *audit.Logger) *jobRunner {
+	return &jobRunner{
+		rootCtx: rootCtx,
+		sid:     sid,
+		router:  router,
+		bridge:  bridge,
+		audit:   auditLog,
+		active:  map[string]context.CancelFunc{},
+	}
+}
+
+func (r *jobRunner) Start(cmd *relay.Command) {
+	var jobCtx context.Context
+	var cancel context.CancelFunc
+	if cmd.TimeoutMS > 0 {
+		jobCtx, cancel = context.WithTimeout(r.rootCtx, time.Duration(cmd.TimeoutMS)*time.Millisecond)
+	} else {
+		jobCtx, cancel = context.WithCancel(r.rootCtx)
+	}
+
+	r.mu.Lock()
+	r.active[cmd.ID] = cancel
+	r.mu.Unlock()
+
+	go func() {
+		defer func() {
+			r.mu.Lock()
+			delete(r.active, cmd.ID)
+			r.mu.Unlock()
+			cancel()
+		}()
+
+		out := r.router.Dispatch(jobCtx, cmd.Kind, cmd.Extras)
+		if err := jobCtx.Err(); err != nil {
+			out.OK = false
+			switch {
+			case errors.Is(err, context.DeadlineExceeded):
+				out.Error = "command timed out"
+			case errors.Is(err, context.Canceled):
+				out.Error = "command cancelled"
+			default:
+				out.Error = err.Error()
 			}
-			_ = log.Write(audit.Entry{
-				SessionID:  sid,
-				Capability: cmd.Kind,
-				Args:       cmd.Extras,
-				Consent:    "session",
-				Result:     res,
-				ElapsedMs:  out.ElapsedMs,
-				Detail:     out.Error,
-			})
 		}
 
-		// Result back to the relay.
-		if err := bridge.SendCommandResult(ctx, cmd.ID, out.OK, out.Result, out.Error, out.ElapsedMs); err != nil {
+		r.writeAudit(cmd, out)
+		if err := r.bridge.SendCommandResult(r.rootCtx, cmd.ID, out.OK, out.Result, out.Error, out.ElapsedMs); err != nil {
+			supportlog.Printf("send result failed sid=%s id=%s: %v", r.sid, cmd.ID, err)
 			fmt.Fprintln(os.Stderr, "could not send result:", err)
 			return
 		}
+		supportlog.Printf("command result sent sid=%s id=%s ok=%v elapsed_ms=%d", r.sid, cmd.ID, out.OK, out.ElapsedMs)
 		fmt.Printf("       -> ok=%v elapsed=%dms\n", out.OK, out.ElapsedMs)
+	}()
+}
+
+func (r *jobRunner) Cancel(id string) bool {
+	r.mu.Lock()
+	cancel := r.active[id]
+	r.mu.Unlock()
+	if cancel == nil {
+		return false
+	}
+	cancel()
+	return true
+}
+
+func (r *jobRunner) CancelAll() {
+	r.mu.Lock()
+	cancels := make([]context.CancelFunc, 0, len(r.active))
+	for _, cancel := range r.active {
+		cancels = append(cancels, cancel)
+	}
+	r.mu.Unlock()
+	for _, cancel := range cancels {
+		cancel()
+	}
+}
+
+func (r *jobRunner) writeAudit(cmd *relay.Command, out dispatch.Outcome) {
+	if r.audit == nil {
+		return
+	}
+	res := "ok"
+	if !out.OK {
+		res = "err"
+	}
+	_ = r.audit.Write(audit.Entry{
+		SessionID:  r.sid,
+		Capability: cmd.Kind,
+		Args:       cmd.Extras,
+		Consent:    "session",
+		Result:     res,
+		ElapsedMs:  out.ElapsedMs,
+		Detail:     out.Error,
+	})
+}
+
+func readStringExtra(extras map[string]json.RawMessage, name string) string {
+	if extras == nil {
+		return ""
+	}
+	var value string
+	if raw, ok := extras[name]; ok {
+		_ = json.Unmarshal(raw, &value)
 	}
+	return value
 }
 
 func defaultRelay() string {

internal/relay/relay.go

@@ -17,6 +17,7 @@ import (
 	"net/url"
 	"sort"
 	"strings"
+	"sync"
 	"time"
 
 	"github.com/coder/websocket"
@@ -34,9 +35,10 @@ type MintResponse struct {
 // `id` and `kind` fields are fixed; any other keys in the JSON payload
 // are accessible via Extras for kind-specific arguments.
 type Command struct {
-	ID     string                     `json:"id"`
-	Kind   string                     `json:"kind"`
-	Extras map[string]json.RawMessage `json:"-"`
+	ID        string                     `json:"id"`
+	Kind      string                     `json:"kind"`
+	TimeoutMS int                        `json:"timeout_ms,omitempty"`
+	Extras    map[string]json.RawMessage `json:"-"`
 }
 
 // UnmarshalJSON decodes a Command keeping the kind-specific fields as
@@ -54,6 +56,10 @@ func (c *Command) UnmarshalJSON(data []byte) error {
 		_ = json.Unmarshal(v, &c.Kind)
 		delete(raw, "kind")
 	}
+	if v, ok := raw["timeout_ms"]; ok {
+		_ = json.Unmarshal(v, &c.TimeoutMS)
+		delete(raw, "timeout_ms")
+	}
 	c.Extras = raw
 	return nil
 }
@@ -91,7 +97,8 @@ func Mint(ctx context.Context, baseURL string) (*MintResponse, error) {
 
 // Bridge is the long-lived WebSocket between this host and the relay.
 type Bridge struct {
-	conn *websocket.Conn
+	conn    *websocket.Conn
+	writeMu sync.Mutex
 }
 
 // Dial opens the WebSocket to the relay's /ws endpoint with the write
@@ -185,6 +192,8 @@ func (b *Bridge) send(ctx context.Context, kind string, payload interface{}) err
 	if err := json.NewEncoder(&buf).Encode(ev); err != nil {
 		return err
 	}
+	b.writeMu.Lock()
+	defer b.writeMu.Unlock()
 	return b.conn.Write(ctx, websocket.MessageText, bytes.TrimRight(buf.Bytes(), "\n"))
 }
 

internal/relay/relay_test.go

@@ -83,3 +83,19 @@ func TestSendHelloIncludesProtocolAndCapabilities(t *testing.T) {
 		t.Fatalf("capabilities = %#v, want %#v", event.Payload.Capabilities, wantCaps)
 	}
 }
+
+func TestCommandUnmarshalCapturesTimeoutOutsideExtras(t *testing.T) {
+	var cmd Command
+	if err := json.Unmarshal([]byte(`{"id":"cmd-1","kind":"ps.exec","timeout_ms":1500,"script":"Start-Sleep 30"}`), &cmd); err != nil {
+		t.Fatalf("Unmarshal: %v", err)
+	}
+	if cmd.ID != "cmd-1" || cmd.Kind != "ps.exec" || cmd.TimeoutMS != 1500 {
+		t.Fatalf("command = %#v", cmd)
+	}
+	if _, ok := cmd.Extras["timeout_ms"]; ok {
+		t.Fatal("timeout_ms should not be forwarded to capability extras")
+	}
+	if _, ok := cmd.Extras["script"]; !ok {
+		t.Fatal("script extra missing")
+	}
+}