fix: reject prompt admission during context reset (#3169)

Co-authored-by: ayattara <alassane.yattara@savoirfairelinux.com>
Co-authored-by: Carlos Florencio <carlosmflorencio@gmail.com>

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.

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

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,