diff --git a/packages/agent/src/server/agent-server.test.ts b/packages/agent/src/server/agent-server.test.ts index 444dfff250..6bb9a27e86 100644 --- a/packages/agent/src/server/agent-server.test.ts +++ b/packages/agent/src/server/agent-server.test.ts @@ -956,6 +956,123 @@ describe("AgentServer HTTP Mode", () => { }, ); + function createRetryTestServer(prompt: ReturnType) { + 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({ diff --git a/packages/agent/src/server/agent-server.ts b/packages/agent/src/server/agent-server.ts index 3b418cf87d..0a6fafe81e 100644 --- a/packages/agent/src/server/agent-server.ts +++ b/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; + }): Promise { + 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 } : {}), diff --git a/packages/agent/src/server/question-relay.test.ts b/packages/agent/src/server/question-relay.test.ts index 6c0a58c22d..3f8648f158 100644 --- a/packages/agent/src/server/question-relay.test.ts +++ b/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",