fix(serve): break approval deadlock + race in wsApprover (#32)

* fix(serve): break approval deadlock + race in wsApprover

Browser approvals were dead. handleWS ran the agent loop synchronously in
the only goroutine reading the socket; when a tool needed approval the
wsApprover blocked in PromptCommand waiting for the browser's
approval_response — which could only be delivered by that same, now-blocked,
receive loop. Every approval hit the 60s timeout and was denied, making the
Web UI safety prompt unusable.

- Move socket reading to a dedicated goroutine. It handles approval_response
  inline (delivering to the blocked PromptCommand via HandleResponse) and
  forwards other messages to a buffered promptCh the processor consumes
  serially. The forward is non-blocking so the reader can never stall and
  re-introduce the deadlock; on socket close it cancels the approver so a
  blocked prompt returns immediately instead of waiting out the timeout.

- Fix the concurrent map read/write in wsApprover: approveAll/trustAll were
  read unlocked at the top of PromptCommand while the "trust" path and
  SetTrustAll wrote them — a data race that can fatally crash serve under
  parallel tool calls. Both accesses now hold the mutex.

Regression tests (both verified to fail against the unpatched code):
- TestServe_E2E_ApprovalRoundTrip: real handleWS + mock LLM tool call;
  approves via WebSocket and asserts completion within 10s (vs the 60s
  deadlock timeout).
- TestWSApprover_ConcurrentTrust_NoRace: concurrent trust round-trips;
  crashes with "concurrent map writes" / DATA RACE before the fix.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

* fix(serve): abort prompts on disconnect via connection-scoped context

vprotocol axis 2.5: the dedicated reader eagerly drains the socket into
promptCh, so prompts queued behind a running one were read before a
disconnect (unlike the old code, which left them unread in the socket
buffer). On disconnect they would then run with a still-live context —
real LLM calls streaming to a dead socket. Derive prompt contexts from a
connection-scoped context the reader cancels on exit, so a disconnect
aborts the in-flight prompt and drains queued ones without LLM calls.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

---------

Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com>

Author: jkyberneees

cmd/odek/serve.go

@@ -485,19 +485,84 @@ func handleWS(store *session.Store, resources *resource.Registry, resolved confi
 	ctx, cancel := signal.NotifyContext(context.Background(), os.Interrupt)
 	defer cancel()
 
+	// Connection-scoped context, cancelled the moment the socket reader exits
+	// (disconnect or shutdown). Prompts run under this context, so a prompt
+	// in flight when the client disconnects is aborted promptly, and any
+	// prompts already buffered in promptCh run against a cancelled context
+	// instead of making real LLM calls whose output streams to a dead socket.
+	connCtx, connCancel := context.WithCancel(ctx)
+	defer connCancel()
+
+	// Dedicated socket reader.
+	//
+	// The agent loop (handlePrompt → RunWithMessages) runs synchronously in the
+	// processor loop below and BLOCKS whenever a tool needs approval — the
+	// wsApprover waits for the browser's approval_response. If the same
+	// goroutine both ran the agent and read the socket, it could never read
+	// that response: every approval would dead-block until the 60s timeout and
+	// be denied, making the Web UI's safety prompt unusable.
+	//
+	// So a separate goroutine owns the socket. It handles approval responses
+	// INLINE (delivering them to the blocked PromptCommand via HandleResponse)
+	// and forwards every other message to promptCh for serial processing. The
+	// forward is non-blocking: the reader must never stall, or a full queue
+	// would re-introduce the deadlock by stopping approval delivery. A buffered
+	// queue absorbs the request/reply UI's normal pacing; an overflow only
+	// happens under a client flooding prompts, which is reported and dropped.
+	promptCh := make(chan []byte, 8)
+	go func() {
+		defer close(promptCh)
+		for {
+			var data []byte
+			if err := golangws.Message.Receive(conn, &data); err != nil {
+				// Socket gone: cancel in-flight/queued prompt work and release
+				// any pending approval so a blocked handlePrompt returns now
+				// instead of waiting out the 60s timeout.
+				connCancel()
+				if approver != nil {
+					approver.Cancel()
+				}
+				return
+			}
+
+			var msgType struct {
+				Type string `json:"type"`
+			}
+			if err := json.Unmarshal(data, &msgType); err != nil {
+				continue
+			}
+
+			// Approval responses are handled here, off the processor goroutine,
+			// so a prompt blocked awaiting approval can be unblocked. This is
+			// the crux of the deadlock fix.
+			if msgType.Type == "approval_response" {
+				var resp approvalResponse
+				if err := json.Unmarshal(data, &resp); err == nil {
+					approver.HandleResponse(resp.ID, resp.Action)
+				}
+				continue
+			}
+
+			select {
+			case promptCh <- data:
+			default:
+				// Processor is busy and the queue is full — only reachable by a
+				// client sending prompts faster than they can run.
+				if msgType.Type == "prompt" {
+					writeWSError(conn, "busy: a prompt is already running")
+				}
+			}
+		}
+	}()
+
 	// Track the current session and model across WebSocket messages
 	var currentSession *session.Session
 	currentModel := resolved.Model // start with configured model
 
 	// Session-level token economics (cumulative across all turns)
 	var sessionInputTokens, sessionOutputTokens int
 
-	for {
-		var data []byte
-		if err := golangws.Message.Receive(conn, &data); err != nil {
-			break
-		}
-
+	for data := range promptCh {
 		// Peek at the message type without full unmarshal
 		var msgType struct {
 			Type string `json:"type"`
@@ -506,15 +571,6 @@ func handleWS(store *session.Store, resources *resource.Registry, resolved confi
 			continue
 		}
 
-		// Handle approval responses separately (non-blocking, from the browser)
-		if msgType.Type == "approval_response" {
-			var resp approvalResponse
-			if err := json.Unmarshal(data, &resp); err == nil {
-				approver.HandleResponse(resp.ID, resp.Action)
-			}
-			continue
-		}
-
 		// Handle skill prompt responses (Save/Skip from skill suggestions)
 		if msgType.Type == "skill_prompt_response" {
 			var resp struct {
@@ -578,8 +634,9 @@ func handleWS(store *session.Store, resources *resource.Registry, resolved confi
 		}
 
 		// Run prompt — passes the persistent agent for buffer continuity
-		// Create a cancelable context for this prompt (so POST /api/cancel can abort it)
-		promptCtx, promptCancel := context.WithCancel(ctx)
+		// Create a cancelable context for this prompt (so POST /api/cancel can abort it).
+		// Derived from connCtx so a disconnect also aborts the running prompt.
+		promptCtx, promptCancel := context.WithCancel(connCtx)
 		currentPromptCancel.Store(promptCancel)
 
 		currentSession = handlePrompt(promptCtx, conn, store, resources, resolved, agent, currentSession, msg.Content, msg.SessionID, &sessionInputTokens, &sessionOutputTokens)

cmd/odek/serve_approval_test.go

@@ -0,0 +1,185 @@
+package main
+
+import (
+	"encoding/json"
+	"fmt"
+	"net"
+	"net/http"
+	"os"
+	"path/filepath"
+	"sync"
+	"testing"
+	"time"
+
+	"github.com/BackendStack21/odek/internal/config"
+	"github.com/BackendStack21/odek/internal/danger"
+	"github.com/BackendStack21/odek/internal/resource"
+	"github.com/BackendStack21/odek/internal/session"
+	golangws "golang.org/x/net/websocket"
+)
+
+// buildServeMuxPromptAll is buildServeMux with the danger policy forced to
+// "prompt" for every risk class, so even a benign `echo` shell command routes
+// through the wsApprover. This lets the test drive a real browser-approval
+// round trip without running a genuinely dangerous command.
+func buildServeMuxPromptAll(t *testing.T, store *session.Store) (net.Listener, *http.ServeMux) {
+	t.Helper()
+	ln, err := net.Listen("tcp", "127.0.0.1:0")
+	if err != nil {
+		t.Fatalf("listen: %v", err)
+	}
+
+	resolved := config.LoadConfig(config.CLIFlags{})
+	systemMessage := resolved.System
+	if systemMessage == "" {
+		systemMessage = defaultSystem
+	}
+	prompt := "prompt"
+	resolved.Dangerous.DefaultAction = &prompt
+
+	cwd, _ := os.Getwd()
+	home, _ := os.UserHomeDir()
+	resourceReg := resource.NewRegistry(
+		resource.NewFileResolver(cwd),
+		resource.NewSessionResolver(filepath.Join(home, ".odek", "sessions")),
+	)
+
+	mux := http.NewServeMux()
+	mux.HandleFunc("/", handleStatic())
+	mux.Handle("/ws", &golangws.Server{
+		Handshake: func(*golangws.Config, *http.Request) error { return nil },
+		Handler: func(conn *golangws.Conn) {
+			handleWS(store, resourceReg, resolved, systemMessage, conn)
+		},
+	})
+	mux.HandleFunc("/api/cancel", handleCancel)
+
+	return ln, mux
+}
+
+// TestServe_E2E_ApprovalRoundTrip is the regression test for the serve-mode
+// approval deadlock. The agent calls a tool that requires approval and blocks
+// in PromptCommand; the browser's approval_response must be read and delivered
+// while that prompt is in flight. Before the fix the socket reader and the
+// agent loop were the same goroutine, so the response could never be read and
+// every approval hit the 60s timeout. The 10s read deadline here is well under
+// that timeout: a regression would blow the deadline, the fix completes in ms.
+func TestServe_E2E_ApprovalRoundTrip(t *testing.T) {
+	llmSrv := mockLLM(t, func(w http.ResponseWriter, callCount int) {
+		w.Header().Set("Content-Type", "application/json")
+		if callCount <= 1 {
+			fmt.Fprint(w, `{"choices":[{"message":{"content":"Running.","tool_calls":[{"id":"c_1","function":{"name":"shell","arguments":"{\"command\":\"echo approved-ok\"}"}}]}}]}`)
+		} else {
+			w.Write([]byte(`{"choices":[{"message":{"content":"All done."}}]}`))
+		}
+	})
+	defer llmSrv.Close()
+
+	envCleanup := setTestEnv(t, llmSrv.URL)
+	defer envCleanup()
+
+	store := newTestSessionStore(t)
+	ln, mux := buildServeMuxPromptAll(t, store)
+	defer ln.Close()
+
+	go func() { _ = serveOnListener(ln, mux) }()
+	waitForHTTP(t, ln.Addr().String())
+
+	wsURL := "ws://" + ln.Addr().String() + "/ws"
+	conn, err := golangws.Dial(wsURL, "", "http://localhost")
+	if err != nil {
+		t.Fatalf("Dial(%q): %v", wsURL, err)
+	}
+	defer conn.Close()
+
+	prompt := map[string]string{"type": "prompt", "content": "run a command"}
+	payload, _ := json.Marshal(prompt)
+	if err := golangws.Message.Send(conn, string(payload)); err != nil {
+		t.Fatalf("Send: %v", err)
+	}
+
+	// 10s ceiling — far below the 60s approval timeout a deadlock would hit.
+	conn.SetReadDeadline(time.Now().Add(10 * time.Second))
+
+	var sawApprovalRequest, sawToolResult, sawDone bool
+	for i := 0; i < 30; i++ {
+		var raw []byte
+		if err := golangws.Message.Receive(conn, &raw); err != nil {
+			t.Fatalf("Receive event %d: %v (approval_request=%v tool_result=%v)", i, err, sawApprovalRequest, sawToolResult)
+		}
+		var evt map[string]any
+		if err := json.Unmarshal(raw, &evt); err != nil {
+			continue
+		}
+		switch evt["type"] {
+		case "approval_request":
+			sawApprovalRequest = true
+			id, _ := evt["id"].(string)
+			resp, _ := json.Marshal(map[string]string{
+				"type":   "approval_response",
+				"id":     id,
+				"action": "approve",
+			})
+			if err := golangws.Message.Send(conn, string(resp)); err != nil {
+				t.Fatalf("send approval_response: %v", err)
+			}
+		case "tool_result":
+			sawToolResult = true
+		case "done":
+			sawDone = true
+			goto finished
+		case "error":
+			t.Fatalf("unexpected error event: %v", evt["message"])
+		}
+	}
+finished:
+	if !sawApprovalRequest {
+		t.Fatal("never received an approval_request — the tool did not prompt")
+	}
+	if !sawToolResult {
+		t.Error("approval was processed but the tool never produced a result")
+	}
+	if !sawDone {
+		t.Fatal("never reached done — approval round trip did not complete (deadlock regression?)")
+	}
+}
+
+// TestWSApprover_ConcurrentTrust_NoRace drives many PromptCommand calls
+// concurrently, each answered with "trust", so the trust-cache write
+// (approveAll[cls]=true) races against the trust-cache read at the top of
+// PromptCommand. Before the fix these map accesses were unsynchronised: this
+// test crashes with "concurrent map writes" (and is flagged under -race).
+func TestWSApprover_ConcurrentTrust_NoRace(t *testing.T) {
+	a := newWSApprover(nil)
+	// Answer every prompt with an immediate "trust" from a separate goroutine,
+	// forcing the approveAll write path under concurrency.
+	a.sendFn = func(v any) error {
+		if req, ok := v.(approvalRequest); ok {
+			go a.HandleResponse(req.ID, "trust")
+		}
+		return nil
+	}
+
+	classes := []danger.RiskClass{
+		danger.Safe, danger.LocalWrite, danger.SystemWrite,
+		danger.NetworkEgress, danger.CodeExecution, danger.Install,
+	}
+
+	var wg sync.WaitGroup
+	for i := 0; i < 200; i++ {
+		for _, cls := range classes {
+			wg.Add(1)
+			go func(c danger.RiskClass) {
+				defer wg.Done()
+				_ = a.PromptCommand(c, "cmd", "desc")
+			}(cls)
+		}
+		// Interleave SetTrustAll to also exercise the trustAll field race.
+		wg.Add(1)
+		go func() {
+			defer wg.Done()
+			a.SetTrustAll(false)
+		}()
+	}
+	wg.Wait()
+}

cmd/odek/wsapprover.go

@@ -114,7 +114,15 @@ func (a *wsApprover) shouldFriction(cls danger.RiskClass) (bool, int) {
 }
 
 func (a *wsApprover) PromptCommand(cls danger.RiskClass, cmd, description string) error {
-	if a.approveAll[cls] || a.trustAll {
+	// Read the trust caches under the lock: PromptCommand runs concurrently
+	// across parallel tool calls, while HandleResponse's "trust" path and
+	// SetTrustAll write these same fields. An unsynchronised read here against
+	// the map write below (or SetTrustAll) is a data race that can fatally
+	// crash serve with "concurrent map read and map write".
+	a.mu.Lock()
+	trusted := a.approveAll[cls] || a.trustAll
+	a.mu.Unlock()
+	if trusted {
 		return nil
 	}
 
@@ -172,7 +180,9 @@ func (a *wsApprover) PromptCommand(cls danger.RiskClass, cmd, description string
 				a.recordApproval(cls)
 				return nil
 			}
+			a.mu.Lock()
 			a.approveAll[cls] = true
+			a.mu.Unlock()
 			return nil
 		default:
 			return fmt.Errorf("operation denied by user: %s", cmd)