diff --git a/src/adapters/base.ts b/src/adapters/base.ts index b743136788..2e376a628a 100644 --- a/src/adapters/base.ts +++ b/src/adapters/base.ts @@ -1,5 +1,6 @@ import type { AdapterEvent, OcxParsedRequest } from "../types"; import type { TranslatorBudget } from "../lib/translator-budget"; +import type { RequestExecutionBudget } from "../lib/request-execution-budget"; import type { AdapterTierMetadata } from "../providers/fastwire"; /** Metadata about the caller's incoming request, for auth-forwarding adapters. */ @@ -139,6 +140,13 @@ export interface AdapterFetchContext { stream?: boolean; /** Custom fetch executor to use for physical upstream network requests (defaults to globalThis.fetch). */ executor?: typeof globalThis.fetch; + /** + * The logical request's send budget (#4546). Optional and unlimited when absent, so an + * adapter unit test that calls a transport context-free keeps its own retry shape. An + * adapter that retries internally must admit EVERY physical send against it: counting one + * adapter entry as one send is how a nested 3x3 ladder stayed invisible to a request cap. + */ + sendBudget?: RequestExecutionBudget; } /** diff --git a/src/adapters/kiro-retry.ts b/src/adapters/kiro-retry.ts index b894338f8e..08ddbbb4d9 100644 --- a/src/adapters/kiro-retry.ts +++ b/src/adapters/kiro-retry.ts @@ -5,6 +5,7 @@ import { readBoundedResponseBody } from "../lib/bounded-body"; import { resolveClientRetryAfter } from "../lib/retry-after"; import { parseRetryAfterMs } from "../combos"; import { + SendBudgetExhaustedError, abortError, cancelResponseBodyBestEffort, fetchWithAttemptDeadline, @@ -162,6 +163,13 @@ async function fetchWithResetRecovery( let lastError: unknown; for (let attempt = 0; attempt < RESET_ATTEMPTS; attempt++) { if (ctx.abortSignal?.aborted) throw abortError(ctx.abortSignal); + // Every physical send is admitted, not just the adapter entry. Kiro nests a throttle loop + // over this ladder and can run the ladder twice per throttle round, so counting one entry + // as one send hid up to eighteen upstream requests from the per-request cap (#4546). + const decision = ctx.sendBudget?.reserveDispatch({ sendClass: "transient", targetKey: url }); + if (decision && (!decision.allowed || !decision.permit.use())) { + throw new SendBudgetExhaustedError(url); + } try { const headers = new Headers(request.headers); const recovered = attempt > 0; diff --git a/src/adapters/kiro/adapter.ts b/src/adapters/kiro/adapter.ts index 8fdaf34c26..1b3a90e80f 100644 --- a/src/adapters/kiro/adapter.ts +++ b/src/adapters/kiro/adapter.ts @@ -13,6 +13,7 @@ import type { } from "../../types"; import type { ProviderAdapter } from "../base"; import type { AdapterFetchContext, AdapterRequest } from "../base"; +import type { RequestExecutionBudget } from "../../lib/request-execution-budget"; import { safeKiroHttpErrorMessage } from "../kiro-errors"; import { calibrateKiroEstimate } from "../kiro-calibration"; import { normalizeKiroImages } from "../kiro-images"; @@ -58,6 +59,9 @@ export function createKiroAdapter(provider: OcxProviderConfig): ProviderAdapter let requestSnapshot: OcxParsedRequest | undefined; let firstRequestBodyBytes = 0; let requestAbortSignal: AbortSignal | undefined; + // Captured the same way as the abort signal, because the text-fallback rebuild below runs + // outside the fetchResponse frame and used to construct a context without either (#4546). + let requestSendBudget: RequestExecutionBudget | undefined; const build = async ( parsed: OcxParsedRequest, @@ -208,6 +212,9 @@ export function createKiroAdapter(provider: OcxProviderConfig): ProviderAdapter abortSignal: requestAbortSignal, returnRawErrors: true, stream: true, + // The text-fallback rebuild used to construct a fresh context and drop the budget, + // so everything after the first send escaped the per-request cap. + ...(requestSendBudget ? { sendBudget: requestSendBudget } : {}), }); return { response, @@ -278,6 +285,7 @@ export function createKiroAdapter(provider: OcxProviderConfig): ProviderAdapter // Keep it for the adapter-owned bounded continuation so cancelling the client turn aborts // both the first Kiro request and its one allowed completion retry. if (ctx?.abortSignal) requestAbortSignal = ctx.abortSignal; + if (ctx?.sendBudget) requestSendBudget = ctx.sendBudget; return fetchKiroWithRetry(request, ctx); }, diff --git a/src/server/responses/compact.ts b/src/server/responses/compact.ts index 28fd4a564f..02060fb000 100644 --- a/src/server/responses/compact.ts +++ b/src/server/responses/compact.ts @@ -76,8 +76,14 @@ import { fetchWithResetRetry, fetchWithTransientRetry, applyUpstreamRecoveryInit, + SendBudgetExhaustedError, + TRANSIENT_RETRY_MAX_ATTEMPTS, type UpstreamSendRecovery, } from "../../lib/upstream-retry"; +import { + createRequestExecutionBudget, + type RequestExecutionBudget, +} from "../../lib/request-execution-budget"; import { classifyTransportFailureKind, transportErrorCode } from "../../lib/upstream-reachability"; import { acquireUpstreamHostAdmission, @@ -224,6 +230,14 @@ export interface HandleResponsesCompactOptions { nativeMainRefreshDependencies?: NativeMainRefreshDependencies; /** Release the listener's idle guard only after the complete request body is accepted. */ onRequestBodyRead?: () => void; + /** + * The logical request's send budget (#4546). Compact used to hold its own: the normal send + * took a fresh transient allowance of three, the 401 replay and the 429 alternate each added + * one -- and the guard that was supposed to make those two mutually exclusive keys on + * `kind === "pool"`, so a main-pool credential could spend all five. The recursive handoff + * child then started over, so one compact could reach ten. + */ + sendBudget?: RequestExecutionBudget; } export function compactResponseTooLargeError(): Response { @@ -760,6 +774,10 @@ export async function handleResponsesCompact( // so routed-model reasoning items (reasoning_text content) don't 400 the ChatGPT backend. const compactBody = sanitizeReasoningInputContent(compactBodyRaw) as typeof compactBodyRaw; const compactUrl = `${base}/responses/compact`; + // One holder for this logical compact, inherited by the handoff child so a second model + // does not start over with a fresh four. + const sendBudget: RequestExecutionBudget = options.sendBudget ?? createRequestExecutionBudget(); + const compactTargetKey = `${route.providerName}|${route.modelId}|compact`; const actualCompactHostKey = upstreamHostHealthKey( route.providerName, safeOriginLabel(compactUrl), @@ -834,6 +852,27 @@ export async function handleResponsesCompact( // wrapping reset retry — because those retries happen before any alternate is even // considered. The alternate is one bounded send: a second ladder would multiply the // work an already-rejecting pool is doing. + // + // Both modes now draw one shared budget. The comment below used to say the 401 replay + // spends the account budget so the 429 alternate is skipped, but that guard is keyed on + // `kind === "pool"` and a main-pool credential left it false -- so 401 then 429 really did + // reach five. The single sends spend the base allowance first and then the one shared + // final-recovery reserve, which is the same rule the Responses path follows. + const sendSingleCompactAttempt = ( + doFetch: () => Promise, + ): Promise => { + if (sendBudget.remainingBaseSends(TRANSIENT_RETRY_MAX_ATTEMPTS) > 0) { + sendBudget.used += 1; + return doFetch(); + } + const decision = sendBudget.reserveDispatch({ + sendClass: "auth-recovery", + targetKey: compactTargetKey, + }); + if (!decision.allowed) return Promise.reject(new SendBudgetExhaustedError(safeHostLabel(compactUrl))); + if (!decision.permit.use()) return Promise.reject(new SendBudgetExhaustedError(safeHostLabel(compactUrl))); + return doFetch(); + }; const sendCompactAttempt = ( sendProvider: OcxProviderConfig, sendHeaders: Headers, @@ -866,8 +905,15 @@ export async function handleResponsesCompact( return res; }); return recovery === "single" - ? doFetch() - : fetchWithTransientRetry(doFetch, { abortSignal: req.signal, label: safeHostLabel(compactUrl) }); + ? sendSingleCompactAttempt(doFetch) + : fetchWithTransientRetry(doFetch, { + abortSignal: req.signal, + label: safeHostLabel(compactUrl), + // Draws the shared remainder instead of a fresh three. Compact is a native endpoint + // of the same logical turn, so its sends belong to the same cap. + attempts: sendBudget.remainingBaseSends(TRANSIENT_RETRY_MAX_ATTEMPTS), + onSendsConsumed: (used: number) => { sendBudget.used += Math.max(0, used); }, + }); }; // The account each outcome belongs to. Reassigned only when the alternate send below @@ -1133,7 +1179,9 @@ export async function handleResponsesCompact( logCtx, turnAdmissionLease, admission, - options, + // The handoff child is the same logical compact on a second model, so it inherits + // the holder. Forwarding `options` alone was not enough: the child minted its own. + { ...options, sendBudget }, ); if (fallback.ok || fallback.status === 499) return fallback; await fallback.body?.cancel().catch(() => undefined); diff --git a/src/server/responses/core.ts b/src/server/responses/core.ts index 06d1937362..2f694105eb 100644 --- a/src/server/responses/core.ts +++ b/src/server/responses/core.ts @@ -7718,6 +7718,7 @@ async function handleResponsesInner( upstreamResponse = await activeAdapter.fetchResponse(builtInitialRequest, { abortSignal: upstream.signal, timeoutMs: connectMs, + sendBudget, stream: parsed.stream, executor: providerFetch(route.provider, options.codexWsRuntimeIdentity, { dispatchOverride: oauthDispatch(builtInitialRequest), @@ -7852,6 +7853,7 @@ async function handleResponsesInner( return await activeAdapter.fetchResponse(retryRequest, { abortSignal: upstream.signal, timeoutMs: connectMs, + sendBudget, stream: parsed.stream, executor: providerFetch(route.provider, options.codexWsRuntimeIdentity, { dispatchOverride: oauthDispatch(retryRequest), @@ -8409,6 +8411,7 @@ async function handleResponsesInner( return await activeAdapter.fetchResponse(builtContinuationRequest, { abortSignal: upstream.signal, timeoutMs: connectMs, + sendBudget, stream: nextParsed.stream, executor: providerFetch(route.provider, options.codexWsRuntimeIdentity, { dispatchOverride: oauthDispatch(builtContinuationRequest, nextParsed),