feat: retry model capacity failures

src/CodexAcpClient.ts

@@ -63,6 +63,7 @@ import {
     createReportedAgentFileChangeReport,
     createUnavailableAgentFileChangeReport,
 } from "./AgentFileChangeReport";
+import {RetryCapacityService, type TurnRetry, type TurnRetryService} from "./RetryCapacityService";
 
 /**
  * Well-known provider id for the client-configurable custom LLM gateway.
@@ -109,6 +110,7 @@ export class CodexAcpClient {
     private pendingLoginCompleted: Promise<AccountLoginCompletedNotification> | null = null;
     private pendingAccountUpdated: Promise<AccountUpdatedNotification> | null = null;
     private readonly sessionNotificationQueues = new Map<string, Promise<void>>();
+    private readonly turnRetryService: TurnRetryService = new RetryCapacityService();
     private skillExtraRoots: string[] = [];
     private configPath: string | null = null;
 
@@ -848,26 +850,54 @@ export class CodexAcpClient {
         disableSummary: boolean,
         cwd: string,
         additionalDirectories: string[],
-        onTurnStarted?: (turnId: string) => void,
+        onTurnStarted?: (turnId: string, retry: TurnRetry | null) => void,
         shouldCancel?: () => boolean,
+        onTurnRetry?: (
+            completed: TurnCompletedNotification,
+            retry: TurnRetry,
+        ) => Promise<void>,
+        retrySignal?: AbortSignal,
+        retryAttempt = 0,
     ): Promise<TurnCompletedNotification | null> {
         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;
         }
-        return await this.codexClient.runTurn({
+        const retry = this.turnRetryService.createRetry(retryAttempt);
+        const completed = await this.codexClient.runTurn({
             threadId: request.sessionId,
-            input: input,
+            input,
             approvalPolicy: agentMode.approvalPolicy,
             approvalsReviewer: agentMode.approvalsReviewer,
             sandboxPolicy: addAdditionalDirectoriesToSandboxPolicy(agentMode.sandboxPolicy, additionalDirectories),
             summary: disableSummary ? "none" : "auto",
-            effort: effort,
+            effort,
             model: modelId.model,
-            serviceTier: serviceTier,
-        }, onTurnStarted);
+            serviceTier,
+        }, (turnId) => onTurnStarted?.(turnId, retry));
+        if (retry === null || !this.turnRetryService.shouldRetry(completed)) {
+            return completed;
+        }
+        await onTurnRetry?.(completed, retry);
+        if (!await this.turnRetryService.wait(retry, retrySignal, shouldCancel)) {
+            return null;
+        }
+        return await this.sendPrompt(
+            this.turnRetryService.createContinuationRequest(request),
+            agentMode,
+            modelId,
+            serviceTier,
+            disableSummary,
+            cwd,
+            additionalDirectories,
+            onTurnStarted,
+            shouldCancel,
+            onTurnRetry,
+            retrySignal,
+            retryAttempt + 1,
+        );
     }
 
     async runAgentFileChangeReport(params: {

src/CodexAcpServer.ts

@@ -19,11 +19,21 @@ import {
     type SessionMetadataWithThread,
     type UrlElicitationRequester
 } from "./CodexAcpClient";
+import {RetryCapacityService, type TurnRetry, type TurnRetryService} from "./RetryCapacityService";
 import {CodexAppServerClient, type McpStartupResult} from "./CodexAppServerClient";
 import {type CodexConnection, startCodexConnection} from "./CodexJsonRpcConnection";
 import {type AcpClientConnection, ACPSessionConnection, type UpdateSessionEvent} from "./ACPSessionConnection";
 import type {InputModality, ReasoningEffort} from "./app-server";
-import type {Account, Model, ReasoningEffortOption, Thread, ThreadGoal, ThreadItem, UserInput} from "./app-server/v2";
+import type {
+    Account,
+    Model,
+    ReasoningEffortOption,
+    Thread,
+    ThreadGoal,
+    ThreadItem,
+    TurnCompletedNotification,
+    UserInput,
+} from "./app-server/v2";
 import type {RateLimitsMap} from "./RateLimitsMap";
 import {ModelId} from "./ModelId";
 import {AgentMode, MODE_CONFIG_ID} from "./AgentMode";
@@ -2288,6 +2298,8 @@ export class CodexAcpServer {
         const disposePromptRequestCancellation = this.observePromptRequestCancellation(signal, sessionState, activePrompt);
         let eventHandler: CodexEventHandler | null = null;
         let promptNotificationsActive = true;
+        const turnRetryService: TurnRetryService = new RetryCapacityService();
+        let turnRetry: TurnRetry | null = null;
         const clearRecoveredSessionFailure = async (handler: CodexEventHandler): Promise<void> => {
             await handler.completeSuccessfulTurn(sessionState.currentTurnId);
             const current = sessionState.sessionFailure;
@@ -2334,11 +2346,16 @@ export class CodexAcpServer {
                         await promptEventHandler.handleSessionScopedNotification(event);
                         return;
                     }
+                    const handledEvent = turnRetryService.transformNotification(
+                        event,
+                        sessionState.currentTurnId,
+                        turnRetry,
+                    );
                     const completesActiveTurn = event.method === "turn/completed"
                         && event.params.turn.id === sessionState.currentTurnId;
-                    permissionContext.handleNotification(event);
-                    await elicitationHandler.handleNotification(event);
-                    await promptEventHandler.handleNotification(event);
+                    permissionContext.handleNotification(handledEvent);
+                    await elicitationHandler.handleNotification(handledEvent);
+                    await promptEventHandler.handleNotification(handledEvent);
                     if (completesActiveTurn) {
                         // The prompt may remain open for plan approval after its turn has ended. Switch at
                         // the causal boundary so a queued late error cannot enter the completed turn's buffer.
@@ -2348,6 +2365,32 @@ export class CodexAcpServer {
                 approvalHandler,
                 elicitationHandler);
 
+            const handleTurnRetry = async (
+                completed: TurnCompletedNotification,
+                retry: TurnRetry,
+            ) => {
+                await this.codexAcpClient.waitForSessionNotifications(params.sessionId);
+                await promptEventHandler.flushPendingErrors();
+                if (!retry.warningPublished) {
+                    await promptEventHandler.handleNotification(turnRetryService.createWarning(
+                        params.sessionId,
+                        completed.turn.id,
+                        retry,
+                    ));
+                }
+                await promptEventHandler.flushPendingPlanUpdates();
+                recoverableSessionFailure = sessionState.sessionFailure;
+                activePrompt.currentTurn = null;
+                sessionState.currentTurnId = null;
+                pendingTurnStart = this.createPendingTurnStart();
+                this.pendingTurnStarts.set(params.sessionId, pendingTurnStart);
+                logger.log("Scheduling turn retry", {
+                    sessionId: params.sessionId,
+                    attempt: retry.attempt,
+                    delaySeconds: retry.delaySeconds,
+                });
+            };
+
             if (activePrompt.signal.aborted) {
                 return cancelledPromptResponse();
             }
@@ -2469,18 +2512,23 @@ export class CodexAcpServer {
                     disableSummary,
                     sessionState.cwd,
                     sessionState.additionalDirectories,
-                    (turnId) => {
+                    (turnId, retry) => {
                         const turn = {threadId: params.sessionId, turnId};
                         activePrompt.currentTurn = turn;
                         if (this.promptShouldStop(params.sessionId, activePrompt)) {
                             this.interruptLateStartedTurn(turn);
                             return;
                         }
                         sessionState.currentTurnId = turnId;
+                        turnRetry = retry;
                         pendingTurnStart?.resolve(turnId);
-                        onTurnStarted?.();
+                        if (retry?.attempt === 1) {
+                            onTurnStarted?.();
+                        }
                     },
                     () => this.promptShouldStop(params.sessionId, activePrompt),
+                    handleTurnRetry,
+                    activePrompt.signal,
                 ));
             void sendPromptPromise.catch((err) => {
                 if (this.activePrompts.get(params.sessionId) !== activePrompt) {
@@ -2559,20 +2607,23 @@ export class CodexAcpServer {
                             disableSummary,
                             sessionState.cwd,
                             sessionState.additionalDirectories,
-                            (turnId) => {
+                            (turnId, retry) => {
                                 const turn = {threadId: params.sessionId, turnId};
                                 activePrompt.currentTurn = turn;
                                 if (this.promptShouldStop(params.sessionId, activePrompt)) {
                                     this.interruptLateStartedTurn(turn);
                                     return;
                                 }
                                 sessionState.currentTurnId = turnId;
+                                turnRetry = retry;
                                 // Keep the approval-to-turn-start gap session-scoped. Once the new turn has
                                 // an identity, snapshot any unchanged session failure as its recovery baseline.
                                 recoverableSessionFailure = sessionState.sessionFailure;
                                 promptNotificationsActive = true;
                             },
                             () => this.promptShouldStop(params.sessionId, activePrompt),
+                            handleTurnRetry,
+                            activePrompt.signal,
                         ),
                     );
                     void implementationPromise.catch((err) => {

src/RetryCapacityService.ts

@@ -0,0 +1,130 @@
+import type {PromptRequest} from "@agentclientprotocol/sdk";
+import type {ServerNotification} from "./app-server";
+import type {TurnCompletedNotification} from "./app-server/v2";
+
+const ERROR_MESSAGE = "Selected model is at capacity. Please try a different model.";
+const CONTINUATION_PROMPT = "Continue from where you left off.";
+const RETRY_WINDOWS_SECONDS = [
+    [1, 10],
+    [1, 30],
+    [30, 60],
+    [60, 120],
+    [120, 300],
+] as const;
+
+export interface TurnRetry {
+    attempt: number;
+    delaySeconds: number;
+    title: string;
+    warningPublished: boolean;
+}
+
+export interface TurnRetryService {
+    createRetry(retryAttempt: number): TurnRetry | null;
+    shouldRetry(completed: TurnCompletedNotification): boolean;
+    wait(
+        retry: TurnRetry,
+        signal: AbortSignal | undefined,
+        shouldCancel: (() => boolean) | undefined,
+    ): Promise<boolean>;
+    createContinuationRequest(request: PromptRequest): PromptRequest;
+    transformNotification(
+        notification: ServerNotification,
+        currentTurnId: string | null,
+        retry: TurnRetry | null,
+    ): ServerNotification;
+    createWarning(threadId: string, turnId: string, retry: TurnRetry): ServerNotification;
+}
+
+export class RetryCapacityService implements TurnRetryService {
+    createRetry(retryAttempt: number): TurnRetry | null {
+        const window = RETRY_WINDOWS_SECONDS[retryAttempt];
+        if (window === undefined) {
+            return null;
+        }
+        const [minimum, maximum] = window;
+        const delaySeconds = minimum + Math.floor(Math.random() * (maximum - minimum + 1));
+        const unit = delaySeconds === 1 ? "second" : "seconds";
+        return {
+            attempt: retryAttempt + 1,
+            delaySeconds,
+            title: `Selected model is at capacity. Retrying in ${delaySeconds} ${unit} `
+                + `(${retryAttempt + 1}/${RETRY_WINDOWS_SECONDS.length}).`,
+            warningPublished: false,
+        };
+    }
+
+    shouldRetry(completed: TurnCompletedNotification): boolean {
+        return completed.turn.status === "failed" && isCapacityError(completed.turn.error);
+    }
+
+    createContinuationRequest(request: PromptRequest): PromptRequest {
+        return {
+            sessionId: request.sessionId,
+            prompt: [{type: "text", text: CONTINUATION_PROMPT}],
+        };
+    }
+
+    transformNotification(
+        notification: ServerNotification,
+        currentTurnId: string | null,
+        retry: TurnRetry | null,
+    ): ServerNotification {
+        if (notification.method !== "error"
+            || notification.params.willRetry
+            || retry === null
+            || notification.params.turnId !== currentTurnId
+            || !isCapacityError(notification.params.error)) {
+            return notification;
+        }
+        retry.warningPublished = true;
+        return {
+            method: "error",
+            params: {
+                ...notification.params,
+                willRetry: true,
+                error: {...notification.params.error, message: retry.title},
+            },
+        };
+    }
+
+    createWarning(threadId: string, turnId: string, retry: TurnRetry): ServerNotification {
+        return {
+            method: "error",
+            params: {
+                threadId,
+                turnId,
+                willRetry: true,
+                error: {
+                    message: retry.title,
+                    codexErrorInfo: "serverOverloaded",
+                    additionalDetails: null,
+                },
+            },
+        };
+    }
+
+    async wait(
+        retry: TurnRetry,
+        signal: AbortSignal | undefined,
+        shouldCancel: (() => boolean) | undefined,
+    ): Promise<boolean> {
+        if (signal?.aborted || shouldCancel?.()) {
+            return false;
+        }
+        return await new Promise<boolean>((resolve) => {
+            const finish = (completed: boolean) => {
+                clearTimeout(timer);
+                signal?.removeEventListener("abort", onAbort);
+                resolve(completed);
+            };
+            const onAbort = () => finish(false);
+            const timer = setTimeout(() => finish(!shouldCancel?.()), retry.delaySeconds * 1000);
+            signal?.addEventListener("abort", onAbort, {once: true});
+        });
+    }
+}
+
+function isCapacityError(error: {message: string; codexErrorInfo: unknown} | null): boolean {
+    return error?.codexErrorInfo === "serverOverloaded" && error.message.trim() === ERROR_MESSAGE;
+}