From abf4e6968ff8754e50f87a42e6b65b4cc0a3e253 Mon Sep 17 00:00:00 2001 From: RHODIZ IT Date: Mon, 14 Sep 2026 15:08:39 -0500 Subject: [PATCH] fix(combos): preserve declared failover targets under send budget Co-authored-by: RHODIZ IT --- src/lib/request-execution-budget.ts | 64 +++++++++++++------ src/server/responses/core-combo.ts | 23 +++---- tests/lib/execution-budget-permits.test.ts | 24 +++++++ .../responses-send-budget-counts.test.ts | 52 +++++++-------- 4 files changed, 98 insertions(+), 65 deletions(-) diff --git a/src/lib/request-execution-budget.ts b/src/lib/request-execution-budget.ts index 80654b0a945..e1363cbce44 100644 --- a/src/lib/request-execution-budget.ts +++ b/src/lib/request-execution-budget.ts @@ -134,33 +134,37 @@ const RESERVE_FUNDED_CLASSES: ReadonlySet = new Set([ let logicalRequestSeq = 0; -export function createRequestExecutionBudget( - policy: RequestExecutionBudgetPolicy = CODEX_TEXT_GUARDED_BUDGET_POLICY, - logicalRequestId?: string, +interface SharedSendCounter { + spent: number; + pendingExternalSends: number; +} + +const sharedSendCounters = new WeakMap(); + +function createRequestExecutionBudgetWithCounter( + policy: RequestExecutionBudgetPolicy, + logicalRequestId: string | undefined, + counter: SharedSendCounter, ): RequestExecutionBudget { - let spent = 0; - // Reservations whose physical send is reported by a retry helper rather than by the permit. - // They are already charged; the reporter's first send settles one instead of charging again. - let pendingExternalSends = 0; let reserveSpent = false; let alternateTargetSends = 0; let targetTransitions = 0; let lastTargetKey: string | undefined; const budget: RequestExecutionBudget = { - get used(): number { return spent; }, + get used(): number { return counter.spent; }, set used(next: number) { // The retry helpers report their real send count by assigning through this field. A // reservation taken with `countedExternally` has already booked one of those sends, so // the report settles the pending booking first and only the surplus is charged. - const delta = next - spent; + const delta = next - counter.spent; if (delta <= 0) { - spent = Math.max(0, next); + counter.spent = Math.max(0, next); return; } - const settled = Math.min(delta, pendingExternalSends); - pendingExternalSends -= settled; - spent += delta - settled; + const settled = Math.min(delta, counter.pendingExternalSends); + counter.pendingExternalSends -= settled; + counter.spent += delta - settled; }, logicalRequestId: logicalRequestId ?? `lr-${Date.now().toString(36)}-${(logicalRequestSeq += 1).toString(36)}`, policyVersion: REQUEST_BUDGET_POLICY_VERSION, @@ -171,11 +175,11 @@ export function createRequestExecutionBudget( get lastTargetKey() { return lastTargetKey; }, remainingBaseSends(cap: number): number { const capped = Number.isFinite(cap) ? Math.trunc(cap) : 0; - return Math.max(0, Math.min(capped, policy.baseSendAllowance - spent)); + return Math.max(0, Math.min(capped, policy.baseSendAllowance - counter.spent)); }, reserveDispatch(intent: DispatchIntent): DispatchDecision { if (intent.replaySafe === false) return { allowed: false, reason: "not-replay-safe" }; - if (spent >= policy.maxTotalModelSends) return { allowed: false, reason: "total-exhausted" }; + if (counter.spent >= policy.maxTotalModelSends) return { allowed: false, reason: "total-exhausted" }; const changesTarget = lastTargetKey !== undefined && lastTargetKey !== intent.targetKey; const isAlternateTarget = changesTarget || intent.sendClass === "account-failover" @@ -190,7 +194,7 @@ export function createRequestExecutionBudget( // The base allowance is spent first. Only once it is gone does a recovery class reach // for the single shared reserve -- an account move and a validated rebuild cannot each // take one. - const drawsReserve = policy.baseSendAllowance - spent <= 0; + const drawsReserve = policy.baseSendAllowance - counter.spent <= 0; if (drawsReserve) { if (!RESERVE_FUNDED_CLASSES.has(intent.sendClass)) { return { allowed: false, reason: "base-allowance-exhausted" }; @@ -205,8 +209,8 @@ export function createRequestExecutionBudget( // one remaining send admitted two physical sends, which is the per-request multiplication // this budget exists to stop. Everything is booked now; `release()` is the way back. const previousTargetKey = lastTargetKey; - spent += 1; - if (intent.countedExternally === true) pendingExternalSends += 1; + counter.spent += 1; + if (intent.countedExternally === true) counter.pendingExternalSends += 1; if (drawsReserve) reserveSpent = true; if (isAlternateTarget) alternateTargetSends += 1; if (changesTarget) targetTransitions += 1; @@ -228,10 +232,10 @@ export function createRequestExecutionBudget( // An externally counted reservation the reporter already settled paid for a send // that physically happened. Refunding it would hand the request a free send back. if (intent.countedExternally === true) { - if (pendingExternalSends === 0) return; - pendingExternalSends -= 1; + if (counter.pendingExternalSends === 0) return; + counter.pendingExternalSends -= 1; } - spent -= 1; + counter.spent -= 1; if (drawsReserve) reserveSpent = false; if (isAlternateTarget) alternateTargetSends -= 1; if (changesTarget) targetTransitions -= 1; @@ -241,9 +245,27 @@ export function createRequestExecutionBudget( }; }, }; + sharedSendCounters.set(budget, counter); return budget; } +export function createRequestExecutionBudget( + policy: RequestExecutionBudgetPolicy = CODEX_TEXT_GUARDED_BUDGET_POLICY, + logicalRequestId?: string, +): RequestExecutionBudget { + return createRequestExecutionBudgetWithCounter(policy, logicalRequestId, { spent: 0, pendingExternalSends: 0 }); +} + +/** Derive a policy scope that shares the exact physical-send ledger with its parent. */ +export function deriveRequestExecutionBudget( + parent: RequestExecutionBudget, + policy: RequestExecutionBudgetPolicy, +): RequestExecutionBudget { + const counter = sharedSendCounters.get(parent); + if (!counter) throw new Error("request execution budget is not factory-backed"); + return createRequestExecutionBudgetWithCounter(policy, parent.logicalRequestId, counter); +} + export function isRequestExecutionBudget( value: TransientSendBudget | undefined, ): value is RequestExecutionBudget { diff --git a/src/server/responses/core-combo.ts b/src/server/responses/core-combo.ts index 43863fc2945..05beb3b396a 100644 --- a/src/server/responses/core-combo.ts +++ b/src/server/responses/core-combo.ts @@ -1,6 +1,7 @@ import { CODEX_TEXT_GUARDED_BUDGET_POLICY, createRequestExecutionBudget, + deriveRequestExecutionBudget, isRequestExecutionBudget, } from "../../lib/request-execution-budget"; import type { @@ -84,8 +85,8 @@ export const COMBO_TARGET_BASE_SENDS = CODEX_TEXT_GUARDED_BUDGET_POLICY.baseSend * combo may make are exactly the targets it declares minus the one it starts on. What stays * capped is the TOTAL: the first target's full ladder, one send for every further declared * target, and the one shared final-recovery reserve. A one-target combo reduces to the guarded - * profile exactly, and a three-target combo whose every target fails hard reaches upstream six - * times instead of the twelve #4546 measured. + * profile exactly. With three hard-failing targets the normal path makes five physical sends + * (3 + 1 + 1); a sixth is available only to one validated final-recovery leg. */ export function comboExecutionBudgetPolicy(declaredTargets: number): RequestExecutionBudgetPolicy { const targets = Math.max(1, Math.trunc(declaredTargets)); @@ -105,25 +106,17 @@ export function comboExecutionBudgetPolicy(declaredTargets: number): RequestExec /** * A budget scope that keeps its own recovery ledgers but spends the SAME request-wide counter. * - * `used` is redefined as an accessor onto the parent because the factory reads it back off this - * object -- `remainingBaseSends` and the total check both do -- so a copied number would let a - * combo target run its ladder against a stale total, which is precisely the per-layer counting - * this work exists to remove. The reserve, alternate-target and transition ledgers stay + * The budget factory binds derived scopes to one shared physical-send counter, including pending + * externally-counted reservations. The reserve, alternate-target and transition ledgers stay * per-scope on purpose: a combo target's account failover is its own recovery decision, while - * the request total still bounds every target together. + * the request total still bounds every target together. Copying only the numeric `used` value + * would re-arm each child against stale state and recreate the multiplication this fixes. */ export function deriveSendBudgetScope( parent: RequestExecutionBudget, policy: RequestExecutionBudgetPolicy, ): RequestExecutionBudget { - const scope = createRequestExecutionBudget(policy, parent.logicalRequestId); - Object.defineProperty(scope, "used", { - get: () => parent.used, - set: (value: number) => { parent.used = value; }, - enumerable: true, - configurable: true, - }); - return scope; + return deriveRequestExecutionBudget(parent, policy); } diff --git a/tests/lib/execution-budget-permits.test.ts b/tests/lib/execution-budget-permits.test.ts index 6aff76df691..33b3de1a4d3 100644 --- a/tests/lib/execution-budget-permits.test.ts +++ b/tests/lib/execution-budget-permits.test.ts @@ -3,6 +3,7 @@ import { readFileSync } from "node:fs"; import { CODEX_TEXT_GUARDED_BUDGET_POLICY, createRequestExecutionBudget, + deriveRequestExecutionBudget, type RequestExecutionBudgetPolicy, } from "../../src/lib/request-execution-budget"; @@ -110,6 +111,29 @@ describe("atomic dispatch permits", () => { expect(budget.used).toBe(3); }); + test("derived scopes share physical sends and counted-externally settlement", () => { + const parent = createRequestExecutionBudget({ + maxTotalModelSends: 8, baseSendAllowance: 7, finalRecoveryAllowance: 1, + maxAlternateTargetSends: 7, maxTargetTransitions: 7, + }); + const combo = deriveRequestExecutionBudget(parent, { + maxTotalModelSends: 8, baseSendAllowance: 7, finalRecoveryAllowance: 1, + maxAlternateTargetSends: 7, maxTargetTransitions: 7, + }); + const target = deriveRequestExecutionBudget(combo, { + maxTotalModelSends: 8, baseSendAllowance: 7, finalRecoveryAllowance: 1, + maxAlternateTargetSends: 1, maxTargetTransitions: 1, + }); + const hop = combo.reserveDispatch({ sendClass: "initial", targetKey: "provider-a/model-a", countedExternally: true }); + expect(hop.allowed).toBe(true); + expect(parent.used).toBe(1); + expect(target.used).toBe(1); + target.used += 1; + expect(parent.used).toBe(1); + target.used += 2; + expect(parent.used).toBe(3); + }); + test("an external report settles the booking, so a late release refunds nothing", () => { const budget = createRequestExecutionBudget(CODEX_TEXT_GUARDED_BUDGET_POLICY); const leg = budget.reserveDispatch({ diff --git a/tests/responses/responses-send-budget-counts.test.ts b/tests/responses/responses-send-budget-counts.test.ts index 78f5a428562..7dd762f2eed 100644 --- a/tests/responses/responses-send-budget-counts.test.ts +++ b/tests/responses/responses-send-budget-counts.test.ts @@ -120,7 +120,7 @@ describe("upstream sends per logical request", () => { expect(sendCounts(logCtx)).toEqual([3]); }); - test("a three-target combo fan-out gives every declared target a send and totals six", async () => { + test("a three-target combo preserves every target while bounding same-target retries", async () => { const upstream = alwaysFailing(502, "upstream busy"); const logCtx: RequestLogContext = { model: "", provider: "" }; @@ -128,34 +128,28 @@ describe("upstream sends per logical request", () => { expect(response.status).toBe(502); await response.text(); - // The measured shape in #4546 was twelve: four sends per target, because each child took a - // fresh full allowance. Sharing one counter alone was not the answer either -- it starved - // the later targets to zero. The first target runs its own ladder, each later target draws - // what is left, and the clamp holds back one send for every target still declared, so the - // last target is still reached. - // Asserted as the INVARIANT the derived policy guarantees rather than as a fixture count. - // An exact per-target vector pins how this harness happens to distribute the ladder, which - // is not what the layer promises and not something this branch can observe: the local suite - // is not run here, so a number guessed from reading is a number nobody checked. - const bearers = upstream.authorizations; - // Every declared target is still reached. Starving the last target is the failure mode that - // sharing one counter WITHOUT a per-target policy produces. - expect(new Set(bearers).size).toBe(3); - expect(bearers).toContain("Bearer sk-t2"); - // The first target keeps its full ladder, so the first sends are all its own. - expect(bearers[0]).toBe("Bearer sk-t0"); - // Bounded by the derived total: the first target's ladder, one send per further declared - // target, and the single shared final-recovery reserve. The measured regression in #4546 was - // twelve, four per target, because each child drew a fresh full allowance. - // The measured bound is NINE, and saying six here would be describing an intention rather - // than the code. #4546 measured twelve -- four sends per target, each child drawing a fresh - // full allowance -- so sharing one counter removes the per-target reserve and takes it to - // nine. The clamp that was meant to hold back one send for every target still declared is - // NOT yet effective; that is stated in the pull request as the open item rather than hidden - // behind an assertion that passes for the wrong reason. - expect(bearers.length).toBeLessThanOrEqual(9); - expect(bearers.length).toBeLessThan(12); - expect(bearers.length).toBeGreaterThanOrEqual(3); + expect(sendCounts(logCtx)).toEqual([3, 1, 1]); + expect(totalSends(logCtx)).toBe(5); + expect(upstream.authorizations).toEqual([ + "Bearer sk-t0", "Bearer sk-t0", "Bearer sk-t0", + "Bearer sk-t1", "Bearer sk-t2", + ]); + }); + + test("a thirteen-target combo reaches every declared fallback before returning failure", async () => { + const upstream = alwaysFailing(502, "upstream busy"); + const logCtx: RequestLogContext = { model: "", provider: "" }; + + const response = await handleResponses(responsesRequest("combo/fan"), comboOverTargets(13), logCtx); + + expect(response.status).toBe(502); + await response.text(); + expect(sendCounts(logCtx)).toEqual([3, ...Array.from({ length: 12 }, () => 1)]); + expect(totalSends(logCtx)).toBe(15); + expect(upstream.authorizations).toHaveLength(15); + for (let index = 0; index < 13; index++) { + expect(upstream.authorizations).toContain(`Bearer sk-t${index}`); + } }); // REMOVED: "a 401 before the 5xx streak spends one of the same three sends".