packages/ai/src/utils/retry.ts

// Adapted from Pi cd34e17ff039502f3664d8363b3c0a23f93a2ca3.
// Copyright (c) 2025 Mario Zechner. MIT: licenses/third-party-notices.txt.
import type { AssistantMessage, ModelRetryStatus } from "../types.ts";

function buildProviderErrorPattern(patterns: readonly string[]): RegExp {
  return new RegExp(patterns.join("|"), "i");
}

const NON_RETRYABLE_PROVIDER_LIMIT_ERROR_PATTERN = buildProviderErrorPattern([
  // OpenCode Go/free-tier limits returned as 429 JSON error types by OpenCode's
  // Zen API. These are subscription/account limits, not transient throttles.
  "GoUsageLimitError",
  "FreeUsageLimitError",

  // OpenCode Go subscription-limit text asks users to enable available-balance
  // usage after rolling/weekly/monthly limits are reached.
  "Monthly usage limit reached",
  "available balance",

  // Generic quota/budget/billing exhaustion. `insufficient_quota` is OpenAI's
  // quota/billing error code; the other strings cover common gateway wording.
  "insufficient_quota",
  "out of budget",
  "quota exceeded",
  "billing",
]);

const RETRYABLE_PROVIDER_ERROR_PATTERN = buildProviderErrorPattern([
  // Generic provider load, HTTP status, and server-side transient failures.
  "overloaded",
  "currently experiencing high demand",
  "rate.?limit",
  "too many requests",
  "429",
  "500",
  "502",
  "503",
  "504",
  "520",
  "524",
  "service.?unavailable",
  "server.?error",
  "internal.?error",

  // Wrapper/provider text for transient upstream failures, including OpenRouter
  // "Provider returned error" responses (#2264).
  "provider.?returned.?error",
  "exceeded request buffer limit while retrying upstream",

  // Network, proxy, and fetch transport failures. This includes OpenAI Codex
  // raw-fetch failures such as "upstream connect", "connection refused", and
  // "reset before headers" (#733), plus OpenRouter connection drops (#3317).
  "network.?error",
  "connection.?error",
  "connection.?refused",
  "connection.?lost",
  "other side closed",
  "fetch failed",
  "getaddrinfo",
  "ENOTFOUND",
  "EAI_AGAIN",
  "upstream.?connect",
  "reset before headers",
  "socket hang up",
  "socket connection was closed",
  "timed? out",
  "timeout",
  "model stream stalled",
  "terminated",

  // Chromium net-stack transport failures (Electron net.fetch), which surface as
  // net::ERR_* strings instead of Node/undici text.
  "ERR_CONNECTION_(CLOSED|RESET|ABORTED|TIMED_?OUT)",
  "ERR_(TIMED_OUT|NETWORK_CHANGED|NAME_NOT_RESOLVED|(HTTP2|QUIC)_PROTOCOL_ERROR|EMPTY_RESPONSE|PROXY_CONNECTION_FAILED)",

  // WebSocket transports can report close/error text instead of HTTP/fetch text.
  "websocket.?closed",
  "websocket.?error",

  // Premature stream endings from SDKs and transports. Anthropic can throw
  // "stream ended without ..." and "Anthropic stream ended before message_stop"
  // (#4433); Bedrock/Smithy can throw an HTTP/2 no-response error (#3594).
  "ended without",
  "stream ended before message_stop",
  "stream ended before a terminal response event",
  "http2 request did not get a response",

  // Provider-requested retry delay cap failures should flow through the outer
  // retry policy so callers can surface/abort the backoff (#1123).
  "retry delay",

  // Explicit retry guidance emitted mid-stream by OpenAI Responses and Bedrock
  // stream exceptions (#6019).
  "you can retry your request",
  "try your request again",
  "please retry your request",

  // gRPC based providers (e.g. NVIDIA NIM)
  "ResourceExhausted",
]);

/**
 * Retry policy: bounded attempts with exponential backoff (`baseDelayMs * 2^(attempt-1)`).
 * `maxAgentDelayMs` caps each computed delay and defaults to 60 seconds.
 * Matches `settings.retry` (`enabled`, `maxRetries`, `baseDelayMs`, `maxAgentDelayMs`) in coding-agent; kept
 * here so the classifier and the policy-driven retry loop live together and stay reusable
 * by the SDK and other callers.
 */
