diff --git a/apps/backend/internal/orchestrator/coordinator_stop_test.go b/apps/backend/internal/orchestrator/coordinator_stop_test.go index 78abe9f57b..b3e76758c4 100644 --- a/apps/backend/internal/orchestrator/coordinator_stop_test.go +++ b/apps/backend/internal/orchestrator/coordinator_stop_test.go @@ -547,7 +547,9 @@ func TestSetSessionStarting_CoordinatorCancellationWinsAfterCurrentStateRead(t * } svc := newCoordinatorStopTestService(repo, newMockTaskRepo(), &mockAgentManager{}) - err = svc.setSessionStarting(ctx, stale.TaskID, stale, true) + err = svc.setSessionStarting( + ctx, stale.TaskID, stale, models.TaskSessionStateRunning, true, + ) require.Error(t, err) stored, getErr := baseRepo.GetTaskSession(ctx, stale.ID) require.NoError(t, getErr) diff --git a/apps/backend/internal/orchestrator/event_handlers_streaming.go b/apps/backend/internal/orchestrator/event_handlers_streaming.go index e714f3461f..e92e299f76 100644 --- a/apps/backend/internal/orchestrator/event_handlers_streaming.go +++ b/apps/backend/internal/orchestrator/event_handlers_streaming.go @@ -1287,7 +1287,24 @@ func (s *Service) deleteCompletedExecutionIfExpired(key string, expiresAt time.T } } -func (s *Service) setSessionStarting(ctx context.Context, taskID string, session *models.TaskSession, promoteTask bool) error { +func allowsSessionStartingRecovery( + nextState, expectedState, currentState models.TaskSessionState, + promoteTask bool, +) bool { + return !promoteTask && + nextState == models.TaskSessionStateStarting && + currentState == expectedState && + (expectedState == models.TaskSessionStateFailed || + expectedState == models.TaskSessionStateCancelled) +} + +func (s *Service) setSessionStarting( + ctx context.Context, + taskID string, + session *models.TaskSession, + expectedState models.TaskSessionState, + promoteTask bool, +) error { if session == nil { return nil } @@ -1303,17 +1320,15 @@ func (s *Service) setSessionStarting(ctx context.Context, taskID string, session if err != nil { return err } - allowedTerminalRecovery := !promoteTask && - session.State == models.TaskSessionStateStarting && - (current.State == models.TaskSessionStateFailed || - (current.State == models.TaskSessionStateCancelled && - models.IsArchiveCancelReason(current.ErrorMessage))) + allowedTerminalRecovery := allowsSessionStartingRecovery( + session.State, expectedState, current.State, promoteTask, + ) if isTerminalSessionState(current.State) && !allowedTerminalRecovery { return &executor.SessionStateSupersededError{SessionID: session.ID, State: current.State} } oldState = current.State - if err := s.persistFullTaskSessionIfCurrent(ctx, session, current.State); err != nil { + if err := s.persistFullTaskSessionIfCurrent(ctx, session, expectedState); err != nil { return err } diff --git a/apps/backend/internal/orchestrator/event_handlers_streaming_test.go b/apps/backend/internal/orchestrator/event_handlers_streaming_test.go index 0b55d913a6..8011ebf2ac 100644 --- a/apps/backend/internal/orchestrator/event_handlers_streaming_test.go +++ b/apps/backend/internal/orchestrator/event_handlers_streaming_test.go @@ -1858,7 +1858,9 @@ func TestSetSessionStartingWritesTaskInProgress(t *testing.T) { taskRepo := newMockTaskRepo() svc := createTestService(repo, newMockStepGetter(), taskRepo) - require.NoError(t, svc.setSessionStarting(ctx, "t1", session, true)) + require.NoError(t, svc.setSessionStarting( + ctx, "t1", session, models.TaskSessionStateRunning, true, + )) updated, err := repo.GetTaskSession(ctx, "s1") require.NoError(t, err) @@ -1881,7 +1883,9 @@ func TestSetSessionStartingCanDeferTaskInProgress(t *testing.T) { taskRepo := newMockTaskRepo() svc := createTestService(repo, newMockStepGetter(), taskRepo) - require.NoError(t, svc.setSessionStarting(ctx, "t1", session, false)) + require.NoError(t, svc.setSessionStarting( + ctx, "t1", session, models.TaskSessionStateRunning, false, + )) updated, err := repo.GetTaskSession(ctx, "s1") require.NoError(t, err) @@ -1907,7 +1911,9 @@ func TestSetSessionStartingRejectsTerminalSession(t *testing.T) { taskRepo := newMockTaskRepo() svc := createTestService(repo, newMockStepGetter(), taskRepo) - require.Error(t, svc.setSessionStarting(ctx, "t1", &next, true)) + require.Error(t, svc.setSessionStarting( + ctx, "t1", &next, models.TaskSessionStateCancelled, true, + )) updated, err := repo.GetTaskSession(ctx, "s1") require.NoError(t, err) @@ -1915,7 +1921,7 @@ func TestSetSessionStartingRejectsTerminalSession(t *testing.T) { require.Empty(t, taskRepo.stateWrites) } -func TestSetSessionStartingRejectsCancelledTerminalResumeWithoutPromotion(t *testing.T) { +func TestSetSessionStartingAllowsExplicitlyCancelledResumeWithoutPromotion(t *testing.T) { ctx := context.Background() repo := setupTestRepo(t) seedSession(t, repo, "t1", "s1", "step1") @@ -1923,6 +1929,7 @@ func TestSetSessionStartingRejectsCancelledTerminalResumeWithoutPromotion(t *tes current, err := repo.GetTaskSession(ctx, "s1") require.NoError(t, err) current.State = models.TaskSessionStateCancelled + current.ErrorMessage = "stopped via API" require.NoError(t, repo.UpdateTaskSession(ctx, current)) next := *current @@ -1933,7 +1940,37 @@ func TestSetSessionStartingRejectsCancelledTerminalResumeWithoutPromotion(t *tes taskRepo := newMockTaskRepo() svc := createTestService(repo, newMockStepGetter(), taskRepo) - require.Error(t, svc.setSessionStarting(ctx, "t1", &next, false)) + require.NoError(t, svc.setSessionStarting( + ctx, "t1", &next, models.TaskSessionStateCancelled, false, + )) + + updated, err := repo.GetTaskSession(ctx, "s1") + require.NoError(t, err) + require.Equal(t, models.TaskSessionStateStarting, updated.State) + require.Empty(t, taskRepo.stateWrites) +} + +func TestSetSessionStartingRejectsCancellationAfterResumeSnapshot(t *testing.T) { + ctx := context.Background() + repo := setupTestRepo(t) + seedSession(t, repo, "t1", "s1", "step1") + + stale, err := repo.GetTaskSession(ctx, "s1") + require.NoError(t, err) + stale.State = models.TaskSessionStateStarting + + current, err := repo.GetTaskSession(ctx, "s1") + require.NoError(t, err) + current.State = models.TaskSessionStateCancelled + current.ErrorMessage = "stopped via API" + require.NoError(t, repo.UpdateTaskSession(ctx, current)) + + taskRepo := newMockTaskRepo() + svc := createTestService(repo, newMockStepGetter(), taskRepo) + + require.Error(t, svc.setSessionStarting( + ctx, "t1", stale, models.TaskSessionStateRunning, false, + )) updated, err := repo.GetTaskSession(ctx, "s1") require.NoError(t, err) @@ -1963,7 +2000,9 @@ func TestSetSessionStartingAllowsTerminalResumeWithoutTaskPromotion(t *testing.T taskRepo := newMockTaskRepo() svc := createTestService(repo, newMockStepGetter(), taskRepo) - require.NoError(t, svc.setSessionStarting(ctx, "t1", &next, false)) + require.NoError(t, svc.setSessionStarting( + ctx, "t1", &next, models.TaskSessionStateFailed, false, + )) updated, err := repo.GetTaskSession(ctx, "s1") require.NoError(t, err) @@ -1999,7 +2038,9 @@ func TestSetSessionStartingAllowsArchiveCancelledResumeWithoutPromotion(t *testi taskRepo := newMockTaskRepo() svc := createTestService(repo, newMockStepGetter(), taskRepo) - require.NoError(t, svc.setSessionStarting(ctx, "t1", &next, false)) + require.NoError(t, svc.setSessionStarting( + ctx, "t1", &next, models.TaskSessionStateCancelled, false, + )) updated, err := repo.GetTaskSession(ctx, "s1") require.NoError(t, err) diff --git a/apps/backend/internal/orchestrator/executor/executor.go b/apps/backend/internal/orchestrator/executor/executor.go index 399c9a0254..0815931156 100644 --- a/apps/backend/internal/orchestrator/executor/executor.go +++ b/apps/backend/internal/orchestrator/executor/executor.go @@ -578,9 +578,17 @@ type SessionStateTransitionFunc func( // SessionStartingFunc is called when the executor has prepared/resumed an // execution and needs to mark the session STARTING while preserving other -// session-row updates such as metadata. promoteTask controls whether the -// callback should also move the parent task to IN_PROGRESS immediately. -type SessionStartingFunc func(ctx context.Context, taskID string, session *models.TaskSession, promoteTask bool) error +// session-row updates such as metadata. expectedState is the state observed +// before runtime launch, so a prior CANCELLED state can resume without +// overwriting a cancellation that raced the launch. promoteTask controls +// whether the callback should also move the parent task to IN_PROGRESS. +type SessionStartingFunc func( + ctx context.Context, + taskID string, + session *models.TaskSession, + expectedState models.TaskSessionState, + promoteTask bool, +) error // ExecutionCleanupClaimFunc atomically claims forced cleanup for one exact // session execution. It returns true when the executor owns cleanup and false diff --git a/apps/backend/internal/orchestrator/executor/executor_execute.go b/apps/backend/internal/orchestrator/executor/executor_execute.go index 5185130227..adcc2cb286 100644 --- a/apps/backend/internal/orchestrator/executor/executor_execute.go +++ b/apps/backend/internal/orchestrator/executor/executor_execute.go @@ -438,12 +438,29 @@ func isStopTerminalSessionState(state models.TaskSessionState) bool { state == models.TaskSessionStateCancelled } +func allowsSessionStartingRecovery( + nextState, expectedState, currentState models.TaskSessionState, + promoteTask bool, +) bool { + return !promoteTask && + nextState == models.TaskSessionStateStarting && + currentState == expectedState && + (expectedState == models.TaskSessionStateFailed || + expectedState == models.TaskSessionStateCancelled) +} + // updateSessionStarting persists a full session-row STARTING transition, using // the orchestrator callback when present so task/session runtime state stays // serialized with guarded REVIEW reconciliation. -func (e *Executor) updateSessionStarting(ctx context.Context, taskID string, session *models.TaskSession, promoteTask bool) error { +func (e *Executor) updateSessionStarting( + ctx context.Context, + taskID string, + session *models.TaskSession, + expectedState models.TaskSessionState, + promoteTask bool, +) error { if e.onSessionStarting != nil { - return e.onSessionStarting(ctx, taskID, session, promoteTask) + return e.onSessionStarting(ctx, taskID, session, expectedState, promoteTask) } current, err := e.repo.GetTaskSession(ctx, session.ID) if err != nil { @@ -452,15 +469,13 @@ func (e *Executor) updateSessionStarting(ctx context.Context, taskID string, ses if current == nil { return fmt.Errorf("%w: agent session not found: %s", models.ErrTaskSessionNotFound, session.ID) } - allowedTerminalRecovery := !promoteTask && - session.State == models.TaskSessionStateStarting && - (current.State == models.TaskSessionStateFailed || - (current.State == models.TaskSessionStateCancelled && - models.IsArchiveCancelReason(current.ErrorMessage))) + allowedTerminalRecovery := allowsSessionStartingRecovery( + session.State, expectedState, current.State, promoteTask, + ) if isStopTerminalSessionState(current.State) && !allowedTerminalRecovery { return &SessionStateSupersededError{SessionID: session.ID, State: current.State} } - return e.persistSessionFullRowIfCurrentState(ctx, session, current.State) + return e.persistSessionFullRowIfCurrentState(ctx, session, expectedState) } func (e *Executor) persistSessionFullRowIfCurrentState( @@ -1480,11 +1495,12 @@ func (e *Executor) startAgentOnExistingWorkspace(ctx context.Context, task *v1.T } // Transition session to STARTING + expectedState := session.State now := time.Now().UTC() session.State = models.TaskSessionStateStarting session.ErrorMessage = "" session.UpdatedAt = now - if err := e.updateSessionStarting(ctx, task.ID, session, true); err != nil { + if err := e.updateSessionStarting(ctx, task.ID, session, expectedState, true); err != nil { e.logger.Error("failed to update session state for agent start", zap.String("session_id", session.ID), zap.Error(err)) diff --git a/apps/backend/internal/orchestrator/executor/executor_interaction.go b/apps/backend/internal/orchestrator/executor/executor_interaction.go index 90824ff046..e7a1a34364 100644 --- a/apps/backend/internal/orchestrator/executor/executor_interaction.go +++ b/apps/backend/internal/orchestrator/executor/executor_interaction.go @@ -673,6 +673,7 @@ func (e *Executor) applyWorktreeToSwitchRequest(req *LaunchAgentRequest, session // container_id / status are written by the lifecycle manager during the launch // itself (lifecycle.persistExecutorRunning) and not touched here. func (e *Executor) persistModelSwitchState(ctx context.Context, taskID, sessionID string, session *models.TaskSession, newModel string) error { + expectedState := session.State session.State = models.TaskSessionStateStarting session.UpdatedAt = time.Now().UTC() @@ -681,7 +682,7 @@ func (e *Executor) persistModelSwitchState(ctx context.Context, taskID, sessionI } session.AgentProfileSnapshot["model"] = newModel - if err := e.updateSessionStarting(ctx, taskID, session, true); err != nil { + if err := e.updateSessionStarting(ctx, taskID, session, expectedState, true); err != nil { e.logger.Error("failed to update session after model switch", zap.String("task_id", taskID), zap.String("session_id", sessionID), diff --git a/apps/backend/internal/orchestrator/executor/executor_interaction_test.go b/apps/backend/internal/orchestrator/executor/executor_interaction_test.go index 6d468f7b0b..e574daa4d6 100644 --- a/apps/backend/internal/orchestrator/executor/executor_interaction_test.go +++ b/apps/backend/internal/orchestrator/executor/executor_interaction_test.go @@ -420,12 +420,13 @@ func TestLaunchModelSwitchAgent_CleansStartedExecutionAfterTerminalRace(t *testi return sessionID == session.ID && executionID == "execution-model-race" }) + callerSession := *session err := exec.launchModelSwitchAgent( ctx, session.TaskID, session.ID, "new-model", - session, + &callerSession, &LaunchAgentRequest{}, nil, ) @@ -440,6 +441,55 @@ func TestLaunchModelSwitchAgent_CleansStartedExecutionAfterTerminalRace(t *testi } } +func TestPersistModelSwitchState_CancellationAfterSnapshotWins(t *testing.T) { + ctx := context.Background() + repo := newMockRepository() + stored := &models.TaskSession{ + ID: "session-model-persist-race", + TaskID: "task-model-persist-race", + State: models.TaskSessionStateRunning, + AgentProfileSnapshot: map[string]interface{}{"model": "old-model"}, + } + repo.sessions[stored.ID] = stored + stale := *stored + stale.AgentProfileSnapshot = map[string]interface{}{"model": "old-model"} + if err := repo.UpdateTaskSessionState( + ctx, stored.ID, models.TaskSessionStateCancelled, "stopped via API", + ); err != nil { + t.Fatalf("cancel session: %v", err) + } + + exec := newTestExecutor(t, &mockAgentManager{}, repo) + var gotExpectedState models.TaskSessionState + exec.SetOnSessionStarting(func( + ctx context.Context, + _ string, + next *models.TaskSession, + expectedState models.TaskSessionState, + _ bool, + ) error { + gotExpectedState = expectedState + return exec.persistSessionFullRowIfCurrentState(ctx, next, expectedState) + }) + + err := exec.persistModelSwitchState( + ctx, stale.TaskID, stale.ID, &stale, "new-model", + ) + if !errors.Is(err, ErrSessionStateSuperseded) { + t.Fatalf("persistModelSwitchState error = %v, want ErrSessionStateSuperseded", err) + } + if gotExpectedState != models.TaskSessionStateRunning { + t.Fatalf("expected state = %q, want RUNNING", gotExpectedState) + } + persisted := repo.sessions[stored.ID] + if persisted.State != models.TaskSessionStateCancelled { + t.Fatalf("stored state = %q, want CANCELLED", persisted.State) + } + if got, _ := persisted.AgentProfileSnapshot["model"].(string); got != "old-model" { + t.Fatalf("stored model = %q, want old-model", got) + } +} + func TestPersistInPlaceModelSwitch_DoesNotRestoreStaleActiveStateAfterCancellation(t *testing.T) { repo := newMockRepository() session := &models.TaskSession{ diff --git a/apps/backend/internal/orchestrator/executor/executor_resume.go b/apps/backend/internal/orchestrator/executor/executor_resume.go index 98bd09ffb2..de0fda7130 100644 --- a/apps/backend/internal/orchestrator/executor/executor_resume.go +++ b/apps/backend/internal/orchestrator/executor/executor_resume.go @@ -304,7 +304,7 @@ func (e *Executor) persistLaunchState(ctx context.Context, taskID, sessionID str var updateErr error if startAgent { - updateErr = e.updateSessionStarting(ctx, taskID, session, true) + updateErr = e.updateSessionStarting(ctx, taskID, session, expectedState, true) } else { updateErr = e.persistSessionFullRowIfCurrentState(ctx, session, expectedState) } @@ -386,6 +386,10 @@ func (e *Executor) persistWorktreeAssociation(ctx context.Context, taskID string // ResumeSession restarts an existing task session using its stored worktree. // When startAgent is false, only the executor runtime is started (agent process is not launched). func (e *Executor) ResumeSession(ctx context.Context, session *models.TaskSession, startAgent bool) (*TaskExecution, error) { + if session != nil { + resumeSnapshot := *session + session = &resumeSnapshot + } task, unlock, err := e.validateAndLockResume(ctx, session) if err != nil { return nil, err @@ -502,6 +506,7 @@ func (e *Executor) validateAndLockResume(ctx context.Context, session *models.Ta if session == nil { return nil, func() {}, ErrExecutionNotFound } + requestedState := session.State // Acquire per-session lock to prevent concurrent resume/launch operations. // This is critical after backend restart when multiple resume requests may arrive @@ -554,6 +559,13 @@ func (e *Executor) validateAndLockResume(ctx context.Context, session *models.Ta return nil, func() {}, fetchErr } if fresh != nil { + if isTerminalSessionState(fresh.State) && fresh.State != requestedState { + unlock() + return nil, func() {}, &SessionStateSupersededError{ + SessionID: session.ID, + State: fresh.State, + } + } session.State = fresh.State } @@ -935,7 +947,7 @@ func (e *Executor) persistResumeState(ctx context.Context, taskID string, sessio var updateErr error if startAgent { - updateErr = e.updateSessionStarting(ctx, taskID, session, false) + updateErr = e.updateSessionStarting(ctx, taskID, session, expectedState, false) } else { updateErr = e.persistSessionFullRowIfCurrentState(ctx, session, expectedState) } diff --git a/apps/backend/internal/orchestrator/executor/executor_resume_test.go b/apps/backend/internal/orchestrator/executor/executor_resume_test.go index 952648af9f..45141a85fe 100644 --- a/apps/backend/internal/orchestrator/executor/executor_resume_test.go +++ b/apps/backend/internal/orchestrator/executor/executor_resume_test.go @@ -299,7 +299,8 @@ func TestResumeSession_CancelledStateForceCleansUpStaleState(t *testing.T) { } exec := newTestExecutor(t, agentMgr, repo) - if _, err := exec.ResumeSession(context.Background(), repo.sessions["sess-1"], true); err != nil { + callerSession := *repo.sessions["sess-1"] + if _, err := exec.ResumeSession(context.Background(), &callerSession, true); err != nil { t.Fatalf("expected success, got: %v", err) } if agentMgr.cleanupStaleExecutionCallCount != 1 { @@ -486,6 +487,41 @@ func TestResumeSession_ConcurrentResumeReReadsFreshState(t *testing.T) { } } +// TestResumeSession_CancellationBeforeResumeLockWins exercises the stop/resume +// race before validateAndLockResume acquires the per-session lock. The resume +// request observed WAITING_FOR_INPUT, then stop committed CANCELLED before the +// in-lock re-read. That cancellation must abort the resume before any runtime +// cleanup or replacement agent launch. +func TestResumeSession_CancellationBeforeResumeLockWins(t *testing.T) { + repo := newMockRepository() + setupLiveResumeTestFixture(repo) + + callerSession := *repo.sessions["sess-1"] + repo.sessions["sess-1"].State = models.TaskSessionStateCancelled + repo.sessions["sess-1"].ErrorMessage = "stopped via API" + + agentMgr := &mockAgentManager{ + launchAgentFunc: func(_ context.Context, _ *LaunchAgentRequest) (*LaunchAgentResponse, error) { + t.Fatal("LaunchAgent must not be called after a concurrent cancellation") + return nil, nil + }, + } + exec := newTestExecutor(t, agentMgr, repo) + + _, err := exec.ResumeSession(context.Background(), &callerSession, true) + if !errors.Is(err, ErrSessionStateSuperseded) { + t.Fatalf("expected ErrSessionStateSuperseded, got: %v", err) + } + if agentMgr.cleanupStaleExecutionCallCount != 0 { + t.Errorf("cleanup must not run after cancellation, got %d calls", + agentMgr.cleanupStaleExecutionCallCount) + } + if agentMgr.launchAgentCallCount != 0 { + t.Errorf("LaunchAgent must not run after cancellation, got %d calls", + agentMgr.launchAgentCallCount) + } +} + // TestResumeSession_AbortsIfSessionReReadFails locks down the abort-on-error // behavior of the in-lock session re-read. If GetTaskSession fails transiently, // silently falling back to the caller's (potentially stale) session.State would diff --git a/apps/backend/internal/orchestrator/executor/executor_test.go b/apps/backend/internal/orchestrator/executor/executor_test.go index 2b5f46b021..31b46b19b7 100644 --- a/apps/backend/internal/orchestrator/executor/executor_test.go +++ b/apps/backend/internal/orchestrator/executor/executor_test.go @@ -451,7 +451,13 @@ func TestLaunchPreparedSession_AbortsWhenStartingPersistenceFails(t *testing.T) persistErr := errors.New("session is terminal") executor := newTestExecutor(t, agentManager, repo) - executor.SetOnSessionStarting(func(ctx context.Context, taskID string, session *models.TaskSession, promoteTask bool) error { + executor.SetOnSessionStarting(func( + ctx context.Context, + taskID string, + session *models.TaskSession, + expectedState models.TaskSessionState, + promoteTask bool, + ) error { return persistErr }) @@ -864,6 +870,66 @@ func TestLaunchPreparedSession_ExistingWorkspace_StartAgent(t *testing.T) { } } +func TestStartAgentOnExistingWorkspace_CancellationAfterSnapshotWins(t *testing.T) { + ctx := context.Background() + repo := newMockRepository() + stored := &models.TaskSession{ + ID: "session-existing-race", + TaskID: "task-existing-race", + AgentProfileID: "profile-123", + State: models.TaskSessionStateCreated, + } + repo.sessions[stored.ID] = stored + callerSession := *stored + + agentManager := &mockAgentManager{ + getExecutionIDForSessionFunc: func(context.Context, string) (string, error) { + return "exec-existing-race", nil + }, + startAgentProcessFunc: func(context.Context, string) error { + t.Fatal("agent process must not start after cancellation") + return nil + }, + } + executor := newTestExecutor(t, agentManager, repo) + + var gotExpectedState models.TaskSessionState + executor.SetOnSessionStarting(func( + ctx context.Context, + _ string, + next *models.TaskSession, + expectedState models.TaskSessionState, + _ bool, + ) error { + gotExpectedState = expectedState + if err := repo.UpdateTaskSessionState( + ctx, stored.ID, models.TaskSessionStateCancelled, "stopped via API", + ); err != nil { + return err + } + return executor.persistSessionFullRowIfCurrentState(ctx, next, expectedState) + }) + + _, err := executor.startAgentOnExistingWorkspace( + ctx, + &v1.Task{ID: stored.TaskID}, + &callerSession, + "", + true, + "", + nil, + ) + if !errors.Is(err, ErrSessionStateSuperseded) { + t.Fatalf("startAgentOnExistingWorkspace error = %v, want ErrSessionStateSuperseded", err) + } + if gotExpectedState != models.TaskSessionStateCreated { + t.Fatalf("expected state = %q, want CREATED", gotExpectedState) + } + if got := repo.sessions[stored.ID].State; got != models.TaskSessionStateCancelled { + t.Fatalf("stored state = %q, want CANCELLED", got) + } +} + // Regression: a workspace-only execution can be created by lazy workspace // restoration before workflow auto-start reaches LaunchPreparedSession. That // execution has no agent command, so it must go through LaunchAgent's promotion @@ -2574,7 +2640,8 @@ func TestPersistResumeState_SetsStartingState(t *testing.T) { CompletedAt: &completedAt, UpdatedAt: now, } - repo.sessions[session.ID] = session + stored := *session + repo.sessions[session.ID] = &stored if err := executor.persistResumeState(context.Background(), "task-1", session, true); err != nil { t.Fatalf("persistResumeState: %v", err) @@ -2597,7 +2664,15 @@ func TestPersistResumeState_SetsStartingState(t *testing.T) { } repo.sessions[session.ID] = session var gotPromoteTask *bool - executor.SetOnSessionStarting(func(ctx context.Context, taskID string, session *models.TaskSession, promoteTask bool) error { + var gotExpectedState models.TaskSessionState + executor.SetOnSessionStarting(func( + ctx context.Context, + taskID string, + session *models.TaskSession, + expectedState models.TaskSessionState, + promoteTask bool, + ) error { + gotExpectedState = expectedState gotPromoteTask = &promoteTask return repo.UpdateTaskSession(ctx, session) }) @@ -2612,6 +2687,9 @@ func TestPersistResumeState_SetsStartingState(t *testing.T) { if gotPromoteTask == nil { t.Fatal("expected onSessionStarting callback") } + if gotExpectedState != models.TaskSessionStateWaitingForInput { + t.Fatalf("expected source state WAITING_FOR_INPUT, got %s", gotExpectedState) + } if *gotPromoteTask { t.Fatal("resume STARTING persistence must defer task promotion until process start succeeds") } @@ -2658,13 +2736,16 @@ func TestPersistResumeState_SetsStartingState(t *testing.T) { }) } -// TestUpdateSessionStarting_ArchiveCancelledSessionRecovers proves +// TestUpdateSessionStarting_CancelledSessionRecovers proves // updateSessionStarting's fallback path (no onSessionStarting callback -// wired) mirrors Service.setSessionStarting: an archive-cancelled session -// (the resume path treats CANCELLED as an expected, resumable terminal -// state) may recover into STARTING like a Failed session. -func TestUpdateSessionStarting_ArchiveCancelledSessionRecovers(t *testing.T) { - for _, reason := range []string{models.SessionArchiveCancelReason, models.SessionArchiveTreeCancelReason} { +// wired) mirrors Service.setSessionStarting: a CANCELLED session may recover +// when CANCELLED was the state observed before the resume launch. +func TestUpdateSessionStarting_CancelledSessionRecovers(t *testing.T) { + for _, reason := range []string{ + "stopped via API", + models.SessionArchiveCancelReason, + models.SessionArchiveTreeCancelReason, + } { t.Run(reason, func(t *testing.T) { repo := newMockRepository() executor := newTestExecutor(t, &mockAgentManager{}, repo) @@ -2680,7 +2761,13 @@ func TestUpdateSessionStarting_ArchiveCancelledSessionRecovers(t *testing.T) { State: models.TaskSessionStateStarting, } - if err := executor.updateSessionStarting(context.Background(), "task-1", next, false); err != nil { + if err := executor.updateSessionStarting( + context.Background(), + "task-1", + next, + models.TaskSessionStateCancelled, + false, + ); err != nil { t.Fatalf("updateSessionStarting: %v", err) } if got := repo.sessions["session-1"].State; got != models.TaskSessionStateStarting { @@ -2690,11 +2777,7 @@ func TestUpdateSessionStarting_ArchiveCancelledSessionRecovers(t *testing.T) { } } -// TestUpdateSessionStarting_OrdinaryCancelledSessionStaysRejected is the -// counterpart to TestUpdateSessionStarting_ArchiveCancelledSessionRecovers: -// an ordinary/explicit cancellation (no archive-cancel reason) must stay -// rejected as superseded, never silently promoted back into STARTING. -func TestUpdateSessionStarting_OrdinaryCancelledSessionStaysRejected(t *testing.T) { +func TestUpdateSessionStarting_CancellationAfterResumeSnapshotStaysRejected(t *testing.T) { repo := newMockRepository() executor := newTestExecutor(t, &mockAgentManager{}, repo) repo.sessions["session-1"] = &models.TaskSession{ @@ -2708,7 +2791,13 @@ func TestUpdateSessionStarting_OrdinaryCancelledSessionStaysRejected(t *testing. State: models.TaskSessionStateStarting, } - err := executor.updateSessionStarting(context.Background(), "task-1", next, false) + err := executor.updateSessionStarting( + context.Background(), + "task-1", + next, + models.TaskSessionStateWaitingForInput, + false, + ) var superseded *SessionStateSupersededError if !errors.As(err, &superseded) { t.Fatalf("updateSessionStarting error = %v, want SessionStateSupersededError", err) diff --git a/apps/backend/internal/orchestrator/service.go b/apps/backend/internal/orchestrator/service.go index fea0d78c7b..5843736b44 100644 --- a/apps/backend/internal/orchestrator/service.go +++ b/apps/backend/internal/orchestrator/service.go @@ -641,8 +641,14 @@ func NewService( return nil }) exec.SetOnSessionStateTransition(s.transitionTaskSessionState) - exec.SetOnSessionStarting(func(ctx context.Context, taskID string, session *models.TaskSession, promoteTask bool) error { - return s.setSessionStarting(ctx, taskID, session, promoteTask) + exec.SetOnSessionStarting(func( + ctx context.Context, + taskID string, + session *models.TaskSession, + expectedState models.TaskSessionState, + promoteTask bool, + ) error { + return s.setSessionStarting(ctx, taskID, session, expectedState, promoteTask) }) exec.SetOnExecutionCleanupClaim(s.claimForcedExecutionCleanup) exec.SetOnExecutionStopOwnerRegistration(s.RegisterExecutionStopOwner) diff --git a/apps/backend/internal/task/models/resume_safety.go b/apps/backend/internal/task/models/resume_safety.go index ce3a8e67a3..ab36dee976 100644 --- a/apps/backend/internal/task/models/resume_safety.go +++ b/apps/backend/internal/task/models/resume_safety.go @@ -74,20 +74,14 @@ func RowMustBePreserved(running *ExecutorRunning, sessionState TaskSessionState) // TaskSession.ErrorMessage values written when a session is auto-cancelled // by archiving — Service.ArchiveTask's CancelActiveTaskSessionsByTaskID call // (single-task archive) and HandoffService.cancelActiveRuns (cascade archive) -// respectively. They exist so the resume path can tell "cancelled because the -// task was archived" apart from an explicit user/coordinator stop. +// respectively. Distinct reasons keep the cancellation source visible in +// session history and diagnostics. const ( SessionArchiveCancelReason = "task archived" SessionArchiveTreeCancelReason = "task tree archived" ) -// IsArchiveCancelReason reports whether reason is a session-cancellation -// message written by an archive path rather than an explicit user or -// coordinator stop. A session cancelled by archiving is expected to resume -// normally once the task is unarchived, so callers gating a STARTING -// transition on "was this session explicitly, intentionally cancelled" must -// treat an archive-cancelled session like a Failed one instead of rejecting -// it outright — see setSessionStarting / updateSessionStarting. +// IsArchiveCancelReason reports whether reason came from an archive path. func IsArchiveCancelReason(reason string) bool { return reason == SessionArchiveCancelReason || reason == SessionArchiveTreeCancelReason }