fix(ws): detach dispatched handler ctx from WS connection lifetime (#959)

apps/backend/internal/gateway/websocket/client.go

@@ -60,8 +60,14 @@ func NewClient(id string, conn *websocket.Conn, hub *Hub, log *logger.Logger) *C
 	}
 }
 
-// ReadPump pumps messages from the WebSocket connection to the hub
-func (c *Client) ReadPump(ctx context.Context) {
+// ReadPump pumps messages from the WebSocket connection to the hub.
+//
+// The ctx argument is retained for API stability but is not consulted by the
+// pump itself — Gorilla's ReadMessage blocks on the conn only, so teardown
+// happens via the conn close path (driven by client disconnect, server
+// shutdown closing all conns, or pong timeout). Dispatched handlers use the
+// hub's lifetime context instead; see handleMessage.
+func (c *Client) ReadPump(_ context.Context) {
 	defer func() {
 		c.hub.Unregister(c)
 		if err := c.conn.Close(); err != nil {
@@ -100,12 +106,16 @@ func (c *Client) ReadPump(ctx context.Context) {
 		// Process the message in a goroutine to avoid blocking the read pump
 		// This allows concurrent message handling so long-running handlers
 		// (like orchestrator.prompt) don't block other requests (like workspace.tree.get)
-		go c.handleMessage(ctx, &msg)
+		go c.handleMessage(&msg)
 	}
 }
 
-// handleMessage processes an incoming message
-func (c *Client) handleMessage(ctx context.Context, msg *ws.Message) {
+// handleMessage processes an incoming message.
+//
+// Intentionally does NOT take the connection context. Dispatched handlers run
+// under the hub's lifetime context so a mid-flight client disconnect doesn't
+// abort in-progress side effects (see the comment on Dispatch below).
+func (c *Client) handleMessage(msg *ws.Message) {
 	c.logger.Debug("Received message",
 		zap.String("action", msg.Action),
 		zap.String("id", msg.ID))
@@ -144,8 +154,17 @@ func (c *Client) handleMessage(ctx context.Context, msg *ws.Message) {
 		return
 	}
 
-	// Dispatch to handler
-	response, err := c.hub.dispatcher.Dispatch(ctx, msg)
+	// Dispatch to handler using the hub's lifetime context, not the per-
+	// connection one. The connection ctx is cancelled when the client
+	// disconnects (page reload, nav, network drop). Using it here would
+	// SIGKILL any exec.CommandContext subprocesses the handler spawned
+	// (e.g. `gh pr`, `git`, agentctl HTTP requests) and abort
+	// side-effecting work like session.launch mid-flight. We can't deliver
+	// the response either way once the connection is gone, but the
+	// handler's work should run to completion so it doesn't leave partial
+	// state behind. The dispatch ctx still cancels on server shutdown.
+	dispatchCtx := c.hub.DispatchContext()
+	response, err := c.hub.dispatcher.Dispatch(dispatchCtx, msg)
 	if err != nil {
 		c.logger.Error("Handler error",
 			zap.String("action", msg.Action),

apps/backend/internal/gateway/websocket/client_dispatch_ctx_test.go

@@ -0,0 +1,191 @@
+package websocket
+
+import (
+	"context"
+	"errors"
+	"testing"
+	"time"
+
+	ws "github.com/kandev/kandev/pkg/websocket"
+)
+
+// TestHub_DispatchContextFallsBackBeforeRun guards the test-setup path:
+// DispatchContext must never return a nil context, even when Run was never
+// started (e.g. tests that just exercise handlers directly).
+func TestHub_DispatchContextFallsBackBeforeRun(t *testing.T) {
+	h := newTestHub(t)
+
+	got := h.DispatchContext()
+	if got == nil {
+		t.Fatal("DispatchContext returned nil before Run; handlers would NPE")
+	}
+	if err := got.Err(); err != nil {
+		t.Fatalf("fallback ctx should not be done, got err=%v", err)
+	}
+}
+
+// TestHub_DispatchContextTracksRunCtx checks that DispatchContext returns the
+// ctx Run was called with, so it cancels on server shutdown (the only
+// cancellation reason that should still kill dispatched handlers).
+func TestHub_DispatchContextTracksRunCtx(t *testing.T) {
+	h := newTestHub(t)
+
+	runCtx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+
+	// Set directly via in-package field access — mirrors what Run does without
+	// the goroutine scheduling race.
+	h.mu.Lock()
+	h.dispatchCtx = runCtx
+	h.mu.Unlock()
+
+	dispatchCtx := h.DispatchContext()
+	if dispatchCtx.Err() != nil {
+		t.Fatalf("dispatchCtx prematurely done: %v", dispatchCtx.Err())
+	}
+
+	cancel()
+
+	// Cancellation is synchronous once the parent ctx is cancelled.
+	if dispatchCtx.Err() == nil {
+		t.Fatal("dispatchCtx did not cancel when hub ctx was cancelled")
+	}
+}
+
+// TestClient_HandleMessageUsesHubCtxNotConnCtx is the regression test for the
+// SIGKILL cascade: a dispatched handler must NOT see a cancelled context when
+// the originating client disconnects. Before the fix, the dispatcher was
+// called with the connection ctx, so any in-flight exec.CommandContext
+// subprocess (gh, git, agentctl HTTP) got SIGKILL'd the moment the user
+// navigated away from the page.
+func TestClient_HandleMessageUsesHubCtxNotConnCtx(t *testing.T) {
+	h := newTestHub(t)
+	dispatcher := ws.NewDispatcher()
+	h.dispatcher = dispatcher
+
+	hubCtx, hubCancel := context.WithCancel(context.Background())
+	defer hubCancel()
+
+	// Set directly before Run so the goroutine has no scheduling race.
+	h.mu.Lock()
+	h.dispatchCtx = hubCtx
+	h.mu.Unlock()
+	go h.Run(hubCtx)
+
+	// Handler that records what ctx it received and waits long enough for
+	// the test to cancel a fake connection ctx before returning. If the
+	// dispatcher were still wired to the connection ctx, gotCtx.Err()
+	// below would be non-nil.
+	handlerEntered := make(chan struct{})
+	releaseHandler := make(chan struct{})
+	var gotCtx context.Context
+	dispatcher.RegisterFunc("test.write", func(ctx context.Context, _ *ws.Message) (*ws.Message, error) {
+		gotCtx = ctx
+		close(handlerEntered)
+		<-releaseHandler
+		return nil, nil
+	})
+
+	c := newTestClient("c-disconnect")
+	c.hub = h
+
+	// Run handleMessage in a goroutine so we can cancel the "connection"
+	// while the handler is mid-flight.
+	done := make(chan struct{})
+	go func() {
+		defer close(done)
+		c.handleMessage(&ws.Message{Action: "test.write", Type: ws.MessageTypeRequest})
+	}()
+
+	select {
+	case <-handlerEntered:
+	case <-time.After(time.Second):
+		t.Fatal("handler never entered; dispatch wiring is broken")
+	}
+
+	// The handler is inside Dispatch. In production this is the moment the
+	// WS client disconnects. Verify the handler's ctx is not derived from
+	// any connection lifetime: we don't even have a connection ctx to
+	// cancel here, which is exactly the point — handleMessage no longer
+	// takes one.
+	if gotCtx == nil {
+		t.Fatal("handler did not receive a context")
+	}
+	if err := gotCtx.Err(); err != nil {
+		t.Fatalf("handler ctx already done at entry: %v", err)
+	}
+
+	// Sanity: the handler ctx must be the hub's lifetime ctx (or derived
+	// from it). Cancel the hub and observe. Cancellation is synchronous
+	// because gotCtx IS hubCtx (DispatchContext returns the field directly).
+	hubCancel()
+	if !errors.Is(gotCtx.Err(), context.Canceled) {
+		t.Fatalf("handler ctx should cancel with hub shutdown, got err=%v", gotCtx.Err())
+	}
+
+	close(releaseHandler)
+	<-done
+}
+
+// TestClient_HandleMessageSurvivesConnectionTeardown is the higher-level
+// regression: even though ReadPump may exit and the client may be torn down,
+// already-dispatched handlers continue running. Models the real bug — a
+// session.launch already inside the handler when the WS closes.
+func TestClient_HandleMessageSurvivesConnectionTeardown(t *testing.T) {
+	h := newTestHub(t)
+	dispatcher := ws.NewDispatcher()
+	h.dispatcher = dispatcher
+
+	hubCtx, hubCancel := context.WithCancel(context.Background())
+	defer hubCancel()
+
+	h.mu.Lock()
+	h.dispatchCtx = hubCtx
+	h.mu.Unlock()
+	go h.Run(hubCtx)
+
+	handlerEntered := make(chan struct{})
+	handlerCompleted := make(chan error, 1)
+	releaseHandler := make(chan struct{})
+	dispatcher.RegisterFunc("session.launch.fake", func(ctx context.Context, _ *ws.Message) (*ws.Message, error) {
+		close(handlerEntered)
+		// Mimic an exec.CommandContext-style subroutine that watches ctx.
+		select {
+		case <-releaseHandler:
+			handlerCompleted <- ctx.Err()
+		case <-ctx.Done():
+			handlerCompleted <- ctx.Err()
+		}
+		return nil, nil
+	})
+
+	c := newTestClient("c-teardown")
+	c.hub = h
+
+	go c.handleMessage(&ws.Message{Action: "session.launch.fake", Type: ws.MessageTypeRequest})
+
+	select {
+	case <-handlerEntered:
+	case <-time.After(time.Second):
+		t.Fatal("handler never entered; dispatch wiring is broken")
+	}
+
+	// Simulate connection teardown: close the client's send channel and
+	// drop our hub reference, the way removeClient would. Crucially, this
+	// does NOT cancel the hub ctx — only client-scoped state goes away.
+	c.mu.Lock()
+	c.closed = true
+	c.mu.Unlock()
+
+	// Handler must still be running. Release it explicitly.
+	close(releaseHandler)
+
+	select {
+	case err := <-handlerCompleted:
+		if err != nil {
+			t.Fatalf("handler ctx was cancelled by connection teardown; want nil, got %v", err)
+		}
+	case <-time.After(time.Second):
+		t.Fatal("handler never completed; deadlock or wrong ctx wiring")
+	}
+}

apps/backend/internal/gateway/websocket/hub.go

@@ -45,6 +45,13 @@ type Hub struct {
 	// effective mode (paused/slow/fast) transitions. See hub_session_mode.go.
 	sessionMode *sessionModeTracker
 
+	// dispatchCtx is the hub's lifetime context, set by Run. Dispatched
+	// message handlers use it instead of the per-connection context so that
+	// a client disconnecting mid-flight does not SIGKILL exec subprocesses
+	// (gh, git, agentctl HTTP calls) or otherwise abort side-effecting work
+	// like session.launch. It still cancels on server shutdown.
+	dispatchCtx context.Context
+
 	mu     sync.RWMutex
 	logger *logger.Logger
 }
@@ -71,6 +78,10 @@ func (h *Hub) Run(ctx context.Context) {
 	h.logger.Info("WebSocket hub started")
 	defer h.logger.Info("WebSocket hub stopped")
 
+	h.mu.Lock()
+	h.dispatchCtx = ctx
+	h.mu.Unlock()
+
 	for {
 		select {
 		case <-ctx.Done():
@@ -437,6 +448,20 @@ func (h *Hub) GetDispatcher() *ws.Dispatcher {
 	return h.dispatcher
 }
 
+// DispatchContext returns a context whose lifetime is tied to the hub (and
+// therefore the server) rather than any single client connection. Dispatched
+// handlers should use this so that a client disconnecting mid-flight does not
+// cancel in-progress writes, exec subprocesses, or downstream HTTP calls.
+// Falls back to context.Background when Run has not been called (test setups).
+func (h *Hub) DispatchContext() context.Context {
+	h.mu.RLock()
+	defer h.mu.RUnlock()
+	if h.dispatchCtx == nil {
+		return context.Background()
+	}
+	return h.dispatchCtx
+}
+
 // SetSessionDataProvider sets the provider for session data on subscription
 func (h *Hub) SetSessionDataProvider(provider SessionDataProvider) {
 	h.sessionDataProvider = provider

apps/backend/internal/orchestrator/task_operations.go

@@ -883,8 +883,13 @@ func (s *Service) ResumeTaskSession(ctx context.Context, taskID, sessionID strin
 		if task, taskErr := s.repo.GetTask(resumeCtx, taskID); taskErr == nil && task != nil && task.ArchivedAt != nil {
 			return nil, executor.ErrTaskArchived
 		}
-		s.updateTaskSessionState(ctx, taskID, sessionID, models.TaskSessionStateFailed, err.Error(), false, session)
-		if stateErr := s.taskRepo.UpdateTaskState(ctx, taskID, v1.TaskStateFailed); stateErr != nil {
+		// Use resumeCtx (WithoutCancel) for the failure-recording writes too —
+		// if the caller's ctx was already cancelled (e.g. WS client navigated
+		// away), the SessionStateFailed and TaskStateFailed updates would
+		// themselves fail with "context canceled" and leave the task stuck
+		// looking "running" forever.
+		s.updateTaskSessionState(resumeCtx, taskID, sessionID, models.TaskSessionStateFailed, err.Error(), false, session)
+		if stateErr := s.taskRepo.UpdateTaskState(resumeCtx, taskID, v1.TaskStateFailed); stateErr != nil {
 			s.logger.Warn("failed to update task state to FAILED after resume error",
 				zap.String("task_id", taskID),
 				zap.String("session_id", sessionID),

apps/backend/internal/orchestrator/task_operations_test.go

@@ -878,6 +878,93 @@ func TestResumeTaskSession_FailedKeepsResumeToken(t *testing.T) {
 	}
 }
 
+// ctxAwareTaskRepo wraps mockTaskRepo and respects ctx cancellation. Used to
+// prove that ResumeTaskSession's failure-recording path is insulated from a
+// pre-cancelled caller ctx (the WS-disconnect scenario).
+type ctxAwareTaskRepo struct {
+	inner *mockTaskRepo
+}
+
+func (c *ctxAwareTaskRepo) GetTask(ctx context.Context, taskID string) (*v1.Task, error) {
+	if err := ctx.Err(); err != nil {
+		return nil, err
+	}
+	return c.inner.GetTask(ctx, taskID)
+}
+
+func (c *ctxAwareTaskRepo) UpdateTaskState(ctx context.Context, taskID string, state v1.TaskState) error {
+	if err := ctx.Err(); err != nil {
+		return err
+	}
+	return c.inner.UpdateTaskState(ctx, taskID, state)
+}
+
+// TestResumeTaskSession_FailedStateWriteSurvivesCancelledCallerCtx verifies the
+// fix for the WS-disconnect cascade: when the caller's ctx was already
+// cancelled (e.g. the user navigated away mid-resume) and the launch then
+// failed, the FAILED state-update writes must still go through using the
+// detached resumeCtx — otherwise the task is stuck looking "running" forever.
+//
+// Before the fix, lines 886-892 used the original ctx, so the failure-state
+// write itself returned "context canceled" and the WARN "failed to update task
+// state to FAILED after resume error: context canceled" appeared in the logs.
+func TestResumeTaskSession_FailedStateWriteSurvivesCancelledCallerCtx(t *testing.T) {
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+	repo := setupTestRepo(t)
+	mockTR := newMockTaskRepo()
+	taskRepo := &ctxAwareTaskRepo{inner: mockTR}
+
+	// Cancel the caller ctx the moment launch is invoked. This mirrors the
+	// WS-disconnect race: the request handler's ctx is alive when the
+	// resume path starts (so it gets through the early-exit checks against
+	// sqlite/etc.) and dies mid-launch. The post-launch failure-recording
+	// writes use resumeCtx (WithoutCancel) and must still succeed.
+	agentMgr := &mockAgentManager{
+		launchAgentFunc: func(_ context.Context, _ *executor.LaunchAgentRequest) (*executor.LaunchAgentResponse, error) {
+			cancel()
+			return nil, errors.New("simulated launch failure")
+		},
+	}
+	svc := createTestServiceWithAgent(repo, newMockStepGetter(), mockTR, agentMgr)
+	svc.executor = executor.NewExecutor(agentMgr, repo, testLogger(), executor.ExecutorConfig{})
+	// Override with the ctx-aware wrapper so we can detect ctx-canceled
+	// writes — the bare mockTaskRepo ignores ctx entirely.
+	svc.taskRepo = taskRepo
+
+	seedTaskAndSession(t, repo, "task1", "session1", models.TaskSessionStateWaitingForInput)
+	session, _ := repo.GetTaskSession(ctx, "session1")
+	session.AgentProfileID = "profile-1"
+	_ = repo.UpdateTaskSession(ctx, session)
+
+	now := time.Now().UTC()
+	_ = repo.UpsertExecutorRunning(ctx, &models.ExecutorRunning{
+		ID: "er1", SessionID: "session1", TaskID: "task1",
+		CreatedAt: now, UpdatedAt: now,
+	})
+
+	_, err := svc.ResumeTaskSession(ctx, "task1", "session1")
+	if err == nil {
+		t.Fatal("expected ResumeTaskSession to return an error from the simulated launch failure")
+	}
+
+	state, ok := mockTR.updatedStates["task1"]
+	if !ok {
+		t.Fatal("task FAILED state was NOT persisted; the failure-recording write was cancelled by the caller ctx")
+	}
+	if state != v1.TaskStateFailed {
+		t.Errorf("expected task1 state=FAILED, got %v", state)
+	}
+
+	persisted, getErr := repo.GetTaskSession(context.Background(), "session1")
+	if getErr != nil {
+		t.Fatalf("failed to reload session: %v", getErr)
+	}
+	if persisted.State != models.TaskSessionStateFailed {
+		t.Errorf("expected session1 state=FAILED, got %v", persisted.State)
+	}
+}
+
 // --- CompleteTask ---
 
 func TestCompleteTask_UpdatesTaskState(t *testing.T) {