fix(orchestrator): preserve step prompt across async agent start failure (#3792)

Co-authored-by: nova28 <17953305+nova28@users.noreply.github.com> Co-authored-by: Carlos Florencio <carlosmflorencio@gmail.com>
2026-09-18 22:19:55 +08:00
parent 3c1bf2da4d
commit 1b02c5fa32
5 changed files with 140 additions and 13 deletions
@@ -154,6 +154,19 @@ func (a *workflowStartPromptAttempt) claimPreservation() bool {
 	)
 }

+// releasePreservationClaim makes a failed queue admission retryable for this
+// launch attempt. The claim is shared by every callback context, so a later
+// callback can retry only while it still owns the same attempt.
+func (a *workflowStartPromptAttempt) releasePreservationClaim() bool {
+	if a == nil || a.claim == nil {
+		return false
+	}
+	return a.claim.state.CompareAndSwap(
+		uint32(workflowStartPromptAttemptPreservationClaimed),
+		uint32(workflowStartPromptAttemptActive),
+	)
+}
+
 func (a *workflowStartPromptAttempt) hasInput() bool {
 	return a != nil && (a.dispatchInputPresent || len(a.attachments) > 0 ||
 		len(a.references) > 0 || strings.TrimSpace(a.handoffText) != "")
@@ -351,6 +364,33 @@ func (s *Service) preserveWorkflowStartPromptAfterFailure(
 	if !attempt.claimPreservation() {
 		return
 	}
+	err := s.persistWorkflowStartPromptAttempt(ctx, attempt)
+	if err != nil && s.workflowStartPromptAttemptIsCurrent(ctx, attempt, executionID) {
+		// Queue admission can fail before the transaction writes a row. Retry
+		// once while this callback still owns the immutable launch attempt. The
+		// attempt claim prevents duplicate callbacks from racing this retry.
+		err = s.persistWorkflowStartPromptAttempt(ctx, attempt)
+	}
+	if err != nil {
+		attempt.releasePreservationClaim()
+		// Keep the startup error as the user-visible recovery signal. This
+		// diagnostic carries only ownership identifiers, never prompt content.
+		s.logger.Warn("failed to preserve workflow prompt after asynchronous startup failure",
+			zap.String("task_id", attempt.taskID),
+			zap.String("session_id", attempt.sessionID),
+			zap.String("agent_execution_id", executionID),
+			zap.String("launch_token", attempt.launchToken),
+			zap.Error(err))
+	}
+}
+
+func (s *Service) persistWorkflowStartPromptAttempt(
+	ctx context.Context,
+	attempt *workflowStartPromptAttempt,
+) error {
+	if attempt == nil {
+		return fmt.Errorf("workflow start prompt attempt is nil")
+	}
 	var err error
 	if attempt.workflowEntryCaptured {
 		_, err = s.persistAutoStartPromptAtWorkflowEntry(
@@ -383,16 +423,7 @@ func (s *Service) preserveWorkflowStartPromptAfterFailure(
 			attempt.handoffText,
 		)
 	}
-	if err != nil {
-		// Keep the startup error as the user-visible recovery signal. This
-		// diagnostic carries only ownership identifiers, never prompt content.
-		s.logger.Warn("failed to preserve workflow prompt after asynchronous startup failure",
-			zap.String("task_id", attempt.taskID),
-			zap.String("session_id", attempt.sessionID),
-			zap.String("agent_execution_id", executionID),
-			zap.String("launch_token", attempt.launchToken),
-			zap.Error(err))
-	}
+	return err
 }

 func (s *Service) workflowStartPromptAttemptIsCurrent(
@@ -263,6 +263,45 @@ func TestWorkflowAsyncStartFailure_QueueFailureKeepsLaunchError(t *testing.T) {
 	}
 }

+func TestWorkflowAsyncStartFailure_RetriesTransientQueueFailure(t *testing.T) {
+	fixture := newWorkflowAsyncStartFailureFixture(t, errors.New("provider startup failed"))
+	queueRepo := &workflowStartPromptQueueRetryRepository{
+		Repository: newAuthoritativeMemoryRepository(fixture.repo),
+		err:        errors.New("temporary queue failure"),
+	}
+	queueRepo.remainingFailures.Store(1)
+	fixture.svc.messageQueue = messagequeue.NewService(
+		queueRepo, messagequeue.DefaultMaxPerSession, testLogger(),
+	)
+
+	if err := fixture.svc.autoStartStepPrompt(
+		context.Background(), fixture.taskID, fixture.session, fixture.step,
+		fixture.prompt, false, true, nil,
+	); err != nil {
+		t.Fatalf("autoStartStepPrompt returned launch error before asynchronous failure: %v", err)
+	}
+	select {
+	case <-fixture.startEntered:
+	case <-time.After(3 * time.Second):
+		t.Fatal("timed out waiting for asynchronous startup")
+	}
+	close(fixture.releaseStart)
+	select {
+	case <-fixture.startReturned:
+	case <-time.After(3 * time.Second):
+		t.Fatal("timed out waiting for startup failure callback")
+	}
+	waitForSessionState(t, fixture.repo, fixture.sessionID, models.TaskSessionStateFailed)
+
+	status := fixture.svc.messageQueue.GetStatus(context.Background(), fixture.sessionID)
+	if status.Count != 1 || status.Entries[0].Content != fixture.prompt {
+		t.Fatalf("queue after transient failure = %#v, want one preserved prompt", status)
+	}
+	if got := queueRepo.insertCalls.Load(); got != 2 {
+		t.Fatalf("queue insert attempts = %d, want one retry after the first failure", got)
+	}
+}
+
 func TestWorkflowAsyncStartFailure_QueueFullKeepsLaunchError(t *testing.T) {
 	fixture := newWorkflowAsyncStartFailureFixture(t, errors.New("provider startup failed"))
 	ctx := context.Background()
@@ -568,6 +607,46 @@ func (r *workflowStartPromptQueueFailureRepository) InsertForSession(
 	return r.err
 }

+type workflowStartPromptQueueRetryRepository struct {
+	messagequeue.Repository
+	err               error
+	remainingFailures atomic.Int32
+	insertCalls       atomic.Int32
+}
+
+func (r *workflowStartPromptQueueRetryRepository) shouldFail() bool {
+	for {
+		remaining := r.remainingFailures.Load()
+		if remaining <= 0 {
+			return false
+		}
+		if r.remainingFailures.CompareAndSwap(remaining, remaining-1) {
+			return true
+		}
+	}
+}
+
+func (r *workflowStartPromptQueueRetryRepository) Insert(
+	ctx context.Context, msg *messagequeue.QueuedMessage, maxPerSession int,
+) error {
+	r.insertCalls.Add(1)
+	if r.shouldFail() {
+		return r.err
+	}
+	return r.Repository.Insert(ctx, msg, maxPerSession)
+}
+
+func (r *workflowStartPromptQueueRetryRepository) InsertForSession(
+	ctx context.Context, identity messagequeue.QueueSessionIdentity,
+	msg *messagequeue.QueuedMessage, maxPerSession int,
+) error {
+	r.insertCalls.Add(1)
+	if r.shouldFail() {
+		return r.err
+	}
+	return r.Repository.InsertForSession(ctx, identity, msg, maxPerSession)
+}
+
 type workflowAsyncStartFailureFixture struct {
 	svc       *Service
 	repo      *sqliterepo.Repository
@@ -32,6 +32,8 @@ Keep existing recovery actions, queue policy, failure projection, and executor c
 - Add the private attempt envelope and retire it on success or synchronous rejection.
 - Bind task, session, execution evidence, turn, workflow entry, and final dynamic attempt ownership.
 - Persist raw queue-form input once through the accepted failure path, without immediate auto-resume scheduling.
+- Retry queue admission once while the same attempt owns the preservation
+  claim; keep the launch error visible after a failed retry.
 - Retain recorded-message identity, attachment ownership, references, handoff, origin, and plan mode.
 - Add all tests named in the plan, including failure-to-recovery service coverage.

@@ -44,6 +46,9 @@ New UI, public APIs, migrations, provider retries, and the existing post-start t
 1. An asynchronous failure after successful launch return preserves the exact input and recovery delivers it once with one transcript row.
 2. Stale, duplicate, terminal, archived, deleted, and superseded-entry callbacks cannot enqueue input or affect successor work.
 3. Error projection, paused queues, synchronous failures, dynamic fallback, and success retain their established behavior.
+4. A transient queue insertion error retries without losing prompt metadata or
+   scheduling an automatic resume; a second failure keeps the original launch
+   error visible.

 ## Regression procedure

@@ -62,7 +67,8 @@ Exercise duplicate callback delivery and a failure before launch return.
 Race replacement execution, same-execution successor turn, cancellation, completion, archive, deletion, and a later workflow entry.
 Exercise a dynamic fallback candidate followed by success and a final failed candidate.
 Cover combined handoff, attachment-only input, entity references, plan mode, and failed original message persistence.
-Cover queue-full/write errors, a paused queue, and reopening the queue repository after preservation.
+Cover queue-full/write errors, one transient insertion retry, a paused queue,
+and reopening the queue repository after preservation.
 Verify no replay from successful startup or post-start prompt failure.

 ## Verification
@@ -143,6 +149,10 @@ metadata reads fail, and recovery uses launch-time effective-input and
 config-mode metadata rather than reclassifying a mode-only prompt from the
 mutable session projection.

+Queue persistence now retries once while the attempt claim remains owned. The
+retry keeps the original launch error as the recovery signal and does not start
+an automatic replacement.
+
 Verification passed:

 ```text
@@ -113,7 +113,10 @@ This draft extension addresses [issue #3753](https://github.com/kdlbs/kandev/iss
 - **AC-TASKS-WORKFLOW-STEP-AGENT-START-OWNERSHIP-005.4:** Prompt preservation shall retain the existing launch-error classification and recovery actions.
  It shall not initiate another launch from the failure callback or turn permanent rejection into an automatic retry loop.
 - **AC-TASKS-WORKFLOW-STEP-AGENT-START-OWNERSHIP-005.5:** After successful queue persistence, the preserved prompt shall survive a backend restart.
-  If persistence fails, the launch error shall remain visible and diagnostics shall identify the preservation failure without exposing prompt content.
+  If queue admission fails, the same launch attempt shall retry once while it
+  still owns the preservation claim. After that retry fails, the launch error
+  shall remain visible and diagnostics shall identify the preservation failure
+  without exposing prompt content.
 - **AC-TASKS-WORKFLOW-STEP-AGENT-START-OWNERSHIP-005.6:** Successful starts and synchronous permanent rejections shall not create asynchronous recovery entries.
  Existing synchronous busy-error recovery shall retain its current behavior.

@@ -327,7 +327,11 @@ Recovery boots the session and lets the existing queue drain send its preserved
 Do not send the same prompt as both a launch description and a queue entry.

 A queue write failure does not mask the original startup failure.
-Log a separate bounded preservation diagnostic with task, session, execution, and attempt identity, without input text or attachment contents.
+The attempt-owned preservation path retries queue admission once while its
+claim remains current. It releases the claim only after that retry fails, so a
+later callback can retry only with the same launch ownership. Log a separate
+bounded preservation diagnostic with task, session, execution, and attempt
+identity, without input text or attachment contents.
 The existing persistent launch error remains the visible recovery signal.
 No frontend layout or copy changes are required by this package.