fix(orchestrator): keep resume readiness waits alive under caller cancellation (#1582)

Co-authored-by: Carlos Florencio <carlosmflorencio@gmail.com>

apps/backend/internal/orchestrator/prompt_idle_session_test.go

@@ -338,6 +338,180 @@ func TestEnsureSessionRunning_WaitsForPromptReadyAfterResume(t *testing.T) {
 	}
 }
 
+// TestEnsureSessionRunning_SurvivesCallerContextCancellationDuringResume
+// pins down the fix for the "context deadline exceeded" stuck-session bug
+// (kdlbs/kandev#1578): the resume call inside ensureSessionRunning is
+// deliberately shielded from the caller's context via
+// context.WithoutCancel (its own comment explains why: a WebSocket/MCP
+// request deadline shouldn't abort an in-flight resume). Before the fix,
+// waitForSessionReady/waitForAgentPromptReady right after the resume still
+// used the original (cancelable) ctx, so cancelling the caller's context
+// mid-resume killed the wait immediately with a misleading
+// "context deadline exceeded" even though the agent was about to become
+// ready. This test cancels ctx once the resume is underway (after
+// LaunchAgent fires) and asserts ensureSessionRunning still succeeds once
+// the agent reports ready shortly after, exactly as a real resume would.
+func TestEnsureSessionRunning_SurvivesCallerContextCancellationDuringResume(t *testing.T) {
+	ctx, cancel := context.WithCancel(context.Background())
+	t.Cleanup(cancel)
+	repo := setupTestRepo(t)
+
+	seedTaskAndSession(t, repo, "task1", "session1", models.TaskSessionStateIdle)
+	session, err := repo.GetTaskSession(ctx, "session1")
+	if err != nil {
+		t.Fatalf("failed to load session: %v", err)
+	}
+	session.AgentExecutionID = "exec-idle-1"
+	session.AgentProfileID = "profile1"
+	if err := repo.UpdateTaskSession(ctx, session); err != nil {
+		t.Fatalf("failed to update session: %v", err)
+	}
+	seedExecutorRunning(t, repo, session.ID, session.TaskID, "exec-idle-1")
+
+	launched := make(chan struct{}, 1)
+	ready := make(chan struct{})
+	agentMgr := &mockAgentManager{
+		isAgentRunning:         false,
+		repoForExecutionLookup: repo,
+		isAgentReadyFn: func(_ context.Context, _ string) bool {
+			select {
+			case <-ready:
+				return true
+			default:
+				return false
+			}
+		},
+		launchAgentFunc: func(_ context.Context, req *executor.LaunchAgentRequest) (*executor.LaunchAgentResponse, error) {
+			select {
+			case launched <- struct{}{}:
+			default:
+			}
+			go func(sessID string) {
+				tick := time.NewTicker(5 * time.Millisecond)
+				defer tick.Stop()
+				timeout := time.After(5 * time.Second)
+				for {
+					select {
+					case <-tick.C:
+						sess, err := repo.GetTaskSession(context.Background(), sessID)
+						if err == nil && sess != nil && sess.State == models.TaskSessionStateStarting {
+							sess.State = models.TaskSessionStateWaitingForInput
+							sess.UpdatedAt = time.Now().UTC()
+							_ = repo.UpdateTaskSession(context.Background(), sess)
+							return
+						}
+					case <-timeout:
+						return
+					}
+				}
+			}(req.SessionID)
+			return &executor.LaunchAgentResponse{AgentExecutionID: "exec-resumed-1"}, nil
+		},
+	}
+
+	taskRepo := newMockTaskRepo()
+	taskRepo.tasks["task1"] = &v1.Task{
+		ID:    "task1",
+		Title: "Test Task",
+		State: v1.TaskStateInProgress,
+	}
+	svc := createTestServiceWithAgent(repo, newMockStepGetter(), taskRepo, agentMgr)
+	svc.executor = executor.NewExecutor(agentMgr, repo, testLogger(), executor.ExecutorConfig{})
+
+	session, err = repo.GetTaskSession(ctx, "session1")
+	if err != nil {
+		t.Fatalf("reload: %v", err)
+	}
+
+	done := make(chan error, 1)
+	go func() {
+		done <- svc.ensureSessionRunning(ctx, "session1", session)
+	}()
+
+	// Simulate the caller's request context expiring mid-resume (e.g. a
+	// WebSocket handler timeout or an MCP tool-call deadline) — this must
+	// NOT abort the resume/ready-wait, since the resume path routes through
+	// context.WithoutCancel internally.
+	select {
+	case <-launched:
+	case <-time.After(3 * time.Second):
+		t.Fatal("expected ResumeSession to call LaunchAgent")
+	}
+	cancel()
+
+	// The agent becomes ready shortly after the caller's context is gone.
+	close(ready)
+
+	select {
+	case err := <-done:
+		if err != nil {
+			t.Fatalf("ensureSessionRunning must survive caller context cancellation mid-resume, got: %v", err)
+		}
+	case <-time.After(2 * time.Second):
+		t.Fatal("ensureSessionRunning did not return after prompt readiness")
+	}
+}
+
+func TestEnsureSessionRunning_FailedStateWriteSurvivesCancelledCallerCtx(t *testing.T) {
+	ctx, cancel := context.WithCancel(context.Background())
+	t.Cleanup(cancel)
+	repo := setupTestRepo(t)
+	mockTR := newMockTaskRepo()
+	taskRepo := &ctxAwareTaskRepo{inner: mockTR}
+
+	agentMgr := &mockAgentManager{
+		isAgentRunning:         false,
+		repoForExecutionLookup: repo,
+		launchAgentFunc: func(_ context.Context, _ *executor.LaunchAgentRequest) (*executor.LaunchAgentResponse, error) {
+			cancel()
+			return nil, errors.New("simulated resume failure")
+		},
+	}
+
+	svc := createTestServiceWithAgent(repo, newMockStepGetter(), mockTR, agentMgr)
+	svc.executor = executor.NewExecutor(agentMgr, repo, testLogger(), executor.ExecutorConfig{})
+	svc.taskRepo = taskRepo
+
+	seedTaskAndSession(t, repo, "task1", "session1", models.TaskSessionStateIdle)
+	session, err := repo.GetTaskSession(ctx, "session1")
+	if err != nil {
+		t.Fatalf("failed to load session: %v", err)
+	}
+	session.AgentExecutionID = "exec-idle-1"
+	session.AgentProfileID = "profile1"
+	if err := repo.UpdateTaskSession(ctx, session); err != nil {
+		t.Fatalf("failed to update session: %v", err)
+	}
+	seedExecutorRunning(t, repo, session.ID, session.TaskID, "exec-idle-1")
+
+	err = svc.ensureSessionRunning(ctx, "session1", session)
+	if err == nil {
+		t.Fatal("expected ensureSessionRunning to return the simulated resume failure")
+	}
+	if !strings.Contains(err.Error(), "simulated resume failure") {
+		t.Fatalf("expected simulated resume failure, got %v", err)
+	}
+	if ctx.Err() == nil {
+		t.Fatal("test did not cancel the caller context")
+	}
+
+	state, ok := mockTR.updatedStates["task1"]
+	if !ok {
+		t.Fatal("task FAILED state was NOT persisted; the failure-recording write used the cancelled 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)
+	}
+}
+
 func TestEnsureSessionRunning_WaitsForPromptReadyWhenExecutionExists(t *testing.T) {
 	ctx := context.Background()
 	repo := setupTestRepo(t)

apps/backend/internal/orchestrator/task_operations.go

@@ -1059,7 +1059,7 @@ func (s *Service) ResumeTaskSession(ctx context.Context, taskID, sessionID strin
 	if err != nil {
 		// If the execution is already running (duplicate resume request), return it as success.
 		if errors.Is(err, executor.ErrExecutionAlreadyRunning) {
-			execution, readySession, err = s.recoverAlreadyRunningResume(ctx, resumeCtx, taskID, sessionID)
+			execution, readySession, err = s.recoverAlreadyRunningResume(resumeCtx, taskID, sessionID)
 			if err != nil && errors.Is(err, ErrAgentNotReadyForPrompt) {
 				return nil, err
 			}
@@ -1093,7 +1093,7 @@ func (s *Service) ResumeTaskSession(ctx context.Context, taskID, sessionID strin
 		}
 	}
 	if readySession == nil {
-		readySession, err = s.waitForResumedSessionReady(ctx, sessionID)
+		readySession, err = s.waitForResumedSessionReady(resumeCtx, sessionID)
 		if err != nil {
 			return nil, err
 		}
@@ -1134,7 +1134,6 @@ func (s *Service) waitForResumedSessionReady(ctx context.Context, sessionID stri
 }
 
 func (s *Service) recoverAlreadyRunningResume(
-	ctx context.Context,
 	resumeCtx context.Context,
 	taskID string,
 	sessionID string,
@@ -1144,7 +1143,7 @@ func (s *Service) recoverAlreadyRunningResume(
 		return nil, nil, executor.ErrExecutionAlreadyRunning
 	}
 
-	readySession, waitErr := s.waitForResumedSessionReady(ctx, sessionID)
+	readySession, waitErr := s.waitForResumedSessionReady(resumeCtx, sessionID)
 	if waitErr == nil {
 		existing.SessionState = v1.TaskSessionState(readySession.State)
 		return existing, readySession, nil
@@ -1169,7 +1168,7 @@ func (s *Service) recoverAlreadyRunningResume(
 	if err != nil {
 		return nil, nil, err
 	}
-	readySession, err = s.waitForResumedSessionReady(ctx, sessionID)
+	readySession, err = s.waitForResumedSessionReady(resumeCtx, sessionID)
 	if err != nil {
 		return nil, nil, err
 	}
@@ -1346,14 +1345,14 @@ func (s *Service) ensureSessionRunning(ctx context.Context, sessionID string, se
 	resumeCtx := context.WithoutCancel(ctx)
 	if _, err = s.executor.ResumeSession(resumeCtx, session, true); err != nil {
 		if errors.Is(err, executor.ErrExecutionAlreadyRunning) {
-			s.recoverAgentPromptStreamIfNeeded(ctx, sessionID)
+			s.recoverAgentPromptStreamIfNeeded(resumeCtx, sessionID)
-			if readyErr := s.waitForAgentPromptReady(ctx, sessionID); readyErr != nil {
+			if readyErr := s.waitForAgentPromptReady(resumeCtx, sessionID); readyErr != nil {
 				return fmt.Errorf("agent not ready after resume race: %w", readyErr)
 			}
 			return nil
 		}
-		s.updateTaskSessionState(ctx, session.TaskID, sessionID, models.TaskSessionStateFailed, err.Error(), false, session)
+		s.updateTaskSessionState(resumeCtx, session.TaskID, sessionID, models.TaskSessionStateFailed, err.Error(), false, session)
-		if stateErr := s.taskRepo.UpdateTaskState(ctx, session.TaskID, v1.TaskStateFailed); stateErr != nil {
+		if stateErr := s.taskRepo.UpdateTaskState(resumeCtx, session.TaskID, v1.TaskStateFailed); stateErr != nil {
 			s.logger.Warn("failed to update task state to FAILED after session ensure resume error",
 				zap.String("task_id", session.TaskID),
 				zap.String("session_id", sessionID),
@@ -1366,10 +1365,19 @@ func (s *Service) ensureSessionRunning(ctx context.Context, sessionID string, se
 
 	// ResumeSession launches the agent asynchronously. Wait for it to finish
 	// initializing before returning, so the caller can send a prompt immediately.
-	if err := s.waitForSessionReady(ctx, sessionID); err != nil {
+	//
+	// Use resumeCtx (context.WithoutCancel) here too, not the original ctx: the
+	// resume itself is already shielded from the caller's request deadline (see
+	// comment above), but a short-lived caller context (WebSocket request,
+	// MCP tool-call timeout, etc.) would otherwise still abort these polling
+	// waits early with a misleading "context deadline exceeded" even though the
+	// resume is progressing fine in the background and would succeed within its
+	// own bounded timeouts (waitForSessionReady's 90s, waitForAgentPromptReady's
+	// 30s below).
+	if err := s.waitForSessionReady(resumeCtx, sessionID); err != nil {
 		return fmt.Errorf("session not ready after resume: %w", err)
 	}
-	if err := s.waitForAgentPromptReady(ctx, sessionID); err != nil {
+	if err := s.waitForAgentPromptReady(resumeCtx, sessionID); err != nil {
 		return fmt.Errorf("agent not ready after resume: %w", err)
 	}
 
@@ -1471,10 +1479,13 @@ func (s *Service) startAgentOnPreparedWorkspace(ctx context.Context, sessionID s
 		return fmt.Errorf("failed to start agent on prepared workspace: %w", err)
 	}
 
-	if err := s.waitForSessionReady(ctx, sessionID); err != nil {
+	// Same reasoning as ensureSessionRunning's resume path: use launchCtx
+	// (already WithoutCancel'd for the launch call above) so a short-lived
+	// caller context doesn't abort these polling waits early.
+	if err := s.waitForSessionReady(launchCtx, sessionID); err != nil {
 		return fmt.Errorf("session not ready after starting agent: %w", err)
 	}
-	if err := s.waitForAgentPromptReady(ctx, sessionID); err != nil {
+	if err := s.waitForAgentPromptReady(launchCtx, sessionID); err != nil {
 		return fmt.Errorf("agent not ready after starting agent: %w", err)
 	}
 	s.logger.Debug("agent started on prepared workspace and ready for prompt")
@@ -1487,17 +1498,27 @@ func (s *Service) waitForSessionReady(ctx context.Context, sessionID string) err
 		pollInterval = 500 * time.Millisecond
 		maxWait      = 90 * time.Second
 	)
-	deadline := time.Now().Add(maxWait)
+	// Derive a bounded context so the overall wait AND each GetTaskSession query
+	// inside the loop honor maxWait. Callers pass context.WithoutCancel(ctx) so
+	// the wait survives the caller's request timeout during resume/launch — but
+	// that context carries no deadline of its own, so without this a single
+	// blocking query could hang well past maxWait (the loop only checks the
+	// wall clock between iterations). Mirrors waitForAgentPromptReady.
+	readyCtx, cancel := context.WithTimeout(ctx, maxWait)
+	defer cancel()
+	// Ticker rather than time.After(pollInterval) in the loop: time.After
+	// allocates a fresh timer each iteration that lives until it fires, so a
+	// long wait would pile up ~maxWait/pollInterval live timers. Mirrors
+	// waitForAgentPromptReady.
+	ticker := time.NewTicker(pollInterval)
+	defer ticker.Stop()
 	for {
-		if time.Now().After(deadline) {
-			return fmt.Errorf("timeout waiting for agent to become ready")
-		}
 		select {
-		case <-ctx.Done():
+		case <-readyCtx.Done():
-			return ctx.Err()
+			return fmt.Errorf("timeout waiting for agent to become ready: %w", readyCtx.Err())
-		case <-time.After(pollInterval):
+		case <-ticker.C:
 		}
-		sess, err := s.repo.GetTaskSession(ctx, sessionID)
+		sess, err := s.repo.GetTaskSession(readyCtx, sessionID)
 		if err != nil {
 			return fmt.Errorf("failed to check session state: %w", err)
 		}

apps/backend/internal/orchestrator/task_operations_resume_test.go

@@ -648,3 +648,127 @@ func TestResumeTaskSession_ReapStopFailureDoesNotFailTask(t *testing.T) {
 		t.Fatalf("expected no task state writes on transient stop failure, got %d", got)
 	}
 }
+
+func TestResumeTaskSession_SurvivesCallerContextCancellationDuringReadyWait(t *testing.T) {
+	ctx, cancel := context.WithCancel(context.Background())
+	t.Cleanup(cancel)
+	repo := setupTestRepo(t)
+	seedTaskAndSession(t, repo, "task1", "session1", models.TaskSessionStateWaitingForInput)
+
+	session, err := repo.GetTaskSession(ctx, "session1")
+	if err != nil {
+		t.Fatalf("failed to load session: %v", err)
+	}
+	session.AgentExecutionID = "exec-old-1"
+	session.AgentProfileID = "profile1"
+	if err := repo.UpdateTaskSession(ctx, session); err != nil {
+		t.Fatalf("failed to update session: %v", err)
+	}
+	seedExecutorRunning(t, repo, session.ID, session.TaskID, "exec-old-1")
+
+	agentMgr := &mockAgentManager{
+		isAgentRunning:         false,
+		repoForExecutionLookup: repo,
+		isAgentReadyFn: func(_ context.Context, _ string) bool {
+			return true
+		},
+		launchAgentFunc: func(_ context.Context, req *executor.LaunchAgentRequest) (*executor.LaunchAgentResponse, error) {
+			cancel()
+			go func(sessID string) {
+				tick := time.NewTicker(5 * time.Millisecond)
+				defer tick.Stop()
+				timeout := time.After(5 * time.Second)
+				for {
+					select {
+					case <-tick.C:
+						sess, err := repo.GetTaskSession(context.Background(), sessID)
+						if err == nil && sess != nil && sess.State == models.TaskSessionStateStarting {
+							sess.State = models.TaskSessionStateWaitingForInput
+							sess.UpdatedAt = time.Now().UTC()
+							_ = repo.UpdateTaskSession(context.Background(), sess)
+							return
+						}
+					case <-timeout:
+						return
+					}
+				}
+			}(req.SessionID)
+			return &executor.LaunchAgentResponse{AgentExecutionID: "exec-resumed-1"}, nil
+		},
+	}
+
+	taskRepo := newMockTaskRepo()
+	taskRepo.tasks["task1"] = &v1.Task{
+		ID:    "task1",
+		Title: "Test Task",
+		State: v1.TaskStateInProgress,
+	}
+	svc := createTestServiceWithAgent(repo, newMockStepGetter(), taskRepo, agentMgr)
+	svc.executor = executor.NewExecutor(agentMgr, repo, testLogger(), executor.ExecutorConfig{})
+
+	exec, err := svc.ResumeTaskSession(ctx, "task1", "session1")
+	if err != nil {
+		t.Fatalf("ResumeTaskSession must survive caller context cancellation mid-resume: %v", err)
+	}
+	if exec == nil {
+		t.Fatal("ResumeTaskSession returned nil execution")
+	}
+	if ctx.Err() == nil {
+		t.Fatal("test did not cancel the caller context")
+	}
+	if exec.SessionState != v1.TaskSessionStateWaitingForInput {
+		t.Fatalf("expected WAITING_FOR_INPUT response state, got %s", exec.SessionState)
+	}
+}
+
+func TestResumeTaskSession_AlreadyRunningReadyWaitSurvivesCancelledCallerContext(t *testing.T) {
+	ctx, cancel := context.WithCancel(context.Background())
+	t.Cleanup(cancel)
+	repo := setupTestRepo(t)
+	seedTaskAndSession(t, repo, "task1", "session1", models.TaskSessionStateWaitingForInput)
+
+	session, err := repo.GetTaskSession(ctx, "session1")
+	if err != nil {
+		t.Fatalf("failed to load session: %v", err)
+	}
+	session.AgentExecutionID = "exec-old-1"
+	session.AgentProfileID = "profile1"
+	if err := repo.UpdateTaskSession(ctx, session); err != nil {
+		t.Fatalf("failed to update session: %v", err)
+	}
+	seedExecutorRunning(t, repo, session.ID, session.TaskID, "exec-old-1")
+
+	agentMgr := &mockAgentManager{
+		repoForExecutionLookup: repo,
+		isAgentRunningFn: func(_ context.Context, _ string) bool {
+			cancel()
+			return true
+		},
+		isAgentReadyFn: func(_ context.Context, _ string) bool {
+			return true
+		},
+	}
+
+	taskRepo := newMockTaskRepo()
+	taskRepo.tasks["task1"] = &v1.Task{
+		ID:    "task1",
+		Title: "Test Task",
+		State: v1.TaskStateInProgress,
+	}
+	svc := createTestServiceWithAgent(repo, newMockStepGetter(), taskRepo, agentMgr)
+	svc.executor = executor.NewExecutor(agentMgr, repo, testLogger(), executor.ExecutorConfig{})
+
+	exec, err := svc.ResumeTaskSession(ctx, "task1", "session1")
+	if err != nil {
+		t.Fatalf("already-running resume must survive caller context cancellation: %v", err)
+	}
+	if exec == nil {
+		t.Fatal("ResumeTaskSession returned nil execution")
+	}
+	if ctx.Err() == nil {
+		t.Fatal("test did not cancel the caller context")
+	}
+	if exec.SessionState != v1.TaskSessionStateWaitingForInput {
+		t.Fatalf("expected WAITING_FOR_INPUT response state, got %s", exec.SessionState)
+	}
+}

apps/backend/internal/orchestrator/task_operations_wait_test.go

@@ -0,0 +1,54 @@
+package orchestrator
+
+import (
+	"context"
+	"strings"
+	"testing"
+
+	"github.com/kandev/kandev/internal/task/models"
+)
+
+// deadlineRecordingRepo wraps the real session store and captures whether the
+// context propagated into GetTaskSession carries a deadline. It forces a
+// terminal session state so waitForSessionReady returns on the first poll
+// instead of running its full budget.
+type deadlineRecordingRepo struct {
+	sessionExecutorStore
+	sawCall     bool
+	hadDeadline bool
+}
+
+func (r *deadlineRecordingRepo) GetTaskSession(ctx context.Context, id string) (*models.TaskSession, error) {
+	r.sawCall = true
+	_, r.hadDeadline = ctx.Deadline()
+	return &models.TaskSession{ID: id, State: models.TaskSessionStateFailed, ErrorMessage: "forced-by-test"}, nil
+}
+
+// TestWaitForSessionReady_BoundsQueriesWithNonCancellableContext guards the fix
+// for cubic's P2: resume/launch pass context.WithoutCancel(ctx) (no deadline) so
+// the wait survives the caller's request timeout, but the GetTaskSession queries
+// inside the poll loop must still be bounded by waitForSessionReady's own
+// deadline — otherwise a single blocking query could hang past the 90s budget,
+// which is only checked between iterations.
+func TestWaitForSessionReady_BoundsQueriesWithNonCancellableContext(t *testing.T) {
+	repo := setupTestRepo(t)
+	svc := createTestService(repo, newMockStepGetter(), newMockTaskRepo())
+	seedTaskAndSession(t, repo, "task1", "session1", models.TaskSessionStateRunning)
+
+	rec := &deadlineRecordingRepo{sessionExecutorStore: svc.repo}
+	svc.repo = rec
+
+	err := svc.waitForSessionReady(context.WithoutCancel(context.Background()), "session1")
+	if err == nil {
+		t.Fatal("expected error from the forced-failed session")
+	}
+	if !strings.Contains(err.Error(), "forced-by-test") {
+		t.Fatalf("expected the forced-failed error, got %v", err)
+	}
+	if !rec.sawCall {
+		t.Fatal("GetTaskSession was never called")
+	}
+	if !rec.hadDeadline {
+		t.Error("GetTaskSession received a context WITHOUT a deadline — queries are unbounded (regression of cubic P2 fix)")
+	}
+}