From cf80887e939d031154d2a672690cad2f958830d4 Mon Sep 17 00:00:00 2001 From: Epinephrine Date: Sat, 31 Oct 2026 07:41:00 +0000 Subject: [PATCH 1/2] fix(workflow-budget): reserve physical sends atomically --- src/lib/workflow-budget.ts | 41 +++++++++++++++++++++ src/server/responses/request-send-budget.ts | 4 -- src/server/responses/request-spend.ts | 25 ++++++++++++- tests/lib/workflow-budget.test.ts | 14 +++++++ 4 files changed, 78 insertions(+), 6 deletions(-) diff --git a/src/lib/workflow-budget.ts b/src/lib/workflow-budget.ts index 90abeaf9ab5..b6265784d88 100644 --- a/src/lib/workflow-budget.ts +++ b/src/lib/workflow-budget.ts @@ -671,6 +671,47 @@ 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. Callers must refund the reservation when + * the dispatch permit is released before use. + */ +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; +} + +/** Return an unused physical-send reservation to its root workflow. */ +export function refundWorkflowSend(rootId: string | undefined, now: number = Date.now()): void { + if (!rootId) return; + const state = roots.get(rootId); + if (!state || windowedSends(state, now) <= 0) return; + const slotMs = windowSlotMs(state.windowMs); + const oldest = Math.floor(now / slotMs) - (WORKFLOW_WINDOW_SLOTS - 1); + let newest = -1; + for (let index = 0; index < WORKFLOW_WINDOW_SLOTS; index += 1) { + if ((state.sendSlotAt[index] ?? Number.NEGATIVE_INFINITY) < oldest) continue; + if ((state.sendSlotCount[index] ?? 0) <= 0) continue; + if (newest < 0 || state.sendSlotAt[index]! > state.sendSlotAt[newest]!) newest = index; + } + if (newest < 0) return; + state.sends = Math.max(0, state.sends - 1); + state.sendSlotCount[newest]!--; + state.lastSeenMs = now; +} + /** * 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/responses/request-send-budget.ts b/src/server/responses/request-send-budget.ts index 0417a6adcf5..4f927ddc525 100644 --- a/src/server/responses/request-send-budget.ts +++ b/src/server/responses/request-send-budget.ts @@ -2,7 +2,6 @@ import { createInferenceSendBudget } from "../inference/context"; import type { ResponsesRequestContext } from "./core-options"; import { claimDispatchSpendProof, createPhysicalSendReporter, createRequestExecutionBudget, isRequestExecutionBudget, reportDispatchSends } from "../../lib/request-execution-budget"; import { - chargeWorkflowSends, workflowSendCeilingReached, workflowSpendCeilingReached, } from "../../lib/workflow-budget"; @@ -79,7 +78,6 @@ export function createResponsesSendBudget( noteTargetSends(charged); reportDispatchSends(sendBudget, charged, permit); options.onCompactionRecoverySendsReported?.(charged); - chargeWorkflowSends(workflowRootId, charged); }; const noteTransientSends = (used: number): void => recordTransientSends(used, pendingHopPermit ?? inheritedPermit); // Capture at helper creation, before another leg can replace the pending handoff. @@ -88,7 +86,6 @@ export function createResponsesSendBudget( identityId: logCtx.accountLogLabel }), permit, charged => { noteTargetSends(Math.max(0, charged)); options.onCompactionRecoverySendsReported?.(charged); - chargeWorkflowSends(workflowRootId, charged); }); // 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 @@ -336,7 +333,6 @@ export function createResponsesSendBudget( } 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; }, diff --git a/src/server/responses/request-spend.ts b/src/server/responses/request-spend.ts index ac217e6899e..7089c77466d 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 { recordWorkflowRefusalEvent, refundWorkflowSend, 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 { @@ -78,6 +79,7 @@ export function createRequestSpendTracker( let seededAccounting = false; let finalRequested = false; let finalUsage: TerminalSpendUsage | undefined; + let workflowReservations = 0; let reporters = 0; const seeds = new Map(); // Abandoned capabilities retain cleanup ownership, never admission lookup. @@ -348,9 +350,24 @@ 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 (alreadySent) { + // Reporting transports cannot be stopped retroactively, but still record the send so + // later reservations see the true total. + if (!reserveWorkflowSend(rootId, { ...currentWorkflowBudgetPolicy(), maxPhysicalSends: Number.MAX_SAFE_INTEGER })) { + throw new Error("Already-sent workflow charge could not be recorded"); + } + } else if (!reserveWorkflowSend(rootId, currentWorkflowBudgetPolicy())) { + return false; + } else if (rootId) { + workflowReservations++; + } if (!alreadySent && tracker.enforced) { const seed = tracker.ensureSeed({ targetKey: options?.targetKey }); - if (!seed) return false; + if (!seed) { + refundWorkflowSend(rootId); + if (rootId) workflowReservations--; + return false; + } options?.onReserved?.({ ledger: ledger(), sendId: seed.sendId }); return true; } @@ -404,6 +421,10 @@ export function createRequestSpendTracker( if (index >= 0) live.push(...live.splice(index, 1)); }, refund(proof): void { + if (workflowReservations > 0) { + refundWorkflowSend(rootId); + workflowReservations--; + } if (proof && proof.ledger === ledgerRef && seededAccounting) { const match = [...seeds.entries()].find(([, state]) => state.seed.sendId === proof.sendId); if (match && !match[1].started) retireSeed(match[0], match[1].seed); diff --git a/tests/lib/workflow-budget.test.ts b/tests/lib/workflow-budget.test.ts index 6a35e44e91b..bd5fff663c0 100644 --- a/tests/lib/workflow-budget.test.ts +++ b/tests/lib/workflow-budget.test.ts @@ -11,7 +11,9 @@ import { DEFAULT_WORKFLOW_BUDGET_POLICY, listTrackedWorkflowRoots, listWorkflowBudgetEvents, + refundWorkflowSend, resetWorkflowBudgetsForTest, + reserveWorkflowSend, settleWorkflowSpend, workflowBudgetSnapshot, workflowDenialSummary, @@ -63,6 +65,18 @@ describe("workflow count caps", () => { if (decision && !decision.admitted) expect(decision.reason).toBe("workflow-sends-exhausted"); }); + test("physical sends reserve atomically and unused permits refund", () => { + 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); + + refundWorkflowSend("atomic-root"); + expect(reserveWorkflowSend("atomic-root", oneSend)).toBe(true); + }); + test("a worker lane may not take the interactive reserve", () => { const workerCeiling = smallPolicy.maxConcurrentChildren - smallPolicy.interactiveReserve; for (let i = 0; i < workerCeiling; i += 1) { From 8f335415c652fc63ffa9819a56da8065cf182a9e Mon Sep 17 00:00:00 2001 From: Epinephrine Date: Wed, 4 Nov 2026 10:44:34 +0900 Subject: [PATCH 2/2] test: keep parallel transfer layout additions mergeable Keep the workflow registration away from the direct-tool-choice registration so both focused fixes merge without a layout conflict. --- scripts/test-layout/layout-responses.json | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/scripts/test-layout/layout-responses.json b/scripts/test-layout/layout-responses.json index f6a92e6ac95..46181d4e113 100644 --- a/scripts/test-layout/layout-responses.json +++ b/scripts/test-layout/layout-responses.json @@ -1,5 +1,4 @@ { - "responses-workflow-send-ceiling.test.ts": "responses", "responses-native-output-identity.test.ts": "responses", "graceful-max-output-tokens.test.ts": "responses", "graceful-max-output-tokens-exclusions.test.ts": "responses", @@ -9,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" }