Skip to content
Closed
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
56 changes: 44 additions & 12 deletions src/providers/request-pacing.ts
Original file line number Diff line number Diff line change
Expand Up @@ -35,7 +35,7 @@ interface Waiter {
modelIntervalMs: number;
queuedAt: number;
signal?: AbortSignal;
resolve: () => void;
start: () => void;
reject: (reason: unknown) => void;
abort?: () => void;
}
Expand Down Expand Up @@ -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<T>(
providerName: string,
provider: OcxProviderConfig,
modelId?: string,
signal?: AbortSignal,
): Promise<void> {
modelId: string | undefined,
signal: AbortSignal | undefined,
start: () => T | PromiseLike<T>,
): Promise<T> {
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<string, number>(),
Expand All @@ -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<void>((resolve, reject) => {
const waiter: Waiter = { modelId, ...intervals, queuedAt: runtime.now(), signal, resolve, reject };
return new Promise<T>((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);
Expand All @@ -252,6 +275,15 @@ export async function waitForProviderRequestSlot(
});
}

export function waitForProviderRequestSlot(
providerName: string,
provider: OcxProviderConfig,
modelId?: string,
signal?: AbortSignal,
): Promise<void> {
return runWithProviderRequestSlot(providerName, provider, modelId, signal, () => undefined);
}

export function providerRequestPacingStatus(
providerName: string,
provider: OcxProviderConfig,
Expand Down
72 changes: 44 additions & 28 deletions src/server/responses/fetch-helpers.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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<Server<WsData>, "timeout"> | undefined): boolean {
Expand Down Expand Up @@ -139,6 +139,7 @@ export function safeOriginLabel(url: string): string {

export interface PaceAwareFetch {
waitForPacing?: (signal?: AbortSignal) => Promise<void>;
runWithPacing?: <T>(signal: AbortSignal | undefined, start: () => T | PromiseLike<T>) => Promise<T>;
unpacedFetch?: typeof globalThis.fetch;
}

Expand Down Expand Up @@ -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<typeof globalThis.fetch>[0], init?: RequestInit) => {
await waitForPacing(init?.signal ?? undefined);
return unpaced(input, init);
const runWithPacing = <T>(signal: AbortSignal | undefined, start: () => T | PromiseLike<T>): Promise<T> => {
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<typeof globalThis.fetch>[0], init?: RequestInit) =>
runWithPacing(init?.signal ?? undefined, () => unpaced(input, init));
const preconnect = (...args: Parameters<typeof globalThis.fetch.preconnect>): void => {
base.preconnect?.(...args);
};
return Object.assign(wrapped, {
preconnect,
waitForPacing,
runWithPacing,
unpacedFetch: Object.assign(unpaced, { preconnect }),
});
}
Expand All @@ -225,29 +236,34 @@ export async function fetchWithHeaderTimeout(
manualRedirect = false,
): Promise<Response> {
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<Response> => {
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();
}
19 changes: 13 additions & 6 deletions tests/request-pacing.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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 = {
Expand All @@ -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");
Expand Down
Loading