From 8467dd94f8fb19b906226a2b238d4efbe0aa9862 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Carlos=20Flor=C3=AAncio?= Date: Tue, 19 May 2026 21:43:45 +0100 Subject: [PATCH] fix(ws): detach dispatched handler ctx from WS connection lifetime (#959) --- .../internal/gateway/websocket/client.go | 33 ++- .../websocket/client_dispatch_ctx_test.go | 191 ++++++++++++++++++ .../backend/internal/gateway/websocket/hub.go | 25 +++ .../internal/orchestrator/task_operations.go | 9 +- .../orchestrator/task_operations_test.go | 87 ++++++++ 5 files changed, 336 insertions(+), 9 deletions(-) create mode 100644 apps/backend/internal/gateway/websocket/client_dispatch_ctx_test.go diff --git a/apps/backend/internal/gateway/websocket/client.go b/apps/backend/internal/gateway/websocket/client.go index d2e8db28c3..f9a039dfb6 100644 --- a/apps/backend/internal/gateway/websocket/client.go +++ b/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), diff --git a/apps/backend/internal/gateway/websocket/client_dispatch_ctx_test.go b/apps/backend/internal/gateway/websocket/client_dispatch_ctx_test.go new file mode 100644 index 0000000000..49b036e98d --- /dev/null +++ b/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") + } +} diff --git a/apps/backend/internal/gateway/websocket/hub.go b/apps/backend/internal/gateway/websocket/hub.go index 41e98539a8..4c2c803ff3 100644 --- a/apps/backend/internal/gateway/websocket/hub.go +++ b/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 diff --git a/apps/backend/internal/orchestrator/task_operations.go b/apps/backend/internal/orchestrator/task_operations.go index cddb33372d..7c934e5df5 100644 --- a/apps/backend/internal/orchestrator/task_operations.go +++ b/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), diff --git a/apps/backend/internal/orchestrator/task_operations_test.go b/apps/backend/internal/orchestrator/task_operations_test.go index 76c81545ce..5f6d85ba5b 100644 --- a/apps/backend/internal/orchestrator/task_operations_test.go +++ b/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) {