fix: reap and retry cold-resume prompt-readiness timeout (#2356)

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)
+	}
+}

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) {