diff --git a/src/CodexAcpClient.ts b/src/CodexAcpClient.ts index bec75265..e4723164 100644 --- a/src/CodexAcpClient.ts +++ b/src/CodexAcpClient.ts @@ -63,6 +63,7 @@ import { createReportedAgentFileChangeReport, createUnavailableAgentFileChangeReport, } from "./AgentFileChangeReport"; +import {RetryCapacityService, type TurnRetry, type TurnRetryService} from "./RetryCapacityService"; /** * Well-known provider id for the client-configurable custom LLM gateway. @@ -109,6 +110,7 @@ export class CodexAcpClient { private pendingLoginCompleted: Promise | null = null; private pendingAccountUpdated: Promise | null = null; private readonly sessionNotificationQueues = new Map>(); + private readonly turnRetryService: TurnRetryService = new RetryCapacityService(); private skillExtraRoots: string[] = []; private configPath: string | null = null; @@ -848,8 +850,14 @@ export class CodexAcpClient { disableSummary: boolean, cwd: string, additionalDirectories: string[], - onTurnStarted?: (turnId: string) => void, + onTurnStarted?: (turnId: string, retry: TurnRetry | null) => void, shouldCancel?: () => boolean, + onTurnRetry?: ( + completed: TurnCompletedNotification, + retry: TurnRetry, + ) => Promise, + retrySignal?: AbortSignal, + retryAttempt = 0, ): Promise { const input = buildPromptItems(request.prompt); const effort = modelId.effort as ReasoningEffort | null; //TODO remove unsafe conversion @@ -857,17 +865,39 @@ export class CodexAcpClient { if (shouldCancel?.()) { return null; } - return await this.codexClient.runTurn({ + const retry = this.turnRetryService.createRetry(retryAttempt); + const completed = await this.codexClient.runTurn({ threadId: request.sessionId, - input: input, + input, approvalPolicy: agentMode.approvalPolicy, approvalsReviewer: agentMode.approvalsReviewer, sandboxPolicy: addAdditionalDirectoriesToSandboxPolicy(agentMode.sandboxPolicy, additionalDirectories), summary: disableSummary ? "none" : "auto", - effort: effort, + effort, model: modelId.model, - serviceTier: serviceTier, - }, onTurnStarted); + serviceTier, + }, (turnId) => onTurnStarted?.(turnId, retry)); + if (retry === null || !this.turnRetryService.shouldRetry(completed)) { + return completed; + } + await onTurnRetry?.(completed, retry); + if (!await this.turnRetryService.wait(retry, retrySignal, shouldCancel)) { + return null; + } + return await this.sendPrompt( + this.turnRetryService.createContinuationRequest(request), + agentMode, + modelId, + serviceTier, + disableSummary, + cwd, + additionalDirectories, + onTurnStarted, + shouldCancel, + onTurnRetry, + retrySignal, + retryAttempt + 1, + ); } async runAgentFileChangeReport(params: { diff --git a/src/CodexAcpServer.ts b/src/CodexAcpServer.ts index f3ebb373..a8aa3397 100644 --- a/src/CodexAcpServer.ts +++ b/src/CodexAcpServer.ts @@ -19,11 +19,21 @@ import { type SessionMetadataWithThread, type UrlElicitationRequester } from "./CodexAcpClient"; +import {RetryCapacityService, type TurnRetry, type TurnRetryService} from "./RetryCapacityService"; import {CodexAppServerClient, type McpStartupResult} from "./CodexAppServerClient"; import {type CodexConnection, startCodexConnection} from "./CodexJsonRpcConnection"; import {type AcpClientConnection, ACPSessionConnection, type UpdateSessionEvent} from "./ACPSessionConnection"; import type {InputModality, ReasoningEffort} from "./app-server"; -import type {Account, Model, ReasoningEffortOption, Thread, ThreadGoal, ThreadItem, UserInput} from "./app-server/v2"; +import type { + Account, + Model, + ReasoningEffortOption, + Thread, + ThreadGoal, + ThreadItem, + TurnCompletedNotification, + UserInput, +} from "./app-server/v2"; import type {RateLimitsMap} from "./RateLimitsMap"; import {ModelId} from "./ModelId"; import {AgentMode, MODE_CONFIG_ID} from "./AgentMode"; @@ -2288,6 +2298,8 @@ export class CodexAcpServer { const disposePromptRequestCancellation = this.observePromptRequestCancellation(signal, sessionState, activePrompt); let eventHandler: CodexEventHandler | null = null; let promptNotificationsActive = true; + const turnRetryService: TurnRetryService = new RetryCapacityService(); + let turnRetry: TurnRetry | null = null; const clearRecoveredSessionFailure = async (handler: CodexEventHandler): Promise => { await handler.completeSuccessfulTurn(sessionState.currentTurnId); const current = sessionState.sessionFailure; @@ -2334,11 +2346,16 @@ export class CodexAcpServer { await promptEventHandler.handleSessionScopedNotification(event); return; } + const handledEvent = turnRetryService.transformNotification( + event, + sessionState.currentTurnId, + turnRetry, + ); const completesActiveTurn = event.method === "turn/completed" && event.params.turn.id === sessionState.currentTurnId; - permissionContext.handleNotification(event); - await elicitationHandler.handleNotification(event); - await promptEventHandler.handleNotification(event); + permissionContext.handleNotification(handledEvent); + await elicitationHandler.handleNotification(handledEvent); + await promptEventHandler.handleNotification(handledEvent); if (completesActiveTurn) { // The prompt may remain open for plan approval after its turn has ended. Switch at // the causal boundary so a queued late error cannot enter the completed turn's buffer. @@ -2348,6 +2365,32 @@ export class CodexAcpServer { approvalHandler, elicitationHandler); + const handleTurnRetry = async ( + completed: TurnCompletedNotification, + retry: TurnRetry, + ) => { + await this.codexAcpClient.waitForSessionNotifications(params.sessionId); + await promptEventHandler.flushPendingErrors(); + if (!retry.warningPublished) { + await promptEventHandler.handleNotification(turnRetryService.createWarning( + params.sessionId, + completed.turn.id, + retry, + )); + } + await promptEventHandler.flushPendingPlanUpdates(); + recoverableSessionFailure = sessionState.sessionFailure; + activePrompt.currentTurn = null; + sessionState.currentTurnId = null; + pendingTurnStart = this.createPendingTurnStart(); + this.pendingTurnStarts.set(params.sessionId, pendingTurnStart); + logger.log("Scheduling turn retry", { + sessionId: params.sessionId, + attempt: retry.attempt, + delaySeconds: retry.delaySeconds, + }); + }; + if (activePrompt.signal.aborted) { return cancelledPromptResponse(); } @@ -2469,7 +2512,7 @@ export class CodexAcpServer { disableSummary, sessionState.cwd, sessionState.additionalDirectories, - (turnId) => { + (turnId, retry) => { const turn = {threadId: params.sessionId, turnId}; activePrompt.currentTurn = turn; if (this.promptShouldStop(params.sessionId, activePrompt)) { @@ -2477,10 +2520,15 @@ export class CodexAcpServer { return; } sessionState.currentTurnId = turnId; + turnRetry = retry; pendingTurnStart?.resolve(turnId); - onTurnStarted?.(); + if (retry?.attempt === 1) { + onTurnStarted?.(); + } }, () => this.promptShouldStop(params.sessionId, activePrompt), + handleTurnRetry, + activePrompt.signal, )); void sendPromptPromise.catch((err) => { if (this.activePrompts.get(params.sessionId) !== activePrompt) { @@ -2559,7 +2607,7 @@ export class CodexAcpServer { disableSummary, sessionState.cwd, sessionState.additionalDirectories, - (turnId) => { + (turnId, retry) => { const turn = {threadId: params.sessionId, turnId}; activePrompt.currentTurn = turn; if (this.promptShouldStop(params.sessionId, activePrompt)) { @@ -2567,12 +2615,15 @@ export class CodexAcpServer { return; } sessionState.currentTurnId = turnId; + turnRetry = retry; // Keep the approval-to-turn-start gap session-scoped. Once the new turn has // an identity, snapshot any unchanged session failure as its recovery baseline. recoverableSessionFailure = sessionState.sessionFailure; promptNotificationsActive = true; }, () => this.promptShouldStop(params.sessionId, activePrompt), + handleTurnRetry, + activePrompt.signal, ), ); void implementationPromise.catch((err) => { diff --git a/src/RetryCapacityService.ts b/src/RetryCapacityService.ts new file mode 100644 index 00000000..5d738959 --- /dev/null +++ b/src/RetryCapacityService.ts @@ -0,0 +1,130 @@ +import type {PromptRequest} from "@agentclientprotocol/sdk"; +import type {ServerNotification} from "./app-server"; +import type {TurnCompletedNotification} from "./app-server/v2"; + +const ERROR_MESSAGE = "Selected model is at capacity. Please try a different model."; +const CONTINUATION_PROMPT = "Continue from where you left off."; +const RETRY_WINDOWS_SECONDS = [ + [1, 10], + [1, 30], + [30, 60], + [60, 120], + [120, 300], +] as const; + +export interface TurnRetry { + attempt: number; + delaySeconds: number; + title: string; + warningPublished: boolean; +} + +export interface TurnRetryService { + createRetry(retryAttempt: number): TurnRetry | null; + shouldRetry(completed: TurnCompletedNotification): boolean; + wait( + retry: TurnRetry, + signal: AbortSignal | undefined, + shouldCancel: (() => boolean) | undefined, + ): Promise; + createContinuationRequest(request: PromptRequest): PromptRequest; + transformNotification( + notification: ServerNotification, + currentTurnId: string | null, + retry: TurnRetry | null, + ): ServerNotification; + createWarning(threadId: string, turnId: string, retry: TurnRetry): ServerNotification; +} + +export class RetryCapacityService implements TurnRetryService { + createRetry(retryAttempt: number): TurnRetry | null { + const window = RETRY_WINDOWS_SECONDS[retryAttempt]; + if (window === undefined) { + return null; + } + const [minimum, maximum] = window; + const delaySeconds = minimum + Math.floor(Math.random() * (maximum - minimum + 1)); + const unit = delaySeconds === 1 ? "second" : "seconds"; + return { + attempt: retryAttempt + 1, + delaySeconds, + title: `Selected model is at capacity. Retrying in ${delaySeconds} ${unit} ` + + `(${retryAttempt + 1}/${RETRY_WINDOWS_SECONDS.length}).`, + warningPublished: false, + }; + } + + shouldRetry(completed: TurnCompletedNotification): boolean { + return completed.turn.status === "failed" && isCapacityError(completed.turn.error); + } + + createContinuationRequest(request: PromptRequest): PromptRequest { + return { + sessionId: request.sessionId, + prompt: [{type: "text", text: CONTINUATION_PROMPT}], + }; + } + + transformNotification( + notification: ServerNotification, + currentTurnId: string | null, + retry: TurnRetry | null, + ): ServerNotification { + if (notification.method !== "error" + || notification.params.willRetry + || retry === null + || notification.params.turnId !== currentTurnId + || !isCapacityError(notification.params.error)) { + return notification; + } + retry.warningPublished = true; + return { + method: "error", + params: { + ...notification.params, + willRetry: true, + error: {...notification.params.error, message: retry.title}, + }, + }; + } + + createWarning(threadId: string, turnId: string, retry: TurnRetry): ServerNotification { + return { + method: "error", + params: { + threadId, + turnId, + willRetry: true, + error: { + message: retry.title, + codexErrorInfo: "serverOverloaded", + additionalDetails: null, + }, + }, + }; + } + + async wait( + retry: TurnRetry, + signal: AbortSignal | undefined, + shouldCancel: (() => boolean) | undefined, + ): Promise { + if (signal?.aborted || shouldCancel?.()) { + return false; + } + return await new Promise((resolve) => { + const finish = (completed: boolean) => { + clearTimeout(timer); + signal?.removeEventListener("abort", onAbort); + resolve(completed); + }; + const onAbort = () => finish(false); + const timer = setTimeout(() => finish(!shouldCancel?.()), retry.delaySeconds * 1000); + signal?.addEventListener("abort", onAbort, {once: true}); + }); + } +} + +function isCapacityError(error: {message: string; codexErrorInfo: unknown} | null): boolean { + return error?.codexErrorInfo === "serverOverloaded" && error.message.trim() === ERROR_MESSAGE; +} diff --git a/src/__tests__/CodexACPAgent/auth-error-events.test.ts b/src/__tests__/CodexACPAgent/auth-error-events.test.ts index ebfc7d13..c79d1078 100644 --- a/src/__tests__/CodexACPAgent/auth-error-events.test.ts +++ b/src/__tests__/CodexACPAgent/auth-error-events.test.ts @@ -110,6 +110,84 @@ describe("CodexEventHandler - auth error events", () => { ); }); + it("retries model capacity failures with continuation turns", async () => { + vi.useFakeTimers(); + const random = vi.spyOn(Math, "random") + .mockReturnValueOnce(0) + .mockReturnValueOnce(0.5) + .mockReturnValue(0); + try { + const {fixture, promptPromise, turnStart} = await startModelCapacityRetryPrompt(2); + await vi.runAllTimersAsync(); + const result = await promptPromise; + + expect(result).toMatchObject({stopReason: "end_turn"}); + expect(turnStart).toHaveBeenCalledTimes(3); + expect(JSON.stringify(turnStart.mock.calls[0]?.[0])).toContain("test capacity retry"); + expect(JSON.stringify(turnStart.mock.calls[1]?.[0])).toContain("Continue from where you left off."); + expect(JSON.stringify(turnStart.mock.calls[2]?.[0])).toContain("Continue from where you left off."); + + const warningTitles = fixture.getAcpConnectionEvents([]) + .map(event => event.args[0].update?._meta?.jetbrains?.air?.sessionFailure) + .filter(failure => failure?.severity === "warning") + .map(failure => failure.title); + expect(warningTitles).toEqual([ + "Selected model is at capacity. Retrying in 1 second (1/5).", + "Selected model is at capacity. Retrying in 16 seconds (2/5).", + ]); + } finally { + random.mockRestore(); + vi.useRealTimers(); + } + }); + + it("returns the terminal model-capacity failure after five retries", async () => { + vi.useFakeTimers(); + const random = vi.spyOn(Math, "random").mockReturnValue(0); + try { + const {fixture, promptPromise, turnStart} = await startModelCapacityRetryPrompt(6); + await vi.runAllTimersAsync(); + const result = await promptPromise; + + expect(turnStart).toHaveBeenCalledTimes(6); + expect(result).toMatchObject({ + stopReason: "end_turn", + _meta: {jetbrains: {air: {sessionFailure: { + category: "service", + severity: "error", + title: "Selected model is at capacity. Please try a different model.", + }}}}, + }); + const warnings = fixture.getAcpConnectionEvents([]) + .map(event => event.args[0].update?._meta?.jetbrains?.air?.sessionFailure) + .filter(failure => failure?.severity === "warning"); + expect(warnings).toHaveLength(5); + } finally { + random.mockRestore(); + vi.useRealTimers(); + } + }); + + it("cancels while waiting to retry model capacity", async () => { + vi.useFakeTimers(); + const random = vi.spyOn(Math, "random").mockReturnValue(0); + const controller = new AbortController(); + try { + const {promptPromise, turnStart} = await startModelCapacityRetryPrompt(1, controller.signal); + await vi.advanceTimersByTimeAsync(0); + expect(turnStart).toHaveBeenCalledTimes(1); + + controller.abort(); + await vi.advanceTimersByTimeAsync(0); + + await expect(promptPromise).resolves.toMatchObject({stopReason: "cancelled"}); + expect(turnStart).toHaveBeenCalledTimes(1); + } finally { + random.mockRestore(); + vi.useRealTimers(); + } + }); + it("does not attach a foreign turn failure to the active prompt", async () => { const {result, updates} = await runPromptWithError(createTestSessionState({ sessionId: "foreign-turn-session", @@ -797,6 +875,54 @@ describe("CodexEventHandler - auth error events", () => { ); }); +async function startModelCapacityRetryPrompt(failuresBeforeSuccess: number, signal?: AbortSignal) { + const fixture = createCodexMockTestFixture(); + const codexAcpAgent = fixture.getCodexAcpAgent(); + const codexAppServerClient = fixture.getCodexAppServerClient(); + const sessionState = createTestSessionState({ + sessionId: "model-capacity-retry-session", + account: {type: "apiKey"}, + }); + await codexAcpAgent.initialize({ + protocolVersion: acp.PROTOCOL_VERSION, + clientCapabilities: typedFailureCapabilities, + }); + vi.spyOn(codexAcpAgent, "getSessionState").mockReturnValue(sessionState); + + let turnNumber = 0; + const turnStart = vi.spyOn(codexAppServerClient, "turnStart").mockImplementation(async () => ({ + turn: createTurn("inProgress", `capacity-turn-${++turnNumber}`), + })); + vi.spyOn(codexAppServerClient, "awaitTurnCompleted").mockImplementation(async (threadId, turnId) => { + const failed = Number(turnId.replace("capacity-turn-", "")) <= failuresBeforeSuccess; + const error: ErrorNotification["error"] | null = failed + ? { + message: "Selected model is at capacity. Please try a different model.", + codexErrorInfo: "serverOverloaded", + additionalDetails: null, + } + : null; + if (error !== null) { + fixture.sendServerNotification({ + method: "error", + params: {threadId, turnId, willRetry: false, error}, + }); + } + const completion: TurnCompletedNotification = { + threadId, + turn: createTurn(failed ? "failed" : "completed", turnId, error), + }; + fixture.sendServerNotification({method: "turn/completed", params: completion}); + return completion; + }); + + const promptPromise = codexAcpAgent.prompt({ + sessionId: sessionState.sessionId, + prompt: [{type: "text", text: "test capacity retry"}], + }, signal); + return {fixture, promptPromise, turnStart}; +} + async function runPromptWithError( sessionState: SessionState, turnError: ErrorNotification["error"],