export interface RetryPolicy {
  enabled: boolean;
  /** Max retry attempts (0 = no retries). The initial call never counts as a retry. */
  maxRetries: number;
  /** Base delay in ms. Per-attempt delay is `baseDelayMs * 2^(attempt-1)`. */
  baseDelayMs: number;
  /** Optional cap for agent-level retry delays in ms. Defaults to 60 seconds. */
  maxAgentDelayMs?: number;
  /** Recovery time allowed after the first failure; <= 0 disables the cap. */
  maxElapsedMs?: number;
}

export const DEFAULT_MAX_AGENT_RETRY_DELAY_MS = 60_000;

/** Keep failure-specific retry limits shared by session and summary recovery. */
export function retryAttemptLimit(message: AssistantMessage, configuredLimit: number): number {
  return message.errorDetails?.code === "EMPTY_MODEL_RESPONSE"
    ? Math.min(1, configuredLimit)
    : configuredLimit;
}

export function retryDelayMs(
  policy: Pick<RetryPolicy, "baseDelayMs" | "maxAgentDelayMs">,
  attempt: number,
  retryAfterMs?: number,
): number {
  const delay = policy.baseDelayMs * 2 ** Math.max(0, attempt - 1);
  const safeDelay = Number.isSafeInteger(delay) ? delay : Number.MAX_SAFE_INTEGER;
  const local = Math.min(safeDelay, policy.maxAgentDelayMs ?? DEFAULT_MAX_AGENT_RETRY_DELAY_MS);
  const server = retryAfterMs !== undefined && Number.isFinite(retryAfterMs) ? Math.max(0, retryAfterMs) : 0;
  // Never let setTimeout overflow into an immediate retry.
  const wait = Math.max(local, server);
  return wait > 2147483647 ? Infinity : wait;
}

/** Optional callbacks emitted by {@link retryAssistantCall} around each retry. */
export interface RetryCallbacks {
  /** Emitted before the backoff sleep of each retry attempt (1-indexed). */
  onRetryScheduled?: (
    attempt: number,
    maxAttempts: number,
    delayMs: number,
    errorMessage: string,
    origin?: ModelRetryStatus["origin"],
  ) => void | Promise<void>;
  /** Emitted after the backoff sleep, immediately before the retried call starts. */
  onRetryAttemptStart?: () => void | Promise<void>;
  /** Optional interruptible backoff. Return true when an explicit user retry resets the recovery budget. */
  waitBeforeRetry?: (delayMs: number, signal?: AbortSignal) => Promise<boolean>;
  /** Emitted once when the loop ends: success if a later call completed normally. */
  onRetryFinished?: (success: boolean, attempt: number, finalError?: string) => void | Promise<void>;
}

class RetrySleepAbortError extends Error {
  constructor() {
    super("Aborted");
  }
}

function sleep(ms: number, signal?: AbortSignal): Promise<void> {
  return new Promise((resolve, reject) => {
    if (signal?.aborted) {
      reject(new RetrySleepAbortError());
      return;
    }
    const abort = () => {
      clearTimeout(timeout);
      reject(new RetrySleepAbortError());
    };
    const timeout = setTimeout(() => {
      signal?.removeEventListener("abort", abort);
      resolve();
    }, ms);
    signal?.addEventListener("abort", abort, { once: true });
  });
}

/**
 * Run a single assistant-producing call with bounded retry on transient errors.
 *
 * Behavior:
 * - A successful response is returned immediately. Aborts are terminal and never
 *   retried, but reported as unsuccessful if they happen after a retry was scheduled.
 *   Aborts during the backoff sleep are normalized to an aborted `AssistantMessage`
 *   too, so callers do not need to care when cancellation happened.
 * - A non-retryable error (per {@link isRetryableAssistantError}, including quota/
 *   billing exhaustion) is returned immediately so deterministic errors fail fast.
 * - Otherwise retries up to `maxRetries` times with exponential backoff, emitting
 *   `onRetryScheduled` before each sleep, `onRetryAttemptStart` after each sleep before
 *   the retried call starts, and `onRetryFinished` once at the end (whether the loop
 *   ends in success, exhausted retries, or an aborted backoff).
 *
 * When `policy` is undefined or disabled, only an explicit manual retry can repeat a request.
 */
