diff --git a/src/providers/request-pacing.ts b/src/providers/request-pacing.ts index e38c62f3c0..5b5f34d3f6 100644 --- a/src/providers/request-pacing.ts +++ b/src/providers/request-pacing.ts @@ -35,7 +35,7 @@ interface Waiter { modelIntervalMs: number; queuedAt: number; signal?: AbortSignal; - resolve: () => void; + start: () => void; reject: (reason: unknown) => void; abort?: () => void; } @@ -190,19 +190,29 @@ function runQueue(providerName: string, state: ProviderPacer): void { if (waiter.modelId && waiter.modelIntervalMs > 0) { state.modelNextStartAt.set(waiter.modelId, startedAt + waiter.modelIntervalMs); } - waiter.resolve(); + // Start the caller-owned transport synchronously in the granted slot. Resolving a + // waiter first would let an async continuation drift later while the next slot + // remained anchored to this earlier bookkeeping timestamp. + waiter.start(); runtime.enqueueMicrotask(() => runQueue(providerName, state)); } -export async function waitForProviderRequestSlot( +export function runWithProviderRequestSlot( providerName: string, provider: OcxProviderConfig, - modelId?: string, - signal?: AbortSignal, -): Promise { + modelId: string | undefined, + signal: AbortSignal | undefined, + start: () => T | PromiseLike, +): Promise { const intervals = requestPacingIntervals(provider, modelId); - if (Math.max(intervals.providerIntervalMs, intervals.modelIntervalMs) <= 0) return; - if (signal?.aborted) throw abortReason(signal); + if (signal?.aborted) return Promise.reject(abortReason(signal)); + if (Math.max(intervals.providerIntervalMs, intervals.modelIntervalMs) <= 0) { + try { + return Promise.resolve(start()); + } catch (error) { + return Promise.reject(error); + } + } const state = pacers.get(providerName) ?? { queue: [], providerNextStartAt: 0, modelNextStartAt: new Map(), @@ -218,15 +228,28 @@ export async function waitForProviderRequestSlot( } runQueue(providerName, state); if (state.queue.length >= maxQueueDepth) { - throw new RequestPacingQueueOverloadError( + return Promise.reject(new RequestPacingQueueOverloadError( providerName, "queue_full", pacingRetryAfterSeconds(state, modelId, runtime.now()), - ); + )); } - await new Promise((resolve, reject) => { - const waiter: Waiter = { modelId, ...intervals, queuedAt: runtime.now(), signal, resolve, reject }; + return new Promise((resolve, reject) => { + const waiter: Waiter = { + modelId, + ...intervals, + queuedAt: runtime.now(), + signal, + reject, + start: () => { + try { + resolve(start()); + } catch (error) { + reject(error); + } + }, + }; waiter.abort = () => { const index = state.queue.indexOf(waiter); if (index >= 0) state.queue.splice(index, 1); @@ -252,6 +275,15 @@ export async function waitForProviderRequestSlot( }); } +export function waitForProviderRequestSlot( + providerName: string, + provider: OcxProviderConfig, + modelId?: string, + signal?: AbortSignal, +): Promise { + return runWithProviderRequestSlot(providerName, provider, modelId, signal, () => undefined); +} + export function providerRequestPacingStatus( providerName: string, provider: OcxProviderConfig, diff --git a/src/server/responses/fetch-helpers.ts b/src/server/responses/fetch-helpers.ts index 21652a5345..01cf3badb9 100644 --- a/src/server/responses/fetch-helpers.ts +++ b/src/server/responses/fetch-helpers.ts @@ -102,7 +102,7 @@ import { } from "../relay"; import { hasResponsesItemIdRepair, relaySseWithResponsesItemIdRepair } from "../responses-item-id-repair"; import type { EffectiveSubagentRoster, SpawnAgentSurface } from "../../codex/catalog"; -import { waitForProviderRequestSlot } from "../../providers/request-pacing"; +import { runWithProviderRequestSlot, waitForProviderRequestSlot } from "../../providers/request-pacing"; export function disableResponsesRequestTimeout(req: Request, server: Pick, "timeout"> | undefined): boolean { @@ -139,6 +139,7 @@ export function safeOriginLabel(url: string): string { export interface PaceAwareFetch { waitForPacing?: (signal?: AbortSignal) => Promise; + runWithPacing?: (signal: AbortSignal | undefined, start: () => T | PromiseLike) => Promise; unpacedFetch?: typeof globalThis.fetch; } @@ -199,16 +200,26 @@ export function providerFetch( const waitForPacing = (signal?: AbortSignal) => options.providerName ? waitForProviderRequestSlot(options.providerName, provider, options.modelId, signal) : Promise.resolve(); - const wrapped = async (input: Parameters[0], init?: RequestInit) => { - await waitForPacing(init?.signal ?? undefined); - return unpaced(input, init); + const runWithPacing = (signal: AbortSignal | undefined, start: () => T | PromiseLike): Promise => { + if (options.providerName) { + return runWithProviderRequestSlot(options.providerName, provider, options.modelId, signal, start); + } + if (signal?.aborted) return Promise.reject(signal.reason ?? new DOMException("The operation was aborted", "AbortError")); + try { + return Promise.resolve(start()); + } catch (error) { + return Promise.reject(error); + } }; + const wrapped = (input: Parameters[0], init?: RequestInit) => + runWithPacing(init?.signal ?? undefined, () => unpaced(input, init)); const preconnect = (...args: Parameters): void => { base.preconnect?.(...args); }; return Object.assign(wrapped, { preconnect, waitForPacing, + runWithPacing, unpacedFetch: Object.assign(unpaced, { preconnect }), }); } @@ -225,29 +236,34 @@ export async function fetchWithHeaderTimeout( manualRedirect = false, ): Promise { const pacing = executor as ProviderFetch; - await pacing.waitForPacing?.(abortSignal); const fetchExecutor = pacing.unpacedFetch ?? executor; - const timeout = new AbortController(); - const timer = setTimeout(() => { - if (!timeout.signal.aborted) timeout.abort(new DOMException("Timeout elapsed", "TimeoutError")); - }, timeoutMs); - const headers = new Headers(init.headers); - // Compressed SSE can be held until the decompressor has a complete block. Streaming calls - // default to identity for low-latency frame delivery, while an explicit caller choice wins. - if (preferIdentityEncoding && !headers.has("accept-encoding")) { - headers.set("accept-encoding", "identity"); - } - try { - return await fetchExecutor(url, { - ...init, - headers, - // Credential-bearing sends opt into manual redirects so a 3xx is relayed - // as a Response instead of being followed into a rejection that is - // indistinguishable from a pre-connection failure (#914). - ...(manualRedirect ? { redirect: "manual" as const } : {}), - signal: AbortSignal.any([abortSignal, timeout.signal]), - }); - } finally { - clearTimeout(timer); - } + const startFetch = async (): Promise => { + const timeout = new AbortController(); + const timer = setTimeout(() => { + if (!timeout.signal.aborted) timeout.abort(new DOMException("Timeout elapsed", "TimeoutError")); + }, timeoutMs); + const headers = new Headers(init.headers); + // Compressed SSE can be held until the decompressor has a complete block. Streaming calls + // default to identity for low-latency frame delivery, while an explicit caller choice wins. + if (preferIdentityEncoding && !headers.has("accept-encoding")) { + headers.set("accept-encoding", "identity"); + } + try { + return await fetchExecutor(url, { + ...init, + headers, + // Credential-bearing sends opt into manual redirects so a 3xx is relayed + // as a Response instead of being followed into a rejection that is + // indistinguishable from a pre-connection failure (#914). + ...(manualRedirect ? { redirect: "manual" as const } : {}), + signal: AbortSignal.any([abortSignal, timeout.signal]), + }); + } finally { + clearTimeout(timer); + } + }; + + if (pacing.runWithPacing) return pacing.runWithPacing(abortSignal, startFetch); + await pacing.waitForPacing?.(abortSignal); + return startFetch(); } diff --git a/tests/request-pacing.test.ts b/tests/request-pacing.test.ts index 4a2dedc60a..23157270aa 100644 --- a/tests/request-pacing.test.ts +++ b/tests/request-pacing.test.ts @@ -84,10 +84,12 @@ describe("requestPacingIntervalMs", () => { }); describe("provider request pacing queue", () => { - test("spaces concurrent starts in one provider FIFO and exposes queue state", async () => { + test("anchors concurrent pacing to actual transport starts and exposes queue state", async () => { + const clock = fakePacingClock(); + setProviderRequestPacingRuntimeForTest(clock.runtime); const starts: number[] = []; const fetchImpl = Object.assign(async () => { - starts.push(Date.now()); + starts.push(clock.now()); return new Response("ok"); }, { preconnect() {} }) as typeof globalThis.fetch; const configured = { @@ -96,12 +98,17 @@ describe("provider request pacing queue", () => { } as OcxProviderConfig & { fetch: typeof globalThis.fetch }; const send = providerFetch(configured, undefined, { providerName: "demo", modelId: "model-a" }); const pending = [send("https://example.test/v1/chat/completions"), send("https://example.test/v1/chat/completions"), send("https://example.test/v1/chat/completions")]; - await Bun.sleep(10); + + expect(starts).toEqual([0]); expect(providerRequestPacingStatus("demo", configured).queued).toBe(2); + clock.advanceBy(99); + expect(starts).toEqual([0]); + clock.advanceBy(1); + expect(starts).toEqual([0, 100]); + clock.advanceBy(100); + expect(starts).toEqual([0, 100, 200]); + await Promise.all(pending); - expect(starts).toHaveLength(3); - expect(starts[1] - starts[0]).toBeGreaterThanOrEqual(85); - expect(starts[2] - starts[1]).toBeGreaterThanOrEqual(85); const status = providerRequestPacingStatus("demo", configured); expect(status.queued).toBe(0); expect(status.lastModelId).toBe("model-a");