diff --git a/docs-site/src/content/docs/reference/configuration/server.md b/docs-site/src/content/docs/reference/configuration/server.md index 6309066f19e..9f280eebcd1 100644 --- a/docs-site/src/content/docs/reference/configuration/server.md +++ b/docs-site/src/content/docs/reference/configuration/server.md @@ -95,7 +95,7 @@ OpenCodex bounds how much one task may fan out. A task is the root conversation, | Field | Default | Allowed | Meaning | | --- | --- | --- | --- | | `maxConcurrentChildren` | `8` | 2–64 | Turns one task may run at the same time. One slot is always reserved for the task's own conversation, so subagents may use one fewer. | -| `maxPhysicalSends` | `256` | 1–10000 | Upstream sends, retries included, after which the task's new requests are refused for the rest of the 10-minute window. The count is checked when a request is admitted, so retries inside a request that was already admitted, and requests admitted concurrently, can take the total past the ceiling. | +| `maxPhysicalSends` | `256` | 1–10000 | Upstream send admissions, retries included, per 10-minute window. Every physical admission checks and claims a slot atomically, so concurrent requests and retries share the same ceiling. An unused dispatch permit consumes no slot; a claimed physical start remains counted even if it later fails. | | `maxDistinctChildren` | `64` | 1–1024 | Distinct subagent threads one task may start inside the same window. | The defaults suit most installs. Raise them when a workflow that legitimately runs many subagents in parallel keeps hitting a ceiling; lower them to stop a runaway fan-out sooner: diff --git a/scripts/test-layout/layout-responses.json b/scripts/test-layout/layout-responses.json index 60f1e7075ea..46181d4e113 100644 --- a/scripts/test-layout/layout-responses.json +++ b/scripts/test-layout/layout-responses.json @@ -8,5 +8,6 @@ "passthrough-headers.test.ts": "responses", "passthrough-abort.test.ts": "responses", "current-turn-tool-catalog-direct.test.ts": "responses", - "current-turn-tool-catalog-namespace.test.ts": "responses" + "current-turn-tool-catalog-namespace.test.ts": "responses", + "responses-workflow-send-ceiling.test.ts": "responses" } diff --git a/src/lib/request-execution-budget.ts b/src/lib/request-execution-budget.ts index c50cc72f322..8812eb80088 100644 --- a/src/lib/request-execution-budget.ts +++ b/src/lib/request-execution-budget.ts @@ -140,6 +140,10 @@ export type DispatchDecision = * send rather than merely describe one. */ export interface RequestSendObserver { + /** Atomically claim the root workflow count at physical admission, before inference I/O. */ + beforePhysicalSend?(): boolean; + /** Record an already-sent legacy receipt without booking its durable spend twice. */ + recordWorkflowSend?(): void; /** * Book one physical send. False refuses the dispatch before the budget charges it. * @@ -149,7 +153,7 @@ export interface RequestSendObserver { * would cross a ceiling never joins the total, the total stays just under, and the ceiling * never fires for any later request either. */ - charge(options?: { alreadySent?: boolean; deferDispatch?: boolean; onReserved?: (proof: SpendReservationProof) => void; targetKey?: string }): boolean; + charge(options?: { alreadySent?: boolean; deferDispatch?: boolean; workflowAlreadyCounted?: boolean; onReserved?: (proof: SpendReservationProof) => void; targetKey?: string }): boolean; /** Confirm the exact reservation once its dispatch is known. */ readonly enforced?: boolean; readonly policyStarted?: boolean; @@ -172,10 +176,12 @@ export interface RequestExecutionBudget extends TransientSendBudget { readonly physicalStarted?: number; readonly physicalLimit?: number; readonly spendEnforced?: boolean; + readonly workflowCountObserved?: boolean; readonly spendPolicyStarted?: boolean; readonly spendAdmissionPolicy?: SpendAdmissionPolicy; startRequest?(target?: SpendScopes & { targetKey?: string }): void; claimPhysicalSend?(): number | undefined; + claimWorkflowSend?(): boolean; beginSpendProducer?(): { close(): void } | undefined; readonly logicalRequestId: string; readonly policyVersion: string; @@ -264,6 +270,10 @@ interface SharedSendLedger { } const sharedSendLedgers = new WeakMap(); +/** Whether physical admission owns this budget's workflow accounting. */ +export function hasWorkflowSendObserver(budget: TransientSendBudget): boolean { + return (budget as RequestExecutionBudget).workflowCountObserved === true; +} const dispatchSpendProofs = new WeakMap 0) counter?.observer?.recordWorkflowSend?.(); + if (workflowAlreadyCounted && counter && counter.observer?.enforced !== true) { + counter.spent += count - prepaid; + for (let index = prepaid; index < count; index++) counter.observer?.charge({ alreadySent: true, workflowAlreadyCounted: true }); + } else budget.used += count - prepaid; } /** A helper owns its starts and can never consume another helper's receipts. */ @@ -319,7 +335,16 @@ export function createPhysicalSendReporter( let currentOrdinal: number | undefined; const report: PhysicalSendReporter = (sends) => { if (closed) return; - if (!producer) { if (observer?.enforced !== true) { reportDispatchSends(budget, sends, permit); telemetry?.(sends); } return; } + if (!producer) { + if (observer?.enforced === true) return; + const count = Number.isFinite(sends) ? Math.max(0, Math.trunc(sends)) : 0; + const admitted = Math.min(count, started - reported); + if (admitted > 0) reportDispatchSends(budget, admitted, permit, true); + if (count > admitted) reportDispatchSends(budget, count - admitted, permit); + reported += admitted; + telemetry?.(count); + return; + } const count = Math.min(Math.max(0, Math.trunc(sends)), started - reported); try { producer.report(count); } finally { @@ -335,7 +360,11 @@ export function createPhysicalSendReporter( const selected = target(); if ((budget as RequestExecutionBudget).startRequest) (budget as RequestExecutionBudget).startRequest!(selected); else observer?.startRequest?.(selected); - if (observer?.enforced !== true) return true; + if (observer?.enforced !== true) { + if (observer?.beforePhysicalSend?.() === false) return false; + started++; + return true; + } producer ??= observer.beginReporter?.(); const receipt = permit && dispatchSpendProofs.get(permit); const proof = receipt?.owner === counter ? receipt?.claim() : undefined; @@ -409,6 +438,7 @@ function createRequestExecutionBudgetWithLedger( get physicalStarted(): number { return counter.physicalStarted; }, get physicalLimit(): number | undefined { return counter.physicalLimit; }, get spendEnforced(): boolean { return observer?.enforced === true; }, + get workflowCountObserved(): boolean { return observer?.beforePhysicalSend !== undefined; }, get spendPolicyStarted(): boolean { return observer?.policyStarted === true; }, get spendAdmissionPolicy(): SpendAdmissionPolicy | undefined { return observer?.spendAdmissionPolicy; }, startRequest(target) { @@ -418,10 +448,14 @@ function createRequestExecutionBudgetWithLedger( } }, beginSpendProducer() { return observer?.enforced ? observer.beginReporter?.() : undefined; }, + claimWorkflowSend(): boolean { + budget.startRequest?.(); + return observer?.beforePhysicalSend?.() ?? true; + }, claimPhysicalSend(): number | undefined { budget.startRequest?.(); if (observer?.enforced !== true || counter.physicalLimit === undefined) return undefined; - if (counter.physicalStarted >= counter.physicalLimit) return undefined; + if (counter.physicalStarted >= counter.physicalLimit || !budget.claimWorkflowSend!()) return undefined; return ++counter.physicalStarted; }, get used(): number { return counter.spent; }, @@ -512,7 +546,7 @@ function createRequestExecutionBudgetWithLedger( let executing = false; let physicalOrdinal: number | undefined; const startPhysical = (): boolean => { - if (observer?.enforced !== true) return true; + if (observer?.enforced !== true) return budget.claimWorkflowSend!(); const seed = observer.ensureSeed?.({ targetKey: intent.targetKey }, spendProof); if (!seed) return false; const ordinal = budget.claimPhysicalSend?.(); diff --git a/src/lib/workflow-budget.ts b/src/lib/workflow-budget.ts index 90abeaf9ab5..3c1621a41b3 100644 --- a/src/lib/workflow-budget.ts +++ b/src/lib/workflow-budget.ts @@ -671,6 +671,28 @@ export function chargeWorkflowSends( state.lastSeenMs = now; } +/** + * Atomically reserve one physical send under a root's live count ceiling. + * + * JavaScript runs this read-and-increment without an await, so concurrent request and retry + * ladders cannot all observe the same remaining slot. Call only at physical admission, after + * unused dispatch permits can be released; a claimed physical start remains counted. + */ +export function reserveWorkflowSend( + rootId: string | undefined, + policy: WorkflowBudgetPolicy = DEFAULT_WORKFLOW_BUDGET_POLICY, + now: number = Date.now(), +): boolean { + if (!rootId) return true; + const state = roots.get(rootId); + if (!state) return true; + if (windowedSends(state, now) >= policy.maxPhysicalSends) return false; + state.sends += 1; + recordWindowedSends(state, now, 1); + state.lastSeenMs = now; + return true; +} + /** * Settle a send's reservation with the usage the response actually reported. Idempotent * per send id -- a second call returns false and books nothing. When the usage frame was diff --git a/src/server/chat-native.ts b/src/server/chat-native.ts index 123659d41d5..d3c25a35d47 100644 --- a/src/server/chat-native.ts +++ b/src/server/chat-native.ts @@ -471,9 +471,9 @@ export async function runNativeChatAttempt( throw new SendBudgetExhaustedError(errorDestinationLabel(url)); } physicalSends += 1; - reportDispatchSends(sendBudget, 1, execution.comboDispatchPermit); + reportDispatchSends(sendBudget, 1, execution.comboDispatchPermit, true); initialSendStarted = true; - } else if (!spendTracker?.charge()) throw new NativeChatSpendRefusal(); + } else if (!spendTracker?.charge({ workflowAlreadyCounted: true })) throw new NativeChatSpendRefusal(); }; let responseOwnsDistributedKey = true; @@ -704,7 +704,8 @@ export async function runNativeChatAttempt( return failLocal(429, sendError.message, SEND_BUDGET_EXHAUSTED_CODE, SEND_BUDGET_EXHAUSTED_CODE); } if (sendError instanceof NativeChatSpendRefusal) { - const refusal = workflowRefusalResponse("workflow-spend-exhausted", logCtx); + const refusal = workflowRefusalResponse(logCtx.errorCode === "workflow_sends_exhausted" + ? "workflow-sends-exhausted" : "workflow-spend-exhausted", logCtx); finishLog(429); return refusal; } diff --git a/src/server/messages-native.ts b/src/server/messages-native.ts index 544552fc0e6..4430c08633c 100644 --- a/src/server/messages-native.ts +++ b/src/server/messages-native.ts @@ -639,7 +639,7 @@ export async function handleNativeMessages(options: HandleNativeMessagesOptions) if (spendTracker.refusals) throw new NativeMessagesSpendRefusal(); throw new SendBudgetExhaustedError(); } - if (!physicalBudget.spendEnforced && !spendTracker.charge()) throw new NativeMessagesSpendRefusal(); + if (!physicalBudget.spendEnforced && !spendTracker.charge({ workflowAlreadyCounted: true })) throw new NativeMessagesSpendRefusal(); noteProviderAttemptSend(logCtx, route.providerName, activeProvider, logCtx.usageLogInputTokens, transportRecovery ?? recovery); // Last synchronous ownership reservation before the physical fetch. Retain this // exact incarnation across its await; a returned response can never reserve one. @@ -844,7 +844,8 @@ export async function handleNativeMessages(options: HandleNativeMessagesOptions) } if (sendError instanceof SendBudgetExhaustedError) return fail(429, sendError.message, "send_budget_exhausted", "send_budget_exhausted"); if (sendError instanceof NativeMessagesSpendRefusal) { - const refusal = workflowRefusalResponse("workflow-spend-exhausted", logCtx); + const refusal = workflowRefusalResponse(logCtx.errorCode === "workflow_sends_exhausted" + ? "workflow-sends-exhausted" : "workflow-spend-exhausted", logCtx); finishLog(429); return refusal; } diff --git a/src/server/responses/adapter-dispatch.ts b/src/server/responses/adapter-dispatch.ts index dad7db0a28e..f72e2b79887 100644 --- a/src/server/responses/adapter-dispatch.ts +++ b/src/server/responses/adapter-dispatch.ts @@ -120,6 +120,7 @@ import { createCodexAuthDispatchGuard, releaseCodexAuthContextProbeLease, unwrap import { isCanonicalOpenAiForwardProvider } from "../../providers/openai-tiers"; import { mapCodexAuthContextErrorToResponse } from "./codex-auth-error"; import { chargeWorkflowSends } from "../../lib/workflow-budget"; +import { hasWorkflowSendObserver } from "../../lib/request-execution-budget"; import { isAntigravityValidationRefusal } from "./antigravity-validation-refusal"; import { createAntigravityModelDenialRecovery, recoverAntigravityAccount } from "./antigravity-model-denial-recovery"; import { unboundPoolSpendRefusalResponse } from "../workflow-refusal"; @@ -403,7 +404,7 @@ export async function prepareAdapterExchange( const autoPauseEnabled = config.providers[route.providerName]?.accountAutoPause?.enabled === true; const billingPoolRequest = isAnthropicBillingKeyPool(route.provider); let billingRecoveryActive = false; - const genericRefetchCountsSends = () => options.comboAttempt === true || sendBudgetState.adapterSendBudget?.spendEnforced === true + const genericRefetchCountsSends = () => !!sendBudgetState.workflowRootId || options.comboAttempt === true || sendBudgetState.adapterSendBudget?.spendEnforced === true || transientRetryPolicyFor(route.provider) !== null || resetReplayPolicyFor(route.provider) !== null || billingPoolRequest || billingRecoveryActive || autoPauseEnabled; let handoffPermit: ResponsesSendBudget["pendingHopPermit"]; @@ -526,7 +527,7 @@ export async function prepareAdapterExchange( abortSignal: upstream.signal, label: safeHostLabel(builtInitialRequest.url), destination: errorDestinationLabel(builtInitialRequest.url), claimAmbiguousResend: claimPreHeaderResend, - ...(sendBudgetState.adapterSendBudget?.spendEnforced || options.comboDispatchPermit || transientPolicy || resetPolicy || compactPrepaid || options.comboAttempt + ...(sendBudgetState.workflowRootId || sendBudgetState.adapterSendBudget?.spendEnforced || options.comboDispatchPermit || transientPolicy || resetPolicy || compactPrepaid || options.comboAttempt || billingPoolRequest || autoPauseEnabled ? { // A pending compaction permit already booked this leg's first physical send. @@ -639,15 +640,14 @@ export async function prepareAdapterExchange( onDispatch?: () => void, preserveFailureResponse?: Response, ): Promise => { - let replacementAdmitted = false; - // The repair permit books the request ledger, but only the transient helper reports - // that send to the root workflow. The other two dispatch paths confirm it here, at - // their physical-send boundary, without booking the request ledger again. + let replacementAdmitted = false, dispatchConfirmed = false; + // Only budgets without a physical observer need this legacy workflow report. let fastDowngradeWorkflowCharged = false; const chargeFastDowngradeWorkflowSend = (): void => { if (recovery !== "anthropic-fast-downgrade" || fastDowngradeWorkflowCharged) return; fastDowngradeWorkflowCharged = true; - chargeWorkflowSends(sendBudgetState.workflowRootId, 1); + if (!sendBudgetState.adapterSendBudget || !hasWorkflowSendObserver(sendBudgetState.adapterSendBudget)) + chargeWorkflowSends(sendBudgetState.workflowRootId, 1); }; let retryRequest: AdapterRequest; if (transportState.sameTargetRequest !== undefined && transportState.sameTargetParsed === parsed && transportState.sameTargetToken === transportState.transportToken) { @@ -713,7 +713,7 @@ export async function prepareAdapterExchange( // an abort, a saturated queue, an expired slot or a removed provider, and none of // those reach the wire. Confirming earlier would hold the charge for a send that the // pacer refused. - onDispatch?.(); + if (!dispatchConfirmed) { onDispatch?.(); dispatchConfirmed = true; } return transportState.activeAdapter.fetchResponse!(retryRequest, { kiroPreferAccountFailover: route.providerName === "kiro" && isGenericOAuthFailoverEnabled(config, "kiro"), abortSignal: upstream.signal, @@ -780,14 +780,14 @@ export async function prepareAdapterExchange( // counters and permits still wait for the physical executor after pacing. if (pendingKeyMetadata) transportState.noteRoutedAttemptSend(retryEstimate, recoveryKind ?? recovery); const confirmDispatch = () => { - if (refetchAllowance?.permit && !refetchPermitUsed) { + if (refetchAllowance?.permit && !refetchPermitUsed && (!onDispatch || refetchAllowance.permit !== prepaid)) { if (!refetchAllowance.permit.use()) throw new SendBudgetExhaustedError(errorDestinationLabel(retryRequest.url)); refetchPermitUsed = true; } if (!pendingKeyMetadata) transportState.noteRoutedAttemptSend(retryEstimate, recoveryKind ?? recovery); // Counted Combo/pool sends confirm at the physical executor; legacy // helper paths retain their existing thunk admission boundary. - onDispatch?.(); + if (!dispatchConfirmed) { onDispatch?.(); dispatchConfirmed = true; } if (preserveFailureResponse) replacementAdmitted = true; if (!refetchAllowance) chargeFastDowngradeWorkflowSend(); }; @@ -1014,7 +1014,7 @@ export async function prepareAdapterExchange( route.provider, transportState.activeAdapter.name); sendBudgetState.pendingHopPermit = hop.permit; const result = await rebuildAndRefetch("oauth-account-429", () => { - if (!adapterOwnsDispatch) hop.permit?.use(); + if (!adapterOwnsDispatch && hop.permit?.use() === false) throw new SendBudgetExhaustedError(); }); if ("failed" in result) return result.failed; upstreamResponse = result; @@ -1136,7 +1136,7 @@ export async function prepareAdapterExchange( }); if (billingHop) sendBudgetState.pendingHopPermit = billingHop.permit; const result = await rebuildAndRefetch("key-401", - billingHop ? () => { billingHop.permit?.use(); } : undefined, failedResponse); + billingHop ? () => { if (billingHop.permit?.use() === false) throw new SendBudgetExhaustedError(); } : undefined, failedResponse); if ("failed" in result) { if (result.failed !== failedResponse) return result.failed; break; @@ -1174,7 +1174,7 @@ export async function prepareAdapterExchange( resolveWireProtocolOverride(route.providerName, route.modelId, route.provider, inboundWire, route.staticPolicy), config.cacheRetention); bindRouteReasoningReplayScope({ parsed, providerName: route.providerName, provider: route.provider, adapterName: transportState.activeAdapter.name }); sendBudgetState.pendingHopPermit = hop.permit; - const result = await rebuildAndRefetch("key-401", () => { hop.permit?.use(); }, failedResponse); + const result = await rebuildAndRefetch("key-401", () => { if (hop.permit?.use() === false) throw new SendBudgetExhaustedError(); }, failedResponse); if ("failed" in result) { if (result.failed !== failedResponse) return result.failed; break; @@ -1217,7 +1217,7 @@ export async function prepareAdapterExchange( let result: Response | { failed: Response }; try { result = await rebuildAndRefetch("anthropic-fast-downgrade", () => { - if (!adapterOwnsDispatch) hop.permit?.use(); + if (!adapterOwnsDispatch && hop.permit?.use() === false) throw new SendBudgetExhaustedError(); }); } finally { sendBudgetState.pendingHopPermit = undefined; @@ -1337,7 +1337,7 @@ export async function prepareAdapterExchange( }); if (billingHop) sendBudgetState.pendingHopPermit = billingHop.permit; const result = await rebuildAndRefetch("key-429", - billingHop ? () => { billingHop.permit?.use(); } : undefined, failedResponse); + billingHop ? () => { if (billingHop.permit?.use() === false) throw new SendBudgetExhaustedError(); } : undefined, failedResponse); if ("failed" in result) { if (result.failed !== failedResponse) return result.failed; break; @@ -1416,7 +1416,7 @@ export async function prepareAdapterExchange( try { result = await rebuildAndRefetch(upstreamResponse.status === 401 ? "oauth-401" : upstreamResponse.status === 403 ? "oauth-account-403" : "anthropic-oauth-429", - () => { if (!adapterOwnsDispatch) hop.permit?.use(); }); + () => { if (!adapterOwnsDispatch && hop.permit?.use() === false) throw new SendBudgetExhaustedError(); }); } finally { sendBudgetState.pendingHopPermit = undefined; } @@ -1523,7 +1523,7 @@ export async function prepareAdapterExchange( // and settling here first would hand it a dead permit, which it reads as an // exhausted request and stops sending on. result = await rebuildAndRefetch("oauth-account-429", () => { - if (!adapterOwnsDispatch) hop.permit?.use(); + if (!adapterOwnsDispatch && hop.permit?.use() === false) throw new SendBudgetExhaustedError(); }, upstreamResponse); } finally { sendBudgetState.pendingHopPermit = undefined; @@ -1628,7 +1628,7 @@ export async function prepareAdapterExchange( // and settling here first would hand it a dead permit, which it reads as an // exhausted request and stops sending on. result = await rebuildAndRefetch("oauth-account-429", () => { - if (!adapterOwnsDispatch) hop.permit?.use(); + if (!adapterOwnsDispatch && hop.permit?.use() === false) throw new SendBudgetExhaustedError(); }); } finally { sendBudgetState.pendingHopPermit = undefined; @@ -1740,7 +1740,7 @@ export async function prepareAdapterExchange( let result: Response | { failed: Response }; try { result = await rebuildAndRefetch("oauth-account-403", () => { - if (!adapterOwnsDispatch) hop.permit?.use(); + if (!adapterOwnsDispatch && hop.permit?.use() === false) throw new SendBudgetExhaustedError(); }, failedResponse); } finally { sendBudgetState.pendingHopPermit = undefined; diff --git a/src/server/responses/antigravity-model-denial-recovery.ts b/src/server/responses/antigravity-model-denial-recovery.ts index 4cb6bf5a285..19a5d9bc03d 100644 --- a/src/server/responses/antigravity-model-denial-recovery.ts +++ b/src/server/responses/antigravity-model-denial-recovery.ts @@ -2,7 +2,7 @@ import { getConfigDir } from "../../config/paths"; import type { OAuthAccessSnapshot } from "../../oauth"; import { failoverAccountSnapshot, rotateAntigravityAccountOnModelNotFound } from "../../oauth/generic-account-failover"; import { captureAntigravityModelObservation, recordAntigravityModelDenial } from "../../providers/antigravity-account-roster"; -import { isNonReplayableResponse } from "../../lib/upstream-retry"; +import { isNonReplayableResponse, SendBudgetExhaustedError } from "../../lib/upstream-retry"; import type { OcxConfig } from "../../types"; import type { ResponsesSendBudget } from "./request-send-budget"; import { isAntigravityModelNotFoundRefusal } from "./antigravity-validation-refusal"; @@ -25,7 +25,9 @@ export async function recoverAntigravityAccount(options: { if ((await options.applyFailoverSnapshot(snapshot))?.accountId !== options.nextId) return null; options.prepare(); options.setPermit(hop.permit); - const result = await options.refetch(() => { if (!options.adapterOwnsDispatch) hop.permit?.use(); }); + const result = await options.refetch(() => { + if (!options.adapterOwnsDispatch && hop.permit?.use() === false) throw new SendBudgetExhaustedError(); + }); if ("failed" in result) return result.failed === options.failedResponse ? null : result.failed; options.complete(); try { void options.failedResponse.body?.cancel().catch(() => {}); } catch { /* already closed */ } diff --git a/src/server/responses/core-codex-account.ts b/src/server/responses/core-codex-account.ts index 736c2c1f1e7..0fd2eb3dde4 100644 --- a/src/server/responses/core-codex-account.ts +++ b/src/server/responses/core-codex-account.ts @@ -73,7 +73,7 @@ import { CodexPoolAccountCreditsOffError, } from "../../codex/auth-context"; import { ACCOUNT_GATED_NATIVE_OPENAI_MODELS } from "../../codex/catalog/native-models"; -import { isRequestExecutionBudget } from "../../lib/request-execution-budget"; +import { hasWorkflowSendObserver, isRequestExecutionBudget } from "../../lib/request-execution-budget"; import type { SingleUseDispatchPermit } from "../../lib/request-execution-budget"; import { codexRouteCredentialOwnership, type CodexCredentialOwnershipOptions } from "./core-auth"; import { bindRouteReasoningReplayScope } from "./core-replay"; @@ -90,7 +90,8 @@ import { noteProviderAttemptSend, } from "../request-log"; import { codexAuthContextLogLabel } from "../../codex/account-label"; -import { chargeWorkflowSends } from "../../lib/workflow-budget"; +import { reserveWorkflowSend, workflowSendCeilingReached } from "../../lib/workflow-budget"; +import { currentWorkflowBudgetPolicy } from "../../lib/workflow-budget-config"; import type { ResponsesTerminalStatus } from "../../bridge"; export function sidecarOutcomeRecorder( @@ -952,23 +953,6 @@ export async function retryCodexPoolOnAlternateAccount( recordAdapterReasoning(logCtx, request); recordAdapterTier(logCtx, request); - await firstResponse.body?.cancel().catch(() => undefined); - options.onCodexAuthContextResolved?.(retryAuthCtx); - route.provider = retryProvider; - logCtx.provider = formatCodexProviderForLog( - route.providerName, - retryAuthCtx.accountId, - config, - ); - logCtx.accountLogLabel = codexAuthContextLogLabel(retryAuthCtx, config); - sealRequestAttemptIdentity( - logCtx.activeAttempt, - logCtx.provider, - retryAdapter.name, - logCtx.accountLogLabel, - ); - recordAttemptCredentialSource(logCtx.activeAttempt, route.providerName, route.provider, retryAdapter.name); - const retrySameConfirmedAccount = outcomeStatus === 400 && ACCOUNT_GATED_NATIVE_OPENAI_MODELS.has(route.modelId) && retryAuthCtx.accountId === firstAuthCtx.accountId; @@ -978,15 +962,6 @@ export async function retryCodexPoolOnAlternateAccount( // exact allow-listed body and fresh entitlement before every later send. Alternate-account and // quota recovery retain their historical one-send bound. // - // Two different bounds, and the effective one is the smaller. `maxRetrySends` answers "how - // many times is it worth re-asking THIS account for a model its roster still grants"; the - // shared budget answers "how many times may this LOGICAL REQUEST reach upstream in total, - // across every layer that can re-send". A ladder of eight layered on sends the request had - // already made is exactly the per-request multiplication #4546 is about, so the ladder is - // capped at what the request has left. The floor of one keeps the single retry this function - // was called to make -- the move already paid for itself with its own permit -- and each rung - // past the first reserves its own send below, so a refusal stops the ladder with the last - // upstream answer intact. // The ladder replays to the SAME account, so it must reserve under the same target key the // other legs use. Folding the account id in made every rung read as a target change, which // spent the one cross-account slot a real move needs on a same-account replay. @@ -996,14 +971,36 @@ export async function retryCodexPoolOnAlternateAccount( // at eight same-account dispatches (tests/server/server-auth.test.ts), and a request that has // already spent sends would silently stop short of it. Reconciling an eight-send same-account // ladder with a four-send request total is a policy decision, not a clamp to add in passing. - // What this diff does fix is that the rungs are now CHARGED instead of free. + // Every physical rung still needs the workflow's allowance, including the first same-account + // retry and a rung admitted by that historical request-local exception. const maxRetrySends = retrySameConfirmedAccount ? 7 : 1; + const claimRetrySend = (permit?: SingleUseDispatchPermit): boolean => { + const workflowObserved = executionBudget && hasWorkflowSendObserver(executionBudget); + const workflowPolicy = currentWorkflowBudgetPolicy(); + // Legacy permits cannot refund after use(). Check the root first, then use the permit + // and reserve the root in this same synchronous turn. A failed permit claims no root + // slot, and a root refusal can still release the unused request reservation. + if (!workflowObserved && workflowSendCeilingReached(options.workflowRootId, workflowPolicy)) { + permit?.release(); + return false; + } + if (permit && !permit.use()) { + permit.release(); + return false; + } + // An observed permit already claimed the root at use(). The exceptional ladder path has + // no permit, so it must ask that same observer directly before another physical dispatch. + if (workflowObserved) { + return permit !== undefined || executionBudget.claimWorkflowSend?.() === true; + } + return reserveWorkflowSend(options.workflowRootId, workflowPolicy); + }; let retrySendCount = 0; // The gated-model 400 has no roster kind; it stays unattributed as before. const moveRecovery: AttemptRecoveryKind | undefined = outcomeStatus === 429 ? "oauth-account-429" : outcomeStatus === 402 ? "oauth-account-402" : isTransientUpstreamStatus(outcomeStatus) ? "transient-5xx" : undefined; - let upstreamResponse: Response; + let upstreamResponse = firstResponse; try { while (true) { // The same-account gated-model 400 ladder below keeps its own `maxRetrySends` bound and @@ -1027,15 +1024,28 @@ export async function retryCodexPoolOnAlternateAccount( recordUnmovedTransientOutcome(); return { kind: "no-alternate" }; } - const charged = accountMovePermit.use(); - accountMovePermit = undefined; - if (!charged) { - releaseCodexAuthContextProbeLease(retryAuthCtx); - recordUnmovedTransientOutcome(); - return { kind: "no-alternate" }; - } - // The move is a physical send like any other, so the root workflow is charged too. - chargeWorkflowSends(args.options.workflowRootId, 1); + } + const rung = retrySendCount > 0 ? executionBudget?.reserveDispatch({ + sendClass: "transient", targetKey: ladderTargetKey, + }) : undefined; + const permit = accountMovePermit ?? (rung?.allowed ? rung.permit : undefined); + accountMovePermit = undefined; + if (!claimRetrySend(permit)) { + if (retrySendCount > 0) break; + releaseCodexAuthContextProbeLease(retryAuthCtx); + recordUnmovedTransientOutcome(); + return { kind: "no-alternate" }; + } + // Keep the last upstream answer readable until its replacement is admitted. In + // particular, root exhaustion on the first retry must leave firstResponse untouched. + await upstreamResponse.body?.cancel().catch(() => undefined); + if (retrySendCount === 0) { + options.onCodexAuthContextResolved?.(retryAuthCtx); + route.provider = retryProvider; + logCtx.provider = formatCodexProviderForLog(route.providerName, retryAuthCtx.accountId, config); + logCtx.accountLogLabel = codexAuthContextLogLabel(retryAuthCtx, config); + sealRequestAttemptIdentity(logCtx.activeAttempt, logCtx.provider, retryAdapter.name, logCtx.accountLogLabel); + recordAttemptCredentialSource(logCtx.activeAttempt, route.providerName, route.provider, retryAdapter.name); } // A move seals the first account's attempt here. Give it the status that account actually // answered, not the 502 the split infers, and name why the second send exists. @@ -1112,25 +1122,6 @@ export async function retryCodexPoolOnAlternateAccount( throw error; } if (!entitledCodexAccountIdsForModel(refreshed, route.modelId)?.has(retryAuthCtx.accountId)) break; - // The next rung is another physical send of this logical request: a same-account, - // same-target replay, charged as an ordinary transient send rather than as a move. - // Reserved here, immediately before looping back, so a refusal stops the ladder with the - // last upstream 400 intact instead of spending a send it cannot make. - // Every rung is CHARGED, and a refusal does not end the ladder. That asymmetry is - // deliberate and it is the one place the shared cap yields. This is a same-account, - // same-target replay of a model-gating 400 whose own bound is eight dispatches, pinned by - // #2097; letting a spent request budget cut it to four would break a recovery that works - // today, which is precisely the mistake 040_send_budget.md warns a flat ceiling makes. - // The request total still governs everything that changes target or credential. - if (executionBudget) { - const rung = executionBudget.reserveDispatch({ - sendClass: "transient", - targetKey: ladderTargetKey, - }); - if (rung.allowed) rung.permit.use(); - chargeWorkflowSends(args.options.workflowRootId, 1); - } - await upstreamResponse.body?.cancel().catch(() => undefined); } } finally { request.releaseBodyObservation?.(); diff --git a/src/server/responses/request-send-budget.ts b/src/server/responses/request-send-budget.ts index 0417a6adcf5..70033dc4f3e 100644 --- a/src/server/responses/request-send-budget.ts +++ b/src/server/responses/request-send-budget.ts @@ -1,8 +1,9 @@ import { createInferenceSendBudget } from "../inference/context"; import type { ResponsesRequestContext } from "./core-options"; -import { claimDispatchSpendProof, createPhysicalSendReporter, createRequestExecutionBudget, isRequestExecutionBudget, reportDispatchSends } from "../../lib/request-execution-budget"; +import { claimDispatchSpendProof, createPhysicalSendReporter, createRequestExecutionBudget, isRequestExecutionBudget, hasWorkflowSendObserver, reportDispatchSends } from "../../lib/request-execution-budget"; import { chargeWorkflowSends, + reserveWorkflowSend, workflowSendCeilingReached, workflowSpendCeilingReached, } from "../../lib/workflow-budget"; @@ -59,6 +60,7 @@ export function createResponsesSendBudget( // sends once per child seven hundred times, so every send charged to the request is charged // to the root as well (#4546). const workflowRootId = req.headers.get("x-codex-parent-thread-id")?.trim() || undefined; + const physicalWorkflowAccounting = hasWorkflowSendObserver(sendBudget); // A child that owns its Combo booking settles that exact permit through the reporter, even // when its caller handed ownership over only as `comboInitialSend`. const inheritedPermit = options.compactionRecoveryPermit ?? options.comboDispatchPermit @@ -74,22 +76,41 @@ export function createResponsesSendBudget( targetSendsUsed += charged; if (charged > 0) initialPermit = undefined; }; - const recordTransientSends = (used: number, permit?: SingleUseDispatchPermit): void => { + const recordTransientSends = (used: number, permit?: SingleUseDispatchPermit, workflowAlreadyCounted = false): void => { const charged = Math.max(0, used); noteTargetSends(charged); - reportDispatchSends(sendBudget, charged, permit); + reportDispatchSends(sendBudget, charged, permit, workflowAlreadyCounted); options.onCompactionRecoverySendsReported?.(charged); - chargeWorkflowSends(workflowRootId, charged); + if (!physicalWorkflowAccounting && !workflowAlreadyCounted) chargeWorkflowSends(workflowRootId, charged); }; const noteTransientSends = (used: number): void => recordTransientSends(used, pendingHopPermit ?? inheritedPermit); // Capture at helper creation, before another leg can replace the pending handoff. - const transientSendReporter = (permit = pendingHopPermit ?? inheritedPermit) => - createPhysicalSendReporter(sendBudget, () => ({ poolId: logCtx.spendPoolId ?? logCtx.provider, + const transientSendReporter = (permit = pendingHopPermit ?? inheritedPermit) => { + let unreportedWorkflowStarts = 0; + const reporter = createPhysicalSendReporter(sendBudget, () => ({ poolId: logCtx.spendPoolId ?? logCtx.provider, identityId: logCtx.accountLogLabel }), permit, charged => { noteTargetSends(Math.max(0, charged)); options.onCompactionRecoverySendsReported?.(charged); - chargeWorkflowSends(workflowRootId, charged); + const prepaid = Math.min(unreportedWorkflowStarts, Math.max(0, charged)); + unreportedWorkflowStarts -= prepaid; + if (!physicalWorkflowAccounting) chargeWorkflowSends(workflowRootId, charged - prepaid); }); + if (workflowRootId && !physicalWorkflowAccounting) { + const beforeSend = reporter.beforeSend, close = reporter.close; + let closed = false; + reporter.beforeSend = () => { + const policy = currentWorkflowBudgetPolicy(); + // Check before a legacy permit's admission can spend it. No await separates this + // read from the atomic claim, and the reporter owns the later accounting receipt. + if (closed || workflowSendCeilingReached(workflowRootId, policy) || beforeSend?.() === false + || !reserveWorkflowSend(workflowRootId, policy)) return false; + unreportedWorkflowStarts++; + return true; + }; + reporter.close = () => { closed = true; close?.(); }; + } + return reporter; + }; // Refused before any dispatch, and deliberately not by evicting the root's ledger entry: // dropping the record to make room would hand the fan-out a fresh allowance, which is the // laundering this ceiling exists to stop. The client is told the task needs a new grant @@ -322,21 +343,29 @@ export function createResponsesSendBudget( remainingTransientSendBudget(cap) + (initialPermit ? 1 : 0)), /** At final executor admission, settle only the owning booking or reserve this send. */ noteInitialDispatch: (prepaid = initialPermit): void => { + const policy = currentWorkflowBudgetPolicy(); + if (!physicalWorkflowAccounting && workflowSendCeilingReached(workflowRootId, policy)) { + prepaid?.release(); + throw new SendBudgetExhaustedError(); + } if (prepaid) { - if (!prepaid.assumeCharge()) throw new SendBudgetExhaustedError(); + if (!prepaid.assumeCharge()) { prepaid.release(); throw new SendBudgetExhaustedError(); } if (prepaid === initialPermit) initialPermit = undefined; } else if (adapterSendBudget && (options.comboInitialSend || options.compactionRecoveryPermit)) { const decision = adapterSendBudget.reserveDispatch({ sendClass: "transient", targetKey: adapterSendBudget.lastTargetKey ?? `${logCtx.provider}/${logCtx.model}` }); - if (!decision.allowed || !decision.permit.use()) throw new SendBudgetExhaustedError(); + if (!decision.allowed) throw new SendBudgetExhaustedError(); + if (!decision.permit.use()) { decision.permit.release(); throw new SendBudgetExhaustedError(); } } else { // Legacy direct callers can prebook without handing over a permit; keep their reporter. - noteTransientSends(1); + if (adapterSendBudget?.claimWorkflowSend && !adapterSendBudget.claimWorkflowSend()) throw new SendBudgetExhaustedError(); + if (!physicalWorkflowAccounting && !reserveWorkflowSend(workflowRootId, policy)) throw new SendBudgetExhaustedError(); + recordTransientSends(1, pendingHopPermit ?? inheritedPermit, true); return; } + if (!physicalWorkflowAccounting && !reserveWorkflowSend(workflowRootId, policy)) throw new SendBudgetExhaustedError(); targetSendsUsed += 1; options.onCompactionRecoverySendsReported?.(1); - chargeWorkflowSends(workflowRootId, 1); }, /** Physical sends reported by this owner; unrelated shared reservations cannot consume it. */ get targetSendsUsed(): number { return targetSendsUsed; }, @@ -391,11 +420,13 @@ function adapterDispatchBudgetView( return { get physicalStarted() { return budget.physicalStarted; }, get physicalLimit() { return budget.physicalLimit; }, + get workflowCountObserved() { return budget.workflowCountObserved; }, get spendEnforced() { return budget.spendEnforced; }, get spendPolicyStarted() { return budget.spendPolicyStarted; }, get spendAdmissionPolicy() { return budget.spendAdmissionPolicy; }, startRequest: target => budget.startRequest?.(target), claimPhysicalSend: () => budget.claimPhysicalSend?.(), + claimWorkflowSend: () => budget.claimWorkflowSend?.() ?? true, beginSpendProducer: () => budget.beginSpendProducer?.(), get used(): number { return budget.used; }, set used(next: number) { budget.used = next; }, diff --git a/src/server/responses/request-spend.ts b/src/server/responses/request-spend.ts index ac217e6899e..5ae06414d7c 100644 --- a/src/server/responses/request-spend.ts +++ b/src/server/responses/request-spend.ts @@ -3,7 +3,8 @@ import type { RequestSendObserver } from "../../lib/request-execution-budget"; import { sharedSpendLedger, sharedSpendPolicy, type SpendReservationLedger, type SpendScopes, type SpendSeed, type SpendReservationProof, type SpendAdmissionPolicy } from "../../lib/spend-reservation-ledger"; import { SpendLedgerOwnerError } from "../../lib/spend-ledger-owner"; import { markLocalRequestLogRefusal, type RequestLogContext } from "../request-log"; -import { recordWorkflowRefusalEvent, workflowDenialSummary, type WorkflowDenial } from "../../lib/workflow-budget"; +import { chargeWorkflowSends, recordWorkflowRefusalEvent, reserveWorkflowSend, workflowDenialSummary, type WorkflowDenial } from "../../lib/workflow-budget"; +import { currentWorkflowBudgetPolicy } from "../../lib/workflow-budget-config"; /** The terminal usage a request reported, in the only two fields the ledger books. */ export interface TerminalSpendUsage { @@ -212,7 +213,20 @@ export function createRequestSpendTracker( recordWorkflowRefusalEvent(rootId, reason, Date.now(), detail); return; }; + const claimWorkflowSend = (): boolean => { + if (reserveWorkflowSend(rootId, currentWorkflowBudgetPolicy())) return true; + refusals++; + const summary = workflowDenialSummary("workflow-sends-exhausted"); + markLocalRequestLogRefusal(logCtx, summary.code); + logCtx.errorCode = summary.code; + recordWorkflowRefusalEvent(rootId, "workflow-sends-exhausted"); + return false; + }; const tracker: RequestSpendTracker = { + recordWorkflowSend() { chargeWorkflowSends(rootId, 1); }, + beforePhysicalSend() { + return claimWorkflowSend(); + }, startRequest(target) { if (admissionPolicy) return; const policy = ledgerRef?.policy ?? sharedSpendPolicy(); @@ -348,6 +362,10 @@ export function createRequestSpendTracker( // fetch -- so without it a root ceiling on the canonical Codex path would sit one send // short of its limit forever and refuse nothing. const alreadySent = options?.alreadySent === true; + if (!options?.deferDispatch && !options?.workflowAlreadyCounted) { + if (alreadySent) chargeWorkflowSends(rootId, 1); + else if (!claimWorkflowSend()) return false; + } if (!alreadySent && tracker.enforced) { const seed = tracker.ensureSeed({ targetKey: options?.targetKey }); if (!seed) return false; diff --git a/structure/transports/responses-spend.md b/structure/transports/responses-spend.md index baac45e580e..6f83b544568 100644 --- a/structure/transports/responses-spend.md +++ b/structure/transports/responses-spend.md @@ -4,6 +4,17 @@ How the Responses data plane reserves credential hops and durable spend before d spent budget tells the client. Dispatch and retry behavior is in [Responses transport](responses.md) and [Responses failover](responses-failover.md). +## Workflow physical-send ceiling + +`src/lib/workflow-budget.ts` atomically checks and claims one root slot at physical admission. +`request-spend.ts` connects this count to both enforced and observe-only request budgets; +reporter retries, concurrent requests and adapter-owned permits share the root's existing window. +Unused dispatch permits consume no root sends, and cleanup cannot refund a claimed physical start. +Observe-only reports reconcile those starts once; already-sent legacy reports remain recorded. +Codex's bounded eight-attempt gated-model recovery may exceed its per-request allowance, but every +retry still claims the shared root ceiling before replacing the previous upstream response. +`tests/responses/responses-workflow-send-ceiling.test.ts` covers both spend modes and retry helpers. + ## Prepaid initial sends `src/server/responses/core-combo.ts` derives each target scope before reserving its initial send. diff --git a/tests/fixtures/test-layout-expected.json b/tests/fixtures/test-layout-expected.json index 37927c69741..9cbc3de9603 100644 --- a/tests/fixtures/test-layout-expected.json +++ b/tests/fixtures/test-layout-expected.json @@ -1,4 +1,4 @@ -{ "kimi-monthly-quota.test.ts": "providers", "sync-fanout.test.ts": "cli", "sync-command-fanout.test.ts": "cli", "link-pinned-sync-argv.test.ts": "clients", "sync-link-status.test.ts": "cli", "sync-fanout-process.test.ts": "cli", "codex-installation-acl.test.ts": "codex-integration", "codex-installation-id.test.ts": "codex-integration", "codex-installation-wire.test.ts": "codex-integration", "codex-installation-failover.test.ts": "responses", "lmstudio-system-policy.test.ts": "config", "openai-chat-single-system.test.ts": "adapters/openai", "update-reclaim-authorizer.test.ts": "update", "update-capability-policy.test.ts": "update", "update-service-home.test.ts": "update", "update-service-home-policy.test.ts": "update", "update-home-recovery.test.ts": "update", "update-recovery-authority.test.ts": "update", "update-recovery-boundaries.test.ts": "update", "update-home-launcher.test.ts": "update", "update-home-worker.test.ts": "update", "update-mutation-boundaries.test.ts": "update", "update-mutation-guard.test.ts": "update", "update-guarded-ops.test.ts": "update", "update-effect-authority.test.ts": "update", "update-windows-wrapper-authority.test.ts": "update", "update-mutation-scanner.test.ts": "update", "refreshed-upstream-deadlines.test.ts": "responses", "local-upstream-deadlines.test.ts": "responses", "service-desktop-recovery.test.ts": "service", "provider-input-modalities.test.ts": "providers", "turn-membership.test.ts": "responses", "turn-membership-ingress.test.ts": "responses", "workbuddy-quota.test.ts": "providers", "workbuddy-login.test.ts": "oauth", "workbuddy-adapter.test.ts": "adapters", "workbuddy-provider.test.ts": "providers", "link-tunnel-history.test.ts": "clients", "link-readiness.test.ts": "clients", "link-diagnostics.test.ts": "clients", "combo-description.test.ts": "routing", "combo-description-catalog.test.ts": "codex-integration", "catalog-auto-review-override.test.ts": "codex-integration", "management-auto-review-settings.test.ts": "server", "management-auto-review-auth.test.ts": "server", "prefer-model-alias-ids.test.ts": "server", "responses-mimo-envelope.test.ts": "responses", "responses-call-id-linkage.test.ts": "responses", "responses-phase-inference.test.ts": "responses", +{ "responses-workflow-send-ceiling.test.ts": "responses", "kimi-monthly-quota.test.ts": "providers", "sync-fanout.test.ts": "cli", "sync-command-fanout.test.ts": "cli", "link-pinned-sync-argv.test.ts": "clients", "sync-link-status.test.ts": "cli", "sync-fanout-process.test.ts": "cli", "codex-installation-acl.test.ts": "codex-integration", "codex-installation-id.test.ts": "codex-integration", "codex-installation-wire.test.ts": "codex-integration", "codex-installation-failover.test.ts": "responses", "lmstudio-system-policy.test.ts": "config", "openai-chat-single-system.test.ts": "adapters/openai", "update-reclaim-authorizer.test.ts": "update", "update-capability-policy.test.ts": "update", "update-service-home.test.ts": "update", "update-service-home-policy.test.ts": "update", "update-home-recovery.test.ts": "update", "update-recovery-authority.test.ts": "update", "update-recovery-boundaries.test.ts": "update", "update-home-launcher.test.ts": "update", "update-home-worker.test.ts": "update", "update-mutation-boundaries.test.ts": "update", "update-mutation-guard.test.ts": "update", "update-guarded-ops.test.ts": "update", "update-effect-authority.test.ts": "update", "update-windows-wrapper-authority.test.ts": "update", "update-mutation-scanner.test.ts": "update", "refreshed-upstream-deadlines.test.ts": "responses", "local-upstream-deadlines.test.ts": "responses", "service-desktop-recovery.test.ts": "service", "provider-input-modalities.test.ts": "providers", "turn-membership.test.ts": "responses", "turn-membership-ingress.test.ts": "responses", "workbuddy-quota.test.ts": "providers", "workbuddy-login.test.ts": "oauth", "workbuddy-adapter.test.ts": "adapters", "workbuddy-provider.test.ts": "providers", "link-tunnel-history.test.ts": "clients", "link-readiness.test.ts": "clients", "link-diagnostics.test.ts": "clients", "combo-description.test.ts": "routing", "combo-description-catalog.test.ts": "codex-integration", "catalog-auto-review-override.test.ts": "codex-integration", "management-auto-review-settings.test.ts": "server", "management-auto-review-auth.test.ts": "server", "prefer-model-alias-ids.test.ts": "server", "responses-mimo-envelope.test.ts": "responses", "responses-call-id-linkage.test.ts": "responses", "responses-phase-inference.test.ts": "responses", "codex-installation-live-content.test.ts": "codex-integration", "codex-installation-invalid-account.test.ts": "codex-integration", "codex-installation-sidecars.test.ts": "codex-integration", diff --git a/tests/lib/workflow-budget.test.ts b/tests/lib/workflow-budget.test.ts index 8f4f28469fe..d8341619617 100644 --- a/tests/lib/workflow-budget.test.ts +++ b/tests/lib/workflow-budget.test.ts @@ -12,6 +12,7 @@ import { listTrackedWorkflowRoots, listWorkflowBudgetEvents, resetWorkflowBudgetsForTest, + reserveWorkflowSend, settleWorkflowSpend, workflowBudgetSnapshot, workflowDenialSummary, @@ -63,6 +64,16 @@ describe("workflow count caps", () => { if (decision && !decision.admitted) expect(decision.reason).toBe("workflow-sends-exhausted"); }); + test("physical starts claim the root allowance atomically", () => { + const oneSend = { ...smallPolicy, maxPhysicalSends: 1 }; + admitWorkflowTurn("atomic-root", "interactive", oneSend)?.lease.release(); + + expect(reserveWorkflowSend("atomic-root", oneSend)).toBe(true); + expect(reserveWorkflowSend("atomic-root", oneSend)).toBe(false); + expect(workflowBudgetSnapshot("atomic-root", oneSend)?.sends).toBe(1); + + }); + test("a worker lane may not take the interactive reserve", () => { const workerCeiling = smallPolicy.maxConcurrentChildren - smallPolicy.interactiveReserve; for (let i = 0; i < workerCeiling; i += 1) { diff --git a/tests/responses/chat-native-spend.test.ts b/tests/responses/chat-native-spend.test.ts index 1d2bd3516ef..64be83ca602 100644 --- a/tests/responses/chat-native-spend.test.ts +++ b/tests/responses/chat-native-spend.test.ts @@ -19,6 +19,8 @@ import * as stateStores from "../../src/lib/state-store-registrations"; import { translatorAggregateCurrentBytesForTests } from "../../src/lib/translator-budget"; import { configureSharedSpendLedger, spendPolicyFromConfig } from "../../src/lib/spend-reservation-ledger"; import { clearKeyCooldowns } from "../../src/providers/key-failover"; +import { resetWorkflowBudgetsForTest, workflowBudgetSnapshot } from "../../src/lib/workflow-budget"; +import { clearComboSelectionState, clearComboTargetCooldowns } from "../../src/combos"; let previousHome: string | undefined; let testDir = ""; @@ -62,6 +64,8 @@ afterEach(async () => { expect(spendLedgerOwnerSnapshot().ownership).toBe("unheld"); resetProviderRequestPacingForTest(); clearKeyCooldowns(); + resetWorkflowBudgetsForTest(); + clearComboSelectionState(); clearComboTargetCooldowns(); if (previousHome === undefined) delete process.env.OPENCODEX_HOME; else process.env.OPENCODEX_HOME = previousHome; isolatedCodexHome?.restore(); @@ -108,6 +112,33 @@ function mockConfig(baseUrl: string, providerOverrides: Partial { + let sends = 0; + activeUpstream = Bun.serve({ hostname: "127.0.0.1", port: 0, fetch() { + sends++; + return Response.json({ error: { message: "temporary fixture refusal" } }, { status: 503, headers: { "retry-after": "0" } }); + } }); + const config = mockConfig(new URL("/v1", activeUpstream.url).href, { transientRetryOn5xx: { attempts: 3 }, liveModels: false, models: ["test-model"] }); + config.workflowBudget = { maxPhysicalSends: 2 }; + if (enforced) config.spend = { pool: { maxTokens: 100_000 } }; + if (combo) config.combos = { fixture: { strategy: "failover", targets: [{ provider: "mock", model: "test-model" }] } }; + saveConfig(config); + activeServer = startServer(0); + const response = await fetch(new URL("/v1/chat/completions", activeServer.url), { + method: "POST", headers: { "content-type": "application/json", "x-codex-parent-thread-id": "native-count-root" }, + body: JSON.stringify({ model: combo ? "combo/fixture" : "mock/test-model", messages: [{ role: "user", content: "fixture" }], max_tokens: 1 }), + }); + await response.text(); + expect(sends).toBe(2); + expect(workflowBudgetSnapshot("native-count-root")?.sends).toBe(2); + if (!combo) { + expect(response.status).toBe(429); + expect(response.headers.get("x-opencodex-local-refusal")).toBe("workflow_sends_exhausted"); + } + }, 10_000); +} + test("native Chat refuses a physical send that exceeds the configured pool spend ceiling", async () => { const upstream = mockChatUpstreamCapturing(); const config = mockConfig(`${upstream.server.url.toString().replace(/\/$/, "")}/v1`); diff --git a/tests/responses/responses-anthropic-fast-downgrade.test.ts b/tests/responses/responses-anthropic-fast-downgrade.test.ts index 3e13b743bd9..eab0e0e7f61 100644 --- a/tests/responses/responses-anthropic-fast-downgrade.test.ts +++ b/tests/responses/responses-anthropic-fast-downgrade.test.ts @@ -1,4 +1,4 @@ -import { afterEach, beforeEach, describe, expect, test } from "bun:test"; +import { afterEach, beforeEach, describe, expect, spyOn, test } from "bun:test"; import { existsSync, mkdtempSync, rmSync } from "node:fs"; import { tmpdir } from "node:os"; import { join } from "node:path"; @@ -15,7 +15,11 @@ import type { RequestLogContext } from "../../src/server/request-log"; import type { OcxConfig } from "../../src/types"; import { saveConfig } from "../../src/config"; import { clearKeyCooldowns } from "../../src/providers/key-failover"; +import { bindWorkflowBudgetConfig } from "../../src/lib/workflow-budget-config"; import { acquireOwnedSpendHome } from "../helpers/owned-spend-home"; +import * as proxyEnv from "../../src/lib/proxy-env"; +import { createResponsesSendBudget } from "../../src/server/responses/request-send-budget"; +import { createTranslatorBudget } from "../../src/lib/translator-budget"; const originalFetch = globalThis.fetch; const originalHome = process.env.OPENCODEX_HOME; @@ -23,6 +27,9 @@ const CREDITS = { type: "error", error: { type: "rate_limit_error", message: "Us const GENERIC = { type: "error", error: { type: "rate_limit_error", message: "Number of request tokens has exceeded your per-minute rate limit" } }; let testDir = ""; let releaseSpendHome: (() => void) | undefined; +const inactiveFetch = (() => { throw new Error("No active in-memory fast-downgrade fixture"); }) as typeof fetch; +let fixtureFetch = inactiveFetch; +let configuredFetch: ReturnType | undefined; function config(fastMode = true): OcxConfig { return { @@ -37,6 +44,7 @@ function config(fastMode = true): OcxConfig { models: ["claude-opus-5-5"], // Anthropic Fast is opt-in; this suite exercises the enabled lane. fastEnabled: true, + fetch: (input: RequestInfo | URL, init?: RequestInit) => fixtureFetch(input, init), }, }, } as OcxConfig; @@ -74,10 +82,11 @@ interface Send { function fakeUpstream(answer: (index: number) => Response): Send[] { const sends: Send[] = []; - globalThis.fetch = (async (_input: RequestInfo | URL, init?: RequestInit) => { + fixtureFetch = (async (_input: RequestInfo | URL, init?: RequestInit) => { sends.push({ body: JSON.parse(String(init?.body)) as Record, headers: new Headers(init?.headers) }); return answer(sends.length); }) as typeof fetch; + globalThis.fetch = fixtureFetch; return sends; } @@ -87,6 +96,9 @@ beforeEach(() => { testDir = mkdtempSync(join(tmpdir(), "ocx-anthropic-fast-downgrade-")); process.env.OPENCODEX_HOME = testDir; releaseSpendHome = acquireOwnedSpendHome(); + fixtureFetch = inactiveFetch; + // SOCKS dispatch may bypass global fetch, so pin the transport boundary as well. + configuredFetch = spyOn(proxyEnv, "configuredOutboundFetch").mockImplementation((input, init) => fixtureFetch(input, init)); }); function cleanupTestState(): void { @@ -94,11 +106,14 @@ function cleanupTestState(): void { releaseSpendHome?.(); } finally { releaseSpendHome = undefined; + fixtureFetch = inactiveFetch; + configuredFetch?.mockRestore(); configuredFetch = undefined; globalThis.fetch = originalFetch; if (originalHome === undefined) delete process.env.OPENCODEX_HOME; else process.env.OPENCODEX_HOME = originalHome; rmSync(testDir, { recursive: true, force: true }); resetWorkflowBudgetsForTest(); + bindWorkflowBudgetConfig(null); clearKeyCooldowns(); } } @@ -164,9 +179,8 @@ describe("Anthropic fast refusal in Responses dispatch", () => { }); test("the standard resend is admitted through the request budget; a refused reservation keeps the original refusal", async () => { - // Anthropic's initial send runs on the reset-only leg, which does not charge the shared - // ledger (pre-existing, every adapter without a transient policy). The resend is the part - // this arm owns, so a zero-send budget must refuse it and hand the real refusal back. + // The root's initial send consumes its one request slot; a refused repair reservation + // preserves the provider's refusal without admitting another physical send. const sends = fakeUpstream(() => refusal(CREDITS)); const logCtx: RequestLogContext = { model: "", provider: "" }; const rootId = "anthropic-fast-refused-repair-root"; @@ -174,16 +188,16 @@ describe("Anthropic fast refusal in Responses dispatch", () => { expect(admitted?.admitted).toBe(true); if (admitted?.admitted) admitted.lease.release(); const budget = createRequestExecutionBudget({ - maxTotalModelSends: 0, baseSendAllowance: 0, finalRecoveryAllowance: 0, + maxTotalModelSends: 1, baseSendAllowance: 1, finalRecoveryAllowance: 0, maxAlternateTargetSends: 0, maxTargetTransitions: 0, - }, "anthropic-fast-zero-send"); + }, "anthropic-fast-one-send"); const response = await handleResponses(request(rootId), config(), logCtx, { sendBudget: budget }); const body = await response.text(); expect(response.status).toBe(429); expect(body).toContain("Usage credits are required for fast mode."); expect(sends).toHaveLength(1); - expect(workflowBudgetSnapshot(rootId)?.sends).toBe(0); + expect(workflowBudgetSnapshot(rootId)?.sends).toBe(1); expect(logCtx.activeAttempt?.recoveryKinds ?? []).not.toContain("anthropic-fast-downgrade"); }); @@ -214,7 +228,7 @@ describe("Anthropic fast refusal in Responses dispatch", () => { const admitted = admitWorkflowTurn(rootId, "interactive"); expect(admitted?.admitted).toBe(true); if (admitted?.admitted) admitted.lease.release(); - chargeWorkflowSends(rootId, DEFAULT_WORKFLOW_BUDGET_POLICY.maxPhysicalSends - 1); + chargeWorkflowSends(rootId, DEFAULT_WORKFLOW_BUDGET_POLICY.maxPhysicalSends - 2); const sends = fakeUpstream(index => index === 1 ? refusal(CREDITS) : success()); const budget = createRequestExecutionBudget(undefined, "anthropic-fast-workflow-child"); const response = await handleResponses(request(rootId), config(), { model: "", provider: "" }, { sendBudget: budget }); @@ -222,7 +236,7 @@ describe("Anthropic fast refusal in Responses dispatch", () => { expect(response.status).toBe(200); expect(sends).toHaveLength(2); - expect(budget.used).toBe(1); + expect(budget.used).toBe(2); expect(workflowBudgetSnapshot(rootId)?.sends).toBe(DEFAULT_WORKFLOW_BUDGET_POLICY.maxPhysicalSends); const next = await handleResponses(request(rootId), config(), { model: "", provider: "" }); @@ -232,6 +246,138 @@ describe("Anthropic fast refusal in Responses dispatch", () => { }); }); +test("a physical workflow denial stops a reserved fast downgrade before fetch", async () => { + const sends = fakeUpstream(index => index === 1 ? refusal(CREDITS) : success()); + const budget = createRequestExecutionBudget(undefined, "fast-denied-at-dispatch", { + beforePhysicalSend: () => false, charge: () => true, refund: () => {}, + }); + const response = await handleResponses(request(), config(), { model: "", provider: "" }, { sendBudget: budget }); + expect(response.status).toBe(429); + expect(await response.text()).toContain("request_send_budget_exhausted"); + expect(sends).toHaveLength(1); + expect(budget.used).toBe(0); +}); + +test("the last prepaid fast downgrade slot is confirmed once", async () => { + const rootId = "fast-last-prepaid-root"; + const admitted = admitWorkflowTurn(rootId, "interactive"); + if (admitted?.admitted) admitted.lease.release(); + const sends = fakeUpstream(index => index === 1 ? refusal(CREDITS) : success()); + const budget = createRequestExecutionBudget({ maxTotalModelSends: 2, baseSendAllowance: 2, + finalRecoveryAllowance: 0, maxAlternateTargetSends: 1, maxTargetTransitions: 1 }); + const response = await handleResponses(request(rootId), config(), { model: "", provider: "" }, { sendBudget: budget }); + expect(response.status).toBe(200); + await response.text(); + expect(sends).toHaveLength(2); + expect(budget.used).toBe(2); + expect(workflowBudgetSnapshot(rootId)?.sends).toBe(2); +}); + +test("a fast downgrade confirms its hop once while an authorized reset replacement keeps its own allowance", async () => { + const configured = config(); + configured.providers["anthropic-apikey"]!.retryOnReset = { replacements: 1 }; + const sends = fakeUpstream(index => { + if (index === 2) throw Object.assign(new Error("synthetic reset"), { code: "ECONNRESET" }); + return index === 1 ? refusal(CREDITS) : success(); + }); + const budget = createRequestExecutionBudget({ maxTotalModelSends: 3, baseSendAllowance: 3, + finalRecoveryAllowance: 0, maxAlternateTargetSends: 1, maxTargetTransitions: 1 }); + const response = await handleResponses(request(), configured, { model: "", provider: "" }, { sendBudget: budget }); + expect(response.status).toBe(200); + await response.text(); + expect(sends).toHaveLength(3); + expect(sends.map(send => send.body.speed)).toEqual(["fast", undefined, undefined]); + expect(budget.used).toBe(3); +}); + +test("a root-only fast downgrade counts each physical send once and preserves the next allowance", async () => { + const rootId = "fast-observed-root"; + const policy = { ...DEFAULT_WORKFLOW_BUDGET_POLICY, maxPhysicalSends: 3 }; + bindWorkflowBudgetConfig({ workflowBudget: { maxPhysicalSends: 3 } }); + const admitted = admitWorkflowTurn(rootId, "interactive", policy); + if (admitted?.admitted) admitted.lease.release(); + const sends = fakeUpstream(index => index === 1 ? refusal(CREDITS) : success()); + const first = await handleResponses(request(rootId), config(), { model: "", provider: "" }); + expect(first.status).toBe(200); + await first.text(); + expect(sends).toHaveLength(2); + expect(workflowBudgetSnapshot(rootId, policy)?.sends).toBe(2); + const next = await handleResponses(request(rootId), config(), { model: "", provider: "" }); + expect(next.status).toBe(200); + await next.text(); + expect(workflowBudgetSnapshot(rootId, policy)?.sends).toBe(3); + const exhausted = await handleResponses(request(rootId), config(), { model: "", provider: "" }); + expect(exhausted.status).toBe(429); + await exhausted.text(); + expect(sends).toHaveLength(3); +}); + +for (const limit of [1, 2]) test(`a legacy request budget obeys the root ceiling before a fast replay (${limit})`, async () => { + const rootId = "fast-legacy-root"; + const policy = { ...DEFAULT_WORKFLOW_BUDGET_POLICY, maxPhysicalSends: limit }; + bindWorkflowBudgetConfig({ workflowBudget: { maxPhysicalSends: limit } }); + const admitted = admitWorkflowTurn(rootId, "interactive", policy); + if (admitted?.admitted) admitted.lease.release(); + const sends = fakeUpstream(index => index === 1 ? refusal(CREDITS) : success()); + const budget = createRequestExecutionBudget(); + const response = await handleResponses(request(rootId), config(), { model: "", provider: "" }, { sendBudget: budget }); + expect(response.status).toBe(limit === 1 ? 429 : 200); + await response.text(); + expect(sends).toHaveLength(limit); + expect(workflowBudgetSnapshot(rootId, policy)?.sends).toBe(limit); + expect(budget.used).toBe(limit); +}); + +function legacyAdmissionFixture() { + const rootId = "legacy-initial-admission-root"; + const policy = { ...DEFAULT_WORKFLOW_BUDGET_POLICY, maxPhysicalSends: 1 }; + bindWorkflowBudgetConfig({ workflowBudget: { maxPhysicalSends: 1 } }); + const admitted = admitWorkflowTurn(rootId, "interactive", policy); + if (admitted?.admitted) admitted.lease.release(); + const budget = createRequestExecutionBudget(); + const translatorBudget = createTranslatorBudget(); + const owner = createResponsesSendBudget({ req: request(rootId), logCtx: { model: "m", provider: "fixture" }, + options: { sendBudget: budget, translatorBudget } }); + if (owner instanceof Response) { translatorBudget.dispose(); throw new Error("fixture admission refused"); } + return { rootId, budget, owner, close: () => translatorBudget.dispose(), + sends: () => workflowBudgetSnapshot(rootId, policy)!.sends }; +} + +test("legacy initial admission releases an unused permit when another child exhausts the root", () => { + const fixture = legacyAdmissionFixture(); + try { + const reservation = fixture.budget.reserveDispatch({ sendClass: "initial", targetKey: "fixture/m", countedExternally: true }); + if (!reservation.allowed) throw new Error("fixture permit refused"); + chargeWorkflowSends(fixture.rootId, 1); + expect(() => fixture.owner.noteInitialDispatch(reservation.permit)).toThrow("request send budget exhausted"); + expect(fixture.budget.used).toBe(0); + expect(fixture.sends()).toBe(1); + } finally { fixture.close(); } +}); + +test("legacy initial admission cannot charge a root for a released permit", () => { + const fixture = legacyAdmissionFixture(); + try { + const reservation = fixture.budget.reserveDispatch({ sendClass: "initial", targetKey: "fixture/m", countedExternally: true }); + if (!reservation.allowed) throw new Error("fixture permit refused"); + reservation.permit.release(); + expect(() => fixture.owner.noteInitialDispatch(reservation.permit)).toThrow("request send budget exhausted"); + expect(fixture.budget.used).toBe(0); + expect(fixture.sends()).toBe(0); + } finally { fixture.close(); } +}); + +test("a closed legacy reporter cannot claim another workflow send", () => { + const fixture = legacyAdmissionFixture(); + try { + const reporter = fixture.owner.transientSendReporter(); + reporter.close!(); + expect(reporter.beforeSend!()).toBe(false); + expect(fixture.budget.used).toBe(0); + expect(fixture.sends()).toBe(0); + } finally { fixture.close(); } +}); + for (const second of ["success", "rate", "billing"] as const) test(`key pool fast downgrade to ${second} counts each physical send once`, async () => { const configured = config(); configured.providers["anthropic-apikey"]!.apiKeyPool = [{ id: "A", key: "test-token" }, { id: "B", key: "test-token-B" }]; diff --git a/tests/responses/responses-pool-401-refresh.test.ts b/tests/responses/responses-pool-401-refresh.test.ts index 75418bb96e7..b4a812cde6a 100644 --- a/tests/responses/responses-pool-401-refresh.test.ts +++ b/tests/responses/responses-pool-401-refresh.test.ts @@ -7,10 +7,14 @@ import { clearAccountNeedsReauth, clearAccountQuota, isAccountNeedsReauth, updat import { codexPoolAffinityKey } from "../../src/codex/auth-context"; import * as authContextModule from "../../src/codex/auth-context"; import * as coreAuthModule from "../../src/server/responses/core-auth"; +import * as proxyEnvModule from "../../src/lib/proxy-env"; import { prepareResponsesRequest } from "../../src/server/responses/request-prepare"; import { handleComboResponses } from "../../src/server/responses/core"; import type { ResponsesAdmissionState } from "../../src/server/responses/core-options"; import { createTranslatorBudget } from "../../src/lib/translator-budget"; +import { createRequestExecutionBudget } from "../../src/lib/request-execution-budget"; +import { admitWorkflowTurn, DEFAULT_WORKFLOW_BUDGET_POLICY, resetWorkflowBudgetsForTest, workflowBudgetSnapshot } from "../../src/lib/workflow-budget"; +import { bindWorkflowBudgetConfig } from "../../src/lib/workflow-budget-config"; import { canAcquireTransientProbe, clearPoolRecoveryState, @@ -75,6 +79,7 @@ let home = ""; let previousOcxHome: string | undefined; let previousCodexHome: string | undefined; let releaseSpendHome: (() => void) | undefined; +let configuredFetchSpy: ReturnType | undefined; function config(options: { secondAccount?: boolean } = {}): OcxConfig { return { @@ -190,7 +195,7 @@ function installHarness(options: { const backupAuths: string[] = []; const backupBodies: string[] = []; let canonicalAliasSends = 0; - globalThis.fetch = (async (input: RequestInfo | URL, init?: RequestInit) => { + const mockFetch = (async (input: RequestInfo | URL, init?: RequestInit) => { const url = new URL(input instanceof Request ? input.url : String(input)); const body = typeof init?.body === "string" ? init.body : ""; const authorization = new Headers(init?.headers).get("authorization") ?? ""; @@ -234,6 +239,13 @@ function installHarness(options: { } return Response.json({ id: "resp_replayed", object: "response", status: "completed", output: [] }); }) as typeof fetch; + globalThis.fetch = mockFetch; + // SOCKS routing can select its own transport before global fetch is reached. Intercept + // that boundary too so this memory fixture never opens a socket under inherited proxies. + configuredFetchSpy?.mockRestore(); + configuredFetchSpy = spyOn(proxyEnvModule, "configuredOutboundFetch").mockImplementation( + (input: RequestInfo | URL, init?: RequestInit) => mockFetch(input, init), + ); return { sends, refreshes, recoveryAuths, backupAuths, backupBodies, get canonicalAliasSends() { return canonicalAliasSends; } }; } @@ -349,6 +361,8 @@ afterEach(() => { // Release before home teardown to prevent Windows removal failures and a live unlinked database. releaseSpendHome?.(); releaseSpendHome = undefined; + configuredFetchSpy?.mockRestore(); + configuredFetchSpy = undefined; globalThis.fetch = originalFetch; clearCompactHandoffRoutesForTests(); clearAccountNeedsReauth(ACCOUNT_ID); @@ -1521,6 +1535,123 @@ describe("ordinary pool 401 refresh and replay (#2887)", () => { }); }); +describe("Codex account retries share the workflow send ceiling", () => { + const rootId = "codex-account-retry-root"; + const gatedModel = "gpt-daybreak-blue-latest"; + + afterEach(() => { + bindWorkflowBudgetConfig(null); + resetWorkflowBudgetsForTest(); + }); + + function admitRoot(limit: number): () => number { + bindWorkflowBudgetConfig({ workflowBudget: { maxPhysicalSends: limit } }); + const policy = { ...DEFAULT_WORKFLOW_BUDGET_POLICY, maxPhysicalSends: limit }; + const admission = admitWorkflowTurn(rootId, "interactive", policy); + if (!admission?.admitted) throw new Error("fixture root admission refused"); + admission.lease.release(); + return () => workflowBudgetSnapshot(rootId, policy)!.sends; + } + + for (const legacyBudget of [false, true]) { + for (const limit of [1, 2, 4, 8]) { + test(`gated-model retries retain the last 400 at root limit ${limit} (legacy=${legacyBudget})`, async () => { + const rootSends = admitRoot(limit); + const budget = legacyBudget ? createRequestExecutionBudget() : undefined; + const responseBody = (sendNumber: number) => ({ + detail: `The '${gatedModel}' model is not supported when using Codex with a ChatGPT account.`, + send_number: sendNumber, + }); + const harness = installHarness({ + responseForSend: (_authorization, sendNumber) => Response.json(responseBody(sendNumber), { status: 400 }), + }); + const response = await handleResponses( + request("/v1/responses", { model: gatedModel, headers: { "x-codex-parent-thread-id": rootId } }), + config(), + { model: "", provider: "" } as RequestLogContext, + { + ...(budget ? { sendBudget: budget } : {}), + resolveCodexModelEntitlements: async () => ({ + modelsByAccount: new Map([[ACCOUNT_ID, new Set([gatedModel])]]), + confirmedAccountIds: new Set([ACCOUNT_ID]), + credentialIdentities: new Map(), + }), + }, + ); + + + expect(response.status).toBe(400); + expect(await response.json()).toEqual(responseBody(limit)); + expect(harness.sends).toEqual(Array(limit).fill("Bearer rejected-access")); + expect(harness.refreshes).toEqual([]); + expect(rootSends()).toBe(limit); + if (budget && limit <= 2) expect(budget.used).toBe(1); + }); + } + + for (const limit of [1, 2]) { + test(`account failover claims one workflow send and retains a refused original (limit=${limit}, legacy=${legacyBudget})`, async () => { + const rootSends = admitRoot(limit); + const budget = legacyBudget ? createRequestExecutionBudget() : undefined; + writeWorkAndOtherAccounts(); + const cfg = config({ secondAccount: true }); + cfg.accountPoolStrategy = "fill-first"; + const originalBody = { error: { code: "usage_limit_reached", message: "fixture account quota" } }; + const harness = installHarness({ + responseForSend: authorization => authorization === "Bearer rejected-access" + ? Response.json(originalBody, { status: 429 }) + : Response.json({ id: "resp_alternate", object: "response", status: "completed", output: [] }), + }); + const response = await handleResponses( + request("/v1/responses", { headers: { "x-codex-parent-thread-id": rootId } }), + cfg, + { model: "", provider: "" } as RequestLogContext, + budget ? { sendBudget: budget } : {}, + ); + + expect(response.status).toBe(limit === 1 ? 429 : 200); + const body = await response.json(); + if (limit === 1) expect(body).toEqual(originalBody); + else expect(body).toMatchObject({ id: "resp_alternate" }); + expect(harness.sends).toEqual(limit === 1 + ? ["Bearer rejected-access"] : ["Bearer rejected-access", "Bearer other-access"]); + expect(harness.refreshes).toEqual([]); + expect(rootSends()).toBe(limit); + if (budget) expect(budget.used).toBe(limit); + }); + } + } + + test("a failed legacy account permit releases its request reservation without claiming a root send", async () => { + const rootSends = admitRoot(2); + writeWorkAndOtherAccounts(); + const cfg = config({ secondAccount: true }); + cfg.accountPoolStrategy = "fill-first"; + const originalBody = { error: { code: "usage_limit_reached", message: "fixture account quota" } }; + const harness = installHarness({ responseForSend: () => Response.json(originalBody, { status: 429 }) }); + const budget = createRequestExecutionBudget(); + const reserve = budget.reserveDispatch.bind(budget); + budget.reserveDispatch = intent => { + const decision = reserve(intent); + return decision.allowed && intent.sendClass === "account-failover" + ? { ...decision, permit: { ...decision.permit, use: () => false } } + : decision; + }; + const response = await handleResponses( + request("/v1/responses", { headers: { "x-codex-parent-thread-id": rootId } }), + cfg, + { model: "", provider: "" } as RequestLogContext, + { sendBudget: budget }, + ); + + expect(response.status).toBe(429); + expect(await response.json()).toEqual(originalBody); + expect(harness.sends).toEqual(["Bearer rejected-access"]); + expect(budget.used).toBe(1); + expect(rootSends()).toBe(1); + }); +}); + describe("stored pool 401 replay then encrypted combo recovery", () => { const assignment = "RECOVERED-POOL-PLAINTEXT-SENTINEL"; diff --git a/tests/responses/responses-spend-contract-c.test.ts b/tests/responses/responses-spend-contract-c.test.ts index fe69f74636b..3c927158ee9 100644 --- a/tests/responses/responses-spend-contract-c.test.ts +++ b/tests/responses/responses-spend-contract-c.test.ts @@ -1,4 +1,4 @@ -import { expect, test } from "bun:test"; +import { expect, spyOn, test } from "bun:test"; import { createHash } from "node:crypto"; import { mkdtempSync, writeFileSync } from "node:fs"; import { tmpdir } from "node:os"; @@ -6,7 +6,7 @@ import { join } from "node:path"; import { createSpendReservationLedger, configureSharedSpendLedger, sharedSpendLedger, DEFAULT_SPEND_RESERVATION_POLICY, type SpendJournal, type SpendReservationPolicy } from "../../src/lib/spend-reservation-ledger"; import { createRequestSpendTracker } from "../../src/server/responses/request-spend"; import { createRequestExecutionBudget, claimDispatchSpendProof, createPhysicalSendReporter } from "../../src/lib/request-execution-budget"; -import { workflowSpendCeilingReached } from "../../src/lib/workflow-budget"; +import { resetWorkflowBudgetsForTest, workflowBudgetSnapshot, workflowSpendCeilingReached } from "../../src/lib/workflow-budget"; import { parseRequest } from "../../src/responses/parser"; import { routeModel } from "../../src/router"; import { applyFinalRouteRequestNormalization } from "../../src/server/responses/core-normalize"; @@ -14,9 +14,19 @@ import { handleNativeMessages } from "../../src/server/messages-native"; import { executeComboResponses } from "../../src/server/responses/core-combo"; import { clearComboSelectionState, clearComboTargetCooldowns } from "../../src/combos"; import { createTranslatorBudget } from "../../src/lib/translator-budget"; +import { saveConfig } from "../../src/config"; +import { flushConfigDirHardeningForTests } from "../../src/config/paths"; +import { flushNativeMainStartupReleases } from "../../src/codex/native-profile-startup"; +import { flushWindowsSecretAclReapsBeforeRemoval } from "../../src/lib/windows-secret-acl"; +import * as proxyEnv from "../../src/lib/proxy-env"; +import { clearKeyCooldowns } from "../../src/providers/key-failover"; +import { resetProviderRequestPacingForTest } from "../../src/providers/request-pacing"; +import { startServer } from "../../src/server"; +import { getRequestLogEntries } from "../../src/server/request-log"; import type { RequestLogContext } from "../../src/server/request-log"; import type { OcxConfig } from "../../src/types"; import { acquireOwnedSpendHome } from "../helpers/owned-spend-home"; +import { installIsolatedCodexHome } from "../helpers/isolated-codex-home"; import { removeTreeWithRetry } from "../helpers/remove-tree"; const salt = "6".repeat(64); @@ -38,6 +48,79 @@ const reserve = (ledger: ReturnType, sendId const log = (provider = "P"): RequestLogContext => ({ model: "model", provider, spendPoolId: provider, usageLogInputTokens: 10, spendOutputCeilingTokens: 0 }); +for (const enforced of [false, true]) for (const status of [503, 429]) { + test(`native Messages workflow ceiling counts each ${status === 503 ? "503 request" : "429 retry"} once (spend=${enforced})`, async () => { + const home = mkdtempSync(join(tmpdir(), "ocx-messages-workflow-")); + const priorHome = process.env.OPENCODEX_HOME; + const codexHome = installIsolatedCodexHome("ocx-messages-workflow-codex-"); + const loopbackFetch = proxyEnv.configuredOutboundFetch; + process.env.OPENCODEX_HOME = home; + resetWorkflowBudgetsForTest(); clearKeyCooldowns(); + let sends = 0; + const upstream = Bun.serve({ hostname: "127.0.0.1", port: 0, fetch(req) { + expect(new URL(req.url).pathname).toBe("/v1/messages"); + sends++; + return Response.json({ type: "error", error: { type: status === 429 ? "rate_limit_error" : "api_error", message: "temporary fixture refusal" } }, + { status, headers: { "retry-after": "0" } }); + } }); + const upstreamOrigin = upstream.url.origin; + // Pin only this fixture's outbound transport; inherited proxies never receive its requests. + const outbound = spyOn(proxyEnv, "configuredOutboundFetch").mockImplementation((input, init) => { + const url = new URL(input instanceof Request ? input.url : String(input)); + if (url.origin !== upstreamOrigin) throw new Error("unexpected fixture outbound destination"); + return loopbackFetch(input, { ...init, proxy: false } as proxyEnv.ProxyCapableRequestInit); + }); + let server: ReturnType | undefined; + try { + const config: OcxConfig = { port: 0, defaultProvider: "fixture", providers: { fixture: { + adapter: "anthropic", authMode: "key", apiKey: "fixture-key", + baseUrl: new URL("/v1", upstream.url).href, allowPrivateNetwork: true, + liveModels: false, models: ["fixture-model"], + ...(status === 429 ? { retryOn429: { attempts: 3, respectRetryAfter: true } } : {}), + } }, protocols: { rollout: { managedMessagesNative: true } }, workflowBudget: { maxPhysicalSends: 2 }, + ...(enforced ? { spend: { pool: { maxTokens: 100_000 } } } : {}) }; + saveConfig(config); + server = startServer(0); + // Messages does not opt into 5xx retries: repeated requests share the root instead. + // Its explicit same-key 429 policy exercises the limit within one native request. + for (let request = 0; request < (status === 503 ? 3 : 1); request++) { + const response = await loopbackFetch(new URL("/v1/messages", server.url), { + method: "POST", headers: { "content-type": "application/json", "x-codex-parent-thread-id": "messages-count-root" }, + body: JSON.stringify({ model: "fixture/fixture-model", messages: [{ role: "user", content: "fixture" }], max_tokens: 1 }), + proxy: false, + } as proxyEnv.ProxyCapableRequestInit); + const body = await response.text(); + if (status === 503 && request < 2) { + // Native Messages normalizes upstream 503 to its 529 overloaded-error wire shape. + expect(response.status, body).toBe(529); + expect(workflowBudgetSnapshot("messages-count-root")?.sends).toBe(request + 1); + } else { + expect(response.status, body).toBe(429); + expect(response.headers.get("x-opencodex-local-refusal")).toBe("workflow_sends_exhausted"); + } + } + expect(sends).toBe(2); + expect(workflowBudgetSnapshot("messages-count-root")?.sends).toBe(2); + expect(getRequestLogEntries().findLast(row => row.inboundProtocol === "messages" && row.protocolTrace?.mode === "native")?.protocolTrace) + .toMatchObject({ mode: "native", requestPath: ["messages", "messages"] }); + } finally { + let drained = false; + try { + try { await server?.stop(true); } finally { await upstream.stop(true); } + await flushNativeMainStartupReleases(); await flushConfigDirHardeningForTests(); + await flushWindowsSecretAclReapsBeforeRemoval(home); + await flushWindowsSecretAclReapsBeforeRemoval(codexHome.path); + drained = true; + } finally { + outbound.mockRestore(); resetWorkflowBudgetsForTest(); clearKeyCooldowns(); resetProviderRequestPacingForTest(); + if (priorHome === undefined) delete process.env.OPENCODEX_HOME; else process.env.OPENCODEX_HOME = priorHome; + codexHome.restoreEnvironment(); + if (drained) { codexHome.removeAfterDrain(); removeTreeWithRetry(home); } + } + } + }, 10_000); +} + test("Responses A B and Messages share canonical ceiling", async () => { const home = mkdtempSync(join(tmpdir(), "ocx-contract-c-messages-")); const priorHome = process.env.OPENCODEX_HOME; diff --git a/tests/responses/responses-workflow-send-ceiling.test.ts b/tests/responses/responses-workflow-send-ceiling.test.ts new file mode 100644 index 00000000000..5451ed0c791 --- /dev/null +++ b/tests/responses/responses-workflow-send-ceiling.test.ts @@ -0,0 +1,131 @@ +import { afterEach, expect, test } from "bun:test"; +import { createPhysicalSendReporter, createRequestExecutionBudget, reportDispatchSends } from "../../src/lib/request-execution-budget"; +import { createSpendReservationLedger, DEFAULT_SPEND_RESERVATION_POLICY } from "../../src/lib/spend-reservation-ledger"; +import { admitWorkflowTurn, DEFAULT_WORKFLOW_BUDGET_POLICY, resetWorkflowBudgetsForTest, workflowBudgetSnapshot } from "../../src/lib/workflow-budget"; +import { bindWorkflowBudgetConfig } from "../../src/lib/workflow-budget-config"; +import { fetchWithTransientRetry } from "../../src/lib/upstream-retry"; +import { createRequestSpendTracker } from "../../src/server/responses/request-spend"; + +afterEach(() => { bindWorkflowBudgetConfig(null); resetWorkflowBudgetsForTest(); }); + +function root(limit: number) { + bindWorkflowBudgetConfig({ workflowBudget: { maxPhysicalSends: limit } }); + const policy = { ...DEFAULT_WORKFLOW_BUDGET_POLICY, maxPhysicalSends: limit }; + const admission = admitWorkflowTurn("physical-root", "interactive", policy); + if (!admission?.admitted) throw new Error("fixture admission refused"); + admission.lease.release(); + return () => workflowBudgetSnapshot("physical-root", policy)!.sends; +} + +function request(enforced: boolean) { + const lines: string[] = []; + const ledger = createSpendReservationLedger({ + journal: { read: () => [...lines], append: line => { lines.push(line); }, rewrite: next => { lines.splice(0, lines.length, ...next); } }, + policy: { ...DEFAULT_SPEND_RESERVATION_POLICY, root: enforced ? { maxTokens: 1_000_000 } : {} }, + }); + const tracker = createRequestSpendTracker({ provider: "fixture", spendInputEstimateTokens: 1 }, "physical-root", ledger); + const budget = createRequestExecutionBudget(undefined, undefined, tracker); + return { tracker, budget }; +} + +test("a prepaid physical report cannot hide an unrelated legacy send", () => { + const sends = root(3); + const { tracker, budget } = request(false); + const prepaid = budget.reserveDispatch({ sendClass: "initial", targetKey: "fixture/m", countedExternally: true }); + if (!prepaid.allowed) throw new Error("fixture reservation refused"); + const reporter = createPhysicalSendReporter(budget, () => ({ poolId: "fixture" }), prepaid.permit); + expect(reporter.beforeSend!()).toBe(true); + reporter(1); reporter.close!(); + expect(sends()).toBe(1); + budget.used += 1; + expect(sends()).toBe(2); + expect(budget.used).toBe(2); + tracker.settle(undefined); +}); + +test("legacy prepaid reports record the first physical send as well as later sends", () => { + const sends = root(3); + const { tracker, budget } = request(false); + const prepaid = budget.reserveDispatch({ sendClass: "initial", targetKey: "fixture/m", countedExternally: true }); + if (!prepaid.allowed) throw new Error("fixture reservation refused"); + reportDispatchSends(budget, 1, prepaid.permit); + expect(sends()).toBe(1); + reportDispatchSends(budget, 1, prepaid.permit); + expect(sends()).toBe(2); + expect(budget.used).toBe(2); + tracker.settle(undefined); +}); + +for (const enforced of [false, true]) { + test(`the retry helper stops before a second upstream call (spend=${enforced})`, async () => { + const sends = root(1); + const { tracker, budget } = request(enforced); + const reporter = createPhysicalSendReporter(budget, () => ({ poolId: "fixture" })); + let calls = 0; + const result = await fetchWithTransientRetry(async () => { + calls++; + return new Response("busy", { status: 502, headers: { "retry-after": "0" } }); + }, { attempts: 3, onSendsConsumed: reporter }).catch(error => error); + expect(result).toBeInstanceOf(Error); + expect(calls).toBe(1); + expect(sends()).toBe(1); + expect(tracker.refusals).toBe(1); + tracker.settle(undefined); + }); + + test(`every reporter start consumes the workflow ceiling before any report (spend=${enforced})`, () => { + const sends = root(2); + const { tracker, budget } = request(enforced); + const reporter = createPhysicalSendReporter(budget, () => ({ poolId: "fixture" })); + expect(reporter.beforeSend!()).toBe(true); + expect(reporter.beforeSend!()).toBe(true); + expect(sends()).toBe(2); + expect(reporter.beforeSend!()).toBe(false); + reporter(2); reporter.close!(); + expect(sends()).toBe(2); + expect(tracker.refusals).toBe(1); + tracker.settle(undefined); + }); + + test(`concurrent request reporters share one root allowance (spend=${enforced})`, () => { + const sends = root(1); + const first = request(enforced), second = request(enforced); + const a = createPhysicalSendReporter(first.budget, () => ({ poolId: "fixture" })); + const b = createPhysicalSendReporter(second.budget, () => ({ poolId: "fixture" })); + expect(a.beforeSend!()).toBe(true); + expect(b.beforeSend!()).toBe(false); + expect(sends()).toBe(1); + a(1); a.close!(); b.close!(); + first.tracker.settle(undefined); second.tracker.settle(undefined); + expect(sends()).toBe(1); + }); + + test(`unused dispatch permits cost no physical sends (spend=${enforced})`, () => { + const sends = root(1); + const { tracker, budget } = request(enforced); + const abandoned = budget.reserveDispatch({ sendClass: "initial", targetKey: "fixture/m" }); + if (!abandoned.allowed) throw new Error("fixture reservation refused"); + expect(sends()).toBe(0); + abandoned.permit.release(); abandoned.permit.release(); + const sent = budget.reserveDispatch({ sendClass: "initial", targetKey: "fixture/m" }); + if (!sent.allowed) throw new Error("fixture reservation refused"); + expect(sent.permit.use()).toBe(true); + sent.permit.release(); + expect(sends()).toBe(1); + tracker.settle(undefined); + }); + + test(`a prepaid reporter is counted once and cleanup cannot refund its start (spend=${enforced})`, () => { + const sends = root(1); + const { tracker, budget } = request(enforced); + const prepaid = budget.reserveDispatch({ sendClass: "initial", targetKey: "fixture/m", countedExternally: true }); + if (!prepaid.allowed) throw new Error("fixture reservation refused"); + const reporter = createPhysicalSendReporter(budget, () => ({ poolId: "fixture" }), prepaid.permit); + expect(reporter.beforeSend!()).toBe(true); + reporter(1); reporter.close!(); prepaid.permit.release(); + expect(sends()).toBe(1); + const next = createPhysicalSendReporter(budget, () => ({ poolId: "fixture" })); + expect(next.beforeSend!()).toBe(false); + next.close!(); tracker.settle(undefined); + }); +} diff --git a/tests/server/server-antigravity-model-roster-routing.test.ts b/tests/server/server-antigravity-model-roster-routing.test.ts index d190bf9b6f5..d78448db30f 100644 --- a/tests/server/server-antigravity-model-roster-routing.test.ts +++ b/tests/server/server-antigravity-model-roster-routing.test.ts @@ -22,6 +22,8 @@ import { removeTreeWithRetry } from "../helpers/remove-tree"; import * as retry from "../../src/lib/upstream-retry"; import * as failover from "../../src/oauth/generic-account-failover"; import { createRequestExecutionBudget } from "../../src/lib/request-execution-budget"; +import { recoverAntigravityAccount } from "../../src/server/responses/antigravity-model-denial-recovery"; +import * as proxyEnv from "../../src/lib/proxy-env"; // #7500: an Antigravity request reaches only accounts whose own model discovery may serve it. const provider = "google-antigravity"; @@ -32,6 +34,9 @@ let previousHome: string | undefined; let codexHome: IsolatedCodexHome; let release: () => void; let originalFetch: typeof fetch; +const inactiveFetch = (() => { throw new Error("No active in-memory Antigravity fixture"); }) as typeof fetch; +let providerFetch = inactiveFetch; +let configuredFetch: ReturnType; let sleep: ReturnType; let ids: string[]; let sends: string[]; @@ -55,7 +60,7 @@ beforeEach(async () => { await setActiveAccount(provider, ids[0]!); sends = []; onSend = undefined; - globalThis.fetch = (async (input, init) => { + providerFetch = (async (input, init) => { const url = new URL(input instanceof Request ? input.url : String(input)); if (url.href === "https://oauth2.googleapis.com/token") return Response.json({ access_token: "synthetic-refreshed", expires_in: 3600, token_type: "Bearer" }); @@ -68,9 +73,13 @@ beforeEach(async () => { usageMetadata: { promptTokenCount: 5, candidatesTokenCount: 3, totalTokenCount: 8 }, } }); }) as typeof fetch; + globalThis.fetch = providerFetch; + configuredFetch = spyOn(proxyEnv, "configuredOutboundFetch").mockImplementation((input, init) => providerFetch(input, init)); }); afterEach(() => { + providerFetch = inactiveFetch; + configuredFetch.mockRestore(); globalThis.fetch = originalFetch; sleep.mockRestore(); release(); @@ -86,6 +95,7 @@ afterEach(() => { function config(enabled?: boolean): OcxConfig { const cfg = { defaultProvider: provider, providers: { [provider]: { adapter: "google", authMode: "oauth", baseUrl: origin, googleMode: "cloud-code-assist", models: [model], + fetch: (input: RequestInfo | URL, init?: RequestInit) => providerFetch(input, init), ...(enabled === undefined ? {} : { oauthAccountFailover: { enabled } }), } } } as unknown as OcxConfig; saveConfig(cfg); @@ -138,6 +148,7 @@ test("an independent request retries a discovery-owned refresh cancelled before if (url.pathname.endsWith(":generateContent")) projects.push(JSON.parse(String(init?.body)).project); return fixtureFetch(input, init); }) as typeof fetch; + providerFetch = globalThis.fetch; const request = turn(cfg); try { await atJoin; @@ -325,3 +336,23 @@ test("replacement admission failure is reported separately from no eligible acco expect(result.body.error?.message).not.toContain("No eligible stored"); } finally { failure?.mockRestore(); } }); + +for (const allowed of [false, true]) test(`generic Antigravity recovery obeys physical permit admission (${allowed})`, async () => { + config(true); + const budget = createRequestExecutionBudget(undefined, undefined, { + beforePhysicalSend: () => allowed, charge: () => true, refund: () => {}, + }); + let calls = 0, completed = false; + const result = await recoverAntigravityAccount({ + nextId: ids[1]!, failedResponse: new Response("model missing", { status: 404 }), + reserve: () => budget.reserveDispatch({ sendClass: "auth-recovery", targetKey: "antigravity/recovery" }), + applyFailoverSnapshot: async snapshot => snapshot, + prepare: () => {}, adapterOwnsDispatch: false, setPermit: () => {}, + refetch: async onDispatch => { onDispatch(); calls++; return new Response("recovered"); }, + complete: () => { completed = true; }, + }); + expect(calls).toBe(allowed ? 1 : 0); + expect(completed).toBe(allowed); + expect(budget.used).toBe(allowed ? 1 : 0); + expect(result?.status ?? null).toBe(allowed ? 200 : null); +});