export async function retryAssistantCall(
  produce: (deadlineMs?: number) => Promise<AssistantMessage>,
  policy: RetryPolicy | undefined,
  signal: AbortSignal | undefined,
  callbacks?: RetryCallbacks,
): Promise<AssistantMessage> {
  const maxAttempts = policy?.enabled ? policy.maxRetries : 0;
  const recoveryLimit = policy?.maxElapsedMs;
  let deadlineMs: number | undefined;

  let attempt = 0;
  let lastRetry: { attempt: number; errorMessage: string } | undefined;
  for (;;) {
    const response = await produce(attempt > 0 ? deadlineMs : undefined);

    // Abort: terminal but not successful. Never retry an aborted message.
    if (response.stopReason === "aborted") {
      if (lastRetry) await callbacks?.onRetryFinished?.(false, lastRetry.attempt);
      return response;
    }

    // Success: non-error, non-abort responses return as-is.
    if (response.stopReason !== "error") {
      if (lastRetry) await callbacks?.onRetryFinished?.(true, lastRetry.attempt);
      return response;
    }

    const manualRetry = response.errorDetails?.code === "MODEL_REQUEST_RETRY";
    if (manualRetry) {
      attempt = 0;
      deadlineMs = recoveryLimit && recoveryLimit > 0 ? Date.now() + recoveryLimit : undefined;
    }
    if (deadlineMs === undefined && recoveryLimit && recoveryLimit > 0)
      deadlineMs = Date.now() + recoveryLimit;
    const delayMs =
      manualRetry || !policy ? 0 : retryDelayMs(policy, attempt + 1, response.errorDetails?.retryAfterMs);
    // An explicit retry remains available when automatic retry is disabled or exhausted.
    if (
      !manualRetry &&
      (!policy ||
        attempt >= retryAttemptLimit(response, maxAttempts) ||
        !isRetryableAssistantError(response) ||
        !Number.isFinite(delayMs) ||
        (deadlineMs !== undefined && Date.now() + delayMs >= deadlineMs))
    ) {
      if (lastRetry) await callbacks?.onRetryFinished?.(false, lastRetry.attempt, response.errorMessage);
      return response;
    }

    attempt++;
    lastRetry = { attempt, errorMessage: response.errorMessage || "Unknown error" };
    await callbacks?.onRetryScheduled?.(
      attempt,
      maxAttempts,
      delayMs,
      lastRetry.errorMessage,
      manualRetry ? "manual" : "auto",
    );

    // Normalize aborts during retry backoff to the same AssistantMessage shape as
    // provider stream aborts, so callers do not need to care when cancellation happened.
    try {
      if (callbacks?.waitBeforeRetry) {
        if (await callbacks.waitBeforeRetry(delayMs, signal)) {
          attempt = 1;
          deadlineMs = recoveryLimit && recoveryLimit > 0 ? Date.now() + recoveryLimit : undefined;
        }
      } else await sleep(delayMs, signal);
    } catch (error) {
      await callbacks?.onRetryFinished?.(false, attempt, lastRetry.errorMessage);
      if (signal?.aborted || error instanceof RetrySleepAbortError) {
        const { errorMessage: _errorMessage, ...rest } = response;
        return { ...rest, stopReason: "aborted" };
      }
      throw error;
    }
    await callbacks?.onRetryAttemptStart?.();
  }
}

/**
 * Classifies whether a failed assistant message looks like a transient provider
 * or transport error, so callers can decide if the last assistant turn should be
 * restarted.
 *
 * This does not implement retry policy. Callers should first handle context
 * overflow separately, then apply their own retry budget, backoff, and reporting
 * before restarting the assistant turn.
 */
export function isRetryableAssistantError(message: AssistantMessage): boolean {
  if (message.stopReason !== "error" || !message.errorMessage) return false;
  if (message.errorDetails?.code === "MODEL_REQUEST_RETRY") return true;
  if (message.errorDetails?.code === "MODEL_RECOVERY_DEADLINE") return false;
  if (message.errorDetails?.retryable === false) return false;
  const errorMessage = message.errorMessage;
  if (NON_RETRYABLE_PROVIDER_LIMIT_ERROR_PATTERN.test(errorMessage)) return false;
  return message.errorDetails?.retryable ?? RETRYABLE_PROVIDER_ERROR_PATTERN.test(errorMessage);
}