From 99469ca8f39cbc522365194979825297b9503c16 Mon Sep 17 00:00:00 2001 From: Nikita Ashikhmin Date: Thu, 27 Aug 2026 10:56:41 +0400 Subject: [PATCH 1/6] feat: retry model capacity failures --- src/CodexAcpServer.ts | 283 +++++++++++++----- src/CodexEventHandler.ts | 16 + .../CodexACPAgent/auth-error-events.test.ts | 146 +++++++++ 3 files changed, 366 insertions(+), 79 deletions(-) diff --git a/src/CodexAcpServer.ts b/src/CodexAcpServer.ts index f3ebb373..b8d5c247 100644 --- a/src/CodexAcpServer.ts +++ b/src/CodexAcpServer.ts @@ -23,7 +23,16 @@ import {CodexAppServerClient, type McpStartupResult} from "./CodexAppServerClien 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"; @@ -172,6 +181,38 @@ export interface SessionFailure { } const CODEX_PROCESS_EXITED_ERROR_CODE = 1001; +const MODEL_CAPACITY_ERROR_MESSAGE = "Selected model is at capacity. Please try a different model."; +const MODEL_CAPACITY_CONTINUATION_PROMPT = "Continue from where you left off."; +const MODEL_CAPACITY_RETRY_WINDOWS_SECONDS = [ + [1, 10], + [1, 30], + [30, 60], + [60, 120], + [120, 300], +] as const; + +interface ModelCapacityRetry { + attempt: number; + delaySeconds: number; + turnId: string | null; + warningPublished: boolean; +} + +function isModelCapacityError(error: {message: string; codexErrorInfo: unknown} | null): boolean { + return error?.codexErrorInfo === "serverOverloaded" + && error.message.trim() === MODEL_CAPACITY_ERROR_MESSAGE; +} + +function modelCapacityRetryDelaySeconds(attempt: number): number { + const [minimum, maximum] = MODEL_CAPACITY_RETRY_WINDOWS_SECONDS[attempt]!; + return minimum + Math.floor(Math.random() * (maximum - minimum + 1)); +} + +function modelCapacityRetryTitle(retry: ModelCapacityRetry): string { + const unit = retry.delaySeconds === 1 ? "second" : "seconds"; + return `Selected model is at capacity. Retrying in ${retry.delaySeconds} ${unit} ` + + `(${retry.attempt}/${MODEL_CAPACITY_RETRY_WINDOWS_SECONDS.length}).`; +} function clientSupportsAirCapability( capabilities: acp.ClientCapabilities | null, @@ -2188,6 +2229,26 @@ export class CodexAcpServer { return activePrompt.signal.aborted || this.activePrompts.get(sessionId) !== activePrompt || this.sessionIsClosing(sessionId); } + private async waitForPromptRetry( + delaySeconds: number, + sessionId: string, + activePrompt: ActivePrompt, + ): Promise { + if (this.promptShouldStop(sessionId, activePrompt)) { + return false; + } + return await new Promise((resolve) => { + const finish = (completed: boolean) => { + clearTimeout(timer); + activePrompt.signal.removeEventListener("abort", onAbort); + resolve(completed); + }; + const onAbort = () => finish(false); + const timer = setTimeout(() => finish(true), delaySeconds * 1000); + activePrompt.signal.addEventListener("abort", onAbort, {once: true}); + }); + } + private async interruptSessionTurn( sessionState: SessionState, requestName: "Cancel" | "Close", @@ -2288,6 +2349,7 @@ export class CodexAcpServer { const disposePromptRequestCancellation = this.observePromptRequestCancellation(signal, sessionState, activePrompt); let eventHandler: CodexEventHandler | null = null; let promptNotificationsActive = true; + let modelCapacityRetry: ModelCapacityRetry | null = null; const clearRecoveredSessionFailure = async (handler: CodexEventHandler): Promise => { await handler.completeSuccessfulTurn(sessionState.currentTurnId); const current = sessionState.sessionFailure; @@ -2334,11 +2396,30 @@ export class CodexAcpServer { await promptEventHandler.handleSessionScopedNotification(event); return; } + let handledEvent = event; + if (event.method === "error" + && !event.params.willRetry + && isModelCapacityError(event.params.error) + && modelCapacityRetry !== null + && event.params.turnId === sessionState.currentTurnId) { + modelCapacityRetry.warningPublished = true; + handledEvent = { + method: "error", + params: { + ...event.params, + willRetry: true, + error: { + ...event.params.error, + message: modelCapacityRetryTitle(modelCapacityRetry), + }, + }, + }; + } 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. @@ -2459,49 +2540,126 @@ export class CodexAcpServer { sessionState.fastModeEnabled, sessionState.currentModelSupportsFast, ); - ensurePendingTurnStart(); - const sendPromptPromise = this.runWithProcessCheck( - () => this.codexAcpClient.sendPrompt( - effectiveParams, - agentMode, - modelId, - serviceTier, - disableSummary, - sessionState.cwd, - sessionState.additionalDirectories, - (turnId) => { - const turn = {threadId: params.sessionId, turnId}; - activePrompt.currentTurn = turn; - if (this.promptShouldStop(params.sessionId, activePrompt)) { - this.interruptLateStartedTurn(turn); - return; + const runTurnWithModelCapacityRetries = async ( + initialRequest: acp.PromptRequest, + notifyOnTurnStarted: boolean, + notificationsActiveBeforeStart: boolean, + ): Promise => { + let turnRequest = initialRequest; + for (let retryAttempt = 0; ; retryAttempt++) { + const retryWindow = MODEL_CAPACITY_RETRY_WINDOWS_SECONDS[retryAttempt]; + modelCapacityRetry = retryWindow === undefined + ? null + : { + attempt: retryAttempt + 1, + delaySeconds: modelCapacityRetryDelaySeconds(retryAttempt), + turnId: null, + warningPublished: false, + }; + promptNotificationsActive = retryAttempt === 0 + ? notificationsActiveBeforeStart + : true; + ensurePendingTurnStart(); + const sendPromptPromise = this.runWithProcessCheck( + () => this.codexAcpClient.sendPrompt( + turnRequest, + agentMode, + modelId, + serviceTier, + disableSummary, + sessionState.cwd, + sessionState.additionalDirectories, + (turnId) => { + const turn = {threadId: params.sessionId, turnId}; + activePrompt.currentTurn = turn; + if (this.promptShouldStop(params.sessionId, activePrompt)) { + this.interruptLateStartedTurn(turn); + return; + } + sessionState.currentTurnId = turnId; + if (!notificationsActiveBeforeStart && retryAttempt === 0) { + recoverableSessionFailure = sessionState.sessionFailure; + } + promptNotificationsActive = true; + if (modelCapacityRetry !== null) { + modelCapacityRetry.turnId = turnId; + } + pendingTurnStart?.resolve(turnId); + if (notifyOnTurnStarted && retryAttempt === 0) { + onTurnStarted?.(); + } + }, + () => this.promptShouldStop(params.sessionId, activePrompt), + )); + void sendPromptPromise.catch((err) => { + if (this.activePrompts.get(params.sessionId) !== activePrompt) { + logger.error(`Prompt for cancelled session ${params.sessionId} failed after prompt returned`, err); } - sessionState.currentTurnId = turnId; - pendingTurnStart?.resolve(turnId); - onTurnStarted?.(); - }, - () => this.promptShouldStop(params.sessionId, activePrompt), - )); - void sendPromptPromise.catch((err) => { - if (this.activePrompts.get(params.sessionId) !== activePrompt) { - logger.error(`Prompt for cancelled session ${params.sessionId} failed after prompt returned`, err); + }); + const completed = await Promise.race([ + sendPromptPromise, + activePrompt.closeSignal, + this.cancelBeforeTurnStarted(activePrompt), + ]); + + if (completed === null) { + modelCapacityRetry = null; + return null; + } + + await this.codexAcpClient.waitForSessionNotifications(params.sessionId); + await promptEventHandler.flushPendingErrors(); + const retry = modelCapacityRetry; + const shouldRetry = retry !== null + && retry.turnId === completed.turn.id + && completed.turn.status === "failed" + && isModelCapacityError(completed.turn.error); + if (!shouldRetry) { + modelCapacityRetry = null; + await promptEventHandler.handleFailedTurn(completed.turn); + promptNotificationsActive = false; + return completed; + } + + if (!retry.warningPublished) { + await promptEventHandler.publishModelCapacityRetryWarning( + completed.turn.id, + modelCapacityRetryTitle(retry), + ); + } + await promptEventHandler.flushPendingPlanUpdates(); + recoverableSessionFailure = sessionState.sessionFailure; + activePrompt.currentTurn = null; + sessionState.currentTurnId = null; + pendingTurnStart = this.createPendingTurnStart(); + this.pendingTurnStarts.set(params.sessionId, pendingTurnStart); + logger.log("Selected model is at capacity; scheduling retry", { + sessionId: params.sessionId, + attempt: retry.attempt, + delaySeconds: retry.delaySeconds, + }); + const retryReady = await this.waitForPromptRetry( + retry.delaySeconds, + params.sessionId, + activePrompt, + ); + if (!retryReady) { + modelCapacityRetry = null; + return null; + } + turnRequest = { + sessionId: params.sessionId, + prompt: [{type: "text", text: MODEL_CAPACITY_CONTINUATION_PROMPT}], + }; } - }); - let turnCompleted = await Promise.race([ - sendPromptPromise, - activePrompt.closeSignal, - this.cancelBeforeTurnStarted(activePrompt), - ]); + }; + + let turnCompleted = await runTurnWithModelCapacityRetries(effectiveParams, true, true); if (turnCompleted === null) { return cancelledPromptResponse(); } - await this.codexAcpClient.waitForSessionNotifications(params.sessionId); - await eventHandler.flushPendingErrors(); - await eventHandler.handleFailedTurn(turnCompleted.turn); - promptNotificationsActive = false; - if (turnCompleted.turn.status === "interrupted") { await eventHandler.flushPendingPlanUpdates(); return cancelledPromptResponse(); @@ -2550,50 +2708,17 @@ export class CodexAcpServer { }; activePrompt.currentTurn = null; sessionState.currentTurnId = null; - const implementationPromise = this.runWithProcessCheck( - () => this.codexAcpClient.sendPrompt( - implementationRequest, - agentMode, - modelId, - serviceTier, - disableSummary, - sessionState.cwd, - sessionState.additionalDirectories, - (turnId) => { - const turn = {threadId: params.sessionId, turnId}; - activePrompt.currentTurn = turn; - if (this.promptShouldStop(params.sessionId, activePrompt)) { - this.interruptLateStartedTurn(turn); - return; - } - sessionState.currentTurnId = turnId; - // 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), - ), - ); - void implementationPromise.catch((err) => { - if (this.activePrompts.get(params.sessionId) !== activePrompt) { - logger.error(`Implementation turn for cancelled prompt ${params.sessionId} failed after prompt returned`, err); - } - }); - turnCompleted = await Promise.race([ - implementationPromise, - activePrompt.closeSignal, - this.cancelBeforeTurnStarted(activePrompt), - ]); + // Keep the approval-to-turn-start gap session-scoped. Snapshot any unchanged failure + // before starting implementation so a successful turn can clear it. + recoverableSessionFailure = sessionState.sessionFailure; + pendingTurnStart = this.createPendingTurnStart(); + this.pendingTurnStarts.set(params.sessionId, pendingTurnStart); + turnCompleted = await runTurnWithModelCapacityRetries(implementationRequest, false, false); if (turnCompleted === null) { return cancelledPromptResponse(); } - await this.codexAcpClient.waitForSessionNotifications(params.sessionId); - await eventHandler.flushPendingErrors(); - await eventHandler.handleFailedTurn(turnCompleted.turn); - promptNotificationsActive = false; if (turnCompleted.turn.status === "interrupted") { await eventHandler.flushPendingPlanUpdates(); return cancelledPromptResponse(); diff --git a/src/CodexEventHandler.ts b/src/CodexEventHandler.ts index 7360b6f9..f9919189 100644 --- a/src/CodexEventHandler.ts +++ b/src/CodexEventHandler.ts @@ -323,6 +323,22 @@ export class CodexEventHandler { delete this.sessionState.sessionFailure; } + async publishModelCapacityRetryWarning(turnId: string, title: string): Promise { + const update = await this.createErrorEvent({ + threadId: this.sessionState.sessionId, + turnId, + willRetry: true, + error: { + message: title, + codexErrorInfo: "serverOverloaded", + additionalDetails: null, + }, + }); + if (update !== null) { + await this.session.update(update); + } + } + async completeSuccessfulTurn(turnId: string | null): Promise { this.lastSessionNotice = undefined; if (!this.supportsTypedSessionFailures || turnId === null) return; diff --git a/src/__tests__/CodexACPAgent/auth-error-events.test.ts b/src/__tests__/CodexACPAgent/auth-error-events.test.ts index ebfc7d13..96634be5 100644 --- a/src/__tests__/CodexACPAgent/auth-error-events.test.ts +++ b/src/__tests__/CodexACPAgent/auth-error-events.test.ts @@ -110,6 +110,104 @@ describe("CodexEventHandler - auth error events", () => { ); }); + it("retries the exact model-capacity failure with jittered delays and 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 updates = fixture.getAcpConnectionEvents([]).map(event => event.args[0].update); + expect(updates).toEqual(expect.arrayContaining([ + expect.objectContaining({ + sessionUpdate: "session_info_update", + _meta: {jetbrains: {air: expect.objectContaining({sessionFailure: expect.objectContaining({ + severity: "warning", + title: "Selected model is at capacity. Retrying in 1 second (1/5).", + })})}}, + }), + expect.objectContaining({ + sessionUpdate: "session_info_update", + _meta: {jetbrains: {air: expect.objectContaining({sessionFailure: expect.objectContaining({ + severity: "warning", + title: "Selected model is at capacity. Retrying in 16 seconds (2/5).", + })})}}, + }), + ])); + } finally { + random.mockRestore(); + vi.useRealTimers(); + } + }); + + it("returns the terminal model-capacity failure after all five custom 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.", + actions: ["retry"], + }}}}, + }); + 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 1 second (2/5).", + "Selected model is at capacity. Retrying in 30 seconds (3/5).", + "Selected model is at capacity. Retrying in 60 seconds (4/5).", + "Selected model is at capacity. Retrying in 120 seconds (5/5).", + ]); + } finally { + random.mockRestore(); + vi.useRealTimers(); + } + }); + + it("cancels immediately while waiting to retry model capacity", async () => { + vi.useFakeTimers(); + const random = vi.spyOn(Math, "random").mockReturnValue(0); + const controller = new AbortController(); + try { + const {fixture, promptPromise, turnStart} = await startModelCapacityRetryPrompt(1, controller.signal); + await vi.advanceTimersByTimeAsync(0); + expect(turnStart).toHaveBeenCalledTimes(1); + expect(fixture.getAcpConnectionEvents([]).some(event => + event.args[0].update?._meta?.jetbrains?.air?.sessionFailure?.severity === "warning", + )).toBe(true); + + 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 +895,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"], From 757f241576e0e2d732b1270e7da0f60106b50812 Mon Sep 17 00:00:00 2001 From: Nikita Ashikhmin Date: Thu, 27 Aug 2026 11:02:39 +0400 Subject: [PATCH 2/6] refactor: isolate model capacity retries --- src/CodexAcpServer.ts | 202 +++++++++----------------------------- src/ModelCapacityRetry.ts | 133 +++++++++++++++++++++++++ 2 files changed, 181 insertions(+), 154 deletions(-) create mode 100644 src/ModelCapacityRetry.ts diff --git a/src/CodexAcpServer.ts b/src/CodexAcpServer.ts index b8d5c247..8724450d 100644 --- a/src/CodexAcpServer.ts +++ b/src/CodexAcpServer.ts @@ -30,7 +30,6 @@ import type { Thread, ThreadGoal, ThreadItem, - TurnCompletedNotification, UserInput, } from "./app-server/v2"; import type {RateLimitsMap} from "./RateLimitsMap"; @@ -127,6 +126,7 @@ import { createUnavailableAgentFileChangeReport, parseAgentFileChangeReportRequest, } from "./AgentFileChangeReport"; +import {ModelCapacityRetryController} from "./ModelCapacityRetry"; export interface SessionState { @@ -181,38 +181,6 @@ export interface SessionFailure { } const CODEX_PROCESS_EXITED_ERROR_CODE = 1001; -const MODEL_CAPACITY_ERROR_MESSAGE = "Selected model is at capacity. Please try a different model."; -const MODEL_CAPACITY_CONTINUATION_PROMPT = "Continue from where you left off."; -const MODEL_CAPACITY_RETRY_WINDOWS_SECONDS = [ - [1, 10], - [1, 30], - [30, 60], - [60, 120], - [120, 300], -] as const; - -interface ModelCapacityRetry { - attempt: number; - delaySeconds: number; - turnId: string | null; - warningPublished: boolean; -} - -function isModelCapacityError(error: {message: string; codexErrorInfo: unknown} | null): boolean { - return error?.codexErrorInfo === "serverOverloaded" - && error.message.trim() === MODEL_CAPACITY_ERROR_MESSAGE; -} - -function modelCapacityRetryDelaySeconds(attempt: number): number { - const [minimum, maximum] = MODEL_CAPACITY_RETRY_WINDOWS_SECONDS[attempt]!; - return minimum + Math.floor(Math.random() * (maximum - minimum + 1)); -} - -function modelCapacityRetryTitle(retry: ModelCapacityRetry): string { - const unit = retry.delaySeconds === 1 ? "second" : "seconds"; - return `Selected model is at capacity. Retrying in ${retry.delaySeconds} ${unit} ` - + `(${retry.attempt}/${MODEL_CAPACITY_RETRY_WINDOWS_SECONDS.length}).`; -} function clientSupportsAirCapability( capabilities: acp.ClientCapabilities | null, @@ -2229,26 +2197,6 @@ export class CodexAcpServer { return activePrompt.signal.aborted || this.activePrompts.get(sessionId) !== activePrompt || this.sessionIsClosing(sessionId); } - private async waitForPromptRetry( - delaySeconds: number, - sessionId: string, - activePrompt: ActivePrompt, - ): Promise { - if (this.promptShouldStop(sessionId, activePrompt)) { - return false; - } - return await new Promise((resolve) => { - const finish = (completed: boolean) => { - clearTimeout(timer); - activePrompt.signal.removeEventListener("abort", onAbort); - resolve(completed); - }; - const onAbort = () => finish(false); - const timer = setTimeout(() => finish(true), delaySeconds * 1000); - activePrompt.signal.addEventListener("abort", onAbort, {once: true}); - }); - } - private async interruptSessionTurn( sessionState: SessionState, requestName: "Cancel" | "Close", @@ -2349,7 +2297,7 @@ export class CodexAcpServer { const disposePromptRequestCancellation = this.observePromptRequestCancellation(signal, sessionState, activePrompt); let eventHandler: CodexEventHandler | null = null; let promptNotificationsActive = true; - let modelCapacityRetry: ModelCapacityRetry | null = null; + const modelCapacityRetry = new ModelCapacityRetryController(); const clearRecoveredSessionFailure = async (handler: CodexEventHandler): Promise => { await handler.completeSuccessfulTurn(sessionState.currentTurnId); const current = sessionState.sessionFailure; @@ -2396,25 +2344,7 @@ export class CodexAcpServer { await promptEventHandler.handleSessionScopedNotification(event); return; } - let handledEvent = event; - if (event.method === "error" - && !event.params.willRetry - && isModelCapacityError(event.params.error) - && modelCapacityRetry !== null - && event.params.turnId === sessionState.currentTurnId) { - modelCapacityRetry.warningPublished = true; - handledEvent = { - method: "error", - params: { - ...event.params, - willRetry: true, - error: { - ...event.params.error, - message: modelCapacityRetryTitle(modelCapacityRetry), - }, - }, - }; - } + const handledEvent = modelCapacityRetry.transformNotification(event, sessionState.currentTurnId); const completesActiveTurn = event.method === "turn/completed" && event.params.turn.id === sessionState.currentTurnId; permissionContext.handleNotification(handledEvent); @@ -2544,87 +2474,64 @@ export class CodexAcpServer { initialRequest: acp.PromptRequest, notifyOnTurnStarted: boolean, notificationsActiveBeforeStart: boolean, - ): Promise => { - let turnRequest = initialRequest; - for (let retryAttempt = 0; ; retryAttempt++) { - const retryWindow = MODEL_CAPACITY_RETRY_WINDOWS_SECONDS[retryAttempt]; - modelCapacityRetry = retryWindow === undefined - ? null - : { - attempt: retryAttempt + 1, - delaySeconds: modelCapacityRetryDelaySeconds(retryAttempt), - turnId: null, - warningPublished: false, - }; - promptNotificationsActive = retryAttempt === 0 - ? notificationsActiveBeforeStart - : true; + ) => modelCapacityRetry.run(initialRequest, { + signal: activePrompt.signal, + shouldStop: () => this.promptShouldStop(params.sessionId, activePrompt), + runTurn: async (turnRequest, retryAttempt) => { + promptNotificationsActive = retryAttempt === 0 ? notificationsActiveBeforeStart : true; ensurePendingTurnStart(); - const sendPromptPromise = this.runWithProcessCheck( - () => this.codexAcpClient.sendPrompt( - turnRequest, - agentMode, - modelId, - serviceTier, - disableSummary, - sessionState.cwd, - sessionState.additionalDirectories, - (turnId) => { - const turn = {threadId: params.sessionId, turnId}; - activePrompt.currentTurn = turn; - if (this.promptShouldStop(params.sessionId, activePrompt)) { - this.interruptLateStartedTurn(turn); - return; - } - sessionState.currentTurnId = turnId; - if (!notificationsActiveBeforeStart && retryAttempt === 0) { - recoverableSessionFailure = sessionState.sessionFailure; - } - promptNotificationsActive = true; - if (modelCapacityRetry !== null) { - modelCapacityRetry.turnId = turnId; - } - pendingTurnStart?.resolve(turnId); - if (notifyOnTurnStarted && retryAttempt === 0) { - onTurnStarted?.(); - } - }, - () => this.promptShouldStop(params.sessionId, activePrompt), - )); + const sendPromptPromise = this.runWithProcessCheck(() => this.codexAcpClient.sendPrompt( + turnRequest, + agentMode, + modelId, + serviceTier, + disableSummary, + sessionState.cwd, + sessionState.additionalDirectories, + (turnId) => { + const turn = {threadId: params.sessionId, turnId}; + activePrompt.currentTurn = turn; + if (this.promptShouldStop(params.sessionId, activePrompt)) { + this.interruptLateStartedTurn(turn); + return; + } + sessionState.currentTurnId = turnId; + if (!notificationsActiveBeforeStart && retryAttempt === 0) { + recoverableSessionFailure = sessionState.sessionFailure; + } + promptNotificationsActive = true; + modelCapacityRetry.markTurnStarted(turnId); + pendingTurnStart?.resolve(turnId); + if (notifyOnTurnStarted && retryAttempt === 0) { + onTurnStarted?.(); + } + }, + () => this.promptShouldStop(params.sessionId, activePrompt), + )); void sendPromptPromise.catch((err) => { if (this.activePrompts.get(params.sessionId) !== activePrompt) { logger.error(`Prompt for cancelled session ${params.sessionId} failed after prompt returned`, err); } }); - const completed = await Promise.race([ + return await Promise.race([ sendPromptPromise, activePrompt.closeSignal, this.cancelBeforeTurnStarted(activePrompt), ]); - - if (completed === null) { - modelCapacityRetry = null; - return null; - } - + }, + afterTurn: async () => { await this.codexAcpClient.waitForSessionNotifications(params.sessionId); await promptEventHandler.flushPendingErrors(); - const retry = modelCapacityRetry; - const shouldRetry = retry !== null - && retry.turnId === completed.turn.id - && completed.turn.status === "failed" - && isModelCapacityError(completed.turn.error); - if (!shouldRetry) { - modelCapacityRetry = null; - await promptEventHandler.handleFailedTurn(completed.turn); - promptNotificationsActive = false; - return completed; - } - + }, + onFinalTurn: async (completed) => { + await promptEventHandler.handleFailedTurn(completed.turn); + promptNotificationsActive = false; + }, + onRetry: async (completed, retry) => { if (!retry.warningPublished) { await promptEventHandler.publishModelCapacityRetryWarning( completed.turn.id, - modelCapacityRetryTitle(retry), + retry.title, ); } await promptEventHandler.flushPendingPlanUpdates(); @@ -2638,21 +2545,8 @@ export class CodexAcpServer { attempt: retry.attempt, delaySeconds: retry.delaySeconds, }); - const retryReady = await this.waitForPromptRetry( - retry.delaySeconds, - params.sessionId, - activePrompt, - ); - if (!retryReady) { - modelCapacityRetry = null; - return null; - } - turnRequest = { - sessionId: params.sessionId, - prompt: [{type: "text", text: MODEL_CAPACITY_CONTINUATION_PROMPT}], - }; - } - }; + }, + }); let turnCompleted = await runTurnWithModelCapacityRetries(effectiveParams, true, true); diff --git a/src/ModelCapacityRetry.ts b/src/ModelCapacityRetry.ts new file mode 100644 index 00000000..6eec886a --- /dev/null +++ b/src/ModelCapacityRetry.ts @@ -0,0 +1,133 @@ +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 ModelCapacityRetryAttempt { + attempt: number; + delaySeconds: number; + title: string; + turnId: string | null; + warningPublished: boolean; +} + +interface ModelCapacityRetryHooks { + signal: AbortSignal; + shouldStop(): boolean; + runTurn(request: PromptRequest, retryAttempt: number): Promise; + afterTurn(completed: TurnCompletedNotification): Promise; + onRetry(completed: TurnCompletedNotification, retry: ModelCapacityRetryAttempt): Promise; + onFinalTurn(completed: TurnCompletedNotification): Promise; +} + +export class ModelCapacityRetryController { + private activeRetry: ModelCapacityRetryAttempt | null = null; + + async run(initialRequest: PromptRequest, hooks: ModelCapacityRetryHooks): Promise { + let request = initialRequest; + for (let retryAttempt = 0; ; retryAttempt++) { + this.activeRetry = this.createRetry(retryAttempt); + const completed = await hooks.runTurn(request, retryAttempt); + if (completed === null) { + this.activeRetry = null; + return null; + } + + await hooks.afterTurn(completed); + const retry = this.activeRetry; + if (retry === null + || retry.turnId !== completed.turn.id + || completed.turn.status !== "failed" + || !isModelCapacityError(completed.turn.error)) { + this.activeRetry = null; + await hooks.onFinalTurn(completed); + return completed; + } + + await hooks.onRetry(completed, retry); + if (!await this.waitForRetry(retry.delaySeconds, hooks)) { + this.activeRetry = null; + return null; + } + request = { + sessionId: initialRequest.sessionId, + prompt: [{type: "text", text: CONTINUATION_PROMPT}], + }; + } + } + + markTurnStarted(turnId: string): void { + if (this.activeRetry !== null) { + this.activeRetry.turnId = turnId; + } + } + + transformNotification(notification: ServerNotification, currentTurnId: string | null): ServerNotification { + if (notification.method !== "error" + || notification.params.willRetry + || !isModelCapacityError(notification.params.error) + || this.activeRetry === null + || notification.params.turnId !== currentTurnId) { + return notification; + } + this.activeRetry.warningPublished = true; + return { + method: "error", + params: { + ...notification.params, + willRetry: true, + error: { + ...notification.params.error, + message: this.activeRetry.title, + }, + }, + }; + } + + private createRetry(retryAttempt: number): ModelCapacityRetryAttempt | 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}).`, + turnId: null, + warningPublished: false, + }; + } + + private async waitForRetry(delaySeconds: number, hooks: ModelCapacityRetryHooks): Promise { + if (hooks.shouldStop()) { + return false; + } + return await new Promise((resolve) => { + const finish = (completed: boolean) => { + clearTimeout(timer); + hooks.signal.removeEventListener("abort", onAbort); + resolve(completed); + }; + const onAbort = () => finish(false); + const timer = setTimeout(() => finish(true), delaySeconds * 1000); + hooks.signal.addEventListener("abort", onAbort, {once: true}); + }); + } +} + +function isModelCapacityError(error: {message: string; codexErrorInfo: unknown} | null): boolean { + return error?.codexErrorInfo === "serverOverloaded" && error.message.trim() === ERROR_MESSAGE; +} From e9a693c634dd410f4af1824582973e464d64df17 Mon Sep 17 00:00:00 2001 From: Nikita Ashikhmin Date: Thu, 27 Aug 2026 11:24:04 +0400 Subject: [PATCH 3/6] refactor: extract capacity retry runner --- src/CodexAcpServer.ts | 110 ++++++++++--------------------- src/ModelCapacityRetry.ts | 135 +++++++++++++++++++++++++++++++++++++- 2 files changed, 169 insertions(+), 76 deletions(-) diff --git a/src/CodexAcpServer.ts b/src/CodexAcpServer.ts index 8724450d..ec7f99de 100644 --- a/src/CodexAcpServer.ts +++ b/src/CodexAcpServer.ts @@ -126,7 +126,7 @@ import { createUnavailableAgentFileChangeReport, parseAgentFileChangeReportRequest, } from "./AgentFileChangeReport"; -import {ModelCapacityRetryController} from "./ModelCapacityRetry"; +import {CodexModelCapacityRetryRunner} from "./ModelCapacityRetry"; export interface SessionState { @@ -2297,7 +2297,7 @@ export class CodexAcpServer { const disposePromptRequestCancellation = this.observePromptRequestCancellation(signal, sessionState, activePrompt); let eventHandler: CodexEventHandler | null = null; let promptNotificationsActive = true; - const modelCapacityRetry = new ModelCapacityRetryController(); + let modelCapacityRetry: CodexModelCapacityRetryRunner | null = null; const clearRecoveredSessionFailure = async (handler: CodexEventHandler): Promise => { await handler.completeSuccessfulTurn(sessionState.currentTurnId); const current = sessionState.sessionFailure; @@ -2344,7 +2344,8 @@ export class CodexAcpServer { await promptEventHandler.handleSessionScopedNotification(event); return; } - const handledEvent = modelCapacityRetry.transformNotification(event, sessionState.currentTurnId); + const handledEvent = modelCapacityRetry?.transformNotification(event, sessionState.currentTurnId) + ?? event; const completesActiveTurn = event.method === "turn/completed" && event.params.turn.id === sessionState.currentTurnId; permissionContext.handleNotification(handledEvent); @@ -2470,85 +2471,43 @@ export class CodexAcpServer { sessionState.fastModeEnabled, sessionState.currentModelSupportsFast, ); - const runTurnWithModelCapacityRetries = async ( - initialRequest: acp.PromptRequest, - notifyOnTurnStarted: boolean, - notificationsActiveBeforeStart: boolean, - ) => modelCapacityRetry.run(initialRequest, { - signal: activePrompt.signal, + const capacityRetryRunner = new CodexModelCapacityRetryRunner({ + sessionId: params.sessionId, + activePrompt, + sessionState, + codexAcpClient: this.codexAcpClient, + eventHandler: promptEventHandler, + agentMode, + modelId, + serviceTier, + disableSummary, + runWithProcessCheck: (operation) => this.runWithProcessCheck(operation), shouldStop: () => this.promptShouldStop(params.sessionId, activePrompt), - runTurn: async (turnRequest, retryAttempt) => { - promptNotificationsActive = retryAttempt === 0 ? notificationsActiveBeforeStart : true; + isActivePrompt: () => this.activePrompts.get(params.sessionId) === activePrompt, + interruptLateStartedTurn: (turn) => this.interruptLateStartedTurn(turn), + cancelBeforeTurnStarted: () => this.cancelBeforeTurnStarted(activePrompt), + ensurePendingTurnStart: () => { ensurePendingTurnStart(); - const sendPromptPromise = this.runWithProcessCheck(() => this.codexAcpClient.sendPrompt( - turnRequest, - agentMode, - modelId, - serviceTier, - disableSummary, - sessionState.cwd, - sessionState.additionalDirectories, - (turnId) => { - const turn = {threadId: params.sessionId, turnId}; - activePrompt.currentTurn = turn; - if (this.promptShouldStop(params.sessionId, activePrompt)) { - this.interruptLateStartedTurn(turn); - return; - } - sessionState.currentTurnId = turnId; - if (!notificationsActiveBeforeStart && retryAttempt === 0) { - recoverableSessionFailure = sessionState.sessionFailure; - } - promptNotificationsActive = true; - modelCapacityRetry.markTurnStarted(turnId); - pendingTurnStart?.resolve(turnId); - if (notifyOnTurnStarted && retryAttempt === 0) { - onTurnStarted?.(); - } - }, - () => this.promptShouldStop(params.sessionId, activePrompt), - )); - void sendPromptPromise.catch((err) => { - if (this.activePrompts.get(params.sessionId) !== activePrompt) { - logger.error(`Prompt for cancelled session ${params.sessionId} failed after prompt returned`, err); - } - }); - return await Promise.race([ - sendPromptPromise, - activePrompt.closeSignal, - this.cancelBeforeTurnStarted(activePrompt), - ]); }, - afterTurn: async () => { - await this.codexAcpClient.waitForSessionNotifications(params.sessionId); - await promptEventHandler.flushPendingErrors(); + resolvePendingTurnStart: (turnId) => pendingTurnStart?.resolve(turnId), + preparePendingTurnStart: () => { + pendingTurnStart = this.createPendingTurnStart(); + this.pendingTurnStarts.set(params.sessionId, pendingTurnStart); }, - onFinalTurn: async (completed) => { - await promptEventHandler.handleFailedTurn(completed.turn); - promptNotificationsActive = false; + setPromptNotificationsActive: (active) => { + promptNotificationsActive = active; }, - onRetry: async (completed, retry) => { - if (!retry.warningPublished) { - await promptEventHandler.publishModelCapacityRetryWarning( - completed.turn.id, - retry.title, - ); - } - await promptEventHandler.flushPendingPlanUpdates(); + snapshotRecoverableSessionFailure: () => { recoverableSessionFailure = sessionState.sessionFailure; - activePrompt.currentTurn = null; - sessionState.currentTurnId = null; - pendingTurnStart = this.createPendingTurnStart(); - this.pendingTurnStarts.set(params.sessionId, pendingTurnStart); - logger.log("Selected model is at capacity; scheduling retry", { - sessionId: params.sessionId, - attempt: retry.attempt, - delaySeconds: retry.delaySeconds, - }); }, + onTurnStarted, }); + modelCapacityRetry = capacityRetryRunner; - let turnCompleted = await runTurnWithModelCapacityRetries(effectiveParams, true, true); + let turnCompleted = await capacityRetryRunner.run(effectiveParams, { + notifyOnTurnStarted: true, + notificationsActiveBeforeStart: true, + }); if (turnCompleted === null) { return cancelledPromptResponse(); @@ -2607,7 +2566,10 @@ export class CodexAcpServer { recoverableSessionFailure = sessionState.sessionFailure; pendingTurnStart = this.createPendingTurnStart(); this.pendingTurnStarts.set(params.sessionId, pendingTurnStart); - turnCompleted = await runTurnWithModelCapacityRetries(implementationRequest, false, false); + turnCompleted = await capacityRetryRunner.run(implementationRequest, { + notifyOnTurnStarted: false, + notificationsActiveBeforeStart: false, + }); if (turnCompleted === null) { return cancelledPromptResponse(); diff --git a/src/ModelCapacityRetry.ts b/src/ModelCapacityRetry.ts index 6eec886a..684fb9bd 100644 --- a/src/ModelCapacityRetry.ts +++ b/src/ModelCapacityRetry.ts @@ -1,6 +1,12 @@ import type {PromptRequest} from "@agentclientprotocol/sdk"; import type {ServerNotification} from "./app-server"; +import type {ServiceTier} from "./app-server/ServiceTier"; import type {TurnCompletedNotification} from "./app-server/v2"; +import type {AgentMode} from "./AgentMode"; +import type {CodexAcpClient} from "./CodexAcpClient"; +import type {CodexEventHandler} from "./CodexEventHandler"; +import {logger} from "./Logger"; +import type {ModelId} from "./ModelId"; const ERROR_MESSAGE = "Selected model is at capacity. Please try a different model."; const CONTINUATION_PROMPT = "Continue from where you left off."; @@ -12,7 +18,7 @@ const RETRY_WINDOWS_SECONDS = [ [120, 300], ] as const; -export interface ModelCapacityRetryAttempt { +interface ModelCapacityRetryAttempt { attempt: number; delaySeconds: number; title: string; @@ -29,7 +35,7 @@ interface ModelCapacityRetryHooks { onFinalTurn(completed: TurnCompletedNotification): Promise; } -export class ModelCapacityRetryController { +class ModelCapacityRetryController { private activeRetry: ModelCapacityRetryAttempt | null = null; async run(initialRequest: PromptRequest, hooks: ModelCapacityRetryHooks): Promise { @@ -128,6 +134,131 @@ export class ModelCapacityRetryController { } } +interface ActivePromptRuntime { + signal: AbortSignal; + closeSignal: Promise; + currentTurn: {threadId: string; turnId: string} | null; +} + +interface ModelCapacitySessionRuntime { + currentTurnId: string | null; + cwd: string; + additionalDirectories: string[]; +} + +interface CodexModelCapacityRetryRunnerOptions { + sessionId: string; + activePrompt: ActivePromptRuntime; + sessionState: ModelCapacitySessionRuntime; + codexAcpClient: CodexAcpClient; + eventHandler: CodexEventHandler; + agentMode: AgentMode; + modelId: ModelId; + serviceTier: ServiceTier | null; + disableSummary: boolean; + runWithProcessCheck(operation: () => Promise): Promise; + shouldStop(): boolean; + isActivePrompt(): boolean; + interruptLateStartedTurn(turn: {threadId: string; turnId: string}): void; + cancelBeforeTurnStarted(): Promise; + ensurePendingTurnStart(): void; + resolvePendingTurnStart(turnId: string): void; + preparePendingTurnStart(): void; + setPromptNotificationsActive(active: boolean): void; + snapshotRecoverableSessionFailure(): void; + onTurnStarted: (() => void) | undefined; +} + +interface RunOptions { + notifyOnTurnStarted: boolean; + notificationsActiveBeforeStart: boolean; +} + +export class CodexModelCapacityRetryRunner { + private readonly retry = new ModelCapacityRetryController(); + + constructor(private readonly options: CodexModelCapacityRetryRunnerOptions) {} + + transformNotification(notification: ServerNotification, currentTurnId: string | null): ServerNotification { + return this.retry.transformNotification(notification, currentTurnId); + } + + async run(initialRequest: PromptRequest, runOptions: RunOptions): Promise { + const options = this.options; + return await this.retry.run(initialRequest, { + signal: options.activePrompt.signal, + shouldStop: options.shouldStop, + runTurn: async (turnRequest, retryAttempt) => { + options.setPromptNotificationsActive( + retryAttempt === 0 ? runOptions.notificationsActiveBeforeStart : true, + ); + options.ensurePendingTurnStart(); + const sendPromptPromise = options.runWithProcessCheck(() => options.codexAcpClient.sendPrompt( + turnRequest, + options.agentMode, + options.modelId, + options.serviceTier, + options.disableSummary, + options.sessionState.cwd, + options.sessionState.additionalDirectories, + (turnId) => { + const turn = {threadId: options.sessionId, turnId}; + options.activePrompt.currentTurn = turn; + if (options.shouldStop()) { + options.interruptLateStartedTurn(turn); + return; + } + options.sessionState.currentTurnId = turnId; + if (!runOptions.notificationsActiveBeforeStart && retryAttempt === 0) { + options.snapshotRecoverableSessionFailure(); + } + options.setPromptNotificationsActive(true); + this.retry.markTurnStarted(turnId); + options.resolvePendingTurnStart(turnId); + if (runOptions.notifyOnTurnStarted && retryAttempt === 0) { + options.onTurnStarted?.(); + } + }, + options.shouldStop, + )); + void sendPromptPromise.catch((error) => { + if (!options.isActivePrompt()) { + logger.error(`Prompt for cancelled session ${options.sessionId} failed after prompt returned`, error); + } + }); + return await Promise.race([ + sendPromptPromise, + options.activePrompt.closeSignal, + options.cancelBeforeTurnStarted(), + ]); + }, + afterTurn: async () => { + await options.codexAcpClient.waitForSessionNotifications(options.sessionId); + await options.eventHandler.flushPendingErrors(); + }, + onFinalTurn: async (completed) => { + await options.eventHandler.handleFailedTurn(completed.turn); + options.setPromptNotificationsActive(false); + }, + onRetry: async (completed, retry) => { + if (!retry.warningPublished) { + await options.eventHandler.publishModelCapacityRetryWarning(completed.turn.id, retry.title); + } + await options.eventHandler.flushPendingPlanUpdates(); + options.snapshotRecoverableSessionFailure(); + options.activePrompt.currentTurn = null; + options.sessionState.currentTurnId = null; + options.preparePendingTurnStart(); + logger.log("Selected model is at capacity; scheduling retry", { + sessionId: options.sessionId, + attempt: retry.attempt, + delaySeconds: retry.delaySeconds, + }); + }, + }); + } +} + function isModelCapacityError(error: {message: string; codexErrorInfo: unknown} | null): boolean { return error?.codexErrorInfo === "serverOverloaded" && error.message.trim() === ERROR_MESSAGE; } From 90118c61a7084bb92418ca630e84a71bcd8a120a Mon Sep 17 00:00:00 2001 From: Nikita Ashikhmin Date: Thu, 27 Aug 2026 11:36:43 +0400 Subject: [PATCH 4/6] refactor: simplify model capacity retries --- src/CodexAcpClient.ts | 129 ++++++++- src/CodexAcpServer.ts | 185 ++++++++---- src/CodexEventHandler.ts | 16 -- src/ModelCapacityRetry.ts | 264 ------------------ .../CodexACPAgent/auth-error-events.test.ts | 50 +--- 5 files changed, 269 insertions(+), 375 deletions(-) delete mode 100644 src/ModelCapacityRetry.ts diff --git a/src/CodexAcpClient.ts b/src/CodexAcpClient.ts index bec75265..e1888ac8 100644 --- a/src/CodexAcpClient.ts +++ b/src/CodexAcpClient.ts @@ -72,6 +72,73 @@ import { export const CUSTOM_GATEWAY_PROVIDER_ID = "custom-gateway"; export const OPENAI_PROVIDER_ID = "openai"; const DEFAULT_OPENAI_BASE_URL = "https://api.openai.com/v1"; +const MODEL_CAPACITY_ERROR_MESSAGE = "Selected model is at capacity. Please try a different model."; +const MODEL_CAPACITY_CONTINUATION_PROMPT = "Continue from where you left off."; +const MODEL_CAPACITY_RETRY_WINDOWS_SECONDS = [ + [1, 10], + [1, 30], + [30, 60], + [60, 120], + [120, 300], +] as const; + +export interface ModelCapacityRetry { + attempt: number; + delaySeconds: number; + title: string; + warningPublished: boolean; +} + +class ModelCapacityRetryError extends Error { + constructor( + readonly completed: TurnCompletedNotification, + readonly retry: ModelCapacityRetry, + ) { + super(retry.title); + } +} + +export function isModelCapacityError(error: {message: string; codexErrorInfo: unknown} | null): boolean { + return error?.codexErrorInfo === "serverOverloaded" + && error.message.trim() === MODEL_CAPACITY_ERROR_MESSAGE; +} + +function createModelCapacityRetry(retryAttempt: number): ModelCapacityRetry | null { + const window = MODEL_CAPACITY_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}/${MODEL_CAPACITY_RETRY_WINDOWS_SECONDS.length}).`, + warningPublished: false, + }; +} + +async function waitForModelCapacityRetry( + delaySeconds: number, + 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?.()), delaySeconds * 1000); + signal?.addEventListener("abort", onAbort, {once: true}); + }); +} /** * The url-mode variant of the ACP `elicitation/create` request params. @@ -848,8 +915,14 @@ export class CodexAcpClient { disableSummary: boolean, cwd: string, additionalDirectories: string[], - onTurnStarted?: (turnId: string) => void, + onTurnStarted?: (turnId: string, retry: ModelCapacityRetry | null) => void, shouldCancel?: () => boolean, + onModelCapacityRetry?: ( + completed: TurnCompletedNotification, + retry: ModelCapacityRetry, + ) => Promise, + retrySignal?: AbortSignal, + retryAttempt = 0, ): Promise { const input = buildPromptItems(request.prompt); const effort = modelId.effort as ReasoningEffort | null; //TODO remove unsafe conversion @@ -857,17 +930,49 @@ export class CodexAcpClient { if (shouldCancel?.()) { return null; } - return await this.codexClient.runTurn({ - threadId: request.sessionId, - input: input, - approvalPolicy: agentMode.approvalPolicy, - approvalsReviewer: agentMode.approvalsReviewer, - sandboxPolicy: addAdditionalDirectoriesToSandboxPolicy(agentMode.sandboxPolicy, additionalDirectories), - summary: disableSummary ? "none" : "auto", - effort: effort, - model: modelId.model, - serviceTier: serviceTier, - }, onTurnStarted); + const retry = createModelCapacityRetry(retryAttempt); + try { + const completed = await this.codexClient.runTurn({ + threadId: request.sessionId, + input: input, + approvalPolicy: agentMode.approvalPolicy, + approvalsReviewer: agentMode.approvalsReviewer, + sandboxPolicy: addAdditionalDirectoriesToSandboxPolicy(agentMode.sandboxPolicy, additionalDirectories), + summary: disableSummary ? "none" : "auto", + effort: effort, + model: modelId.model, + serviceTier: serviceTier, + }, (turnId) => onTurnStarted?.(turnId, retry)); + if (retry !== null && completed.turn.status === "failed" && isModelCapacityError(completed.turn.error)) { + throw new ModelCapacityRetryError(completed, retry); + } + return completed; + } catch (error) { + if (!(error instanceof ModelCapacityRetryError)) { + throw error; + } + await onModelCapacityRetry?.(error.completed, error.retry); + if (!await waitForModelCapacityRetry(error.retry.delaySeconds, retrySignal, shouldCancel)) { + return null; + } + return await this.sendPrompt( + { + sessionId: request.sessionId, + prompt: [{type: "text", text: MODEL_CAPACITY_CONTINUATION_PROMPT}], + }, + agentMode, + modelId, + serviceTier, + disableSummary, + cwd, + additionalDirectories, + onTurnStarted, + shouldCancel, + onModelCapacityRetry, + retrySignal, + retryAttempt + 1, + ); + } } async runAgentFileChangeReport(params: { diff --git a/src/CodexAcpServer.ts b/src/CodexAcpServer.ts index ec7f99de..7da0b130 100644 --- a/src/CodexAcpServer.ts +++ b/src/CodexAcpServer.ts @@ -13,7 +13,9 @@ import {type CodexAuthRequest, getCodexAuthMethods, isCodexAuthRequest} from "./ import {clientSupportsUrlElicitation} from "./ElicitationCapabilities"; import { CodexAcpClient, + isModelCapacityError, type JsonObject, + type ModelCapacityRetry, OPENAI_PROVIDER_ID, type SessionMetadata, type SessionMetadataWithThread, @@ -30,6 +32,7 @@ import type { Thread, ThreadGoal, ThreadItem, + TurnCompletedNotification, UserInput, } from "./app-server/v2"; import type {RateLimitsMap} from "./RateLimitsMap"; @@ -126,7 +129,6 @@ import { createUnavailableAgentFileChangeReport, parseAgentFileChangeReportRequest, } from "./AgentFileChangeReport"; -import {CodexModelCapacityRetryRunner} from "./ModelCapacityRetry"; export interface SessionState { @@ -2297,7 +2299,7 @@ export class CodexAcpServer { const disposePromptRequestCancellation = this.observePromptRequestCancellation(signal, sessionState, activePrompt); let eventHandler: CodexEventHandler | null = null; let promptNotificationsActive = true; - let modelCapacityRetry: CodexModelCapacityRetryRunner | null = null; + let modelCapacityRetry: ModelCapacityRetry | null = null; const clearRecoveredSessionFailure = async (handler: CodexEventHandler): Promise => { await handler.completeSuccessfulTurn(sessionState.currentTurnId); const current = sessionState.sessionFailure; @@ -2344,8 +2346,22 @@ export class CodexAcpServer { await promptEventHandler.handleSessionScopedNotification(event); return; } - const handledEvent = modelCapacityRetry?.transformNotification(event, sessionState.currentTurnId) - ?? event; + let handledEvent = event; + if (event.method === "error" + && !event.params.willRetry + && modelCapacityRetry !== null + && event.params.turnId === sessionState.currentTurnId + && isModelCapacityError(event.params.error)) { + modelCapacityRetry.warningPublished = true; + handledEvent = { + method: "error", + params: { + ...event.params, + willRetry: true, + error: {...event.params.error, message: modelCapacityRetry.title}, + }, + }; + } const completesActiveTurn = event.method === "turn/completed" && event.params.turn.id === sessionState.currentTurnId; permissionContext.handleNotification(handledEvent); @@ -2360,6 +2376,40 @@ export class CodexAcpServer { approvalHandler, elicitationHandler); + const handleModelCapacityRetry = async ( + completed: TurnCompletedNotification, + retry: ModelCapacityRetry, + ) => { + await this.codexAcpClient.waitForSessionNotifications(params.sessionId); + await promptEventHandler.flushPendingErrors(); + if (!retry.warningPublished) { + await promptEventHandler.handleNotification({ + method: "error", + params: { + threadId: params.sessionId, + turnId: completed.turn.id, + willRetry: true, + error: { + message: retry.title, + codexErrorInfo: "serverOverloaded", + additionalDetails: null, + }, + }, + }); + } + await promptEventHandler.flushPendingPlanUpdates(); + recoverableSessionFailure = sessionState.sessionFailure; + activePrompt.currentTurn = null; + sessionState.currentTurnId = null; + pendingTurnStart = this.createPendingTurnStart(); + this.pendingTurnStarts.set(params.sessionId, pendingTurnStart); + logger.log("Selected model is at capacity; scheduling retry", { + sessionId: params.sessionId, + attempt: retry.attempt, + delaySeconds: retry.delaySeconds, + }); + }; + if (activePrompt.signal.aborted) { return cancelledPromptResponse(); } @@ -2471,48 +2521,54 @@ export class CodexAcpServer { sessionState.fastModeEnabled, sessionState.currentModelSupportsFast, ); - const capacityRetryRunner = new CodexModelCapacityRetryRunner({ - sessionId: params.sessionId, - activePrompt, - sessionState, - codexAcpClient: this.codexAcpClient, - eventHandler: promptEventHandler, - agentMode, - modelId, - serviceTier, - disableSummary, - runWithProcessCheck: (operation) => this.runWithProcessCheck(operation), - shouldStop: () => this.promptShouldStop(params.sessionId, activePrompt), - isActivePrompt: () => this.activePrompts.get(params.sessionId) === activePrompt, - interruptLateStartedTurn: (turn) => this.interruptLateStartedTurn(turn), - cancelBeforeTurnStarted: () => this.cancelBeforeTurnStarted(activePrompt), - ensurePendingTurnStart: () => { - ensurePendingTurnStart(); - }, - resolvePendingTurnStart: (turnId) => pendingTurnStart?.resolve(turnId), - preparePendingTurnStart: () => { - pendingTurnStart = this.createPendingTurnStart(); - this.pendingTurnStarts.set(params.sessionId, pendingTurnStart); - }, - setPromptNotificationsActive: (active) => { - promptNotificationsActive = active; - }, - snapshotRecoverableSessionFailure: () => { - recoverableSessionFailure = sessionState.sessionFailure; - }, - onTurnStarted, - }); - modelCapacityRetry = capacityRetryRunner; - - let turnCompleted = await capacityRetryRunner.run(effectiveParams, { - notifyOnTurnStarted: true, - notificationsActiveBeforeStart: true, + ensurePendingTurnStart(); + const sendPromptPromise = this.runWithProcessCheck( + () => this.codexAcpClient.sendPrompt( + effectiveParams, + agentMode, + modelId, + serviceTier, + disableSummary, + sessionState.cwd, + sessionState.additionalDirectories, + (turnId, retry) => { + const turn = {threadId: params.sessionId, turnId}; + activePrompt.currentTurn = turn; + if (this.promptShouldStop(params.sessionId, activePrompt)) { + this.interruptLateStartedTurn(turn); + return; + } + sessionState.currentTurnId = turnId; + modelCapacityRetry = retry; + pendingTurnStart?.resolve(turnId); + if (retry?.attempt === 1) { + onTurnStarted?.(); + } + }, + () => this.promptShouldStop(params.sessionId, activePrompt), + handleModelCapacityRetry, + activePrompt.signal, + )); + void sendPromptPromise.catch((err) => { + if (this.activePrompts.get(params.sessionId) !== activePrompt) { + logger.error(`Prompt for cancelled session ${params.sessionId} failed after prompt returned`, err); + } }); + let turnCompleted = await Promise.race([ + sendPromptPromise, + activePrompt.closeSignal, + this.cancelBeforeTurnStarted(activePrompt), + ]); if (turnCompleted === null) { return cancelledPromptResponse(); } + await this.codexAcpClient.waitForSessionNotifications(params.sessionId); + await eventHandler.flushPendingErrors(); + await eventHandler.handleFailedTurn(turnCompleted.turn); + promptNotificationsActive = false; + if (turnCompleted.turn.status === "interrupted") { await eventHandler.flushPendingPlanUpdates(); return cancelledPromptResponse(); @@ -2561,20 +2617,53 @@ export class CodexAcpServer { }; activePrompt.currentTurn = null; sessionState.currentTurnId = null; - // Keep the approval-to-turn-start gap session-scoped. Snapshot any unchanged failure - // before starting implementation so a successful turn can clear it. - recoverableSessionFailure = sessionState.sessionFailure; - pendingTurnStart = this.createPendingTurnStart(); - this.pendingTurnStarts.set(params.sessionId, pendingTurnStart); - turnCompleted = await capacityRetryRunner.run(implementationRequest, { - notifyOnTurnStarted: false, - notificationsActiveBeforeStart: false, + const implementationPromise = this.runWithProcessCheck( + () => this.codexAcpClient.sendPrompt( + implementationRequest, + agentMode, + modelId, + serviceTier, + disableSummary, + sessionState.cwd, + sessionState.additionalDirectories, + (turnId, retry) => { + const turn = {threadId: params.sessionId, turnId}; + activePrompt.currentTurn = turn; + if (this.promptShouldStop(params.sessionId, activePrompt)) { + this.interruptLateStartedTurn(turn); + return; + } + sessionState.currentTurnId = turnId; + modelCapacityRetry = 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), + handleModelCapacityRetry, + activePrompt.signal, + ), + ); + void implementationPromise.catch((err) => { + if (this.activePrompts.get(params.sessionId) !== activePrompt) { + logger.error(`Implementation turn for cancelled prompt ${params.sessionId} failed after prompt returned`, err); + } }); + turnCompleted = await Promise.race([ + implementationPromise, + activePrompt.closeSignal, + this.cancelBeforeTurnStarted(activePrompt), + ]); if (turnCompleted === null) { return cancelledPromptResponse(); } + await this.codexAcpClient.waitForSessionNotifications(params.sessionId); + await eventHandler.flushPendingErrors(); + await eventHandler.handleFailedTurn(turnCompleted.turn); + promptNotificationsActive = false; if (turnCompleted.turn.status === "interrupted") { await eventHandler.flushPendingPlanUpdates(); return cancelledPromptResponse(); diff --git a/src/CodexEventHandler.ts b/src/CodexEventHandler.ts index f9919189..7360b6f9 100644 --- a/src/CodexEventHandler.ts +++ b/src/CodexEventHandler.ts @@ -323,22 +323,6 @@ export class CodexEventHandler { delete this.sessionState.sessionFailure; } - async publishModelCapacityRetryWarning(turnId: string, title: string): Promise { - const update = await this.createErrorEvent({ - threadId: this.sessionState.sessionId, - turnId, - willRetry: true, - error: { - message: title, - codexErrorInfo: "serverOverloaded", - additionalDetails: null, - }, - }); - if (update !== null) { - await this.session.update(update); - } - } - async completeSuccessfulTurn(turnId: string | null): Promise { this.lastSessionNotice = undefined; if (!this.supportsTypedSessionFailures || turnId === null) return; diff --git a/src/ModelCapacityRetry.ts b/src/ModelCapacityRetry.ts deleted file mode 100644 index 684fb9bd..00000000 --- a/src/ModelCapacityRetry.ts +++ /dev/null @@ -1,264 +0,0 @@ -import type {PromptRequest} from "@agentclientprotocol/sdk"; -import type {ServerNotification} from "./app-server"; -import type {ServiceTier} from "./app-server/ServiceTier"; -import type {TurnCompletedNotification} from "./app-server/v2"; -import type {AgentMode} from "./AgentMode"; -import type {CodexAcpClient} from "./CodexAcpClient"; -import type {CodexEventHandler} from "./CodexEventHandler"; -import {logger} from "./Logger"; -import type {ModelId} from "./ModelId"; - -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; - -interface ModelCapacityRetryAttempt { - attempt: number; - delaySeconds: number; - title: string; - turnId: string | null; - warningPublished: boolean; -} - -interface ModelCapacityRetryHooks { - signal: AbortSignal; - shouldStop(): boolean; - runTurn(request: PromptRequest, retryAttempt: number): Promise; - afterTurn(completed: TurnCompletedNotification): Promise; - onRetry(completed: TurnCompletedNotification, retry: ModelCapacityRetryAttempt): Promise; - onFinalTurn(completed: TurnCompletedNotification): Promise; -} - -class ModelCapacityRetryController { - private activeRetry: ModelCapacityRetryAttempt | null = null; - - async run(initialRequest: PromptRequest, hooks: ModelCapacityRetryHooks): Promise { - let request = initialRequest; - for (let retryAttempt = 0; ; retryAttempt++) { - this.activeRetry = this.createRetry(retryAttempt); - const completed = await hooks.runTurn(request, retryAttempt); - if (completed === null) { - this.activeRetry = null; - return null; - } - - await hooks.afterTurn(completed); - const retry = this.activeRetry; - if (retry === null - || retry.turnId !== completed.turn.id - || completed.turn.status !== "failed" - || !isModelCapacityError(completed.turn.error)) { - this.activeRetry = null; - await hooks.onFinalTurn(completed); - return completed; - } - - await hooks.onRetry(completed, retry); - if (!await this.waitForRetry(retry.delaySeconds, hooks)) { - this.activeRetry = null; - return null; - } - request = { - sessionId: initialRequest.sessionId, - prompt: [{type: "text", text: CONTINUATION_PROMPT}], - }; - } - } - - markTurnStarted(turnId: string): void { - if (this.activeRetry !== null) { - this.activeRetry.turnId = turnId; - } - } - - transformNotification(notification: ServerNotification, currentTurnId: string | null): ServerNotification { - if (notification.method !== "error" - || notification.params.willRetry - || !isModelCapacityError(notification.params.error) - || this.activeRetry === null - || notification.params.turnId !== currentTurnId) { - return notification; - } - this.activeRetry.warningPublished = true; - return { - method: "error", - params: { - ...notification.params, - willRetry: true, - error: { - ...notification.params.error, - message: this.activeRetry.title, - }, - }, - }; - } - - private createRetry(retryAttempt: number): ModelCapacityRetryAttempt | 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}).`, - turnId: null, - warningPublished: false, - }; - } - - private async waitForRetry(delaySeconds: number, hooks: ModelCapacityRetryHooks): Promise { - if (hooks.shouldStop()) { - return false; - } - return await new Promise((resolve) => { - const finish = (completed: boolean) => { - clearTimeout(timer); - hooks.signal.removeEventListener("abort", onAbort); - resolve(completed); - }; - const onAbort = () => finish(false); - const timer = setTimeout(() => finish(true), delaySeconds * 1000); - hooks.signal.addEventListener("abort", onAbort, {once: true}); - }); - } -} - -interface ActivePromptRuntime { - signal: AbortSignal; - closeSignal: Promise; - currentTurn: {threadId: string; turnId: string} | null; -} - -interface ModelCapacitySessionRuntime { - currentTurnId: string | null; - cwd: string; - additionalDirectories: string[]; -} - -interface CodexModelCapacityRetryRunnerOptions { - sessionId: string; - activePrompt: ActivePromptRuntime; - sessionState: ModelCapacitySessionRuntime; - codexAcpClient: CodexAcpClient; - eventHandler: CodexEventHandler; - agentMode: AgentMode; - modelId: ModelId; - serviceTier: ServiceTier | null; - disableSummary: boolean; - runWithProcessCheck(operation: () => Promise): Promise; - shouldStop(): boolean; - isActivePrompt(): boolean; - interruptLateStartedTurn(turn: {threadId: string; turnId: string}): void; - cancelBeforeTurnStarted(): Promise; - ensurePendingTurnStart(): void; - resolvePendingTurnStart(turnId: string): void; - preparePendingTurnStart(): void; - setPromptNotificationsActive(active: boolean): void; - snapshotRecoverableSessionFailure(): void; - onTurnStarted: (() => void) | undefined; -} - -interface RunOptions { - notifyOnTurnStarted: boolean; - notificationsActiveBeforeStart: boolean; -} - -export class CodexModelCapacityRetryRunner { - private readonly retry = new ModelCapacityRetryController(); - - constructor(private readonly options: CodexModelCapacityRetryRunnerOptions) {} - - transformNotification(notification: ServerNotification, currentTurnId: string | null): ServerNotification { - return this.retry.transformNotification(notification, currentTurnId); - } - - async run(initialRequest: PromptRequest, runOptions: RunOptions): Promise { - const options = this.options; - return await this.retry.run(initialRequest, { - signal: options.activePrompt.signal, - shouldStop: options.shouldStop, - runTurn: async (turnRequest, retryAttempt) => { - options.setPromptNotificationsActive( - retryAttempt === 0 ? runOptions.notificationsActiveBeforeStart : true, - ); - options.ensurePendingTurnStart(); - const sendPromptPromise = options.runWithProcessCheck(() => options.codexAcpClient.sendPrompt( - turnRequest, - options.agentMode, - options.modelId, - options.serviceTier, - options.disableSummary, - options.sessionState.cwd, - options.sessionState.additionalDirectories, - (turnId) => { - const turn = {threadId: options.sessionId, turnId}; - options.activePrompt.currentTurn = turn; - if (options.shouldStop()) { - options.interruptLateStartedTurn(turn); - return; - } - options.sessionState.currentTurnId = turnId; - if (!runOptions.notificationsActiveBeforeStart && retryAttempt === 0) { - options.snapshotRecoverableSessionFailure(); - } - options.setPromptNotificationsActive(true); - this.retry.markTurnStarted(turnId); - options.resolvePendingTurnStart(turnId); - if (runOptions.notifyOnTurnStarted && retryAttempt === 0) { - options.onTurnStarted?.(); - } - }, - options.shouldStop, - )); - void sendPromptPromise.catch((error) => { - if (!options.isActivePrompt()) { - logger.error(`Prompt for cancelled session ${options.sessionId} failed after prompt returned`, error); - } - }); - return await Promise.race([ - sendPromptPromise, - options.activePrompt.closeSignal, - options.cancelBeforeTurnStarted(), - ]); - }, - afterTurn: async () => { - await options.codexAcpClient.waitForSessionNotifications(options.sessionId); - await options.eventHandler.flushPendingErrors(); - }, - onFinalTurn: async (completed) => { - await options.eventHandler.handleFailedTurn(completed.turn); - options.setPromptNotificationsActive(false); - }, - onRetry: async (completed, retry) => { - if (!retry.warningPublished) { - await options.eventHandler.publishModelCapacityRetryWarning(completed.turn.id, retry.title); - } - await options.eventHandler.flushPendingPlanUpdates(); - options.snapshotRecoverableSessionFailure(); - options.activePrompt.currentTurn = null; - options.sessionState.currentTurnId = null; - options.preparePendingTurnStart(); - logger.log("Selected model is at capacity; scheduling retry", { - sessionId: options.sessionId, - attempt: retry.attempt, - delaySeconds: retry.delaySeconds, - }); - }, - }); - } -} - -function isModelCapacityError(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 96634be5..c79d1078 100644 --- a/src/__tests__/CodexACPAgent/auth-error-events.test.ts +++ b/src/__tests__/CodexACPAgent/auth-error-events.test.ts @@ -110,7 +110,7 @@ describe("CodexEventHandler - auth error events", () => { ); }); - it("retries the exact model-capacity failure with jittered delays and continuation turns", async () => { + it("retries model capacity failures with continuation turns", async () => { vi.useFakeTimers(); const random = vi.spyOn(Math, "random") .mockReturnValueOnce(0) @@ -127,30 +127,21 @@ describe("CodexEventHandler - auth error events", () => { 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 updates = fixture.getAcpConnectionEvents([]).map(event => event.args[0].update); - expect(updates).toEqual(expect.arrayContaining([ - expect.objectContaining({ - sessionUpdate: "session_info_update", - _meta: {jetbrains: {air: expect.objectContaining({sessionFailure: expect.objectContaining({ - severity: "warning", - title: "Selected model is at capacity. Retrying in 1 second (1/5).", - })})}}, - }), - expect.objectContaining({ - sessionUpdate: "session_info_update", - _meta: {jetbrains: {air: expect.objectContaining({sessionFailure: expect.objectContaining({ - severity: "warning", - title: "Selected model is at capacity. Retrying in 16 seconds (2/5).", - })})}}, - }), - ])); + 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 all five custom retries", async () => { + it("returns the terminal model-capacity failure after five retries", async () => { vi.useFakeTimers(); const random = vi.spyOn(Math, "random").mockReturnValue(0); try { @@ -165,37 +156,26 @@ describe("CodexEventHandler - auth error events", () => { category: "service", severity: "error", title: "Selected model is at capacity. Please try a different model.", - actions: ["retry"], }}}}, }); - const warningTitles = fixture.getAcpConnectionEvents([]) + const warnings = 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 1 second (2/5).", - "Selected model is at capacity. Retrying in 30 seconds (3/5).", - "Selected model is at capacity. Retrying in 60 seconds (4/5).", - "Selected model is at capacity. Retrying in 120 seconds (5/5).", - ]); + .filter(failure => failure?.severity === "warning"); + expect(warnings).toHaveLength(5); } finally { random.mockRestore(); vi.useRealTimers(); } }); - it("cancels immediately while waiting to retry model capacity", async () => { + 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 {fixture, promptPromise, turnStart} = await startModelCapacityRetryPrompt(1, controller.signal); + const {promptPromise, turnStart} = await startModelCapacityRetryPrompt(1, controller.signal); await vi.advanceTimersByTimeAsync(0); expect(turnStart).toHaveBeenCalledTimes(1); - expect(fixture.getAcpConnectionEvents([]).some(event => - event.args[0].update?._meta?.jetbrains?.air?.sessionFailure?.severity === "warning", - )).toBe(true); controller.abort(); await vi.advanceTimersByTimeAsync(0); From 65930f28e5a90444b5a39b56cf2bffeb62d65044 Mon Sep 17 00:00:00 2001 From: Nikita Ashikhmin Date: Thu, 27 Aug 2026 11:44:18 +0400 Subject: [PATCH 5/6] refactor: extract retry capacity service --- src/CodexAcpClient.ts | 144 +++++++---------------------------- src/CodexAcpServer.ts | 47 ++++-------- src/RetryCapacityService.ts | 147 ++++++++++++++++++++++++++++++++++++ 3 files changed, 187 insertions(+), 151 deletions(-) create mode 100644 src/RetryCapacityService.ts diff --git a/src/CodexAcpClient.ts b/src/CodexAcpClient.ts index e1888ac8..57f97d29 100644 --- a/src/CodexAcpClient.ts +++ b/src/CodexAcpClient.ts @@ -63,6 +63,7 @@ import { createReportedAgentFileChangeReport, createUnavailableAgentFileChangeReport, } from "./AgentFileChangeReport"; +import {type CapacityRetry, RetryCapacityService} from "./RetryCapacityService"; /** * Well-known provider id for the client-configurable custom LLM gateway. @@ -72,73 +73,6 @@ import { export const CUSTOM_GATEWAY_PROVIDER_ID = "custom-gateway"; export const OPENAI_PROVIDER_ID = "openai"; const DEFAULT_OPENAI_BASE_URL = "https://api.openai.com/v1"; -const MODEL_CAPACITY_ERROR_MESSAGE = "Selected model is at capacity. Please try a different model."; -const MODEL_CAPACITY_CONTINUATION_PROMPT = "Continue from where you left off."; -const MODEL_CAPACITY_RETRY_WINDOWS_SECONDS = [ - [1, 10], - [1, 30], - [30, 60], - [60, 120], - [120, 300], -] as const; - -export interface ModelCapacityRetry { - attempt: number; - delaySeconds: number; - title: string; - warningPublished: boolean; -} - -class ModelCapacityRetryError extends Error { - constructor( - readonly completed: TurnCompletedNotification, - readonly retry: ModelCapacityRetry, - ) { - super(retry.title); - } -} - -export function isModelCapacityError(error: {message: string; codexErrorInfo: unknown} | null): boolean { - return error?.codexErrorInfo === "serverOverloaded" - && error.message.trim() === MODEL_CAPACITY_ERROR_MESSAGE; -} - -function createModelCapacityRetry(retryAttempt: number): ModelCapacityRetry | null { - const window = MODEL_CAPACITY_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}/${MODEL_CAPACITY_RETRY_WINDOWS_SECONDS.length}).`, - warningPublished: false, - }; -} - -async function waitForModelCapacityRetry( - delaySeconds: number, - 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?.()), delaySeconds * 1000); - signal?.addEventListener("abort", onAbort, {once: true}); - }); -} /** * The url-mode variant of the ACP `elicitation/create` request params. @@ -176,6 +110,7 @@ export class CodexAcpClient { private pendingLoginCompleted: Promise | null = null; private pendingAccountUpdated: Promise | null = null; private readonly sessionNotificationQueues = new Map>(); + private readonly retryCapacityService = new RetryCapacityService(); private skillExtraRoots: string[] = []; private configPath: string | null = null; @@ -915,64 +850,37 @@ export class CodexAcpClient { disableSummary: boolean, cwd: string, additionalDirectories: string[], - onTurnStarted?: (turnId: string, retry: ModelCapacityRetry | null) => void, + onTurnStarted?: (turnId: string, retry: CapacityRetry | null) => void, shouldCancel?: () => boolean, onModelCapacityRetry?: ( completed: TurnCompletedNotification, - retry: ModelCapacityRetry, + retry: CapacityRetry, ) => Promise, retrySignal?: AbortSignal, - retryAttempt = 0, ): Promise { - const input = buildPromptItems(request.prompt); const effort = modelId.effort as ReasoningEffort | null; //TODO remove unsafe conversion - await this.refreshSkills(cwd, additionalDirectories); - if (shouldCancel?.()) { - return null; - } - const retry = createModelCapacityRetry(retryAttempt); - try { - const completed = await this.codexClient.runTurn({ - threadId: request.sessionId, - input: input, - approvalPolicy: agentMode.approvalPolicy, - approvalsReviewer: agentMode.approvalsReviewer, - sandboxPolicy: addAdditionalDirectoriesToSandboxPolicy(agentMode.sandboxPolicy, additionalDirectories), - summary: disableSummary ? "none" : "auto", - effort: effort, - model: modelId.model, - serviceTier: serviceTier, - }, (turnId) => onTurnStarted?.(turnId, retry)); - if (retry !== null && completed.turn.status === "failed" && isModelCapacityError(completed.turn.error)) { - throw new ModelCapacityRetryError(completed, retry); - } - return completed; - } catch (error) { - if (!(error instanceof ModelCapacityRetryError)) { - throw error; - } - await onModelCapacityRetry?.(error.completed, error.retry); - if (!await waitForModelCapacityRetry(error.retry.delaySeconds, retrySignal, shouldCancel)) { - return null; - } - return await this.sendPrompt( - { - sessionId: request.sessionId, - prompt: [{type: "text", text: MODEL_CAPACITY_CONTINUATION_PROMPT}], - }, - agentMode, - modelId, - serviceTier, - disableSummary, - cwd, - additionalDirectories, - onTurnStarted, - shouldCancel, - onModelCapacityRetry, - retrySignal, - retryAttempt + 1, - ); - } + return await this.retryCapacityService.run(request, { + signal: retrySignal, + shouldCancel, + onRetry: onModelCapacityRetry, + runTurn: async (turnRequest, retry) => { + await this.refreshSkills(cwd, additionalDirectories); + if (shouldCancel?.()) { + return null; + } + return await this.codexClient.runTurn({ + threadId: turnRequest.sessionId, + input: buildPromptItems(turnRequest.prompt), + approvalPolicy: agentMode.approvalPolicy, + approvalsReviewer: agentMode.approvalsReviewer, + sandboxPolicy: addAdditionalDirectoriesToSandboxPolicy(agentMode.sandboxPolicy, additionalDirectories), + summary: disableSummary ? "none" : "auto", + effort, + model: modelId.model, + serviceTier, + }, (turnId) => onTurnStarted?.(turnId, retry)); + }, + }); } async runAgentFileChangeReport(params: { diff --git a/src/CodexAcpServer.ts b/src/CodexAcpServer.ts index 7da0b130..c1311c63 100644 --- a/src/CodexAcpServer.ts +++ b/src/CodexAcpServer.ts @@ -13,14 +13,13 @@ import {type CodexAuthRequest, getCodexAuthMethods, isCodexAuthRequest} from "./ import {clientSupportsUrlElicitation} from "./ElicitationCapabilities"; import { CodexAcpClient, - isModelCapacityError, type JsonObject, - type ModelCapacityRetry, OPENAI_PROVIDER_ID, type SessionMetadata, type SessionMetadataWithThread, type UrlElicitationRequester } from "./CodexAcpClient"; +import {type CapacityRetry, RetryCapacityService} from "./RetryCapacityService"; import {CodexAppServerClient, type McpStartupResult} from "./CodexAppServerClient"; import {type CodexConnection, startCodexConnection} from "./CodexJsonRpcConnection"; import {type AcpClientConnection, ACPSessionConnection, type UpdateSessionEvent} from "./ACPSessionConnection"; @@ -2299,7 +2298,8 @@ export class CodexAcpServer { const disposePromptRequestCancellation = this.observePromptRequestCancellation(signal, sessionState, activePrompt); let eventHandler: CodexEventHandler | null = null; let promptNotificationsActive = true; - let modelCapacityRetry: ModelCapacityRetry | null = null; + const retryCapacityService = new RetryCapacityService(); + let modelCapacityRetry: CapacityRetry | null = null; const clearRecoveredSessionFailure = async (handler: CodexEventHandler): Promise => { await handler.completeSuccessfulTurn(sessionState.currentTurnId); const current = sessionState.sessionFailure; @@ -2346,22 +2346,11 @@ export class CodexAcpServer { await promptEventHandler.handleSessionScopedNotification(event); return; } - let handledEvent = event; - if (event.method === "error" - && !event.params.willRetry - && modelCapacityRetry !== null - && event.params.turnId === sessionState.currentTurnId - && isModelCapacityError(event.params.error)) { - modelCapacityRetry.warningPublished = true; - handledEvent = { - method: "error", - params: { - ...event.params, - willRetry: true, - error: {...event.params.error, message: modelCapacityRetry.title}, - }, - }; - } + const handledEvent = retryCapacityService.transformNotification( + event, + sessionState.currentTurnId, + modelCapacityRetry, + ); const completesActiveTurn = event.method === "turn/completed" && event.params.turn.id === sessionState.currentTurnId; permissionContext.handleNotification(handledEvent); @@ -2378,24 +2367,16 @@ export class CodexAcpServer { const handleModelCapacityRetry = async ( completed: TurnCompletedNotification, - retry: ModelCapacityRetry, + retry: CapacityRetry, ) => { await this.codexAcpClient.waitForSessionNotifications(params.sessionId); await promptEventHandler.flushPendingErrors(); if (!retry.warningPublished) { - await promptEventHandler.handleNotification({ - method: "error", - params: { - threadId: params.sessionId, - turnId: completed.turn.id, - willRetry: true, - error: { - message: retry.title, - codexErrorInfo: "serverOverloaded", - additionalDetails: null, - }, - }, - }); + await promptEventHandler.handleNotification(retryCapacityService.createWarning( + params.sessionId, + completed.turn.id, + retry, + )); } await promptEventHandler.flushPendingPlanUpdates(); recoverableSessionFailure = sessionState.sessionFailure; diff --git a/src/RetryCapacityService.ts b/src/RetryCapacityService.ts new file mode 100644 index 00000000..bc70c105 --- /dev/null +++ b/src/RetryCapacityService.ts @@ -0,0 +1,147 @@ +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 CapacityRetry { + attempt: number; + delaySeconds: number; + title: string; + warningPublished: boolean; +} + +interface RetryCapacityOptions { + signal: AbortSignal | undefined; + shouldCancel: (() => boolean) | undefined; + runTurn(request: PromptRequest, retry: CapacityRetry | null): Promise; + onRetry: ((completed: TurnCompletedNotification, retry: CapacityRetry) => Promise) | undefined; +} + +class CapacityRetryError extends Error { + constructor( + readonly completed: TurnCompletedNotification, + readonly retry: CapacityRetry, + ) { + super(retry.title); + } +} + +export class RetryCapacityService { + async run( + request: PromptRequest, + options: RetryCapacityOptions, + retryAttempt = 0, + ): Promise { + if (options.signal?.aborted || options.shouldCancel?.()) { + return null; + } + const retry = this.createRetry(retryAttempt); + try { + const completed = await options.runTurn(request, retry); + if (completed === null) { + return null; + } + if (retry !== null && completed.turn.status === "failed" && isCapacityError(completed.turn.error)) { + throw new CapacityRetryError(completed, retry); + } + return completed; + } catch (error) { + if (!(error instanceof CapacityRetryError)) { + throw error; + } + await options.onRetry?.(error.completed, error.retry); + if (!await this.wait(error.retry.delaySeconds, options)) { + return null; + } + return await this.run({ + sessionId: request.sessionId, + prompt: [{type: "text", text: CONTINUATION_PROMPT}], + }, options, retryAttempt + 1); + } + } + + transformNotification( + notification: ServerNotification, + currentTurnId: string | null, + retry: CapacityRetry | 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: CapacityRetry): ServerNotification { + return { + method: "error", + params: { + threadId, + turnId, + willRetry: true, + error: { + message: retry.title, + codexErrorInfo: "serverOverloaded", + additionalDetails: null, + }, + }, + }; + } + + private createRetry(retryAttempt: number): CapacityRetry | 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, + }; + } + + private async wait(delaySeconds: number, options: RetryCapacityOptions): Promise { + if (options.signal?.aborted || options.shouldCancel?.()) { + return false; + } + return await new Promise((resolve) => { + const finish = (completed: boolean) => { + clearTimeout(timer); + options.signal?.removeEventListener("abort", onAbort); + resolve(completed); + }; + const onAbort = () => finish(false); + const timer = setTimeout(() => finish(!options.shouldCancel?.()), delaySeconds * 1000); + options.signal?.addEventListener("abort", onAbort, {once: true}); + }); + } +} + +function isCapacityError(error: {message: string; codexErrorInfo: unknown} | null): boolean { + return error?.codexErrorInfo === "serverOverloaded" && error.message.trim() === ERROR_MESSAGE; +} From c84b2369df8382244fee4022852a1b06576016fc Mon Sep 17 00:00:00 2001 From: Nikita Ashikhmin Date: Thu, 27 Aug 2026 11:47:43 +0400 Subject: [PATCH 6/6] refactor: keep turn retry orchestration generic --- src/CodexAcpClient.ts | 69 ++++++++++++-------- src/CodexAcpServer.ts | 26 ++++---- src/RetryCapacityService.ts | 121 ++++++++++++++++-------------------- 3 files changed, 108 insertions(+), 108 deletions(-) diff --git a/src/CodexAcpClient.ts b/src/CodexAcpClient.ts index 57f97d29..e4723164 100644 --- a/src/CodexAcpClient.ts +++ b/src/CodexAcpClient.ts @@ -63,7 +63,7 @@ import { createReportedAgentFileChangeReport, createUnavailableAgentFileChangeReport, } from "./AgentFileChangeReport"; -import {type CapacityRetry, RetryCapacityService} from "./RetryCapacityService"; +import {RetryCapacityService, type TurnRetry, type TurnRetryService} from "./RetryCapacityService"; /** * Well-known provider id for the client-configurable custom LLM gateway. @@ -110,7 +110,7 @@ export class CodexAcpClient { private pendingLoginCompleted: Promise | null = null; private pendingAccountUpdated: Promise | null = null; private readonly sessionNotificationQueues = new Map>(); - private readonly retryCapacityService = new RetryCapacityService(); + private readonly turnRetryService: TurnRetryService = new RetryCapacityService(); private skillExtraRoots: string[] = []; private configPath: string | null = null; @@ -850,37 +850,54 @@ export class CodexAcpClient { disableSummary: boolean, cwd: string, additionalDirectories: string[], - onTurnStarted?: (turnId: string, retry: CapacityRetry | null) => void, + onTurnStarted?: (turnId: string, retry: TurnRetry | null) => void, shouldCancel?: () => boolean, - onModelCapacityRetry?: ( + onTurnRetry?: ( completed: TurnCompletedNotification, - retry: CapacityRetry, + retry: TurnRetry, ) => Promise, retrySignal?: AbortSignal, + retryAttempt = 0, ): Promise { + const input = buildPromptItems(request.prompt); const effort = modelId.effort as ReasoningEffort | null; //TODO remove unsafe conversion - return await this.retryCapacityService.run(request, { - signal: retrySignal, + await this.refreshSkills(cwd, additionalDirectories); + if (shouldCancel?.()) { + return null; + } + const retry = this.turnRetryService.createRetry(retryAttempt); + const completed = await this.codexClient.runTurn({ + threadId: request.sessionId, + input, + approvalPolicy: agentMode.approvalPolicy, + approvalsReviewer: agentMode.approvalsReviewer, + sandboxPolicy: addAdditionalDirectoriesToSandboxPolicy(agentMode.sandboxPolicy, additionalDirectories), + summary: disableSummary ? "none" : "auto", + effort, + model: modelId.model, + 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, - onRetry: onModelCapacityRetry, - runTurn: async (turnRequest, retry) => { - await this.refreshSkills(cwd, additionalDirectories); - if (shouldCancel?.()) { - return null; - } - return await this.codexClient.runTurn({ - threadId: turnRequest.sessionId, - input: buildPromptItems(turnRequest.prompt), - approvalPolicy: agentMode.approvalPolicy, - approvalsReviewer: agentMode.approvalsReviewer, - sandboxPolicy: addAdditionalDirectoriesToSandboxPolicy(agentMode.sandboxPolicy, additionalDirectories), - summary: disableSummary ? "none" : "auto", - effort, - model: modelId.model, - serviceTier, - }, (turnId) => onTurnStarted?.(turnId, retry)); - }, - }); + onTurnRetry, + retrySignal, + retryAttempt + 1, + ); } async runAgentFileChangeReport(params: { diff --git a/src/CodexAcpServer.ts b/src/CodexAcpServer.ts index c1311c63..a8aa3397 100644 --- a/src/CodexAcpServer.ts +++ b/src/CodexAcpServer.ts @@ -19,7 +19,7 @@ import { type SessionMetadataWithThread, type UrlElicitationRequester } from "./CodexAcpClient"; -import {type CapacityRetry, RetryCapacityService} from "./RetryCapacityService"; +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"; @@ -2298,8 +2298,8 @@ export class CodexAcpServer { const disposePromptRequestCancellation = this.observePromptRequestCancellation(signal, sessionState, activePrompt); let eventHandler: CodexEventHandler | null = null; let promptNotificationsActive = true; - const retryCapacityService = new RetryCapacityService(); - let modelCapacityRetry: CapacityRetry | null = null; + 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; @@ -2346,10 +2346,10 @@ export class CodexAcpServer { await promptEventHandler.handleSessionScopedNotification(event); return; } - const handledEvent = retryCapacityService.transformNotification( + const handledEvent = turnRetryService.transformNotification( event, sessionState.currentTurnId, - modelCapacityRetry, + turnRetry, ); const completesActiveTurn = event.method === "turn/completed" && event.params.turn.id === sessionState.currentTurnId; @@ -2365,14 +2365,14 @@ export class CodexAcpServer { approvalHandler, elicitationHandler); - const handleModelCapacityRetry = async ( + const handleTurnRetry = async ( completed: TurnCompletedNotification, - retry: CapacityRetry, + retry: TurnRetry, ) => { await this.codexAcpClient.waitForSessionNotifications(params.sessionId); await promptEventHandler.flushPendingErrors(); if (!retry.warningPublished) { - await promptEventHandler.handleNotification(retryCapacityService.createWarning( + await promptEventHandler.handleNotification(turnRetryService.createWarning( params.sessionId, completed.turn.id, retry, @@ -2384,7 +2384,7 @@ export class CodexAcpServer { sessionState.currentTurnId = null; pendingTurnStart = this.createPendingTurnStart(); this.pendingTurnStarts.set(params.sessionId, pendingTurnStart); - logger.log("Selected model is at capacity; scheduling retry", { + logger.log("Scheduling turn retry", { sessionId: params.sessionId, attempt: retry.attempt, delaySeconds: retry.delaySeconds, @@ -2520,14 +2520,14 @@ export class CodexAcpServer { return; } sessionState.currentTurnId = turnId; - modelCapacityRetry = retry; + turnRetry = retry; pendingTurnStart?.resolve(turnId); if (retry?.attempt === 1) { onTurnStarted?.(); } }, () => this.promptShouldStop(params.sessionId, activePrompt), - handleModelCapacityRetry, + handleTurnRetry, activePrompt.signal, )); void sendPromptPromise.catch((err) => { @@ -2615,14 +2615,14 @@ export class CodexAcpServer { return; } sessionState.currentTurnId = turnId; - modelCapacityRetry = retry; + 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), - handleModelCapacityRetry, + handleTurnRetry, activePrompt.signal, ), ); diff --git a/src/RetryCapacityService.ts b/src/RetryCapacityService.ts index bc70c105..5d738959 100644 --- a/src/RetryCapacityService.ts +++ b/src/RetryCapacityService.ts @@ -12,67 +12,63 @@ const RETRY_WINDOWS_SECONDS = [ [120, 300], ] as const; -export interface CapacityRetry { +export interface TurnRetry { attempt: number; delaySeconds: number; title: string; warningPublished: boolean; } -interface RetryCapacityOptions { - signal: AbortSignal | undefined; - shouldCancel: (() => boolean) | undefined; - runTurn(request: PromptRequest, retry: CapacityRetry | null): Promise; - onRetry: ((completed: TurnCompletedNotification, retry: CapacityRetry) => Promise) | undefined; -} - -class CapacityRetryError extends Error { - constructor( - readonly completed: TurnCompletedNotification, - readonly retry: CapacityRetry, - ) { - super(retry.title); - } +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 { - async run( - request: PromptRequest, - options: RetryCapacityOptions, - retryAttempt = 0, - ): Promise { - if (options.signal?.aborted || options.shouldCancel?.()) { +export class RetryCapacityService implements TurnRetryService { + createRetry(retryAttempt: number): TurnRetry | null { + const window = RETRY_WINDOWS_SECONDS[retryAttempt]; + if (window === undefined) { return null; } - const retry = this.createRetry(retryAttempt); - try { - const completed = await options.runTurn(request, retry); - if (completed === null) { - return null; - } - if (retry !== null && completed.turn.status === "failed" && isCapacityError(completed.turn.error)) { - throw new CapacityRetryError(completed, retry); - } - return completed; - } catch (error) { - if (!(error instanceof CapacityRetryError)) { - throw error; - } - await options.onRetry?.(error.completed, error.retry); - if (!await this.wait(error.retry.delaySeconds, options)) { - return null; - } - return await this.run({ - sessionId: request.sessionId, - prompt: [{type: "text", text: CONTINUATION_PROMPT}], - }, options, retryAttempt + 1); - } + 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: CapacityRetry | null, + retry: TurnRetry | null, ): ServerNotification { if (notification.method !== "error" || notification.params.willRetry @@ -92,7 +88,7 @@ export class RetryCapacityService { }; } - createWarning(threadId: string, turnId: string, retry: CapacityRetry): ServerNotification { + createWarning(threadId: string, turnId: string, retry: TurnRetry): ServerNotification { return { method: "error", params: { @@ -108,36 +104,23 @@ export class RetryCapacityService { }; } - private createRetry(retryAttempt: number): CapacityRetry | 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, - }; - } - - private async wait(delaySeconds: number, options: RetryCapacityOptions): Promise { - if (options.signal?.aborted || options.shouldCancel?.()) { + 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); - options.signal?.removeEventListener("abort", onAbort); + signal?.removeEventListener("abort", onAbort); resolve(completed); }; const onAbort = () => finish(false); - const timer = setTimeout(() => finish(!options.shouldCancel?.()), delaySeconds * 1000); - options.signal?.addEventListener("abort", onAbort, {once: true}); + const timer = setTimeout(() => finish(!shouldCancel?.()), retry.delaySeconds * 1000); + signal?.addEventListener("abort", onAbort, {once: true}); }); } }