From 534f84e4552b06fd96ba11bd99bef5990271a599 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Cl=C3=A9ment=20Fiere?= Date: Fri, 7 Aug 2026 16:26:06 +0200 Subject: [PATCH] fix: reap and retry cold-resume prompt-readiness timeout (#2356) --- ...ensure_session_running_cold_resume_test.go | 159 ++++++++++++++++++ .../internal/orchestrator/task_operations.go | 101 ++++++++--- 2 files changed, 235 insertions(+), 25 deletions(-) create mode 100644 apps/backend/internal/orchestrator/ensure_session_running_cold_resume_test.go diff --git a/apps/backend/internal/orchestrator/ensure_session_running_cold_resume_test.go b/apps/backend/internal/orchestrator/ensure_session_running_cold_resume_test.go new file mode 100644 index 0000000000..bd24da13d9 --- /dev/null +++ b/apps/backend/internal/orchestrator/ensure_session_running_cold_resume_test.go @@ -0,0 +1,159 @@ +package orchestrator + +import ( + "context" + "sync" + "sync/atomic" + "testing" + "time" + + "github.com/kandev/kandev/internal/orchestrator/executor" + "github.com/kandev/kandev/internal/task/models" + v1 "github.com/kandev/kandev/pkg/api/v1" +) + +// TestEnsureSessionRunning_ColdResumeReapsAndRetriesOnPromptReadyTimeout pins +// down a self-healing gap specific to the "kandev restart" cold-resume path: +// when the backend restarts, the in-memory execution store is empty even +// though a resumable executors_running row survives in the DB. The very +// first ensureSessionRunning call for such a session (the task-open +// auto-resume in session_ensure.go's tryEnsureExecution, which has no +// external retry — failures are only logged) takes the "no execution +// tracked yet" branch: it calls ResumeSession once and waits once for +// prompt-readiness. +// +// The sibling branch just above it (an execution IS already tracked but +// isn't prompt-ready) self-heals on the same timeout by reaping the stuck +// execution and relaunching — see reapPromptUnreadyExecution's call site a +// few lines up. The cold-resume branch had no equivalent: a launch that +// hangs past agentPromptReadyTimeout on its first attempt returned the bare +// timeout error with no reap-and-retry, even though the exact same recovery +// this test drives (stop the wedged execution, relaunch, wait again) is +// already proven safe elsewhere (TestResumeTaskSession_ReapsPromptDeadExecutionAndRelaunches). +func TestEnsureSessionRunning_ColdResumeReapsAndRetriesOnPromptReadyTimeout(t *testing.T) { + oldReadyTimeout := agentPromptReadyTimeout + oldReadyInterval := agentPromptReadyInterval + agentPromptReadyTimeout = 20 * time.Millisecond + agentPromptReadyInterval = time.Millisecond + t.Cleanup(func() { + agentPromptReadyTimeout = oldReadyTimeout + agentPromptReadyInterval = oldReadyInterval + }) + + ctx := context.Background() + 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.AgentProfileID = "profile1" + if err := repo.UpdateTaskSession(ctx, session); err != nil { + t.Fatalf("failed to update session: %v", err) + } + + // Surviving executors_running row, no in-memory execution — the shape + // left behind by a backend restart. + now := time.Now().UTC() + if err := repo.UpsertExecutorRunning(ctx, &models.ExecutorRunning{ + ID: "er1", + SessionID: "session1", + TaskID: "task1", + AgentExecutionID: "exec-wedged", + Status: "running", + Resumable: true, + ResumeToken: "resume-token-123", + CreatedAt: now, + UpdatedAt: now, + }); err != nil { + t.Fatalf("failed to seed executor running: %v", err) + } + + // Both launch attempts flip the session's DB state to WAITING_FOR_INPUT + // promptly (so waitForSessionReady never falls back to the 6-minute + // constants.AgentLaunchTimeout in this test) — only agent-level + // prompt-readiness (isAgentReadyFn) distinguishes the wedged first + // attempt from the healthy replacement. + // + // The launch goroutines below are bound to the test lifetime via + // launchCtx/launchWG: without that, a goroutine can still be polling + // (and calling into the test repo) after setupTestRepo's own t.Cleanup + // closes the database, which panics or logs after test completion. + // t.Cleanup runs LIFO, so this cancel-and-wait cleanup (registered after + // setupTestRepo's) drains every goroutine before the repo closes. + var launchCalls atomic.Int32 + launchCtx, cancelLaunches := context.WithCancel(ctx) + var launchWG sync.WaitGroup + t.Cleanup(func() { + cancelLaunches() + launchWG.Wait() + }) + agentMgr := &mockAgentManager{ + repoForExecutionLookup: repo, + isAgentRunningFn: func(_ context.Context, _ string) bool { + return false + }, + isAgentReadyFn: func(_ context.Context, _ string) bool { + // Only the reap-and-retry relaunch (attempt 2+) ever reports ready. + return launchCalls.Load() >= 2 + }, + stopAgentWithReasonFunc: func(_ context.Context, _ string, _ string, _ bool) error { + return nil + }, + launchAgentFunc: func(_ context.Context, req *executor.LaunchAgentRequest) (*executor.LaunchAgentResponse, error) { + n := launchCalls.Add(1) + launchWG.Add(1) + go func(sessID string) { + defer launchWG.Done() + tick := time.NewTicker(time.Millisecond) + defer tick.Stop() + timeout := time.After(5 * time.Second) + for { + select { + case <-launchCtx.Done(): + return + case <-tick.C: + sess, err := repo.GetTaskSession(launchCtx, sessID) + if err == nil && sess != nil && sess.State == models.TaskSessionStateStarting { + sess.State = models.TaskSessionStateWaitingForInput + sess.UpdatedAt = time.Now().UTC() + _ = repo.UpdateTaskSession(launchCtx, sess) + return + } + case <-timeout: + return + } + } + }(req.SessionID) + if n == 1 { + // First attempt: session state flips fine, but the agent + // process never reports prompt-ready — the wedged-resume + // shape this test targets. + return &executor.LaunchAgentResponse{AgentExecutionID: "exec-wedged-relaunch"}, nil + } + return &executor.LaunchAgentResponse{AgentExecutionID: "exec-replacement"}, 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) + } + + if err := svc.ensureSessionRunning(ctx, "session1", session); err != nil { + t.Fatalf("expected cold resume to reap the wedged launch and retry, got: %v", err) + } + if got := launchCalls.Load(); got != 2 { + t.Fatalf("expected a reap-and-retry relaunch (2 LaunchAgent calls), got %d", got) + } +} diff --git a/apps/backend/internal/orchestrator/task_operations.go b/apps/backend/internal/orchestrator/task_operations.go index 3da4871256..4e2749c2d5 100644 --- a/apps/backend/internal/orchestrator/task_operations.go +++ b/apps/backend/internal/orchestrator/task_operations.go @@ -1865,13 +1865,9 @@ func (s *Service) ensureSessionRunning(ctx context.Context, sessionID string, se if isOfficeTask { return errOfficeTaskResumeRequiresScheduler } - recoveryCtx := context.WithoutCancel(ctx) - if stopErr := s.reapPromptUnreadyExecution(recoveryCtx, sessionID, err); stopErr != nil { - return fmt.Errorf("failed to stop prompt-unready execution: %w", stopErr) - } - refreshed, refreshErr := s.repo.GetTaskSession(recoveryCtx, sessionID) - if refreshErr != nil { - return fmt.Errorf("failed to reload session after prompt-readiness recovery: %w", refreshErr) + refreshed, reapErr := s.reapAndReloadSession(ctx, sessionID, err) + if reapErr != nil { + return reapErr } session = refreshed } else { @@ -1883,6 +1879,61 @@ func (s *Service) ensureSessionRunning(ctx context.Context, sessionID string, se zap.String("session_id", sessionID), zap.String("session_state", string(session.State))) + // Bounded to two attempts: a fresh cold resume, and — if the launched + // agent never reports prompt-ready — one reap-and-retry, mirroring the + // already-tracked-execution branch above. Without this, a wedged launch + // right after a backend restart (the common "kandev restart" cold-resume + // shape: in-memory execution store empty, executors_running row intact) + // surfaced a bare "agent not ready after resume: ... context deadline + // exceeded" with no self-heal, unlike every other resume path in this + // file. + for attempt := 1; ; attempt++ { + retryable, err := s.attemptColdResume(ctx, sessionID, session, isOfficeTask) + if err == nil { + return nil + } + if attempt >= 2 || !retryable { + return err + } + refreshed, reapErr := s.reapAndReloadSession(ctx, sessionID, err) + if reapErr != nil { + return reapErr + } + session = refreshed + } +} + +// reapAndReloadSession stops a prompt-unready execution and reloads the +// session row afterward. Shared by ensureSessionRunning's already-tracked- +// execution branch and its cold-resume retry loop so both recovery paths +// stay in lockstep. +func (s *Service) reapAndReloadSession(ctx context.Context, sessionID string, cause error) (*models.TaskSession, error) { + recoveryCtx := context.WithoutCancel(ctx) + if stopErr := s.reapPromptUnreadyExecution(recoveryCtx, sessionID, cause); stopErr != nil { + return nil, fmt.Errorf("failed to stop prompt-unready execution: %w", stopErr) + } + refreshed, refreshErr := s.repo.GetTaskSession(recoveryCtx, sessionID) + if refreshErr != nil { + return nil, fmt.Errorf("failed to reload session after prompt-readiness recovery: %w", refreshErr) + } + return refreshed, nil +} + +// attemptColdResume performs one resume attempt for ensureSessionRunning's +// cold-resume loop (no execution currently tracked in memory). retryable is +// true only when the caller should reap the stuck execution and retry — +// specifically the main-path prompt-readiness timeout below, which is the +// exact "kandev restart" wedged-launch shape this loop exists to self-heal. +// The concurrent-resume-race branch (ErrExecutionAlreadyRunning) reports its +// own timeout as non-retryable: that failure is against an execution this +// attempt never launched, not a wedged cold launch, so retrying it here would +// blur two distinct failure modes the caller's tests pin down separately. +func (s *Service) attemptColdResume( + ctx context.Context, + sessionID string, + session *models.TaskSession, + isOfficeTask bool, +) (retryable bool, err error) { // If the session is in CREATED state with an existing workspace (executors_running // row exists), the workspace was prepared but the agent was never started. Use // LaunchPreparedSession which routes to startAgentOnExistingWorkspace to reuse @@ -1891,20 +1942,20 @@ func (s *Service) ensureSessionRunning(ctx context.Context, sessionID string, se if session.State == models.TaskSessionStateCreated { hasRunning, _ := s.repo.HasExecutorRunningRow(ctx, sessionID) if hasRunning { - return s.startAgentOnPreparedWorkspace(ctx, sessionID, session) + return false, s.startAgentOnPreparedWorkspace(ctx, sessionID, session) } } if isOfficeTask { - return errOfficeTaskResumeRequiresScheduler + return false, errOfficeTaskResumeRequiresScheduler } - running, err := s.repo.GetExecutorRunningBySessionID(ctx, sessionID) - if err != nil || running == nil { - return fmt.Errorf("session is not resumable: no executor record (state: %s)", session.State) + running, lookupErr := s.repo.GetExecutorRunningBySessionID(ctx, sessionID) + if lookupErr != nil || running == nil { + return false, fmt.Errorf("session is not resumable: no executor record (state: %s)", session.State) } if err := validateSessionWorktrees(session); err != nil { - return err + return false, err } // Use context.WithoutCancel to prevent WebSocket request timeout from canceling the resume. @@ -1912,19 +1963,19 @@ func (s *Service) ensureSessionRunning(ctx context.Context, sessionID string, se // when the agent's ACP session initializes — that's what unblocks waitForSessionReady, // no flag-tracking needed. resumeCtx := context.WithoutCancel(ctx) - if _, err = s.executor.ResumeSession(resumeCtx, session, true); err != nil { - if errors.Is(err, executor.ErrExecutionAlreadyRunning) { + if _, launchErr := s.executor.ResumeSession(resumeCtx, session, true); launchErr != nil { + if errors.Is(launchErr, executor.ErrExecutionAlreadyRunning) { s.recoverAgentPromptStreamIfNeeded(resumeCtx, sessionID) if readyErr := s.waitForAgentPromptReady(resumeCtx, sessionID); readyErr != nil { - return fmt.Errorf("agent not ready after resume race: %w", readyErr) + return false, fmt.Errorf("agent not ready after resume race: %w", readyErr) } - return nil + return false, nil } - return s.handleSessionLaunchFailure( + return false, s.handleSessionLaunchFailure( resumeCtx, session.TaskID, sessionID, - fmt.Errorf("failed to resume session: %w", err), + fmt.Errorf("failed to resume session: %w", launchErr), session, ) } @@ -1941,14 +1992,14 @@ func (s *Service) ensureSessionRunning(ctx context.Context, sessionID string, se // own bounded timeouts (waitForSessionReady's AgentLaunchTimeout launch // budget, and waitForAgentPromptReady's 30s below). if err := s.waitForSessionReady(resumeCtx, sessionID); err != nil { - return fmt.Errorf("session not ready after resume: %w", err) + return false, fmt.Errorf("session not ready after resume: %w", err) } - if err := s.waitForAgentPromptReady(resumeCtx, sessionID); err != nil { - return fmt.Errorf("agent not ready after resume: %w", err) + readyErr := s.waitForAgentPromptReady(resumeCtx, sessionID) + if readyErr == nil { + s.logger.Debug("session resumed and ready for prompt") + return false, nil } - - s.logger.Debug("session resumed and ready for prompt") - return nil + return errors.Is(readyErr, ErrAgentNotReadyForPrompt), fmt.Errorf("agent not ready after resume: %w", readyErr) } func (s *Service) recoverAgentPromptStreamIfNeeded(ctx context.Context, sessionID string) {