diff --git a/apps/backend/internal/orchestrator/prompt_idle_session_test.go b/apps/backend/internal/orchestrator/prompt_idle_session_test.go index c37cbd904f..1fa80b14dc 100644 --- a/apps/backend/internal/orchestrator/prompt_idle_session_test.go +++ b/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) diff --git a/apps/backend/internal/orchestrator/task_operations.go b/apps/backend/internal/orchestrator/task_operations.go index f12bed200d..0f546d8f95 100644 --- a/apps/backend/internal/orchestrator/task_operations.go +++ b/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) - if readyErr := s.waitForAgentPromptReady(ctx, sessionID); readyErr != nil { + s.recoverAgentPromptStreamIfNeeded(resumeCtx, sessionID) + 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) - if stateErr := s.taskRepo.UpdateTaskState(ctx, session.TaskID, v1.TaskStateFailed); stateErr != nil { + s.updateTaskSessionState(resumeCtx, session.TaskID, sessionID, models.TaskSessionStateFailed, err.Error(), false, session) + 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(): - return ctx.Err() - case <-time.After(pollInterval): + case <-readyCtx.Done(): + return fmt.Errorf("timeout waiting for agent to become ready: %w", readyCtx.Err()) + 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) } diff --git a/apps/backend/internal/orchestrator/task_operations_resume_test.go b/apps/backend/internal/orchestrator/task_operations_resume_test.go index 8b60fd3e6a..23ccfe83ed 100644 --- a/apps/backend/internal/orchestrator/task_operations_resume_test.go +++ b/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) + } +} diff --git a/apps/backend/internal/orchestrator/task_operations_wait_test.go b/apps/backend/internal/orchestrator/task_operations_wait_test.go new file mode 100644 index 0000000000..9da20d8eb8 --- /dev/null +++ b/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)") + } +}