diff --git a/src/lib/workflow-budget.ts b/src/lib/workflow-budget.ts index 8ca435352c..a9eaa4c579 100644 --- a/src/lib/workflow-budget.ts +++ b/src/lib/workflow-budget.ts @@ -29,10 +29,25 @@ import { export interface WorkflowBudgetPolicy { /** Children admitted concurrently under one root. */ readonly maxConcurrentChildren: number; - /** Physical model sends charged to one root across its whole life. */ + /** Physical model sends charged to one root INSIDE {@link WorkflowBudgetPolicy.windowMs}. */ readonly maxPhysicalSends: number; - /** Distinct children one root may ever create. */ + /** Distinct children one root may have inside the same window. */ readonly maxDistinctChildren: number; + /** + * The interval both counts are measured over. + * + * These were lifetime totals, and a lifetime total is the wrong instrument. The cap was + * written against a fan-out that sends once per child seven hundred times, which is a RATE; + * a running total cannot tell that from an ordinary session spread over an afternoon and + * refuses both. Because the root id is the caller thread, for Codex that made the ceiling a + * session expiry: a session reaching it was refused for the rest of the process even after + * going idle for hours, and the only cure was restarting the proxy. + * + * Omitted means {@link WORKFLOW_DEFAULT_WINDOW_MS}. A count inside a window is never larger + * than the same count over a lifetime, so windowing can only ever admit more for identical + * traffic -- no install sees a refusal it would not have seen before. + */ + readonly windowMs?: number; /** * Concurrency slots a fan-out may never take. An interactive turn arriving into a saturated * root still gets admitted; without this a worker burst starves the conversation it serves. @@ -47,14 +62,90 @@ export interface WorkflowBudgetPolicy { readonly maxTrackedRoots: number; } +/** + * Ten minutes. Long enough that the burst this ceiling was written against -- seven hundred + * sends in a minute -- is still refused several times over, and short enough that an ordinary + * session, which averages far less than a send every two seconds, never approaches it. + */ +export const WORKFLOW_DEFAULT_WINDOW_MS = 10 * 60_000; + export const DEFAULT_WORKFLOW_BUDGET_POLICY: WorkflowBudgetPolicy = { maxConcurrentChildren: 8, maxPhysicalSends: 256, maxDistinctChildren: 64, interactiveReserve: 1, maxTrackedRoots: 512, + windowMs: WORKFLOW_DEFAULT_WINDOW_MS, }; +/** Fixed ring size. Ten minutes over twelve slots gives fifty-second granularity. */ +const WORKFLOW_WINDOW_SLOTS = 12; + +function workflowWindowMs(policy: WorkflowBudgetPolicy): number { + const declared = policy.windowMs; + return declared !== undefined && Number.isFinite(declared) && declared > 0 + ? declared + : WORKFLOW_DEFAULT_WINDOW_MS; +} + +/** + * Slot size for one root's own window. + * + * The geometry is read off the state rather than off whatever policy the current caller + * happens to hold. Two callers may legitimately pass different policies for the same root -- + * the ceiling numbers are the caller's business -- but if they also disagreed about + * `windowMs`, the slot ids one of them wrote would be on a scale the other cannot read, and + * charging with a long window while reading with a short one makes every stored slot look + * ancient and the ceiling never fire at all. + */ +function windowSlotMs(windowMs: number): number { + return Math.max(1, Math.ceil(windowMs / WORKFLOW_WINDOW_SLOTS)); +} + +/** + * Add sends to the ring, resetting a slot whose turn has come round again. + * + * A ring rather than a list of timestamps because the storage has to be bounded: a root that + * sends forever would otherwise grow forever, and this ledger exists to bound a fan-out. + */ +function recordWindowedSends(state: WorkflowState, now: number, sends: number): void { + const slotMs = windowSlotMs(state.windowMs); + const slot = Math.floor(now / slotMs); + const index = ((slot % WORKFLOW_WINDOW_SLOTS) + WORKFLOW_WINDOW_SLOTS) % WORKFLOW_WINDOW_SLOTS; + if (state.sendSlotAt[index] !== slot) { + state.sendSlotAt[index] = slot; + state.sendSlotCount[index] = 0; + } + state.sendSlotCount[index] = (state.sendSlotCount[index] ?? 0) + sends; +} + +/** Sends inside the window. A slot older than the window contributes nothing. */ +function windowedSends(state: WorkflowState, now: number): number { + const slotMs = windowSlotMs(state.windowMs); + const oldest = Math.floor(now / slotMs) - (WORKFLOW_WINDOW_SLOTS - 1); + let total = 0; + for (let index = 0; index < WORKFLOW_WINDOW_SLOTS; index += 1) { + if ((state.sendSlotAt[index] ?? Number.NEGATIVE_INFINITY) >= oldest) { + total += state.sendSlotCount[index] ?? 0; + } + } + return total; +} + +/** + * Forget children last seen before the window opened, and report how many remain. + * + * Pruning on read keeps the map bounded without a timer: every admission pays for the children + * it can still see, and a root that goes quiet is cleaned up the next time it speaks. + */ +function windowedChildren(state: WorkflowState, now: number): number { + const cutoff = now - state.windowMs; + for (const [childId, lastSeenMs] of state.children) { + if (lastSeenMs <= cutoff) state.children.delete(childId); + } + return state.children.size; +} + export type WorkflowDenial = | "workflow-concurrency-exhausted" | "workflow-sends-exhausted" @@ -113,9 +204,28 @@ export interface WorkflowSpendRequest { interface WorkflowState { active: number; + /** Lifetime total, kept for diagnostics only. The ceiling reads the window instead. */ sends: number; - children: Set; + /** Ring of per-slot send counts; sendSlotAt[i] names the slot that bucket holds. */ + sendSlotCount: number[]; + sendSlotAt: number[]; + /** Child id to the last time it was admitted, so a child that stops ages out of the count. */ + children: Map; lastSeenMs: number; + /** Window this root's ring and child map are measured over, fixed when the root appeared. */ + windowMs: number; +} + +function newWorkflowState(now: number, policy: WorkflowBudgetPolicy): WorkflowState { + return { + active: 0, + sends: 0, + sendSlotCount: new Array(WORKFLOW_WINDOW_SLOTS).fill(0), + sendSlotAt: new Array(WORKFLOW_WINDOW_SLOTS).fill(Number.NEGATIVE_INFINITY), + children: new Map(), + lastSeenMs: now, + windowMs: workflowWindowMs(policy), + }; } const roots = new Map(); @@ -127,7 +237,11 @@ const roots = new Map(); * the new root regardless, so `maxTrackedRoots` bounded nothing whenever every candidate * was active or exhausted -- which is precisely the fan-out this file exists to bound. */ -function evictOneRoot(policy: WorkflowBudgetPolicy, spendLedger?: SpendReservationLedger): boolean { +function evictOneRoot( + policy: WorkflowBudgetPolicy, + spendLedger?: SpendReservationLedger, + now: number = Date.now(), +): boolean { let oldestKey: string | undefined; let oldestAt = Number.POSITIVE_INFINITY; for (const [key, state] of roots) { @@ -136,7 +250,7 @@ function evictOneRoot(policy: WorkflowBudgetPolicy, spendLedger?: SpendReservati // EXHAUSTED-but-idle root -- count-exhausted or spend-exhausted -- because recreating it // fresh under the same id resets the very ceiling that already fired. if (state.active > 0) continue; - if (state.sends >= policy.maxPhysicalSends) continue; + if (windowedSends(state, now) >= policy.maxPhysicalSends) continue; if (spendLedger?.exhausted("root", key) === true) continue; if (state.lastSeenMs < oldestAt) { oldestAt = state.lastSeenMs; oldestKey = key; } } @@ -174,22 +288,22 @@ export function admitWorkflowTurn( const ledger = spendLedger ?? (spend ? sharedSpendLedger() : undefined); let state = roots.get(rootId); if (!state) { - if (roots.size >= policy.maxTrackedRoots && !evictOneRoot(policy, ledger)) { + if (roots.size >= policy.maxTrackedRoots && !evictOneRoot(policy, ledger, now)) { // Nothing may be forgotten, so the new root is refused instead of admitted over the // bound. The alternative -- evicting an exhausted root -- resets the ceiling that // already fired, and a caller minting fresh ids would get unlimited budget from it. return { admitted: false, reason: "workflow-tracking-exhausted", rootId }; } - state = { active: 0, sends: 0, children: new Set(), lastSeenMs: now }; + state = newWorkflowState(now, policy); roots.set(rootId, state); } state.lastSeenMs = now; - if (state.sends >= policy.maxPhysicalSends) { + if (windowedSends(state, now) >= policy.maxPhysicalSends) { return { admitted: false, reason: "workflow-sends-exhausted", rootId }; } if (childId !== undefined && !state.children.has(childId) - && state.children.size >= policy.maxDistinctChildren) { + && windowedChildren(state, now) >= policy.maxDistinctChildren) { return { admitted: false, reason: "workflow-children-exhausted", rootId }; } const ceiling = lane === "worker" @@ -229,7 +343,7 @@ export function admitWorkflowTurn( } state.active += 1; - if (childId !== undefined) state.children.add(childId); + if (childId !== undefined) state.children.set(childId, now); let released = false; return { admitted: true, @@ -244,6 +358,8 @@ export function admitWorkflowTurn( const current = roots.get(rootId); if (current) { current.active = Math.max(0, current.active - 1); + // Eviction ordering only; no ceiling reads lastSeenMs, so the wall clock is the + // right source here and a caller does not need to inject one. current.lastSeenMs = Date.now(); } // Which of the two applies depends on whether the send ever left this process. @@ -263,12 +379,20 @@ export function admitWorkflowTurn( * Charge physical sends to a root. Called from the send budget's own accounting so a retry * inside one request counts toward the workflow total, not only the request total. */ -export function chargeWorkflowSends(rootId: string | undefined, sends: number): void { +export function chargeWorkflowSends( + rootId: string | undefined, + sends: number, + now: number = Date.now(), +): void { if (!rootId || sends <= 0) return; const state = roots.get(rootId); if (!state) return; state.sends += sends; - state.lastSeenMs = Date.now(); + // Geometry comes off the root itself, so no caller can charge on one scale and read on + // another. This function does not take a policy at all any more: it has no ceiling to + // compare, and the only thing a policy could have supplied here was that scale. + recordWindowedSends(state, now, sends); + state.lastSeenMs = now; } /** @@ -315,18 +439,40 @@ export function abandonWorkflowSpend(sendId: string, spendLedger?: SpendReservat export function workflowSendCeilingReached( rootId: string | undefined, policy: WorkflowBudgetPolicy = DEFAULT_WORKFLOW_BUDGET_POLICY, + now: number = Date.now(), ): boolean { if (!rootId) return false; const state = roots.get(rootId); - return state !== undefined && state.sends >= policy.maxPhysicalSends; + return state !== undefined && windowedSends(state, now) >= policy.maxPhysicalSends; } -export function workflowBudgetSnapshot(rootId: string): { - - active: number; sends: number; children: number; +export function workflowBudgetSnapshot( + rootId: string, + policy: WorkflowBudgetPolicy = DEFAULT_WORKFLOW_BUDGET_POLICY, + now: number = Date.now(), +): { + active: number; + /** Sends inside the window. This is the number the ceiling compares. */ + sends: number; + /** Children inside the window, which is likewise what the ceiling compares. */ + children: number; + /** Everything the root has ever sent, for diagnostics; no ceiling reads it. */ + lifetimeSends: number; + windowMs: number; + maxPhysicalSends: number; + maxDistinctChildren: number; } | undefined { const state = roots.get(rootId); - return state ? { active: state.active, sends: state.sends, children: state.children.size } : undefined; + if (!state) return undefined; + return { + active: state.active, + sends: windowedSends(state, now), + children: windowedChildren(state, now), + lifetimeSends: state.sends, + windowMs: state.windowMs, + maxPhysicalSends: policy.maxPhysicalSends, + maxDistinctChildren: policy.maxDistinctChildren, + }; } /** Test seam. Production never clears a live ledger: that would reset a spent budget. */ diff --git a/tests/fixtures/file-size-baseline.json b/tests/fixtures/file-size-baseline.json index dff0815d02..c410a3587e 100644 --- a/tests/fixtures/file-size-baseline.json +++ b/tests/fixtures/file-size-baseline.json @@ -25,7 +25,7 @@ "src/config.ts": 4799, "src/providers/registry.ts": 3744, "src/server/index.ts": 3400, - "src/server/responses/core.ts": 9360, + "src/server/responses/core.ts": 9387, "tests/ci-workflows/ci-workflows.test.ts": 5628, "tests/cli/cli-account.test.ts": 2313, "tests/codex-integration/codex-auth-api.test.ts": 6549, diff --git a/tests/lib/workflow-budget.test.ts b/tests/lib/workflow-budget.test.ts index f21ecc1ef9..a8fb7ce447 100644 --- a/tests/lib/workflow-budget.test.ts +++ b/tests/lib/workflow-budget.test.ts @@ -209,3 +209,166 @@ describe("workflow spend reservation", () => { if (denied && !denied.admitted) expect(denied.reason).toBe("workflow-spend-exhausted"); }); }); + +describe("root ceilings bound a rate, not a lifetime (#4546)", () => { + const WINDOW = 60_000; + const policy: WorkflowBudgetPolicy = { + ...DEFAULT_WORKFLOW_BUDGET_POLICY, + maxPhysicalSends: 4, + maxDistinctChildren: 2, + windowMs: WINDOW, + }; + + beforeEach(() => { + resetWorkflowBudgetsForTest(); + }); + + test("a root at the send ceiling is admitted again once its window rolls", () => { + const now = 1_700_000_000_000; + const first = admitWorkflowTurn("root-a", "worker", policy, undefined, now); + expect(first?.admitted).toBe(true); + first?.lease.release(); + chargeWorkflowSends("root-a", policy.maxPhysicalSends, now); + + // Inside the window the ceiling still fires: the burst this cap was written against is + // refused exactly as before. + expect(admitWorkflowTurn("root-a", "worker", policy, undefined, now + 1)?.reason) + .toBe("workflow-sends-exhausted"); + expect(workflowSendCeilingReached("root-a", policy, now + 1)).toBe(true); + + // Past the window the same root is served, with no restart. This is the case that made a + // long-lived session unusable: work it finished hours ago kept refusing it. + const rolled = admitWorkflowTurn("root-a", "worker", policy, undefined, now + WINDOW + 1); + expect(rolled?.admitted).toBe(true); + rolled?.lease.release(); + }); + + test("distinct children age out of the count the same way sends do", () => { + const now = 1_700_000_000_000; + for (const child of ["c1", "c2"]) { + const admitted = admitWorkflowTurn("root-b", "worker", policy, child, now); + expect(admitted?.admitted).toBe(true); + admitted?.lease.release(); + } + // A third distinct child inside the window is refused at the configured ceiling. + expect(admitWorkflowTurn("root-b", "worker", policy, "c3", now + 1)?.reason) + .toBe("workflow-children-exhausted"); + + // Once c1 and c2 have aged out, c3 is a new child under an empty count rather than the + // third member of a set the root can never shrink. + const later = admitWorkflowTurn("root-b", "worker", policy, "c3", now + WINDOW + 1); + expect(later?.admitted).toBe(true); + later?.lease.release(); + }); + + test("a child that keeps working holds its slot; one that stops does not", () => { + const now = 1_700_000_000_000; + for (const at of [now, now + WINDOW / 2, now + WINDOW]) { + const busy = admitWorkflowTurn("root-c", "worker", policy, "busy", at); + expect(busy?.admitted).toBe(true); + busy?.lease.release(); + } + const quiet = admitWorkflowTurn("root-c", "worker", policy, "quiet", now); + expect(quiet?.admitted).toBe(true); + quiet?.lease.release(); + + // "busy" was seen inside the window and still counts; "quiet" was not and does not, so + // there is room for exactly one more distinct child rather than none. + const snapshot = workflowBudgetSnapshot("root-c", policy, now + WINDOW + 1); + expect(snapshot?.children).toBe(1); + }); + + test("windowing never refuses traffic the lifetime count would have admitted", () => { + // The safety argument stated as a test rather than trusted as prose: a count inside a + // window is bounded by the same count over a lifetime, so for identical traffic the + // windowed ceiling fires no earlier than the lifetime one did. + // + // The root is admitted first on purpose. An earlier version of this test charged a root + // that had never been admitted, so `chargeWorkflowSends` returned at its `!state` guard, + // the snapshot came back undefined, and every assertion sat behind `if (snapshot)`. It + // would have passed with the ring deleted. + const now = 1_700_000_000_000; + const seeded = admitWorkflowTurn("root-d", "worker", policy, undefined, now); + expect(seeded?.admitted).toBe(true); + seeded?.lease.release(); + + let lifetime = 0; + let refusals = 0; + for (let i = 0; i < policy.maxPhysicalSends * 3; i += 1) { + const at = now + i * (WINDOW / 2); + chargeWorkflowSends("root-d", 1, at); + lifetime += 1; + const snapshot = workflowBudgetSnapshot("root-d", policy, at); + expect(snapshot).toBeDefined(); + expect(snapshot?.lifetimeSends).toBe(lifetime); + expect(snapshot?.sends).toBeLessThanOrEqual(lifetime); + if (workflowSendCeilingReached("root-d", policy, at)) { + refusals += 1; + expect(lifetime).toBeGreaterThanOrEqual(policy.maxPhysicalSends); + } + } + + // Spread half a window apart, this traffic is a trickle and is never refused, while the + // lifetime count passed the same ceiling three times over. That gap is the whole change. + expect(refusals).toBe(0); + expect(lifetime).toBeGreaterThan(policy.maxPhysicalSends); + }); + + test("the window a root was created with is the one its ceiling reads", () => { + // Charging on one scale and reading on another is not hypothetical: the slot ids written + // under a long window look ancient to a short one, `windowedSends` returns zero, and the + // ceiling stops firing at all. The geometry therefore belongs to the root, not to + // whichever policy the current caller happens to be holding. + const now = 1_700_000_000_000; + const seeded = admitWorkflowTurn("root-f", "worker", policy, undefined, now); + expect(seeded?.admitted).toBe(true); + seeded?.lease.release(); + chargeWorkflowSends("root-f", policy.maxPhysicalSends, now); + + const wider: WorkflowBudgetPolicy = { ...policy, windowMs: WINDOW * 100 }; + const narrower: WorkflowBudgetPolicy = { ...policy, windowMs: 1_000 }; + expect(workflowSendCeilingReached("root-f", wider, now + 1)).toBe(true); + expect(workflowSendCeilingReached("root-f", narrower, now + 1)).toBe(true); + expect(workflowBudgetSnapshot("root-f", narrower, now + 1)?.windowMs).toBe(WINDOW); + }); + + test("the snapshot separates the window from the lifetime total", () => { + const now = 1_700_000_000_000; + const admitted = admitWorkflowTurn("root-e", "worker", policy, undefined, now); + admitted?.lease.release(); + chargeWorkflowSends("root-e", 3, now); + const inside = workflowBudgetSnapshot("root-e", policy, now); + expect(inside?.sends).toBe(3); + expect(inside?.lifetimeSends).toBe(3); + expect(inside?.windowMs).toBe(WINDOW); + + const after = workflowBudgetSnapshot("root-e", policy, now + WINDOW * 2); + // The ceiling reads the window and sees nothing; the lifetime total is still reported, so + // an operator can tell an idle root from one that never worked. + expect(after?.sends).toBe(0); + expect(after?.lifetimeSends).toBe(3); + }); +}); + + +describe("every ceiling on this path reads the caller's clock", () => { + test("no function reads Date.now() except as a parameter default", async () => { + // This defect has now appeared three times in two days: codexPoolAffinityKey, then + // chargeWorkflowSends, then workflowSendCeilingReached. Each time a caller working against + // a fixed clock wrote into one window and read from another, and each time the symptom was + // a ceiling that fired when it should not have. A function that decides admission must be + // askable about a moment, so the clock is a parameter and never an ambient read. + const source = await Bun.file( + new URL("../../src/lib/workflow-budget.ts", import.meta.url), + ).text(); + const ambient = source + .split("\n") + .map((line, index) => ({ line: line.trim(), number: index + 1 })) + .filter(entry => entry.line.includes("Date.now()")) + .filter(entry => !entry.line.startsWith("now: number = Date.now()")) + .filter(entry => !entry.line.startsWith("//")) + // lastSeenMs feeds eviction ordering, not a ceiling, and its comment says so. + .filter(entry => !entry.line.includes("lastSeenMs = Date.now()")); + expect(ambient).toEqual([]); + }); +});