diff --git a/apps/backend/internal/orchestrator/event_handlers_queue_lifecycle_test.go b/apps/backend/internal/orchestrator/event_handlers_queue_lifecycle_test.go index 8b7c63be09..608523b814 100644 --- a/apps/backend/internal/orchestrator/event_handlers_queue_lifecycle_test.go +++ b/apps/backend/internal/orchestrator/event_handlers_queue_lifecycle_test.go @@ -12,6 +12,19 @@ import ( v1 "github.com/kandev/kandev/pkg/api/v1" ) +type turnStartSignalService struct { + TurnService + started chan<- struct{} +} + +func (s *turnStartSignalService) StartTurn(ctx context.Context, sessionID string) (*models.Turn, error) { + turn, err := s.TurnService.StartTurn(ctx, sessionID) + if err == nil { + s.started <- struct{}{} + } + return turn, err +} + func TestExecuteQueuedMessage_LifecycleReselectsReplacementSession(t *testing.T) { ctx := context.Background() repo := setupTestRepo(t) @@ -125,6 +138,97 @@ func TestExecuteQueuedMessage_LifecycleNonPromptableSessionRequeues(t *testing.T } } +// A reset can cancel a lifecycle turn while task-state reconciliation is +// blocked. The queue must retry that entry instead of dispatching the stale +// turn or recording a duplicate visible message. +func TestExecuteQueuedMessage_LifecycleResetAfterTurnCreationRequeues(t *testing.T) { + ctx := context.Background() + svc, repo, manager, session := newActiveResetTestService(t) + setupTurnValue, ok := svc.activeTurns.Load(session.ID) + if !ok { + t.Fatal("test setup did not create an active turn") + } + setupTurnID, ok := setupTurnValue.(string) + if !ok || setupTurnID == "" { + t.Fatalf("active turn cache value = %#v, want turn ID", setupTurnValue) + } + if err := svc.turnService.CompleteTurn(ctx, setupTurnID); err != nil { + t.Fatalf("complete setup turn: %v", err) + } + svc.activeTurns.Delete(session.ID) + if err := repo.UpdateTaskSessionState(ctx, session.ID, models.TaskSessionStateWaitingForInput, ""); err != nil { + t.Fatalf("set session waiting: %v", err) + } + session.State = models.TaskSessionStateWaitingForInput + manager.isAgentRunning = true + svc.messageCreator = &mockMessageCreator{} + svc.executor = executor.NewExecutor(manager, repo, testLogger(), executor.ExecutorConfig{}) + + turnStarted := make(chan struct{}, 1) + svc.turnService = &turnStartSignalService{ + TurnService: svc.turnService, + started: turnStarted, + } + _, _, accepted, err := svc.messageQueue.QueueLifecycleMessageWithCoalesceKey( + ctx, session.ID, session.TaskID, "merged lifecycle prompt", "", messagequeue.QueuedByWorkflow, + false, nil, map[string]interface{}{"origin": githubPRAutomationOrigin}, + "github-pr:repo:1:merged", true, + ) + if err != nil || !accepted { + t.Fatalf("queue lifecycle entry: accepted=%v err=%v", accepted, err) + } + reserved, ok := svc.messageQueue.ReserveQueued(ctx, session.ID) + if !ok { + t.Fatal("reserve lifecycle entry") + } + svc.markQueuedDispatchInFlight(session.ID, reserved.ID) + + svc.taskRuntimeStateMu.Lock() + locked := true + defer func() { + if locked { + svc.taskRuntimeStateMu.Unlock() + } + }() + dispatchDone := make(chan struct{}) + go func() { + svc.executeQueuedMessage(session.ID, reserved) + close(dispatchDone) + }() + select { + case <-turnStarted: + case <-time.After(time.Second): + t.Fatal("timed out waiting for lifecycle turn creation") + } + + resetDone := make(chan bool, 1) + go func() { + resetDone <- svc.resetAgentContext(ctx, session.TaskID, session, "Successor") + }() + requireResetEvents(t, manager.events, "cancel", "reset") + if resetOK := <-resetDone; !resetOK { + t.Fatal("resetAgentContext returned false") + } + + svc.taskRuntimeStateMu.Unlock() + locked = false + select { + case <-dispatchDone: + case <-time.After(2 * time.Second): + t.Fatal("timed out waiting for lifecycle dispatch") + } + + if got := len(manager.capturedPromptCalls); got != 0 { + t.Fatalf("prompts after reset cancelled lifecycle turn = %d, want 0", got) + } + if got := len(svc.messageCreator.(*mockMessageCreator).userMessages); got != 0 { + t.Fatalf("visible lifecycle messages after reset = %d, want 0", got) + } + if got := svc.messageQueue.GetStatus(ctx, session.ID).Count; got != 1 { + t.Fatalf("lifecycle retries after reset = %d, want 1", got) + } +} + // A lifecycle prompt cannot be dispatched unless its visible automation chat // message was persisted. The failed write must restore the pre-claim state, // complete the just-started turn, and retain exactly one coalesced retry. diff --git a/apps/backend/internal/orchestrator/event_handlers_workflow_reset_quiescence_test.go b/apps/backend/internal/orchestrator/event_handlers_workflow_reset_quiescence_test.go index e9f77b2cb1..f219d35e58 100644 --- a/apps/backend/internal/orchestrator/event_handlers_workflow_reset_quiescence_test.go +++ b/apps/backend/internal/orchestrator/event_handlers_workflow_reset_quiescence_test.go @@ -14,6 +14,22 @@ import ( v1 "github.com/kandev/kandev/pkg/api/v1" ) +type lifecycleClaimSignalRepo struct { + repoStore + claimed chan<- struct{} +} + +func (r lifecycleClaimSignalRepo) ClaimPromptableTaskSessionIfActive( + ctx context.Context, + sessionID string, +) (models.PromptableTaskSessionClaim, error) { + claim, err := r.repoStore.ClaimPromptableTaskSessionIfActive(ctx, sessionID) + if err == nil && claim.Status == models.PromptableTaskSessionClaimed { + r.claimed <- struct{}{} + } + return claim, err +} + type orderedResetAgentManager struct { *mockAgentManager events chan string @@ -205,6 +221,74 @@ func TestResetAgentContext_CancelFailureStopsProviderReset(t *testing.T) { } } +// TestResetAgentContext_ResetMarkerPrecedesCancellationWait fixes the ordering +// that made TestResetAgentContext_SerializesPromptAdmission flaky. Once reset +// has published its marker and yielded the session guard for cancellation, a +// prompt must reject the reset before it considers waiting on that cancellation. +func TestResetAgentContext_ResetMarkerPrecedesCancellationWait(t *testing.T) { + ctx := context.Background() + svc, _, manager, session := newActiveResetTestService(t) + cancelEntered := make(chan struct{}, 1) + cancelRelease := make(chan struct{}) + manager.cancelAgentEntered = cancelEntered + manager.cancelAgentBlock = cancelRelease + var releaseCancellation sync.Once + releaseCancel := func() { releaseCancellation.Do(func() { close(cancelRelease) }) } + t.Cleanup(releaseCancel) + + resetDone := make(chan bool, 1) + go func() { + resetDone <- svc.resetAgentContext(ctx, session.TaskID, session, "Successor") + }() + <-cancelEntered + + cancelledCtx, cancel := context.WithCancel(ctx) + cancel() + _, _, _, _, _, err := svc.claimSessionRunningForPrompt( + cancelledCtx, session.TaskID, session.ID, "", false, nil, nil, "", false, + ) + if !errors.Is(err, ErrSessionResetInProgress) { + t.Fatalf("prompt admission error = %v, want %v", err, ErrSessionResetInProgress) + } + + releaseCancel() + if resetOK := <-resetDone; !resetOK { + t.Fatal("resetAgentContext returned false") + } +} + +func TestResetAgentContext_ResetMarkerPrecedesLifecycleCancellationWait(t *testing.T) { + ctx := context.Background() + svc, _, manager, session := newActiveResetTestService(t) + cancelEntered := make(chan struct{}, 1) + cancelRelease := make(chan struct{}) + manager.cancelAgentEntered = cancelEntered + manager.cancelAgentBlock = cancelRelease + var releaseCancellation sync.Once + releaseCancel := func() { releaseCancellation.Do(func() { close(cancelRelease) }) } + t.Cleanup(releaseCancel) + + resetDone := make(chan bool, 1) + go func() { + resetDone <- svc.resetAgentContext(ctx, session.TaskID, session, "Successor") + }() + <-cancelEntered + + cancelledCtx, cancel := context.WithCancel(ctx) + cancel() + _, _, _, _, err := svc.claimLifecycleSessionRunning( + cancelledCtx, session.TaskID, session.ID, "", + ) + if !errors.Is(err, ErrSessionResetInProgress) { + t.Fatalf("lifecycle prompt admission error = %v, want %v", err, ErrSessionResetInProgress) + } + + releaseCancel() + if resetOK := <-resetDone; !resetOK { + t.Fatal("resetAgentContext returned false") + } +} + // TestResetAgentContext_SerializesPromptAdmission stages a prompt immediately // before its final guarded claim, then starts reset while that claim is paused. // If reset wins the shared guard, the marker rejects the prompt while reset is @@ -283,3 +367,58 @@ func TestResetAgentContext_SerializesPromptAdmission(t *testing.T) { // reset. The provider replacement is therefore ordered after cancellation. requireResetEvents(t, manager.events, "cancel", "reset") } + +func TestResetAgentContext_SeesLifecycleTurnBeforeProviderReset(t *testing.T) { + ctx := context.Background() + svc, repo, manager, session := newActiveResetTestService(t) + turnIDValue, ok := svc.activeTurns.Load(session.ID) + if !ok { + t.Fatal("test setup did not create an active turn") + } + turnID, ok := turnIDValue.(string) + if !ok || turnID == "" { + t.Fatalf("active turn cache value = %#v, want turn ID", turnIDValue) + } + if err := svc.turnService.CompleteTurn(ctx, turnID); err != nil { + t.Fatalf("complete setup turn: %v", err) + } + svc.activeTurns.Delete(session.ID) + if err := repo.UpdateTaskSessionState(ctx, session.ID, models.TaskSessionStateWaitingForInput, ""); err != nil { + t.Fatalf("set session waiting: %v", err) + } + session.State = models.TaskSessionStateWaitingForInput + + claimed := make(chan struct{}, 1) + svc.repo = lifecycleClaimSignalRepo{repoStore: repo, claimed: claimed} + svc.taskRuntimeStateMu.Lock() + + promptDone := make(chan error, 1) + go func() { + _, _, err := svc.claimLifecyclePromptDispatch(ctx, session.TaskID, session.ID, "", nil) + promptDone <- err + }() + <-claimed + + resetDone := make(chan bool, 1) + go func() { + resetDone <- svc.resetAgentContext(ctx, session.TaskID, session, "Successor") + }() + + // The lifecycle claim is now blocked before its task-state reconciliation. + // Reset must still see and cancel its turn before replacing provider context; + // the lifecycle claim must then reject the cancelled turn. + requireResetEvents(t, manager.events, "cancel", "reset") + if resetOK := <-resetDone; !resetOK { + t.Fatal("resetAgentContext returned false") + } + + svc.taskRuntimeStateMu.Unlock() + select { + case err := <-promptDone: + if !errors.Is(err, ErrSessionResetInProgress) { + t.Fatalf("lifecycle prompt claim error = %v, want %v", err, ErrSessionResetInProgress) + } + case <-time.After(time.Second): + t.Fatal("timed out waiting for lifecycle prompt claim") + } +} diff --git a/apps/backend/internal/orchestrator/task_operations.go b/apps/backend/internal/orchestrator/task_operations.go index 74c7bdac03..4a235e53a3 100644 --- a/apps/backend/internal/orchestrator/task_operations.go +++ b/apps/backend/internal/orchestrator/task_operations.go @@ -3943,12 +3943,11 @@ func (promptTaskOptions) failureContext(ctx context.Context) (context.Context, c // promptTask is PromptTask's implementation. Its options carry queued-dispatch // ownership and the bounded-context exception. When options.claimEntryID is -// non-empty, this -// acquires sessionID's cancelInFlight guard around a second ownership check -// and the "mark the session RUNNING" step immediately before startTurnForSession -// and the (potentially long-blocking) executor.Prompt call below it. The worker -// claims the handoff before visible side effects; this check keeps the final -// session transition tied to that same ownership record. +// non-empty, this acquires sessionID's cancelInFlight guard around a second +// ownership check and the "mark the session RUNNING" step immediately before +// startTurnForSession and the (potentially long-blocking) executor.Prompt call +// below it. The worker claims the handoff before visible side effects; this +// check keeps the final session transition tied to that same ownership record. // // A bare (unguarded) "check the token, then separately call // setSessionRunning" would still let two settling dispatches for the same @@ -3958,10 +3957,10 @@ func (promptTaskOptions) failureContext(ctx context.Context) (context.Context, c // *and* the mark through the same per-session guard used for every other // cancel/take-and-dispatch decision (see the Service.cancelInFlight field // doc comment) makes "exactly one dispatch wins" a property of mutual -// exclusion instead of a racy read. The guard is held only for this -// fast, DB-only step — released well before the long-running -// executor.Prompt call — matching every other guard holder's -// bounded-hold-time contract in this file. +// exclusion instead of a racy read. Ordinary prompts release the guard after +// this fast, DB-only step. Lifecycle prompts hold the same guard until agentctl +// accepts the prompt, so a context reset cannot cancel the claimed turn in the +// gap before provider dispatch. // // Every return path *before* this step (invalid session, not promptable, // ensureSessionRunning failure, a model switch) is covered by @@ -4048,7 +4047,11 @@ func (s *Service) promptTask(ctx context.Context, taskID, sessionID string, prom s.publishForegroundActivityChanged(ctx, taskID, sessionID) } var releaseDispatchGuard func() - if options.requireNonterminalSession { + if rollback.dispatchGuardRelease != nil { + releaseDispatchGuard = rollback.dispatchGuardRelease + defer releaseDispatchGuard() + } + if options.requireNonterminalSession && releaseDispatchGuard == nil { var guard *sync.Mutex guard, releaseDispatchGuard = s.acquireCancelInFlightGuard(sessionID) guard.Lock() @@ -4333,6 +4336,7 @@ type promptClaimRollback struct { createdTurn bool reservedTurn *models.Turn reservedTurnAccepted bool + dispatchGuardRelease func() } func (s *Service) claimPromptDispatch( @@ -4372,29 +4376,44 @@ func (s *Service) claimLifecyclePromptDispatch( taskID, sessionID, claimEntryID string, afterClaim func() error, ) (*models.TaskSession, promptClaimRollback, error) { - session, previousSessionState, err := s.claimLifecycleSessionRunning( + session, previousSessionState, turnID, createdTurn, err := s.claimLifecycleSessionRunning( ctx, taskID, sessionID, claimEntryID, ) if err != nil { return nil, promptClaimRollback{}, err } - rollback := promptClaimRollback{previousSessionState: previousSessionState} + rollback := promptClaimRollback{ + previousSessionState: previousSessionState, + turnID: turnID, + createdTurn: createdTurn, + } rollback.previousTaskState, rollback.taskStateClaimed, err = s.reconcileLifecycleClaim( ctx, taskID, sessionID, previousSessionState, session, ) if err != nil { + s.rollbackPromptClaimAfterAdmissionFailure(ctx, taskID, sessionID, rollback) return nil, promptClaimRollback{}, err } - rollback.turnID, rollback.createdTurn, _, err = s.startTurnForSessionWithOwnershipChecked( - ctx, sessionID, false, nil, - ) - if err != nil { + dispatchGuard := s.lockCancelInFlightGuard(sessionID) + var releaseGuardOnce sync.Once + releaseDispatchGuard := func() { + releaseGuardOnce.Do(dispatchGuard.release) + } + keepDispatchGuard := false + defer func() { + if !keepDispatchGuard { + releaseDispatchGuard() + } + }() + rollback.dispatchGuardRelease = releaseDispatchGuard + if err := s.validateLifecyclePromptTurn(ctx, sessionID, rollback.turnID); err != nil { s.rollbackPromptClaimAfterAdmissionFailure(ctx, taskID, sessionID, rollback) - return nil, promptClaimRollback{}, fmt.Errorf("persist lifecycle prompt turn: %w", err) + return nil, promptClaimRollback{}, err } if err := s.acknowledgePromptClaim(ctx, taskID, sessionID, afterClaim, rollback); err != nil { return nil, promptClaimRollback{}, err } + keepDispatchGuard = true return session, rollback, nil } @@ -4432,6 +4451,9 @@ func (s *Service) rollbackPromptClaim( taskID, sessionID string, rollback promptClaimRollback, ) { + if rollback.dispatchGuardRelease != nil { + rollback.dispatchGuardRelease() + } s.restoreLifecycleClaim(ctx, taskID, sessionID, rollback.previousSessionState) if rollback.taskStateClaimed { s.restoreLifecycleTaskState(ctx, taskID, rollback.previousTaskState) @@ -4675,6 +4697,12 @@ func (s *Service) claimSessionRunningForPrompt( defer release() lock.Lock() defer lock.Unlock() + // A reset-owned cancellation must not make prompt admission wait for the + // reset that already forbids it. Keep the second check below for a reset + // that starts while this prompt is waiting on an unrelated cancellation. + if s.isSessionResetInProgress(sessionID) { + return nil, "", "", false, nil, ErrSessionResetInProgress + } if err := s.waitForCancellationWithGuard(ctx, sessionID, lock.Unlock, lock.Lock); err != nil { return nil, "", "", false, nil, err } @@ -4773,52 +4801,67 @@ func (s *Service) claimSessionRunningForPrompt( func (s *Service) claimLifecycleSessionRunning( ctx context.Context, taskID, sessionID, claimEntryID string, -) (*models.TaskSession, models.TaskSessionState, error) { +) (*models.TaskSession, models.TaskSessionState, string, bool, error) { lock, release := s.acquireCancelInFlightGuard(sessionID) defer release() lock.Lock() defer lock.Unlock() + // Match ordinary prompt admission: a reset-owned cancellation must reject + // lifecycle admission before its cancellation wait can observe ctx.Err(). + if s.isSessionResetInProgress(sessionID) { + return nil, "", "", false, ErrSessionResetInProgress + } if err := s.waitForCancellationWithGuard(ctx, sessionID, lock.Unlock, lock.Lock); err != nil { - return nil, "", err + return nil, "", "", false, err } if s.isSessionResetInProgress(sessionID) { - return nil, "", ErrSessionResetInProgress + return nil, "", "", false, ErrSessionResetInProgress } if claimEntryID != "" && !s.isCurrentQueuedDispatch(sessionID, claimEntryID) { - return nil, "", errQueuedDispatchSuperseded + return nil, "", "", false, errQueuedDispatchSuperseded } if err := s.validateLifecyclePromptSelection(ctx, taskID, sessionID); err != nil { - return nil, "", err + return nil, "", "", false, err } claim, err := s.repo.ClaimPromptableTaskSessionIfActive(ctx, sessionID) if err != nil { - return nil, "", fmt.Errorf("%w: %v", errLifecyclePromptClaim, err) + return nil, "", "", false, fmt.Errorf("%w: %v", errLifecyclePromptClaim, err) } switch claim.Status { case models.PromptableTaskSessionInactive: - return nil, "", errLifecyclePromptInactive + return nil, "", "", false, errLifecyclePromptInactive case models.PromptableTaskSessionBusy: - return nil, "", fmt.Errorf("%w: lifecycle prompt session is busy", ErrAgentPromptInProgress) + return nil, "", "", false, fmt.Errorf("%w: lifecycle prompt session is busy", ErrAgentPromptInProgress) } if claimEntryID != "" && !s.isCurrentQueuedDispatch(sessionID, claimEntryID) { s.restoreLifecycleClaim(ctx, taskID, sessionID, claim.PreviousState) - return nil, "", errQueuedDispatchSuperseded + return nil, "", "", false, errQueuedDispatchSuperseded } if s.isSessionResetInProgress(sessionID) { s.restoreLifecycleClaim(ctx, taskID, sessionID, claim.PreviousState) - return nil, "", ErrSessionResetInProgress + return nil, "", "", false, ErrSessionResetInProgress } freshSession, err := s.repo.GetTaskSession(ctx, sessionID) if err != nil || freshSession == nil { s.restoreLifecycleClaim(ctx, taskID, sessionID, claim.PreviousState) - return nil, "", errLifecyclePromptInactive + return nil, "", "", false, errLifecyclePromptInactive + } + // Keep the lifecycle admission guard through durable turn creation. A reset + // that follows this claim must observe this turn and quiesce it before it + // replaces the provider context. + turnID, createdTurn, _, err := s.startTurnForSessionWithOwnershipChecked(ctx, sessionID, false, nil) + if err != nil { + failureCtx, cancel := (promptTaskOptions{}).failureContext(ctx) + defer cancel() + s.restoreLifecycleClaim(failureCtx, taskID, sessionID, claim.PreviousState) + return nil, "", "", false, fmt.Errorf("persist lifecycle prompt turn: %w", err) } s.releaseQueuedDispatchPendingIfCurrent( sessionID, s.queuedDispatchReservationForEntry(sessionID, claimEntryID), ) - return freshSession, claim.PreviousState, nil + return freshSession, claim.PreviousState, turnID, createdTurn, nil } func (s *Service) validateLifecyclePromptSelection( @@ -4839,6 +4882,26 @@ func (s *Service) validateLifecyclePromptSelection( return nil } +func (s *Service) validateLifecyclePromptTurn( + ctx context.Context, + sessionID, turnID string, +) error { + if s.isSessionResetInProgress(sessionID) { + return ErrSessionResetInProgress + } + if s.turnService == nil || turnID == "" { + return nil + } + activeTurn, err := s.turnService.GetActiveTurn(ctx, sessionID) + if err != nil { + return fmt.Errorf("verify lifecycle prompt turn: %w", err) + } + if activeTurn == nil || activeTurn.ID != turnID { + return ErrSessionResetInProgress + } + return nil +} + func (s *Service) restoreLifecycleClaim( ctx context.Context, taskID, sessionID string,