fix: preserve resumable auth and web batches

Author: dcramer

packages/junior/src/chat/runtime/slack-turn.ts

@@ -981,6 +981,7 @@ export function createSlackTurn(deps: SlackTurnDeps) {
         // Once model output is settled, later commit errors must not trigger a
         // second visible failure reply.
         let runResultHandled = false;
+        let authPauseAccepted = false;
         let acceptedDeliveryId: string | undefined;
         let turnCompletionNotified = false;
         const recordDispatchOutcome = async (
@@ -1516,9 +1517,12 @@ export function createSlackTurn(deps: SlackTurnDeps) {
               conversation: preparedState.conversation,
               sessionId: turnId,
             });
-            await persistThreadState(thread, {
-              conversation: preparedState.conversation,
-            });
+            authPauseAccepted = true;
+            await persistWithRetry(() =>
+              persistThreadState(thread, {
+                conversation: preparedState.conversation,
+              }),
+            );
             persistedAtLeastOnce = true;
             shouldPersistFailureState = false;
             return;
@@ -1602,6 +1606,24 @@ export function createSlackTurn(deps: SlackTurnDeps) {
             }
           }
         } catch (error) {
+          if (authPauseAccepted && !persistedAtLeastOnce) {
+            // The private auth link is already visible. Reconcile the paused
+            // state without turning a resumable Turn into a terminal failure.
+            logException(error, "slack.auth_pause_state_persist.failed");
+            try {
+              await persistThreadState(thread, {
+                conversation: preparedState.conversation,
+              });
+              persistedAtLeastOnce = true;
+              shouldPersistFailureState = false;
+              return;
+            } catch (recoveryError) {
+              logException(
+                recoveryError,
+                "slack.auth_pause_state_recovery.failed",
+              );
+            }
+          }
           if (runResultHandled) {
             // Errors after the completed run produced output or intentional
             // silence (redundant-ack cleanup, completion callbacks) must not

packages/junior/src/chat/task-execution/web-cancellation.ts

@@ -92,18 +92,20 @@ export async function completeCancelledWebTurn(args: {
   sandboxRef?: SandboxRef;
   signal?: AbortSignal;
   turnId: string;
-  userMessageId: string;
+  userMessageIds: readonly string[];
 }): Promise<void> {
   await abandonTurnRecord({
     conversationId: args.conversationId,
     turnId: args.turnId,
     errorMessage: "Web Turn cancelled",
   });
   clearPendingAuth(args.conversation, args.turnId);
-  markConversationMessage(args.conversation, args.userMessageId, {
-    replied: false,
-    skippedReason: "turn cancelled",
-  });
+  for (const userMessageId of args.userMessageIds) {
+    markConversationMessage(args.conversation, userMessageId, {
+      replied: false,
+      skippedReason: "turn cancelled",
+    });
+  }
   markTurnClosed({
     conversation: args.conversation,
     nowMs: Date.now(),

packages/junior/src/chat/task-execution/web-work.ts

@@ -206,6 +206,7 @@ export function createWebWorker(options: {
     let text: string;
     let turnId: string;
     let userMessageId: string;
+    let userMessageIds: string[];
     let startedAtMs: number;
 
     const storedConversation = await getConversationStore().get({
@@ -238,6 +239,7 @@ export function createWebWorker(options: {
 
     if (resolved.kind === "mailbox") {
       const batch = mailboxBatch?.length ? mailboxBatch : resolved.batch;
+      mailboxBatch = batch;
       const first = batch[0]!;
       text = batch
         .map((entry) => entry.message.input.text.trim())
@@ -246,6 +248,7 @@ export function createWebWorker(options: {
       actor = actorFromMetadata(first.metadata);
       turnId = webTurnIdForMessage(first.metadata.messageId);
       userMessageId = first.metadata.messageId;
+      userMessageIds = batch.map((entry) => entry.metadata.messageId);
       startedAtMs = first.message.createdAtMs;
     } else {
       turnId = resolved.turnId;
@@ -266,6 +269,7 @@ export function createWebWorker(options: {
       });
       text = userMessage.text;
       userMessageId = userMessage.id;
+      userMessageIds = [userMessage.id];
       startedAtMs = userMessage.createdAtMs;
     }
 
@@ -292,6 +296,23 @@ export function createWebWorker(options: {
         ...(actor.userName ? { userName: actor.userName } : {}),
       },
       async () => {
+        const markUserMessages = (
+          target: typeof conversation,
+          patch: { replied?: boolean; skippedReason?: string },
+        ): void => {
+          for (const messageId of userMessageIds) {
+            markConversationMessage(target, messageId, patch);
+          }
+        };
+        const buildInstructionContext = (): string | undefined => {
+          const currentMessageIds = new Set(userMessageIds);
+          return buildConversationContext({
+            ...conversation,
+            messages: conversation.messages.filter(
+              (message) => !currentMessageIds.has(message.id),
+            ),
+          });
+        };
         let acknowledged = isResume || context.attempt.messages.length === 0;
         const acknowledge = async (): Promise<void> => {
           if (acknowledged) {
@@ -370,33 +391,38 @@ export function createWebWorker(options: {
               conversationId: context.conversationId,
             });
           }
-          upsertConversationMessage(conversation, {
-            id: userMessageId,
-            role: "user",
-            text: normalizeConversationText(text),
-            createdAtMs: startedAtMs,
-            author: {
-              ...(actor.email ? { email: actor.email } : {}),
-              ...(actor.fullName ? { fullName: actor.fullName } : {}),
-              userId: actor.userId,
-              ...(actor.userName ? { userName: actor.userName } : {}),
-            },
-            meta: {
-              explicitMention: true,
-              replied: false,
-              source: "web",
-            },
-          });
+          for (const entry of mailboxBatch!) {
+            const messageActor = actorFromMetadata(entry.metadata);
+            upsertConversationMessage(conversation, {
+              id: entry.metadata.messageId,
+              role: "user",
+              text: normalizeConversationText(entry.message.input.text),
+              createdAtMs: entry.message.createdAtMs,
+              author: {
+                ...(messageActor.email ? { email: messageActor.email } : {}),
+                ...(messageActor.fullName
+                  ? { fullName: messageActor.fullName }
+                  : {}),
+                userId: messageActor.userId,
+                ...(messageActor.userName
+                  ? { userName: messageActor.userName }
+                  : {}),
+              },
+              meta: {
+                explicitMention: true,
+                replied: false,
+                source: "web",
+              },
+            });
+          }
           await persistConversationMessages({
             conversation,
             conversationId: context.conversationId,
           });
           await lifecycle.start({
             conversationId: context.conversationId,
             createdAtMs: Date.now(),
-            inputMessageIds: mailboxBatch!.map(
-              (entry) => entry.metadata.messageId,
-            ),
+            inputMessageIds: userMessageIds,
             surface: "api",
             turnId,
           });
@@ -409,6 +435,7 @@ export function createWebWorker(options: {
         let currentRunId: string | undefined;
         let failureCode: ConversationTurnFailureCode = "persistence_failed";
         let assistantMessageAccepted = false;
+        let authPauseAccepted = false;
         let cancellationCleanupStarted = false;
         let terminalConversationStateCommitted = false;
         const completeCancelledTurn =
@@ -425,7 +452,7 @@ export function createWebWorker(options: {
                 sandboxRef,
                 signal: cancellationSignal,
                 turnId,
-                userMessageId,
+                userMessageIds,
               });
             } catch (error) {
               if (hasLostTurnInputCommit(error)) {
@@ -458,6 +485,10 @@ export function createWebWorker(options: {
             text: replyText,
             userMessageId,
           });
+          markUserMessages(conversation, {
+            replied: true,
+            skippedReason: undefined,
+          });
           assistantMessageAccepted = true;
           try {
             await persistWithRetry(() =>
@@ -515,6 +546,10 @@ export function createWebWorker(options: {
                   sessionId: turnId,
                   userMessageId,
                 });
+                markUserMessages(completedState.conversation, {
+                  replied: true,
+                  skippedReason: undefined,
+                });
                 await persistThreadStateById(context.conversationId, {
                   conversation: completedState.conversation,
                   sandboxRef: result.sandboxRef ?? sandboxRef,
@@ -558,6 +593,12 @@ export function createWebWorker(options: {
                       markConversationMessage,
                     });
                   }
+                  markUserMessages(
+                    conversation,
+                    turnHasReply(conversation, turnId)
+                      ? { replied: true, skippedReason: undefined }
+                      : { replied: false, skippedReason: "reply failed" },
+                  );
                   try {
                     await persistThreadStateById(context.conversationId, {
                       conversation,
@@ -594,9 +635,7 @@ export function createWebWorker(options: {
               runId: currentRunId,
               instruction: {
                 text,
-                context: buildConversationContext(conversation, {
-                  excludeMessageId: userMessageId,
-                }),
+                context: buildInstructionContext(),
               },
               history: piMessages,
               actor,
@@ -651,10 +690,17 @@ export function createWebWorker(options: {
               conversation,
               sessionId: turnId,
             });
-            await persistThreadStateById(context.conversationId, {
-              conversation,
-              sandboxRef,
+            markUserMessages(conversation, {
+              replied: true,
+              skippedReason: undefined,
             });
+            authPauseAccepted = true;
+            await persistWithRetry(() =>
+              persistThreadStateById(context.conversationId, {
+                conversation,
+                sandboxRef,
+              }),
+            );
             if (cancellationSignal) {
               options.cancellation?.park(
                 context.conversationId,
@@ -722,6 +768,46 @@ export function createWebWorker(options: {
             }
             return await completeCancelledTurn();
           }
+          if (authPauseAccepted) {
+            // The auth request is already visible. Reconcile the paused state
+            // without turning a resumable Turn into a terminal failure.
+            captureWebBoundaryFailure({
+              conversationId: context.conversationId,
+              error,
+              failureCode: "persistence_failed",
+              runId: currentRunId,
+              turnId,
+            });
+            try {
+              await persistThreadStateById(context.conversationId, {
+                conversation,
+                sandboxRef,
+              });
+            } catch (persistenceError) {
+              if (context.attempt.isFinalAttempt) {
+                finishCancellation();
+              }
+              throw new AggregateError(
+                [error, persistenceError],
+                "Web auth pause state could not be persisted",
+              );
+            }
+            if (cancellationSignal) {
+              options.cancellation?.park(
+                context.conversationId,
+                cancellationSignal,
+              );
+            }
+            try {
+              await acknowledge();
+            } catch (acknowledgeError) {
+              if (hasLostTurnInputCommit(acknowledgeError)) {
+                return { status: "lost_lease" };
+              }
+              throw acknowledgeError;
+            }
+            return { status: "completed" };
+          }
           if (!context.attempt.isFinalAttempt) {
             throw error;
           }
@@ -753,6 +839,12 @@ export function createWebWorker(options: {
                   markConversationMessage,
                 });
               }
+              markUserMessages(
+                conversation,
+                turnHasReply(conversation, turnId) || userMessageHandled
+                  ? { replied: true, skippedReason: undefined }
+                  : { replied: false, skippedReason: "reply failed" },
+              );
               await persistThreadStateById(context.conversationId, {
                 conversation,
                 sandboxRef: initialSandboxRef ?? null,

packages/junior/tests/component/auth/mcp-auth-runtime-slack.test.ts

@@ -424,7 +424,7 @@ describe("mcp auth runtime slack integration", () => {
     process.env = { ...ORIGINAL_ENV };
   }, 45_000);
 
-  it("parks an MCP auth challenge from the real Slack runtime and resumes after OAuth callback", async () => {
+  it("keeps a delivered MCP auth challenge resumable after a transient state failure", async () => {
     const threadId = "slack:C123:1700000000.001";
     const turnId = "turn_user-1";
     const { createTestChatRuntime } = chatRuntimeModule;
@@ -461,6 +461,22 @@ describe("mcp auth runtime slack integration", () => {
       },
     });
     await mirrorThreadStateToAdapter(thread);
+    const stateAdapter = stateAdapterModule.getStateAdapter();
+    const setState = stateAdapter.set.bind(stateAdapter);
+    let authPauseStateFailures = 0;
+    vi.spyOn(stateAdapter, "set").mockImplementation(
+      async (key, value, ttlMs) => {
+        if (
+          authPauseStateFailures < 3 &&
+          key === `thread-state:${threadId}` &&
+          getCapturedSlackApiCalls("chat.postEphemeral").length > 0
+        ) {
+          authPauseStateFailures += 1;
+          throw new Error("thread state unavailable");
+        }
+        await setState(key, value, ttlMs);
+      },
+    );
 
     await slackRuntime.handleNewMention(
       thread,
@@ -485,6 +501,7 @@ describe("mcp auth runtime slack integration", () => {
 
     expect(agentProbe.promptCallCount).toBe(1);
     expect(agentProbe.continueCallCount).toBe(0);
+    expect(authPauseStateFailures).toBe(3);
 
     expect(getCapturedSlackApiCalls("chat.postEphemeral")).toEqual([
       expect.objectContaining({