fix(agent): continue unattended turns after transient upstream stream failures

packages/agent/src/server/agent-server.test.ts

@@ -956,6 +956,123 @@ describe("AgentServer HTTP Mode", () => {
       },
     );
 
+    function createRetryTestServer(prompt: ReturnType<typeof vi.fn>) {
+      const testServer = createFailureTestServer();
+      testServer.session = {
+        acpSessionId: "acp-1",
+        payload: { run_id: "run-1" },
+        logWriter: { appendRawLine: vi.fn(), flush: vi.fn(async () => {}) },
+        clientConnection: { prompt },
+      };
+      return testServer as unknown as {
+        promptWithUpstreamRetry(request: {
+          sessionId: string;
+          prompt: ContentBlock[];
+        }): Promise<{ stopReason: string }>;
+      };
+    }
+
+    it("continues an unattended turn after a transient upstream stream death", async () => {
+      vi.useFakeTimers();
+      try {
+        const prompt = vi
+          .fn()
+          .mockRejectedValueOnce(new Error("API Error: terminated"))
+          .mockResolvedValueOnce({ stopReason: "end_turn" });
+        const testServer = createRetryTestServer(prompt);
+
+        const resultPromise = testServer.promptWithUpstreamRetry({
+          sessionId: "acp-1",
+          prompt: [{ type: "text", text: "do the task" }],
+        });
+        await vi.advanceTimersByTimeAsync(5_000);
+
+        await expect(resultPromise).resolves.toEqual({
+          stopReason: "end_turn",
+        });
+        expect(prompt).toHaveBeenCalledTimes(2);
+        const retryRequest = prompt.mock.calls[1][0] as {
+          sessionId: string;
+          prompt: Array<{ type: string; text: string }>;
+        };
+        expect(retryRequest.sessionId).toBe("acp-1");
+        expect(retryRequest.prompt[0].text).toContain(
+          "interrupted by a transient connection error",
+        );
+      } finally {
+        vi.useRealTimers();
+      }
+    });
+
+    it("re-sends the original prompt when the failure happened before the stream started", async () => {
+      vi.useFakeTimers();
+      try {
+        const prompt = vi
+          .fn()
+          .mockRejectedValueOnce(new Error("API Error: Connection error."))
+          .mockResolvedValueOnce({ stopReason: "end_turn" });
+        const testServer = createRetryTestServer(prompt);
+
+        const resultPromise = testServer.promptWithUpstreamRetry({
+          sessionId: "acp-1",
+          prompt: [{ type: "text", text: "do the task" }],
+        });
+        await vi.advanceTimersByTimeAsync(5_000);
+
+        await expect(resultPromise).resolves.toEqual({
+          stopReason: "end_turn",
+        });
+        expect(prompt).toHaveBeenCalledTimes(2);
+        const retryRequest = prompt.mock.calls[1][0] as {
+          sessionId: string;
+          prompt: Array<{ type: string; text: string }>;
+        };
+        expect(retryRequest.sessionId).toBe("acp-1");
+        expect(retryRequest.prompt).toEqual([
+          { type: "text", text: "do the task" },
+        ]);
+      } finally {
+        vi.useRealTimers();
+      }
+    });
+
+    it("does not retry a genuine agent error", async () => {
+      const prompt = vi.fn().mockRejectedValue(new Error("boom"));
+      const testServer = createRetryTestServer(prompt);
+
+      await expect(
+        testServer.promptWithUpstreamRetry({
+          sessionId: "acp-1",
+          prompt: [{ type: "text", text: "do the task" }],
+        }),
+      ).rejects.toThrow("boom");
+      expect(prompt).toHaveBeenCalledTimes(1);
+    });
+
+    it("stops continuing once the bounded retry budget is exhausted", async () => {
+      vi.useFakeTimers();
+      try {
+        const prompt = vi
+          .fn()
+          .mockRejectedValue(new Error("API Error: terminated"));
+        const testServer = createRetryTestServer(prompt);
+
+        const resultPromise = testServer.promptWithUpstreamRetry({
+          sessionId: "acp-1",
+          prompt: [{ type: "text", text: "do the task" }],
+        });
+        const assertion = expect(resultPromise).rejects.toThrow(
+          "API Error: terminated",
+        );
+        await vi.advanceTimersByTimeAsync(10_000);
+
+        await assertion;
+        expect(prompt).toHaveBeenCalledTimes(3);
+      } finally {
+        vi.useRealTimers();
+      }
+    });
+
     it("persists structured turn completion notifications", () => {
       const appendRawLine = vi.fn();
       const testServer = new AgentServer({

packages/agent/src/server/agent-server.ts

@@ -125,6 +125,12 @@ type MessageCallback = (message: unknown) => void;
 
 export const SSE_KEEPALIVE_INTERVAL_MS = 25_000;
 
+// Bounded per-turn retries for unattended (initial/resume) turns that hit a
+// transient upstream failure. Two covers a retry whose own attempt also gets
+// cut once, without letting a hard upstream outage loop forever.
+const MAX_UPSTREAM_TURN_RETRIES = 2;
+const UPSTREAM_TURN_RETRY_DELAY_MS = 5_000;
+
 class NdJsonTap {
   private decoder = new TextDecoder();
   private buffer = "";
@@ -1516,6 +1522,72 @@ export class AgentServer {
     return { classification: classifyAgentError(message), message };
   }
 
+  /**
+   * Send an initial/resume turn prompt, absorbing transient upstream
+   * failures with a bounded number of retries. These turns run unattended
+   * (no user watching who could retry), so without this a single transient
+   * transport cut fails the whole run. A stream that died mid-response has
+   * already delivered the original prompt into the session history, so that
+   * case retries with a hidden continuation; failures where the request may
+   * never have been processed re-send the original prompt instead.
+   */
+  private async promptWithUpstreamRetry(request: {
+    sessionId: string;
+    prompt: ContentBlock[];
+    _meta?: Record<string, unknown>;
+  }): Promise<PromptResponse> {
+    let retries = 0;
+    let continueInterruptedTurn = false;
+    for (;;) {
+      // Re-read the session on every attempt: it can be torn down or
+      // replaced while the retry delay is pending.
+      const session = this.session;
+      if (!session) {
+        throw new Error("Agent session ended before the turn could be sent");
+      }
+      const attempt = continueInterruptedTurn
+        ? {
+            sessionId: session.acpSessionId,
+            prompt: [
+              hiddenTextBlock(
+                "The previous response was interrupted by a transient connection error. " +
+                  "Continue from where you left off — do not repeat work that already completed.",
+              ),
+            ],
+          }
+        : { ...request, sessionId: session.acpSessionId };
+      try {
+        return await session.clientConnection.prompt(attempt);
+      } catch (error) {
+        const { classification, message } =
+          this.extractErrorClassification(error);
+        if (
+          !upstreamProviderFailureClassifications.has(classification) ||
+          retries >= MAX_UPSTREAM_TURN_RETRIES
+        ) {
+          throw error;
+        }
+        retries += 1;
+        // Only a mid-response stream death guarantees the prompt reached the
+        // model; connection/timeout/status failures re-send the original.
+        continueInterruptedTurn =
+          classification === "upstream_stream_terminated";
+        this.logger.warn(
+          "Turn hit a transient upstream failure; retrying after a short delay",
+          {
+            classification,
+            message,
+            attempt: retries,
+            continueInterruptedTurn,
+          },
+        );
+        await new Promise((resolve) =>
+          setTimeout(resolve, UPSTREAM_TURN_RETRY_DELAY_MS),
+        );
+      }
+    }
+  }
+
   private async handleTurnFailure(
     payload: JwtPayload,
     phase: "initial" | "resume" | "followup",
@@ -1665,7 +1737,7 @@ export class AgentServer {
 
       this.session.logWriter.resetTurnMessages(payload.run_id);
 
-      const result = await this.session.clientConnection.prompt({
+      const result = await this.promptWithUpstreamRetry({
         sessionId: this.session.acpSessionId,
         prompt: initialPrompt,
         ...(initialPromptMeta ? { _meta: initialPromptMeta } : {}),
@@ -1877,7 +1949,7 @@ export class AgentServer {
 
       this.session.logWriter.resetTurnMessages(payload.run_id);
 
-      const result = await this.session.clientConnection.prompt({
+      const result = await this.promptWithUpstreamRetry({
         sessionId: this.session.acpSessionId,
         prompt: builtPrompt.prompt,
         ...(builtPrompt.meta ? { _meta: builtPrompt.meta } : {}),

packages/agent/src/server/question-relay.test.ts

@@ -730,7 +730,7 @@ describe("Question relay", () => {
       expect(promptSpy).not.toHaveBeenCalled();
     });
 
-    it("does not replay a transient upstream termination before any session activity", async () => {
+    it("replays a transient upstream termination with a continuation prompt", async () => {
       vi.spyOn(server.posthogAPI, "getTask").mockResolvedValue({
         id: "test-task-id",
         title: "t",
@@ -744,7 +744,8 @@ describe("Question relay", () => {
 
       const promptSpy = vi
         .fn()
-        .mockRejectedValueOnce(createTransientPromptError());
+        .mockRejectedValueOnce(createTransientPromptError())
+        .mockResolvedValueOnce({ stopReason: "cancelled" });
       const updateTaskRunSpy = vi
         .spyOn(server.posthogAPI, "updateTaskRun")
         .mockResolvedValue({} as TaskRun);
@@ -762,20 +763,26 @@ describe("Question relay", () => {
         },
       };
 
-      await server.sendInitialTaskMessage(TEST_PAYLOAD);
-
-      expect(promptSpy).toHaveBeenCalledTimes(1);
-      expect(updateTaskRunSpy).toHaveBeenCalledWith(
-        "test-task-id",
-        "test-run-id",
-        {
-          status: "failed",
-          error_message: UPSTREAM_PROVIDER_FAILURE_MESSAGE,
-        },
+      vi.useFakeTimers();
+      try {
+        const sendPromise = server.sendInitialTaskMessage(TEST_PAYLOAD);
+        await vi.advanceTimersByTimeAsync(5_000);
+        await sendPromise;
+      } finally {
+        vi.useRealTimers();
+      }
+
+      expect(promptSpy).toHaveBeenCalledTimes(2);
+      const continuation = promptSpy.mock.calls[1][0] as {
+        prompt: Array<{ type: string; text: string }>;
+      };
+      expect(continuation.prompt[0].text).toContain(
+        "interrupted by a transient connection error",
       );
+      expect(updateTaskRunSpy).not.toHaveBeenCalled();
     });
 
-    it("surfaces upstream provider failures with a retryable message", async () => {
+    it("re-sends the original prompt after an upstream provider failure", async () => {
       vi.spyOn(server.posthogAPI, "getTask").mockResolvedValue({
         id: "test-task-id",
         title: "t",
@@ -789,7 +796,8 @@ describe("Question relay", () => {
 
       const promptSpy = vi
         .fn()
-        .mockRejectedValueOnce(createUpstreamProviderFailureError());
+        .mockRejectedValueOnce(createUpstreamProviderFailureError())
+        .mockResolvedValueOnce({ stopReason: "cancelled" });
       const updateTaskRunSpy = vi
         .spyOn(server.posthogAPI, "updateTaskRun")
         .mockResolvedValue({} as TaskRun);
@@ -807,20 +815,24 @@ describe("Question relay", () => {
         },
       };
 
-      await server.sendInitialTaskMessage(TEST_PAYLOAD);
-
-      expect(promptSpy).toHaveBeenCalledTimes(1);
-      expect(updateTaskRunSpy).toHaveBeenCalledWith(
-        "test-task-id",
-        "test-run-id",
-        {
-          status: "failed",
-          error_message: UPSTREAM_PROVIDER_FAILURE_MESSAGE,
-        },
-      );
+      vi.useFakeTimers();
+      try {
+        const sendPromise = server.sendInitialTaskMessage(TEST_PAYLOAD);
+        await vi.advanceTimersByTimeAsync(5_000);
+        await sendPromise;
+      } finally {
+        vi.useRealTimers();
+      }
+
+      expect(promptSpy).toHaveBeenCalledTimes(2);
+      const retryRequest = promptSpy.mock.calls[1][0] as {
+        prompt: Array<{ type: string; text: string }>;
+      };
+      expect(retryRequest.prompt[0].text).toBe("original task description");
+      expect(updateTaskRunSpy).not.toHaveBeenCalled();
     });
 
-    it("surfaces upstream connection errors with the shared provider failure message", async () => {
+    it("surfaces the shared provider failure message once upstream retries are exhausted", async () => {
       vi.spyOn(server.posthogAPI, "getTask").mockResolvedValue({
         id: "test-task-id",
         title: "t",
@@ -832,7 +844,7 @@ describe("Question relay", () => {
         state: {},
       } as unknown as TaskRun);
 
-      const promptSpy = vi.fn().mockImplementationOnce(async () => {
+      const promptSpy = vi.fn().mockImplementation(async () => {
         throw createTransientConnectionError();
       });
       const updateTaskRunSpy = vi
@@ -852,9 +864,16 @@ describe("Question relay", () => {
         },
       };
 
-      await server.sendInitialTaskMessage(TEST_PAYLOAD);
+      vi.useFakeTimers();
+      try {
+        const sendPromise = server.sendInitialTaskMessage(TEST_PAYLOAD);
+        await vi.advanceTimersByTimeAsync(10_000);
+        await sendPromise;
+      } finally {
+        vi.useRealTimers();
+      }
 
-      expect(promptSpy).toHaveBeenCalledTimes(1);
+      expect(promptSpy).toHaveBeenCalledTimes(3);
       expect(updateTaskRunSpy).toHaveBeenCalledWith(
         "test-task-id",
         "test-run-id",