diff --git a/.claude/docs/oauth-continuation.md b/.claude/docs/oauth-continuation.md index 8d359760..2616cda3 100644 --- a/.claude/docs/oauth-continuation.md +++ b/.claude/docs/oauth-continuation.md @@ -27,6 +27,147 @@ explicit programmatic option outranks the environment so tests are never perturb `evictions` array on every `ws_head_decision` diagnostic — sustained `*_lru_cap` counts mean a cap is too small. +### Pacing new connections + +`src/oauth/ws-upgrade-pacer.ts`. OpenAI's edge rejects a WebSocket upgrade with HTTP 403, and in the +traffic sampled below those rejections clustered in the minutes that opened the most new +connections. The rejection is handled (see below) but the rate was previously unlimited. That the +rate is what the edge reacts to is this module's working assumption, not a demonstrated cause. A process-wide token bucket now gates **primary connection creation** — a request that +reuses an established or nursery head never consults it, so pacing can never add latency to a +continuation, and the two replacement paths below are exempt by design. + +Defaults: 60 new connections per minute sustained, burst 10. Override the rate with +`CLODEX_WS_MAX_NEW_CONNECTIONS_PER_MIN` (integer 1-600; `0` disables pacing, values above 600 clamp, +malformed values are reported once and ignored). The bucket is shared process-wide for the same +reason the pools are: the server holds a separate transport per model, so a per-transport bucket +would multiply the rate by the number of models in play. What the throttle is scoped to is not +known — one account on one machine cannot tell an account-, IP-, model- or edge-level limit +apart — so one shared bucket is the conservative reading, not a modelled one. + +**Overflow is refused, not delayed indefinitely.** A request the rate cannot serve within the wait +bound gets the same retryable 429 frame shape the upgrade 403 produces — `code: '429'` plus the +load-bearing `retry after Ns` prose — and the AI SDK backs off and retries it. Admitting anyway past +the bound was tried first and does not work: with the bound doubling as the debt floor, sustained +output settles at exactly the offered rate delayed by the bound, so an 82/min fan-out still went out +at 82/min. Refusing sheds the overflow instead, so the rate of *admissions* is capped. It is not +free: the refused request returns through the SDK's retry ladder, and that backoff runs *inside* the +same no-data deadline a queue wait spends — which is why the bound below budgets the whole ladder +rather than one wait. + +**With `CLODEX_UPSTREAM_MAX_RETRIES=0` the pacer cannot refuse** — the SDK rethrows before consulting +`shouldRetry`, so a refusal would be an immediate hard failure. In that mode it shapes the opening +burst (`burst + bound x refill`, 25 connections at the defaults) and then stops delaying anything at +all. **That is not a safety guarantee, it is limiting switched off past the floor**: sustained +traffic is unshaped, exactly as it would be with pacing disabled. Delaying every request by the +bound instead was measured to shape nothing — sustained output simply equals sustained input, +late — so it taxes the user for no benefit. The burst is kept because the burst is the part that +correlates with rejection. + +**What pacing costs.** One new connection per second is an aggregate ceiling, not a per-request +delay. By Little's law, N agents that each need a new connection per turn settle at roughly N +seconds per turn once the burst is spent: about 20s per turn at 20 agents against ~3s unpaced. The +trade is throughput for a lower chance of tripping the throttle, and it is the point of the feature +rather than a side effect. It is a reduction in risk, not a guarantee: the causal link is assumed +(see the scope note below), and a fan-out large enough to exhaust the bound is refused by the pacer +itself, which the client sees as a rate limit. + +**A refusal debits nothing.** That is what makes the retry ladder safe — a refused request opens no +connection and will be retried, so charging it a token would let each retry deepen the deficit that +caused the refusal. Because only an admitted request debits, and only within the bound, `tokens` +cannot fall below `-bound x refill`: the queue is bounded by construction and admissions in any +window stay within `burst + bound x refill + rate x T + cancellations` however many retries arrive. + +**That bounds admissions, not sockets.** Both replacement paths — a transport retry and a +`previous_response_not_found` retry — build their connection through `createReplacement`, which does +not consult the pacer, so connections opened can exceed admissions granted. A cancellation likewise +refunds its token without rescheduling the reservations queued behind it, so each one permits one +extra admission at that instant. + +**The wait bound is derived, not chosen.** Every attempt of one request shares ONE no-data deadline +(the timer starts before the SDK call and only a stream part resets it), so the whole ladder must +fit: `(maxRetries + 1) x bound + totalBackoff < idleTimeout`. At the default 120s deadline and five +retries the backoff ladder alone is 62s and the bound works out at ~4.8s; a flat 15s would instead +let six attempts plus backoff reach 152s against 120s. + +**Both terms are read, not assumed.** Every term of that inequality is user-configurable +(`CLODEX_UPSTREAM_IDLE_TIMEOUT_MS`, `CLODEX_UPSTREAM_TOTAL_TIMEOUT_MS`, +`CLODEX_UPSTREAM_MAX_RETRIES`) and they interact — a shorter deadline lowers the retry ceiling — so +the pacer resolves them together through the same `upstreamRequestBudget()` call every SDK +generation entry point makes, and sizes its bound against the deadline the paced request will +actually spend. No production caller overrides `idleTimeoutMs` on that call, so the two resolve +identically. Hardcoding either term would leave the bound correct only at the default +configuration. + +The inequality holds strictly for every resolvable configuration rather than by coincidence at one +of them: pacing takes at most half of what the ladder leaves, so `attempts x bound + backoff <= +(idle + backoff) / 2 < idle` whenever `backoff < idle`, and `upstreamRequestBudget` guarantees that +side condition by capping `maxRetries` at the largest ladder fitting the resolved deadline. Where a +configuration leaves too little room — the extreme being a deadline barely wider than its own +backoff ladder, e.g. `CLODEX_UPSTREAM_IDLE_TIMEOUT_MS=14001` — the bound floors to zero and the +pacer disables itself with a notice, since refusing everything past the burst would be worse than +not pacing. It degrades to less pacing, never to a request pushed past its deadline. + +**The backoff ladder is NOT an upper bound on the pacing case.** `getRetryDelayInMs` SUBSTITUTES a +supplied `retry-after` for its own rung rather than taking the larger of the two, and a refusal +carries one. So the gap between paced attempts is the hint, and since the hint is capped at the +bound — which can exceed an early rung, 15s against a 2s first rung — a paced gap can be longer +than the rung it replaced. The conservative term is the per-gap maximum: + + (maxRetries + 1) x bound + SUM_i max(cappedHint, rung_i) < idleTimeout + +That is the property the tests assert across the resolvable space. `wsNewConnectionMaxWaitMs` +budgets the ladder alone; the halving is the slack that keeps the stronger inequality true, and +that is measured rather than argued. + +**An uncapped hint was a real defect, and capping it created a second one.** The hint used to be +the raw token deficit, so a 30s hint could be spent inside a 10s deadline: the request died having +made one attempt, with its remaining retries never run. Capping the hint at the bound fixed that +and broke low rates in the other direction — at 1/minute the first token is 60s away while six +attempts ~4s apart are all spent inside 20s, so every refused request exhausted its retries before +a token could exist. Measured at 1/min and 2/min: 10 of 10 refused requests terminal. + +**That is a trade-off, not an impossibility** — an earlier draft of this section claimed no hint +strategy could fix it and was wrong. A separately budgeted 12s hint does reach the 1/minute refill +(attempts at 0, 12, 24, 36, 48, 60s) and still fits the conservative mixed-gap bound: +`6 x 4833 + max(12,2) + max(12,4) + max(12,8) + max(12,16) + max(12,32) = 112,998ms < 120,000ms`. +What is true is narrower: no strategy admits ALL the overflow inside the deadline while preserving +the configured ceiling, because at 1/minute ten simultaneous overflow requests need ten minutes of +capacity. + +So **the pacer refuses only when its retry schedule can outlast the wait for a refill** +(`canRefuseAtRate`); below that it shapes the opening burst and then admits the remaining overflow +rather than failing it, with a notice. That is the same rule the zero-bound case already used. It +avoids guaranteed local failures while retaining bounded opening-burst shaping, **at the cost of +relaxing the configured ceiling** — which is what is actually given up here. The user configured a +connection rate, not a latency, and the fallback mostly adds no latency: at 1/minute with 20 +simultaneous requests, 19 are admitted immediately, one waits out the bound and none is refused. +A separately budgeted hint would be a reasonable follow-up. + +Because the head scan runs before the wait, an admitted request re-reads the clock, reaps whatever +expired while it was queued, and demotes itself to `parallel_isolated` if a same-partition request +went in flight meanwhile — otherwise two requests would each register a persistent nursery head for +one key and a fan-out would evict other conversations' heads. + +Diagnostics: a `ws_new_connection_paced` event (`outcome` of `admitted`, `refused`, or `aborted`, +with `waitedMs` / `requiredWaitMs` / `retryAfterSeconds`) and `pacingWaitedMs` on the same request's +`ws_head_decision`. Requests admitted on arrival record nothing. A request cancelled while queued +returns its reservation and opens no connection. + +The numbers come from re-reading one machine's `ws_head_decision` log (103,698 records over about a +day and a half), bucketing records that carry a `createdConnectionId` by wall-clock minute. Every +upgrade 403 fell in three minutes, and 39 of the 40 fell in two minutes that each opened 82 new +connections; across the 1,158 minutes that opened any connection the median was 6, the 90th +percentile 22 and the 99th 48, and four exceeded 60. + +**Scope of that measurement**, so it is not over-read: one account on one machine, one contiguous +window of roughly a day and a half, counted by the `ws_head_decision` predicate +`createdConnectionId != null`, which counts PRIMARY connections only — replacements never emit a +head decision, so they are absent from every figure above. It is a correlation, not a published +limit and not a demonstrated cause, which is why the default is conservative and tunable. The replacement connections +are deliberately **not** paced — both the transport retry and the `previous_response_not_found` +retry: each recovers a request that was already admitted, each is capped at one per request, and +both are built inside socket callbacks where an await would restructure the retry path. + ### Upstream timeouts and retries Every AI SDK generation entry point, for both Anthropic- and OpenAI-format routes, resolves one diff --git a/.claude/harnesses/README.md b/.claude/harnesses/README.md index c4b07086..c7133fd8 100644 --- a/.claude/harnesses/README.md +++ b/.claude/harnesses/README.md @@ -58,4 +58,4 @@ Several need real Claude Code bundles. Extract them once with | --- | --- | | `pr92-selfconnect-guard-matrix` | The self-connection guard across address forms — exact, loopback alias, wildcard bind. | | `pr92-selfconnect-loop-repro` | Reproduces the recursive self-tunnel the guard exists to prevent. | -| `fix-parent-notice-tui-and-epipe-round2` | Parent notices under a real Claude Code TUI, plus async EPIPE containment. Needs a real binary; set `CLODEX_CLAUDE_PATH` and `MAINBASE_DIR`. | +| `fix-parent-notice-tui-and-epipe-round2` | **BROKEN — does not run.** It imports `tests/helpers/register-ts-resolve-hook.mjs`, which was never committed, so both probes exit 1 at module resolution. Supply that hook before trusting anything here. Claim it was written to settle: parent notices under a real Claude Code TUI, plus async EPIPE containment. Needs a real binary; set `CLODEX_CLAUDE_PATH` and `MAINBASE_DIR`. | diff --git a/.claude/harnesses/fix-parent-notice-tui-and-epipe-round2.harness.ts b/.claude/harnesses/fix-parent-notice-tui-and-epipe-round2.harness.ts index b76a0d39..60fe1913 100644 --- a/.claude/harnesses/fix-parent-notice-tui-and-epipe-round2.harness.ts +++ b/.claude/harnesses/fix-parent-notice-tui-and-epipe-round2.harness.ts @@ -7,6 +7,12 @@ import { afterAll, describe, expect, it } from 'vitest'; const ROOT = fileURLToPath(new URL('..', import.meta.url)); const MAINBASE = process.env['MAINBASE_DIR'] ?? '../clodex-review/mainbase'; +// NOTE: this file does not exist in the repository and never has — it was not +// committed with the harness. Both probes below therefore exit 1 at module +// resolution before running anything. Recorded rather than silently left: the +// call this harness makes into src/upstream-retry.ts was updated when +// `upstreamMaxRetries` was removed, but that repair is UNVERIFIED: nothing here +// has been executed, and it cannot be until this hook is supplied. const REGISTER_HOOK = join(ROOT, 'tests/helpers/register-ts-resolve-hook.mjs'); const LAUNCH_URL = pathToFileURL(join(ROOT, 'src/launch.ts')).href; const NOTICE_URL = pathToFileURL(join(ROOT, 'src/parent-notice.ts')).href; @@ -115,7 +121,7 @@ describe('round-two parent notice runtime review', () => { const retryUrl = pathToFileURL(join(root, 'src/upstream-retry.ts')).href; writeFileSync(child, `#!/bin/sh\nprintf 'child-first-line\\n'\ntouch ${JSON.stringify(started)}\nsleep 1\n`); chmodSync(child, 0o755); - writeFileSync(probe, `import { existsSync } from 'node:fs';\nimport { launchClaude } from ${JSON.stringify(launchUrl)};\nimport { upstreamMaxRetries } from ${JSON.stringify(retryUrl)};\nconst sleep = ms => new Promise(r => setTimeout(r, ms));\nprocess.env.CLODEX_CLAUDE_PATH = ${JSON.stringify(child)};\nconst running = launchClaude({ ...process.env }, undefined, []);\nwhile (!existsSync(${JSON.stringify(started)})) await sleep(10);\nawait sleep(300);\nupstreamMaxRetries({ CLODEX_UPSTREAM_MAX_RETRIES: '6' });\nawait running;\nawait sleep(200);\n`); + writeFileSync(probe, `import { existsSync } from 'node:fs';\nimport { launchClaude } from ${JSON.stringify(launchUrl)};\nimport { upstreamRequestBudget } from ${JSON.stringify(retryUrl)};\nconst sleep = ms => new Promise(r => setTimeout(r, ms));\nprocess.env.CLODEX_CLAUDE_PATH = ${JSON.stringify(child)};\nconst running = launchClaude({ ...process.env }, undefined, []);\nwhile (!existsSync(${JSON.stringify(started)})) await sleep(10);\nawait sleep(300);\nupstreamRequestBudget({ env: { CLODEX_UPSTREAM_MAX_RETRIES: '6' } });\nawait running;\nawait sleep(200);\n`); writeFileSync(runner, `#!/bin/bash\nset +e\nset -o pipefail\n${JSON.stringify(process.execPath)} --experimental-strip-types --no-warnings --import ${JSON.stringify(REGISTER_HOOK)} ${JSON.stringify(probe)} 2>&1 | head -n 1\ncodes=(\"\${PIPESTATUS[@]}\")\nprintf '${label}_node=%s head=%s\\n' \"\${codes[0]}\" \"\${codes[1]}\"\n`); chmodSync(runner, 0o755); return execFileSync(runner, { encoding: 'utf8', env: { ...process.env, CLODEX_HOME: join(dir, 'home') } }); diff --git a/README.md b/README.md index 852f7ece..26da13b6 100644 --- a/README.md +++ b/README.md @@ -398,6 +398,40 @@ clodex --version # version so telling the client never to resend a request is not quietly undone one layer down. Recovered requests appear in the inference log as `response_retried`. +- **Connection pacing (ChatGPT/Codex plans):** when many agents run at once, + clodex spaces out the new connections it opens to OpenAI, which should make a + burst of parallel work less likely to trip OpenAI's own rate limit. (In the + traffic we sampled, the rejections clustered in the busiest minutes; that the + rate is what triggers them is a reasonable reading of that, not something we + can prove.) A follow-up turn that can reuse the connection it already has is + never delayed by this; + what goes through the limiter is work that needs a *new* connection — a first + turn, a conversation that branched, or several agents running at once. The + default is 60 new connections a minute, with an allowance of 10 opened back + to back after a quiet spell. + **This is a real throughput ceiling, not a brief pause.** One new connection + per second means that if you run many agents at once and each needs its own + connection, they end up sharing that budget: roughly 20 agents settle at + about 20 seconds per turn instead of a few seconds. That is the trade — you + wait longer, in exchange for a lower chance of losing turns to rate-limit + errors. It reduces that risk rather than removing it: clodex cannot see + OpenAI's actual limit, and under a heavy enough fan-out pacing can itself + answer a turn with a rate-limit response. Work over the + rate is queued for a few seconds, and anything still over is answered with + the same "try again shortly" response OpenAI itself would return, which + clodex retries for you with backoff. Set + `CLODEX_WS_MAX_NEW_CONNECTIONS_PER_MIN` to another whole number between `1` + and `600` to change the rate, or to `0` to turn pacing off. Higher values + clamp to `600` with a one-time warning, and an unreadable value is reported + once and ignored. If you have turned retries off with + `CLODEX_UPSTREAM_MAX_RETRIES=0`, pacing never refuses a request — but it also + stops limiting once its initial allowance is used up, because there would be + nothing left to retry a refused request. The same applies at very low rates: + if clodex cannot retry a turned-away request for long enough to reach the + next free connection slot — which is the case around 1 or 2 connections a + minute at the default timeouts — it admits the excess late rather than + failing it, and says so once. Turning a rate that low into hard failures + would manufacture the errors this feature exists to reduce. ## Known limitations diff --git a/src/oauth/responses-websocket.ts b/src/oauth/responses-websocket.ts index 4656f4f8..8de3617f 100644 --- a/src/oauth/responses-websocket.ts +++ b/src/oauth/responses-websocket.ts @@ -15,6 +15,12 @@ import { outboundWsProxyAgent } from '../outbound-proxy.js'; import { emitParentNotice } from '../parent-notice.js'; import { anthropicErrorType, clampRetryAfterSeconds, frameStatusCode } from '../upstream-error.js'; import { sanitizeToolInput } from '../tool-input-sanitize.js'; +import { + resetWsUpgradePacerForTests, + sharedWsUpgradePacer, + type ConnectionPacer, + type UpgradeAdmission, +} from './ws-upgrade-pacer.js'; const RESPONSES_LITE_HEADER = 'x-openai-internal-codex-responses-lite'; const TERMINAL_EVENT_TYPES = new Set(['response.completed', 'response.failed', 'response.incomplete']); @@ -35,6 +41,8 @@ export interface ResponsesWebSocketFetchOptions { nurseryIdleTtlMs?: number; maxConnections?: number; maxNurseryConnections?: number; + /** Test override; production shares the process-wide new-connection pacer. */ + pacer?: ConnectionPacer; now?: () => number; /** Opt-in structured transport diagnostics; never receives conversation content. */ onDiagnostic?: (event: ResponsesWebSocketDiagnosticEvent) => void; @@ -201,6 +209,9 @@ export function resetResponsesWebSocketConnectionsForTests(): void { } connections.clear(); nextConnectionDebugId = 1; + // The shared pacer counts connections, so it has to be dropped with them: + // otherwise one test file's sockets pace the next test's first request. + resetWsUpgradePacerForTests(); } /** Normalize the SDK's HeadersInit into a plain record for `ws`. */ @@ -1801,6 +1812,46 @@ function handleSocketMessage(entry: ConnectionEntry, data: RawData): void { } } +/** + * What the pacer hands back instead of opening a connection it cannot afford. + * + * Shaped like the frame `failContext` writes for an upgrade 403, but built + * standalone: a refusal happens before any connection or request context + * exists, so there is no stream to write into, nothing registered to delete and + * no context to close. `code` is the stringified status `frameStatusCode` reads + * preferentially. The `Retry-After` header is what the SDK's own backoff reads; + * the "retry after Ns" prose is the separate channel that carries the hint to + * the CLIENT, because the SDK's chunk schema strips `retry_after_seconds` + * before `sdkUpstreamErrorDetails` ever sees it. + */ +function pacedRefusalResponse(retryAfterSeconds: number): Response { + const frame = { + type: 'error', + sequence_number: 0, + error: { + type: anthropicErrorType(429), + code: '429', + message: 'clodex is limiting how fast it opens new OpenAI connections to reduce the chance ' + + `of an upstream rate limit; retry after ${retryAfterSeconds}s`, + param: null, + retry_after_seconds: retryAfterSeconds, + }, + }; + return new Response(`data: ${JSON.stringify(frame)}\n\n`, { + status: 200, + headers: { + 'content-type': 'text/event-stream; charset=utf-8', + // A real header, not just the prose: `getRetryDelayInMs` reads headers and + // ignores the body, so without one the SDK falls back to its fixed 2s/4s + // ladder. This does NOT de-correlate the group — refusals debit nothing, + // so everyone refused at the same instant sees the same deficit and gets + // the same hint — it defers the whole group by long enough for the bucket + // to refill, which is what turns a retry storm into a successful retry. + 'retry-after': String(retryAfterSeconds), + }, + }); +} + function numericRetryAfterHeader(value: string | string[] | undefined): number | undefined { const single = Array.isArray(value) ? value[0] : value; return typeof single === 'string' && /^\d+$/.test(single.trim()) @@ -1992,7 +2043,9 @@ export function createResponsesWebSocketFetch( const promptFieldHashes = responsesWebSocketPromptFieldHashes(payload); const instructionsSnapshot = instructionsFromPayload(payload); const diagnosticCorrelation = diagnosticContext.getStore(); - const now = resolvedOptions.now(); + // Re-read after a pacing wait, so head ages stay comparable with the pool + // counts reported alongside them. + let now = resolvedOptions.now(); const evictions = cleanupExpiredConnections(now); const candidates = partitionKey ? connectionEntries(partitionKey) : []; @@ -2082,6 +2135,76 @@ export function createResponsesWebSocketFetch( decision = 'unpartitioned_socket'; } + // Pace only the path that opens a NEW connection. Reusing a head — nursery + // or established — sends on a socket that already exists, costs the account + // no upgrade, and must never wait behind the bucket. The wait sits ahead of + // the nursery eviction below so a queued request does not retire a head it + // may still be seconds away from needing. + let pacingWaitedMs: number | undefined; + if (!selected) { + const pacer = options.pacer ?? sharedWsUpgradePacer(); + const pacingStartedAt = resolvedOptions.now(); + let admission: UpgradeAdmission; + try { + admission = await pacer.admit(init?.signal ?? undefined); + } catch (error) { + // Cancelled while queued: never open the connection it was waiting for. + emitDiagnostic(options, { + event: 'ws_new_connection_paced', + outcome: 'aborted', + decision, + waitedMs: Math.max(0, resolvedOptions.now() - pacingStartedAt), + }, diagnosticCorrelation); + throw error; + } + if (admission.kind === 'refused') { + debug( + `refused a new connection to hold the pacing rate; retry after ${admission.retryAfterSeconds}s`, + ); + emitDiagnostic(options, { + event: 'ws_new_connection_paced', + outcome: 'refused', + decision, + requiredWaitMs: admission.requiredWaitMs, + retryAfterSeconds: admission.retryAfterSeconds, + }, diagnosticCorrelation); + return pacedRefusalResponse(admission.retryAfterSeconds); + } + if (admission.waitedMs > 0) { + pacingWaitedMs = admission.waitedMs; + debug(`paced new connection by ${admission.waitedMs}ms`); + emitDiagnostic(options, { + event: 'ws_new_connection_paced', + outcome: 'admitted', + decision, + waitedMs: admission.waitedMs, + }, diagnosticCorrelation); + } + // Unconditional, NOT gated on having waited: `admit` is async, so even an + // immediate admission resumes a microtask later, and two same-partition + // requests arriving in one tick both resume having waited zero. The head + // scan above ran before that yield either way. + // + // Re-read the clock and reap what expired meanwhile, so head ages are + // measured from now rather than from arrival. The candidate SET is still + // the one scanned on arrival: a head that appeared or was reaped during + // the wait is not reflected in `heads`, only in the pool counts. + now = resolvedOptions.now(); + evictions.push(...cleanupExpiredConnections(now)); + // A same-partition request may have gone in flight across the yield. It + // was classified when no head existed, so without this both would open a + // persistent nursery head for one key — and a fan-out would fill the + // nursery with duplicates, evicting other conversations' heads and + // forcing the full-context resends that open still more connections. An + // overlap like this takes an isolated socket today; keep that. + if (persistent && partitionKey + && connectionEntries(partitionKey).some(entry => entry.inFlight)) { + persistent = false; + decision = 'parallel_isolated'; + debug('parallel request using an isolated socket after pacing'); + } + } + if (!selected && persistent) { evictions.push(...evictOldestIdleGeneration( 'nursery', @@ -2128,6 +2251,7 @@ export function createResponsesWebSocketFetch( continuationMatchMode: selectedMatch?.mode, promotedConnectionId, createdConnectionId: selected ? undefined : nextConnectionDebugId, + ...(pacingWaitedMs !== undefined ? { pacingWaitedMs } : {}), createdGeneration: selected ? undefined : persistent ? 'nursery' : 'isolated', incrementalInputItems: selectedDelta?.length, heads: candidates.map(entry => ({ diff --git a/src/oauth/ws-upgrade-pacer.ts b/src/oauth/ws-upgrade-pacer.ts new file mode 100644 index 00000000..2339753c --- /dev/null +++ b/src/oauth/ws-upgrade-pacer.ts @@ -0,0 +1,619 @@ +// ws-upgrade-pacer.ts — client-side pacing for NEW ChatGPT/Codex Responses +// WebSocket connections. +// +// OpenAI's edge rejects a Responses WebSocket upgrade with HTTP 403, and in the +// traffic sampled below those rejections clustered in the minutes that opened +// the most new connections. `responses-websocket.ts` already handles the +// rejection — every upgrade 403 becomes a retryable 429 carrying a backoff hint +// — but nothing limited how fast clodex asked for new connections. This module +// shapes that rate on the assumption, not the proof, that it is what the edge +// is reacting to. A request that reuses an existing chain head never comes here. +// +// Why these numbers. Measured by re-reading one machine's own +// `ws_head_decision` diagnostics log (103,698 records spanning roughly a day +// and a half), bucketing new connections — records carrying a +// `createdConnectionId` — by wall-clock minute: +// +// * 1,158 minutes opened at least one connection. Median 6, 90th percentile +// 22, 99th percentile 48, maximum 82. Four minutes exceeded 60. +// * All 40 upgrade rejections fell in three of those minutes. 39 of them fell +// in the two minutes that each opened 82, and the fortieth in a minute that +// opened 41. +// * On a 200k-line slice of the same log, 11,417 of 26,430 head decisions +// opened a connection and 10,229 of those were the non-reusable parallel +// fan-out kind, so new connections move with the number of concurrent +// agents rather than with conversation length. +// +// Scope: one account, one machine, one contiguous window. The predicate is +// `createdConnectionId != null` on `ws_head_decision`, which counts PRIMARY +// connections only — the two replacement paths emit no head decision, so they +// appear in none of these figures. +// +// Read this as a correlation in one account's traffic over one window, not as a +// published limit or a demonstrated cause: the rejections cluster in the +// highest-rate minutes, and the instantaneous rate was 3-5/second in those +// minutes and in quiet ones alike, which is why the limiter shapes a sustained +// rate rather than a burst. A single 41/minute rejection sits outside that +// pattern and is unexplained. + +import { emitParentNotice } from '../parent-notice.js'; +import { clampRetryAfterSeconds } from '../upstream-error.js'; +import { upstreamRequestBudget } from '../upstream-retry.js'; + +export const WS_NEW_CONNECTIONS_PER_MIN_ENV = 'CLODEX_WS_MAX_NEW_CONNECTIONS_PER_MIN'; + +/** + * Sustained ceiling on new connections, per minute. It sits between the 99th + * percentile of observed per-minute demand (48) and the two minutes in which 39 + * of the 40 rejections were observed (82 each). Those minutes are where the + * rejections fell; that they were caused by the rate is the working assumption + * this module is built on, not a finding. + * + * Do not read that as "rarely engages". Only four of the 1,158 observed active + * minutes exceeded 60 overall, but the limiter also holds requests whenever the + * burst allowance is spent WITHIN a minute, which heavy fan-out does routinely. + * One connection per second is a real aggregate ceiling: by Little's law, N + * agents that each need a new connection per turn settle at about N seconds per + * turn once the burst is gone — roughly 20s per turn at 20 agents, against ~3s + * unpaced. That trade is the feature: throughput for a LOWER CHANCE of tripping + * the throttle. Not for immunity from it — the causal link is this module's + * assumption (see the header), and a fan-out big enough to exhaust the bound is + * refused here, which the client sees as a rate limit. + */ +export const DEFAULT_WS_NEW_CONNECTIONS_PER_MIN = 60; + +/** + * Sanity bound on the configured rate, not a measured threshold. Ten per second + * is far above anything in the sample; a value that large leaves the bucket + * effectively open, which is what `0` is for. + */ +export const MAX_WS_NEW_CONNECTIONS_PER_MIN = 600; + +/** + * Connections opened with no delay at all after an idle stretch. Sized to the + * observed concurrency rather than to the instantaneous rate: connections peaked + * at 13 with no capacity evictions, while the 3-5/second instantaneous rate was + * the same in rejecting and quiet minutes alike and so distinguishes nothing. A + * fan-out of up to ten subagents is therefore admitted without waiting, provided + * the bucket has had ten seconds to refill. + */ +export const WS_NEW_CONNECTION_BURST = 10; + +/** No single request is queued longer than this, whatever the arithmetic says. */ +export const WS_NEW_CONNECTION_MAX_WAIT_CEILING_MS = 15_000; + +/** + * The request budget this pacer sizes its wait bound against. + * + * Both terms of that arithmetic are READ, never assumed. `maxRetries` and the + * no-data deadline are user-configurable (`CLODEX_UPSTREAM_MAX_RETRIES`, + * `CLODEX_UPSTREAM_IDLE_TIMEOUT_MS`, `CLODEX_UPSTREAM_TOTAL_TIMEOUT_MS`) and + * interact — a shorter deadline lowers the retry ceiling — so they are resolved + * together by the same `upstreamRequestBudget` call every SDK generation entry + * point makes. Hardcoding either would make this feature's correctness depend + * on a constant a user can change out from under it: too small a deadline or + * too large a retry count and the ladder overruns the deadline it shares. + * + * The no-argument call is the one the paced request itself resolves: no + * production caller passes `upstreamRequestBudget` an `idleTimeoutMs` + * override (that seam exists for direct adapter callers and is unused), so + * both read the same environment. The environment cannot change under a live + * process, so reading it once when the process-wide bucket is built is enough. + */ +export function resolvedPacingBudget(): { idleTimeoutMs: number; maxRetries: number } { + const { idleTimeoutMs, maxRetries } = upstreamRequestBudget(); + return { idleTimeoutMs, maxRetries }; +} + +/** First step of the AI SDK's backoff ladder; each later step doubles. */ +const SDK_INITIAL_BACKOFF_MS = 2_000; + +/** + * Longest a single request may be queued, given the budget it is spending. + * + * Every attempt of one request shares ONE no-data deadline — the timer starts + * before the SDK call and is only reset by a stream part — so the whole retry + * ladder has to fit inside it: + * + * (maxRetries + 1) x bound + totalBackoff < idleTimeout + * + * At the default 120s deadline and five retries the ladder alone is + * 2+4+8+16+32 = 62s, leaving 58s for six attempts. Half of that is reserved for + * the provider's own first byte, so the bound is ~4.8s. A flat 15s bound would + * instead allow six attempts plus backoff to reach 152s against a 120s + * deadline — not a certain timeout, since an attempt need not spend its whole + * bound, but a ceiling that no longer fits the budget. Reserving time cannot + * GUARANTEE the ladder completes either: first-byte latency is unbounded within + * the deadline. + * + * Halving is what makes the inequality STRICT for every input, rather than an + * arithmetic coincidence at one configuration: attempts x bound <= + * (idle - backoff) / 2, so attempts x bound + backoff <= (idle + backoff) / 2, + * which is below `idle` whenever `backoff < idle`. `upstreamRequestBudget` + * guarantees that side condition for every budget it resolves, because it caps + * `maxRetries` at the largest ladder that fits the resolved deadline. On any + * other input — an injected retry count above that cap, a deadline barely wider + * than its own ladder — the shareable term floors to zero and the bound with + * it, which the constructor reads as "do not pace". + * + * `totalBackoffMs` is the SDK's own exponential ladder, used when a failure + * carries no `Retry-After`. A refusal from this module DOES carry one, and + * `getRetryDelayInMs` SUBSTITUTES it for the rung rather than taking the larger + * of the two, so the real gap between paced attempts is the hint. + * + * The ladder is therefore NOT an upper bound on the pacing case — do not read + * it as one. `pacedRetryAfterSeconds` caps the hint at this bound, and this + * bound can exceed an early rung (a 15s cap against a 2s first rung), so a + * paced gap can be longer than the rung it replaced. The conservative term is + * the per-gap maximum: + * + * (maxRetries + 1) x bound + SUM_i max(cappedHint, rung_i) < idleTimeout + * + * which is what the tests assert across the resolvable configuration space. + * This function budgets the ladder alone; the halving above is the slack that + * keeps the stronger inequality true as well, and it is measured rather than + * assumed. + * + * Both `idleTimeoutMs` and `maxRetries` are read from the resolved request + * budget rather than assumed; see `resolvedPacingBudget`. + */ +export function wsNewConnectionMaxWaitMs( + idleTimeoutMs: number, + maxRetries: number, +): number { + // Total, on every input. A NaN reaching the bound would make every + // `waitMs > maxWaitMs` comparison false, which fails open into unbounded + // waits — the one failure this bound exists to prevent. + if (!Number.isFinite(idleTimeoutMs) || !Number.isFinite(maxRetries) + || idleTimeoutMs <= 0 || maxRetries < 0) { + return 0; + } + const attempts = maxRetries + 1; + const totalBackoffMs = SDK_INITIAL_BACKOFF_MS * (2 ** maxRetries - 1); + const shareableMs = Math.max(0, idleTimeoutMs - totalBackoffMs) / 2; + const bound = Math.floor(shareableMs / attempts); + return Number.isFinite(bound) + ? Math.min(WS_NEW_CONNECTION_MAX_WAIT_CEILING_MS, Math.max(0, bound)) + : 0; +} + +/** + * Backoff hint a refusal sends the client, in whole seconds. + * + * THIS IS PART OF THE BOUND, NOT A COURTESY. `wsNewConnectionMaxWaitMs` budgets + * the SDK's own 2s/4s/8s/16s/32s ladder as the delay between attempts, but + * `getRetryDelayInMs` (ai@7.0.22, `util/retry-with-exponential-backoff.ts`) + * RETURNS a supplied hint in place of that rung whenever the hint is under 60 + * seconds — it does not take the larger of the two. So whatever this function + * emits is what the request actually spends between attempts, and an uncapped + * hint taken from the token deficit would silently replace the very term the + * bound was derived against. + * + * It did. At 2 connections/minute the deficit is 30s while a 10s deadline funds + * a 666ms bound: the refusal asked for 30s, the request hit its deadline having + * made one attempt, and the two retries its budget had paid for never ran. + * + * So the hint is capped at the bound. `requiredWaitMs` is still reported + * honestly to diagnostics; only the number the client is asked to honour is + * capped. The floor is one second because `Retry-After` is whole seconds and + * zero would mean "immediately"; one second is below the SDK's smallest rung + * (2s), so it always fits inside the ladder term the bound already budgets. + * + * The cost is real and is the right trade: a refused request comes back sooner + * than the deficit needs and may be refused again. It spends its retries inside + * its deadline instead of spending its whole deadline on one oversized sleep. + */ +export function pacedRetryAfterSeconds(requiredWaitMs: number, maxWaitMs: number): number { + const capSeconds = Number.isFinite(maxWaitMs) ? Math.max(1, Math.floor(maxWaitMs / 1_000)) : 1; + const wantSeconds = Number.isFinite(requiredWaitMs) + ? Math.max(1, Math.ceil(requiredWaitMs / 1_000)) + : 1; + return Math.min(wantSeconds, capSeconds); +} + +/** + * How long a refused request's retry schedule can span, in milliseconds. + * + * A refusal only helps if the request can come back after a token exists. Every + * gap in that schedule is the capped hint — a refused request's deficit exceeds + * the bound by definition, so the cap, not the deficit, is what it waits — and + * there are `maxRetries` gaps. Zero when nothing would retry. + */ +export function refusalScheduleMs(maxWaitMs: number, maxRetries: number): number { + if (!Number.isFinite(maxRetries) || maxRetries <= 0) return 0; + return maxRetries * pacedRetryAfterSeconds(Number.MAX_SAFE_INTEGER, maxWaitMs) * 1_000; +} + +/** + * Whether refusing can actually make progress at this rate. + * + * Capping the hint at the bound fixed requests overrunning their deadline, but + * it created the opposite failure at low rates: at 1/minute the first token is + * 60s away while six attempts four seconds apart are all spent inside 20s, so + * every refused request exhausted its retries before a token could exist and + * died as a rate-limit error — manufactured by the very thing meant to avoid + * them. Measured at 1/min and 2/min: 10 of 10 refused requests terminal. + * + * This is a TRADE-OFF, not an impossibility. An earlier version of this comment + * claimed no hint strategy could fix it; that was wrong, and a review disproved + * it. A separately budgeted 12s hint reaches the 1/minute refill — attempts at + * 0, 12, 24, 36, 48 and 60s — and still fits the conservative mixed-gap bound + * at the shipped budget: 6 x 4833 + max(12,2) + max(12,4) + max(12,8) + + * max(12,16) + max(12,32) = 112,998ms < 120,000ms. + * + * What IS true is narrower: no strategy can admit all the overflow inside the + * deadline while preserving the configured ceiling. At 1/minute, ten + * simultaneous overflow requests need ten minutes of capacity, so a longer hint + * rescues one of them and still fails the other nine. + * + * Given that, shaping is chosen over refusing because it avoids guaranteed + * local failures while retaining bounded opening-burst shaping, at the cost of + * RELAXING THE CONFIGURED CEILING. Be exact about what is spent: the user + * configured a connection rate, not a latency, and this fallback mostly does + * not add latency — at 1/minute with 20 simultaneous requests, 19 are admitted + * immediately, one waits out the bound and none is refused. What gives way is + * the ceiling itself. It also keeps one hint rule instead of two. A separately + * budgeted hint is a reasonable follow-up, not a correction. This is the same + * rule already applied to a zero bound, where refusing everything past the + * burst is worse than not pacing. + */ +export function canRefuseAtRate( + maxWaitMs: number, + maxRetries: number, + ratePerMinute: number, +): boolean { + if (!Number.isFinite(ratePerMinute) || ratePerMinute <= 0) return false; + const refillIntervalMs = 60_000 / ratePerMinute; + return refusalScheduleMs(maxWaitMs, maxRetries) >= refillIntervalMs; +} + +/** Outcome of asking the pacer for permission to open a new connection. */ +export type UpgradeAdmission = + | { + kind: 'admitted'; + /** Milliseconds spent queued. 0 means admitted on arrival. */ + waitedMs: number; + } + | { + kind: 'refused'; + /** What the rate would have required, before the bound refused it. */ + requiredWaitMs: number; + /** Backoff hint for the client, in seconds. */ + retryAfterSeconds: number; + }; + +export interface ConnectionPacer { + admit(signal?: AbortSignal): Promise; +} + +export interface WsUpgradePacerOptions { + /** New connections per minute; 0 disables pacing entirely. */ + ratePerMinute?: number; + burst?: number; + /** No-data deadline the queued request is spending. */ + idleTimeoutMs?: number; + /** Retry attempts that deadline has to cover. */ + maxRetries?: number; + now?: () => number; + /** Runs `fire` after `ms`; returns a cancel function. Injected by tests. */ + schedule?: (ms: number, fire: () => void) => () => void; +} + +type Reservation = + | { admitted: true; waitMs: number; consumed: number } + | { admitted: false; requiredWaitMs: number }; + +function defaultSchedule(ms: number, fire: () => void): () => void { + // Deliberately not unref'd: an in-flight request must keep the process alive + // for as long as it is queued, exactly as it does while it is on the wire. + const timer = setTimeout(fire, ms); + return () => clearTimeout(timer); +} + +/** + * Mirrors `streamAbortError` in sdk-adapter.ts: an Error reason is the caller's + * own error and is surfaced unchanged; anything else becomes a named + * AbortError, which the AI SDK recognizes and never retries. + */ +function abortError(signal: AbortSignal | undefined): Error { + if (signal?.reason instanceof Error) return signal.reason; + const error = new Error( + typeof signal?.reason === 'string' ? signal.reason : 'WebSocket connection pacing aborted', + ); + error.name = 'AbortError'; + return error; +} + +const reportedNotices = new Set(); + +/** + * `key` must carry its own namespace prefix. Rate-validation notices key off an + * arbitrary environment string, so an unprefixed key let a hostile or unlucky + * value collide with the pacing-disabled key and suppress the one notice that + * must never go missing — that pacing turned itself off. + */ +function reportOnce(key: string, message: string, warn: (message: string) => void): void { + if (reportedNotices.has(key)) return; + reportedNotices.add(key); + try { + warn(message); + } catch { + // A diagnostic must never turn a pacing setting into a request failure. + } +} + +/** + * Token bucket over new-connection creation. + * + * Admission is decided by one synchronous reservation — refill, read, debit — + * so concurrent callers can never interleave inside it and each leaves holding + * its own deadline. Ordering therefore follows arrival order without a queue to + * scan, and nothing is held across the `await`. + * + * A request the rate cannot serve within the bound is REFUSED rather than + * queued longer or admitted anyway. Admitting anyway was tried first and does + * not work: with the bound also acting as the debt floor, sustained output + * settles at exactly the offered rate delayed by the bound, so a 82/min fan-out + * still went out at 82/min and simply arrived one bound later. Refusing sheds + * that overflow instead, in the same retryable 429 shape the upgrade 403 + * already produces today — the shape the client is known to handle. + * + * Refusing is not free. The retry it provokes backs off INSIDE the same no-data + * deadline a queue wait would have spent, which is why `wsNewConnectionMaxWaitMs` + * budgets the whole ladder rather than one wait. + * + * A refusal deliberately debits NOTHING. That is what makes the retry ladder + * safe: a refused request opens no connection, and while retries remain it will + * be retried — so charging it a token would let every retry deepen the deficit + * that caused the refusal and the ladder could never recover. A refusal on the + * FINAL attempt is terminal and surfaces to the user as a rate limit, which is + * why the bound is sized so the whole ladder fits the deadline. Because only an + * admitted request debits, and a request is only admitted when it needs at most + * `maxWaitMs`, `tokens` can never fall below `-maxWaitMs x refillPerMs`. The + * queue is therefore bounded by construction, every queued request drains within + * the bound, and — while refusals are available, i.e. with retries enabled — + * admissions in any window T are at most `burst + maxWaitMs x refillPerMs + + * rate x T + cancellations` however many retries arrive. With retries disabled + * nothing is refused, so only the opening burst is shaped and that bound does + * not hold. + * + * Two caveats on that bound, both real: + * + * * It counts ADMISSIONS, not sockets. A transport retry and a + * `previous_response_not_found` retry each build a replacement connection + * through `createReplacement`, which does not consult the pacer, so + * connections opened can exceed admissions granted. + * * A cancellation refunds its token but does not reschedule the reservations + * already queued behind it, so a later arrival can take the vacated slot + * alongside them. Each cancellation therefore permits one extra admission + * at that instant. It cannot reorder admissions, only coalesce them. + * + * The bound is also on reservations, not on wall-clock departures: a stalled + * event loop releases overdue timers together, and the pacer neither observes + * that nor corrects for it. + */ +export class WsUpgradePacer implements ConnectionPacer { + /** Longest any one request will be queued. Derived; exposed for diagnostics. */ + readonly maxWaitMs: number; + private readonly enabled: boolean; + private readonly refillPerMs: number; + private readonly capacity: number; + private readonly canRefuse: boolean; + private readonly maxDebt: number; + private readonly now: () => number; + private readonly schedule: (ms: number, fire: () => void) => () => void; + private tokens: number; + private lastRefillAt: number; + + constructor(options: WsUpgradePacerOptions = {}) { + const ratePerMinute = options.ratePerMinute ?? DEFAULT_WS_NEW_CONNECTIONS_PER_MIN; + const burst = options.burst ?? WS_NEW_CONNECTION_BURST; + // Resolve the environment only when a term is missing, so a fully injected + // test pacer never depends on ambient configuration. + const budget = options.idleTimeoutMs === undefined || options.maxRetries === undefined + ? resolvedPacingBudget() + : { idleTimeoutMs: options.idleTimeoutMs, maxRetries: options.maxRetries }; + const idleTimeoutMs = options.idleTimeoutMs ?? budget.idleTimeoutMs; + const maxRetries = options.maxRetries ?? budget.maxRetries; + this.maxWaitMs = wsNewConnectionMaxWaitMs(idleTimeoutMs, maxRetries); + // A bound of zero means the deadline has no room to queue anything, so + // pacing would degenerate into refusing everything past the burst. Not + // pacing at all is the safer reading of that configuration. + const rateRequested = Number.isFinite(ratePerMinute) && ratePerMinute > 0; + this.enabled = rateRequested && this.maxWaitMs > 0; + if (rateRequested && !this.enabled) { + // Turning itself off on a seam is the one failure mode nobody would + // notice, so it is never silent. + reportOnce( + `disabled:${maxRetries}:${idleTimeoutMs}`, + `not pacing new OpenAI connections: a ${maxRetries}-retry budget leaves no room to ` + + `queue inside the resolved ${idleTimeoutMs}ms request deadline`, + message => emitParentNotice(`clodex: ${message}`), + ); + } + // A refusal is only safe when something will retry it, AND only useful when + // that retry can outlast the wait for a token. With retries turned off the + // SDK rethrows before it ever consults `shouldRetry`, so a refusal would be + // an immediate hard failure; at a rate whose refill the retry schedule + // cannot reach, it is a slower hard failure. Both admit instead. + this.canRefuse = maxRetries > 0 + && canRefuseAtRate(this.maxWaitMs, maxRetries, ratePerMinute); + if (this.enabled && maxRetries > 0 && !this.canRefuse) { + // The user asked for a hard ceiling and is not getting one. Never silent. + reportOnce( + `norefuse:${ratePerMinute}:${maxRetries}:${this.maxWaitMs}`, + `pacing new OpenAI connections at ${ratePerMinute}/minute without refusing overflow: a ` + + `${maxRetries}-retry schedule spans only ` + + `${refusalScheduleMs(this.maxWaitMs, maxRetries)}ms, which cannot outlast the ` + + `${Math.round(60_000 / ratePerMinute)}ms wait for a free connection slot, so clodex ` + + 'shapes the opening burst and then admits remaining overflow instead of failing it', + message => emitParentNotice(`clodex: ${message}`), + ); + } + this.refillPerMs = this.enabled ? ratePerMinute / 60_000 : 0; + this.capacity = Number.isFinite(burst) ? Math.max(1, burst) : WS_NEW_CONNECTION_BURST; + this.maxDebt = this.maxWaitMs * this.refillPerMs; + this.now = options.now ?? Date.now; + this.schedule = options.schedule ?? defaultSchedule; + this.tokens = this.capacity; + this.lastRefillAt = this.now(); + } + + /** + * Resolves when this request may open a new connection, or resolves to a + * refusal the caller must report as a retryable rate limit. Callers that + * reuse an existing connection must not call this at all. + */ + async admit(signal?: AbortSignal): Promise { + if (!this.enabled) return { kind: 'admitted', waitedMs: 0 }; + // A request whose consumer is already gone must not spend a token, and must + // not be parked on a listener an aborted signal will never fire. + if (signal?.aborted) throw abortError(signal); + + const reservation = this.reserve(this.now()); + if (!reservation.admitted) { + return { + kind: 'refused', + requiredWaitMs: reservation.requiredWaitMs, + // Capped at the bound: the SDK substitutes this hint for its own + // backoff rung, so an uncapped one would overrun the deadline the + // bound was derived to fit. See `pacedRetryAfterSeconds`. + retryAfterSeconds: clampRetryAfterSeconds( + pacedRetryAfterSeconds(reservation.requiredWaitMs, this.maxWaitMs), + ), + }; + } + if (reservation.waitMs <= 0) return { kind: 'admitted', waitedMs: 0 }; + + const startedAt = this.now(); + try { + await this.sleep(reservation.waitMs, signal); + } catch (error) { + // A cancelled request opens no connection, so its reservation goes back + // rather than pacing someone else against an upgrade that never happened. + this.refund(reservation.consumed); + throw error; + } + return { kind: 'admitted', waitedMs: Math.max(0, this.now() - startedAt) }; + } + + private reserve(now: number): Reservation { + // A clock that moves backwards refills nothing rather than draining. + const elapsed = Math.max(0, now - this.lastRefillAt); + this.tokens = Math.min(this.capacity, this.tokens + elapsed * this.refillPerMs); + this.lastRefillAt = now; + const waitMs = this.tokens >= 1 ? 0 : Math.ceil((1 - this.tokens) / this.refillPerMs); + if (waitMs > this.maxWaitMs) { + if (this.canRefuse) return { admitted: false, requiredWaitMs: waitMs }; + // Nothing would retry a refusal here, so this must admit. Shape the + // opening burst — the part that correlates with rejection — and then stop + // once the debt floor is reached: past it, delaying every request by the + // bound shapes NOTHING (sustained output would equal sustained input, + // merely late) and only taxes the user. `consumed` is zero exactly at the + // floor, which is the signal that shaping has run out. + const before = this.tokens; + this.tokens = Math.max(-this.maxDebt, this.tokens - 1); + const consumed = before - this.tokens; + return { admitted: true, waitMs: consumed > 0 ? this.maxWaitMs : 0, consumed }; + } + this.tokens -= 1; + return { admitted: true, waitMs, consumed: 1 }; + } + + private refund(consumed: number): void { + this.tokens = Math.min(this.capacity, this.tokens + consumed); + } + + /** + * INVARIANT REQUIRED OF FUTURE EDITS: `admit` must reject an aborted signal + * before reaching here, and nothing may be awaited between that check and + * this call. An abort arriving in such a window would leave the request + * parked on a listener an already-aborted signal never fires, and it would + * wait out the full duration. There is deliberately no second check here to + * catch that, because an untested guard is not a guarantee. + */ + private sleep(ms: number, signal?: AbortSignal): Promise { + return new Promise((resolve, reject) => { + let settled = false; + let cancelTimer: (() => void) | undefined; + const onAbort = () => { + if (settled) return; + settled = true; + cancelTimer?.(); + reject(abortError(signal)); + }; + const fire = () => { + if (settled) return; + settled = true; + signal?.removeEventListener('abort', onAbort); + resolve(); + }; + signal?.addEventListener('abort', onAbort, { once: true }); + cancelTimer = this.schedule(ms, fire); + // The schedule implementation may have fired synchronously. + if (settled) cancelTimer(); + }); + } +} + +/** + * Optional override for the sustained new-connection rate. `0` turns pacing + * off; a malformed value leaves the default in control rather than failing the + * request that happened to read it. + */ +export function wsNewConnectionsPerMinute( + env: NodeJS.ProcessEnv = process.env, + // emitParentNotice rather than console.error: this fires from a request while + // `clodex claude` has the parent's stdout/stderr muted for Claude Code's TUI. + warn: (message: string) => void = message => emitParentNotice(`clodex: ${message}`), +): number { + const raw = env[WS_NEW_CONNECTIONS_PER_MIN_ENV]?.trim(); + if (raw === undefined || raw === '') return DEFAULT_WS_NEW_CONNECTIONS_PER_MIN; + const value = Number(raw); + if (!Number.isInteger(value) || value < 0) { + reportOnce( + `rate:${raw}`, + `ignoring ${WS_NEW_CONNECTIONS_PER_MIN_ENV}=${raw} ` + + `(expected a non-negative integer; using ${DEFAULT_WS_NEW_CONNECTIONS_PER_MIN})`, + warn, + ); + return DEFAULT_WS_NEW_CONNECTIONS_PER_MIN; + } + if (value > MAX_WS_NEW_CONNECTIONS_PER_MIN) { + reportOnce( + `rate:${raw}`, + `clamping ${WS_NEW_CONNECTIONS_PER_MIN_ENV}=${raw} to ${MAX_WS_NEW_CONNECTIONS_PER_MIN} ` + + '(a higher rate shapes nothing OpenAI throttles on)', + warn, + ); + return MAX_WS_NEW_CONNECTIONS_PER_MIN; + } + return value; +} + +let sharedPacer: ConnectionPacer | undefined; + +/** + * The process-wide pacer. + * + * Shared rather than per-transport for the same reason the connection pools + * are: the server keeps a separate transport per model, so a per-transport + * bucket would multiply the rate by the number of models in play. + * + * What the throttle is actually scoped to is NOT known. One account on one + * machine cannot distinguish an account-, IP-, model- or edge-level limit, and + * the sample behind this module is exactly that. Sharing one bucket is the + * conservative choice under that uncertainty: it paces a multi-account process + * harder than it may need to be paced, which is the safe direction to be wrong + * in. + */ +export function sharedWsUpgradePacer(): ConnectionPacer { + sharedPacer ??= new WsUpgradePacer({ ratePerMinute: wsNewConnectionsPerMinute() }); + return sharedPacer; +} + +/** Test-only: drop the shared bucket's accumulated state and re-read the env. */ +export function resetWsUpgradePacerForTests(): void { + sharedPacer = undefined; + reportedNotices.clear(); +} diff --git a/tests/responses-websocket.test.ts b/tests/responses-websocket.test.ts index 3abe2b77..07037549 100644 --- a/tests/responses-websocket.test.ts +++ b/tests/responses-websocket.test.ts @@ -33,6 +33,7 @@ import { type ResponsesWebSocketDiagnosticEvent, } from '../src/oauth/responses-websocket.js'; import { sdkUpstreamErrorDetails } from '../src/upstream-error.js'; +import type { UpgradeAdmission } from '../src/oauth/ws-upgrade-pacer.js'; const WS_URL = 'wss://chatgpt.com/backend-api/codex/responses'; @@ -3837,3 +3838,593 @@ describe('createResponsesWebSocketFetch', () => { .toBe(responsesWebSocketPromptFingerprint({ tools: [{ parameters: { a: 1, b: 2 }, name: 'x' }], model: 'm', input: ['different'] })); }); }); + +describe('new-connection pacing', () => { + beforeEach(() => { + resetResponsesWebSocketConnectionsForTests(); + fakeSockets.length = 0; + }); + + /** Records every admission request; production shares one process-wide pacer. */ + function recordingPacer(admission: UpgradeAdmission = { kind: 'admitted', waitedMs: 0 }) { + return { admit: vi.fn(async () => admission) }; + } + + it('never asks the pacer for a request that reuses an existing connection', async () => { + const pacer = recordingPacer(); + const diagnostics: ResponsesWebSocketDiagnosticEvent[] = []; + const wsFetch = createResponsesWebSocketFetch(WS_URL, undefined, { + accountId: 'acct-pacing-reuse', + pacer, + onDiagnostic: event => diagnostics.push(event), + }); + const input = [{ role: 'user', content: [{ type: 'input_text', text: 'first turn' }] }]; + + const first = await wsFetch('https://x', { + method: 'POST', headers: {}, body: JSON.stringify(sessionPayload(input)), + }); + const socket = lastSocket(); + socket.emit('open'); + socket.emit('message', Buffer.from(JSON.stringify({ + type: 'response.created', response: { id: 'resp_pace_1' }, + }))); + socket.emit('message', Buffer.from(JSON.stringify({ + type: 'response.output_item.done', output_index: 0, + item: { + type: 'function_call', id: 'fc_1', call_id: 'call_1', name: 'Read', + arguments: '{"path":"file.ts"}', status: 'completed', + }, + }))); + socket.emit('message', Buffer.from(JSON.stringify({ + type: 'response.completed', response: { id: 'resp_pace_1' }, + }))); + await readAll(first); + expect(pacer.admit).toHaveBeenCalledTimes(1); + + // Second turn continues the same head: no new connection, so no pacing. + const echoedCall = { + type: 'function_call', call_id: 'call_1', name: 'Read', arguments: '{"path":"file.ts"}', + }; + const toolOutput = { type: 'function_call_output', call_id: 'call_1', output: 'contents' }; + const second = await wsFetch('https://x', { + method: 'POST', headers: {}, + body: JSON.stringify(sessionPayload([...input, echoedCall, toolOutput])), + }); + expect(fakeSockets).toHaveLength(1); + expect(diagnostics.at(-1)).toMatchObject({ event: 'ws_head_decision', decision: 'continuation' }); + expect(pacer.admit).toHaveBeenCalledTimes(1); + emitTextResponse(socket, 'resp_pace_2', 'done'); + await readAll(second); + expect(pacer.admit).toHaveBeenCalledTimes(1); + }); + + it('asks the pacer for each primary connection it opens', async () => { + const pacer = recordingPacer(); + const wsFetch = createResponsesWebSocketFetch(WS_URL, undefined, { + accountId: 'acct-pacing-open', + pacer, + }); + // Two unrelated conversations: each needs its own head. + for (const text of ['alpha', 'beta']) { + const response = await wsFetch('https://x', { + method: 'POST', headers: {}, + body: JSON.stringify(sessionPayload( + [{ role: 'user', content: [{ type: 'input_text', text }] }], + { prompt_cache_key: `relay-session-${text}` }, + )), + }); + const socket = lastSocket(); + socket.emit('open'); + emitTextResponse(socket, `resp_${text}`, 'ok'); + await readAll(response); + } + expect(fakeSockets).toHaveLength(2); + expect(pacer.admit).toHaveBeenCalledTimes(2); + }); + + it('records how long a delayed connection waited, on both diagnostics', async () => { + const debug: string[] = []; + const diagnostics: ResponsesWebSocketDiagnosticEvent[] = []; + const wsFetch = createResponsesWebSocketFetch(WS_URL, message => debug.push(message), { + accountId: 'acct-pacing-delay', + pacer: recordingPacer({ kind: 'admitted', waitedMs: 2_500 }), + onDiagnostic: event => diagnostics.push(event), + }); + + const response = await withResponsesWebSocketDiagnosticContext( + { requestId: 'req-paced' }, + () => wsFetch('https://x', { + method: 'POST', headers: {}, body: JSON.stringify(sessionPayload([])), + }), + ); + const socket = lastSocket(); + socket.emit('open'); + emitTextResponse(socket, 'resp_paced', 'ok'); + await readAll(response); + + expect(diagnostics).toContainEqual(expect.objectContaining({ + event: 'ws_new_connection_paced', + outcome: 'admitted', + waitedMs: 2_500, + requestId: 'req-paced', + })); + expect(diagnostics).toContainEqual(expect.objectContaining({ + event: 'ws_head_decision', + pacingWaitedMs: 2_500, + })); + expect(debug).toContain('ws: paced new connection by 2500ms'); + }); + + it('refuses without opening a socket, in a shape that classifies as a retryable 429', async () => { + const debug: string[] = []; + const diagnostics: ResponsesWebSocketDiagnosticEvent[] = []; + const wsFetch = createResponsesWebSocketFetch(WS_URL, message => debug.push(message), { + accountId: 'acct-pacing-refused', + pacer: recordingPacer({ kind: 'refused', requiredWaitMs: 7_400, retryAfterSeconds: 8 }), + onDiagnostic: event => diagnostics.push(event), + }); + + const response = await wsFetch('https://x', { + method: 'POST', headers: {}, body: JSON.stringify(sessionPayload([])), + }); + + // The whole point: the connection this would have opened is not opened. + expect(fakeSockets).toHaveLength(0); + expect(response.status).toBe(200); + expect(response.headers.get('content-type')).toBe('text/event-stream; charset=utf-8'); + // A real header, not only the prose: the SDK's backoff reads headers and + // ignores the body. It does not de-correlate a fan-out — everyone refused at + // the same instant gets the same hint — it defers the group long enough for + // the bucket to refill. + expect(response.headers.get('retry-after')).toBe('8'); + + // A refusal happens before any connection or request context exists, so it + // cannot go through failContext. Prove the standalone frame still CLASSIFIES + // downstream as a retryable rate limit carrying the backoff hint — the same + // classification the real upgrade 403 gets. `classifyThroughSdk` runs with + // maxRetries: 0, so this pins the classification, not the retrying; the + // 403 test above is what exercises an actual SDK retry. + const body = await readAll(response); + // The frame itself, matching the shape failContext writes for a 403. + expect(JSON.parse(body.replace(/^data: /, '').trim())).toMatchObject({ + type: 'error', + error: { type: 'rate_limit_error', code: '429', retry_after_seconds: 8 }, + }); + // And through the real provider. Note the hint survives via the PROSE: the + // SDK's chunk schema strips `retry_after_seconds`, so the frame field alone + // would not reach the client. + expect(await classifyThroughSdk(body)).toMatchObject({ + statusCode: 429, + isRetryable: true, + retryAfterSeconds: 8, + }); + + expect(diagnostics).toContainEqual(expect.objectContaining({ + event: 'ws_new_connection_paced', + outcome: 'refused', + requiredWaitMs: 7_400, + retryAfterSeconds: 8, + })); + expect(diagnostics.some(event => event.event === 'ws_head_decision')).toBe(false); + expect(debug).toContain( + 'ws: refused a new connection to hold the pacing rate; retry after 8s', + ); + }); + + it('does not open a second persistent head for a partition that filled up during the wait', async () => { + // The head scan runs before the wait. A same-partition request that starts + // while this one is queued would otherwise leave BOTH registering + // persistent nursery heads for one key, filling the nursery with + // duplicates and evicting other conversations' reusable heads. + const diagnostics: ResponsesWebSocketDiagnosticEvent[] = []; + let releaseQueued: (() => void) | undefined; + const queued = new Promise(resolve => { releaseQueued = resolve; }); + let markQueued: (() => void) | undefined; + const isQueued = new Promise(resolve => { markQueued = resolve; }); + let admissions = 0; + const wsFetch = createResponsesWebSocketFetch(WS_URL, undefined, { + accountId: 'acct-pacing-overlap', + pacer: { + admit: async (): Promise => { + admissions += 1; + if (admissions > 1) return { kind: 'admitted', waitedMs: 0 }; + markQueued!(); + await queued; + return { kind: 'admitted', waitedMs: 25 }; + }, + }, + onDiagnostic: event => diagnostics.push(event), + }); + const body = (text: string) => JSON.stringify(sessionPayload( + [{ role: 'user', content: [{ type: 'input_text', text }] }], + )); + + // Classified with an empty partition, then held inside the pacer. + const held = wsFetch('https://x', { method: 'POST', headers: {}, body: body('held') }); + await isQueued; + + // A second request for the same partition gets in and goes in flight. + const overtaking = await wsFetch('https://x', { method: 'POST', headers: {}, body: body('overtaking') }); + lastSocket().emit('open'); + + releaseQueued!(); + const heldResponse = await held; + + const decisions = diagnostics.filter(event => event.event === 'ws_head_decision'); + expect(decisions).toHaveLength(2); + expect(decisions[1]).toMatchObject({ + decision: 'parallel_isolated', + createdGeneration: 'isolated', + pacingWaitedMs: 25, + }); + + for (const socket of fakeSockets) socket.emit('open'); + emitTextResponse(fakeSockets[1]!, 'resp_held', 'ok'); + emitTextResponse(fakeSockets[0]!, 'resp_overtaking', 'ok'); + await readAll(heldResponse); + await readAll(overtaking); + }); + + it('asks the pacer for every shape of primary new connection', async () => { + // Four shapes reach the creation path. The parallel one dominates real + // traffic (2,844 of 2,987 observed upgrades), so none of them may be left + // on prose. + const pacer = recordingPacer(); + const diagnostics: ResponsesWebSocketDiagnosticEvent[] = []; + const wsFetch = createResponsesWebSocketFetch(WS_URL, undefined, { + accountId: 'acct-pacing-shapes', + pacer, + onDiagnostic: event => diagnostics.push(event), + }); + const lastDecision = () => diagnostics.filter(event => event.event === 'ws_head_decision').at(-1); + const send = (payload: unknown) => wsFetch('https://x', { + method: 'POST', headers: {}, body: JSON.stringify(payload), + }); + const turn = (text: string) => [{ role: 'user', content: [{ type: 'input_text', text }] }]; + + // 1. unpartitioned_socket — no prompt_cache_key, so there is no partition. + const unpartitioned = await send({ model: 'gpt-5.6-sol', input: turn('no session') }); + lastSocket().emit('open'); + emitTextResponse(lastSocket(), 'resp_unpartitioned', 'ok'); + await readAll(unpartitioned); + expect(lastDecision()).toMatchObject({ decision: 'unpartitioned_socket' }); + expect(pacer.admit).toHaveBeenCalledTimes(1); + + // 2. new_partition_head — first turn of a session. + const root = await send(sessionPayload(turn('root'))); + lastSocket().emit('open'); + emitTextResponse(lastSocket(), 'resp_root', 'ok'); + await readAll(root); + expect(lastDecision()).toMatchObject({ decision: 'new_partition_head' }); + expect(pacer.admit).toHaveBeenCalledTimes(2); + + // 3. history_mismatch_new_head — same partition, divergent history. + const diverged = await send(sessionPayload(turn('a different root'))); + expect(lastDecision()).toMatchObject({ decision: 'history_mismatch_new_head' }); + expect(pacer.admit).toHaveBeenCalledTimes(3); + + // 4. parallel_isolated — same partition while that one is still in flight. + const parallel = await send(sessionPayload(turn('a third root'))); + expect(lastDecision()).toMatchObject({ decision: 'parallel_isolated' }); + expect(pacer.admit).toHaveBeenCalledTimes(4); + + for (const socket of fakeSockets) if (socket.listenerCount('open') > 0) socket.emit('open'); + emitTextResponse(fakeSockets[2]!, 'resp_diverged', 'ok'); + emitTextResponse(fakeSockets[3]!, 'resp_parallel', 'ok'); + await readAll(diverged); + await readAll(parallel); + }); + + it('does not pace the replacement a transport retry opens', async () => { + // Deliberate exemption: it recovers a request that was already admitted, + // it is capped at one per request, and it is built inside a socket + // callback. Pinned so the exemption cannot drift into an accident. + const pacer = recordingPacer(); + const wsFetch = createResponsesWebSocketFetch(WS_URL, undefined, { + accountId: 'acct-pacing-transport-replacement', + pacer, + }); + const response = await wsFetch('https://x', { + method: 'POST', headers: {}, + body: JSON.stringify(sessionPayload([{ role: 'user', content: [{ type: 'input_text', text: 'hi' }] }])), + }); + fakeSockets[0]!.emit('open'); + fakeSockets[0]!.emit('error', new Error('connection reset')); + + expect(fakeSockets).toHaveLength(2); + expect(pacer.admit).toHaveBeenCalledTimes(1); + + fakeSockets[1]!.emit('open'); + emitTextResponse(fakeSockets[1]!, 'resp_replaced', 'ok'); + await readAll(response); + }); + + it('does not pace the replacement a missing previous response opens', async () => { + // The second unpaced replacement path, and the one the deep doc used to + // omit entirely. + const pacer = recordingPacer(); + const wsFetch = createResponsesWebSocketFetch(WS_URL, undefined, { + accountId: 'acct-pacing-prev-missing', + pacer, + }); + const input = [{ role: 'user', content: [{ type: 'input_text', text: 'hi' }] }]; + const first = await wsFetch('https://x', { + method: 'POST', headers: {}, body: JSON.stringify(sessionPayload(input)), + }); + fakeSockets[0]!.emit('open'); + emitTextResponse(fakeSockets[0]!, 'resp_prev_1', 'ok'); + await readAll(first); + expect(pacer.admit).toHaveBeenCalledTimes(1); + + const echoed = [ + ...input, + { role: 'assistant', content: [{ type: 'output_text', text: 'ok' }] }, + { role: 'user', content: [{ type: 'input_text', text: 'next' }] }, + ]; + const second = await wsFetch('https://x', { + method: 'POST', headers: {}, body: JSON.stringify(sessionPayload(echoed)), + }); + const continued = fakeSockets.length; + fakeSockets[0]!.emit('message', Buffer.from(JSON.stringify({ + type: 'error', status: 400, + error: { code: 'previous_response_not_found', message: 'gone' }, + }))); + + // A replacement socket was opened without consulting the pacer. + expect(fakeSockets.length).toBeGreaterThan(continued); + expect(pacer.admit).toHaveBeenCalledTimes(1); + + const replacement = lastSocket(); + replacement.emit('open'); + emitTextResponse(replacement, 'resp_prev_2', 'ok'); + await readAll(second); + }); + + it('demotes a concurrent sibling even when neither request waited', async () => { + // `admit` is async, so even an immediate admission resumes a microtask + // later and BOTH requests classify themselves against an empty partition + // before either registers. Gating the re-check on having waited would let + // both open a persistent head for one key. + // + // The barrier holds both inside the pacer until both have been classified, + // which is the interleaving that makes this deterministic rather than a + // race the scheduler happens to win. + const diagnostics: ResponsesWebSocketDiagnosticEvent[] = []; + let admitted = 0; + let barrierArmed = false; + let openBarrier: (() => void) | undefined; + const barrier = new Promise(resolve => { openBarrier = resolve; }); + const debug: string[] = []; + const wsFetch = createResponsesWebSocketFetch(WS_URL, message => debug.push(message), { + accountId: 'acct-pacing-concurrent', + pacer: { + admit: async (): Promise => { + if (!barrierArmed) return { kind: 'admitted', waitedMs: 0 }; + admitted += 1; + if (admitted === 2) openBarrier!(); + await barrier; + return { kind: 'admitted', waitedMs: 0 }; + }, + }, + onDiagnostic: event => diagnostics.push(event), + }); + const send = (text: string) => wsFetch('https://x', { + method: 'POST', headers: {}, + body: JSON.stringify(sessionPayload([{ role: 'user', content: [{ type: 'input_text', text }] }])), + }); + + // Resolve the mocked `ws` module BEFORE the concurrent phase: two dynamic + // imports of a mocked module in flight at once can race in vitest and hand + // one caller the unmocked module. + const warmup = await wsFetch('https://x', { + method: 'POST', headers: {}, + body: JSON.stringify(sessionPayload( + [{ role: 'user', content: [{ type: 'input_text', text: 'warmup' }] }], + { prompt_cache_key: 'relay-session-warmup' }, + )), + }); + lastSocket().emit('open'); + emitTextResponse(lastSocket(), 'resp_warmup', 'ok'); + await readAll(warmup); + barrierArmed = true; + + const [alpha, beta] = await Promise.all([send('alpha'), send('beta')]); + + // The warmup produced a decision of its own; the pair is the last two. + const decisions = diagnostics.filter(event => event.event === 'ws_head_decision').slice(-2); + expect(decisions).toHaveLength(2); + // Both were classified against an empty partition — the race is real... + expect(decisions.map(event => event.candidateCount)).toEqual([0, 0]); + // ...but only one of them may end up holding a persistent head for it. + // Without the demotion both are 'nursery': two persistent heads, one key. + expect(decisions.map(event => event.createdGeneration).sort()) + .toEqual(['isolated', 'nursery']); + expect(debug).toContain('ws: parallel request using an isolated socket after pacing'); + + for (const socket of fakeSockets.slice(-2)) { + socket.emit('open'); + emitTextResponse(socket, `resp_${fakeSockets.indexOf(socket)}`, 'ok'); + } + await Promise.all([readAll(alpha), readAll(beta)]); + }); + + it('ages the heads it reports from after the wait, not from arrival', async () => { + // The decision record must not mix head ages read on arrival with pool + // counts read after the wait; a queued request can be seconds old by then. + let clockMs = 0; + let waits = 0; + const diagnostics: ResponsesWebSocketDiagnosticEvent[] = []; + const wsFetch = createResponsesWebSocketFetch(WS_URL, undefined, { + accountId: 'acct-pacing-clock', + now: () => clockMs, + pacer: { + admit: async (): Promise => { + waits += 1; + if (waits === 1) return { kind: 'admitted', waitedMs: 0 }; + clockMs += 4_000; + return { kind: 'admitted', waitedMs: 4_000 }; + }, + }, + onDiagnostic: event => diagnostics.push(event), + }); + const body = (text: string) => JSON.stringify(sessionPayload( + [{ role: 'user', content: [{ type: 'input_text', text }] }], + )); + + const first = await wsFetch('https://x', { method: 'POST', headers: {}, body: body('first') }); + lastSocket().emit('open'); + emitTextResponse(lastSocket(), 'resp_clock', 'ok'); + await readAll(first); + + // Same partition, different history: an idle head is reported, not reused. + const second = await wsFetch('https://x', { method: 'POST', headers: {}, body: body('second') }); + const decision = diagnostics.filter(event => event.event === 'ws_head_decision').at(-1); + expect(decision).toMatchObject({ pacingWaitedMs: 4_000 }); + // Read on arrival this head looks freshly used; it is 4s idle by admission. + expect((decision as { heads: Array<{ idleMs: number }> }).heads[0]!.idleMs).toBe(4_000); + + lastSocket().emit('open'); + emitTextResponse(lastSocket(), 'resp_clock_2', 'ok'); + await readAll(second); + }); + + it('uses the shared bucket and its env var when no pacer is injected', async () => { + // Nothing here injects a pacer, so this fails if production stops + // consulting the shared one — or stops reading the environment. + // 3/minute, not 1: three is the lowest rate at which refusing can still + // reach a refill inside a request's retry schedule at the default + // timeouts, and it is far from the default of 60, so this still proves the + // environment is read. Rate 1 is covered by the test below. + process.env.CLODEX_WS_MAX_NEW_CONNECTIONS_PER_MIN = '3'; + try { + resetResponsesWebSocketConnectionsForTests(); + const diagnostics: ResponsesWebSocketDiagnosticEvent[] = []; + const wsFetch = createResponsesWebSocketFetch(WS_URL, undefined, { + accountId: 'acct-pacing-shared', + onDiagnostic: event => diagnostics.push(event), + }); + const open = async (text: string) => { + const response = await wsFetch('https://x', { + method: 'POST', + headers: {}, + body: JSON.stringify(sessionPayload( + [{ role: 'user', content: [{ type: 'input_text', text }] }], + { prompt_cache_key: `relay-session-${text}` }, + )), + }); + const socket = fakeSockets[fakeSockets.length - 1]; + if (socket && socket.listenerCount('open') > 0) { + socket.emit('open'); + emitTextResponse(socket, `resp_${text}`, 'ok'); + } + return readAll(response); + }; + + // At three new connections per minute the burst of ten is free; the + // eleventh needs twenty seconds, far past the wait bound, so it is + // refused. + for (let index = 0; index < 10; index += 1) await open(`burst-${index}`); + expect(fakeSockets).toHaveLength(10); + + const refusedBody = await open('past-the-burst'); + expect(fakeSockets).toHaveLength(10); + expect(await classifyThroughSdk(refusedBody)).toMatchObject({ + statusCode: 429, + isRetryable: true, + }); + expect(diagnostics).toContainEqual(expect.objectContaining({ + event: 'ws_new_connection_paced', + outcome: 'refused', + })); + } finally { + delete process.env.CLODEX_WS_MAX_NEW_CONNECTIONS_PER_MIN; + resetResponsesWebSocketConnectionsForTests(); + } + }); + + it('admits late instead of refusing at a rate a retry cannot outlast', async () => { + // REGRESSION, end to end through the production singleton. Capping the + // refusal hint at the wait bound stopped requests overrunning their + // deadline, but at 1/minute the first free slot is 60s away while the whole + // retry schedule is spent inside 20s — so every refused request used up its + // retries and died as a rate-limit error, manufactured by the feature meant + // to reduce them. Nothing may be refused here. + process.env.CLODEX_WS_MAX_NEW_CONNECTIONS_PER_MIN = '1'; + try { + resetResponsesWebSocketConnectionsForTests(); + const diagnostics: ResponsesWebSocketDiagnosticEvent[] = []; + const wsFetch = createResponsesWebSocketFetch(WS_URL, undefined, { + accountId: 'acct-pacing-lowrate', + onDiagnostic: event => diagnostics.push(event), + }); + const open = async (text: string) => { + const response = await wsFetch('https://x', { + method: 'POST', + headers: {}, + body: JSON.stringify(sessionPayload( + [{ role: 'user', content: [{ type: 'input_text', text }] }], + { prompt_cache_key: `relay-session-${text}` }, + )), + }); + const socket = fakeSockets[fakeSockets.length - 1]; + if (socket && socket.listenerCount('open') > 0) { + socket.emit('open'); + emitTextResponse(socket, `resp_${text}`, 'ok'); + } + return readAll(response); + }; + + // Eleven, opened sequentially: ten spend the burst and the eleventh has + // to wait. Arrivals here are sequential rather than simultaneous, so each + // one refills roughly what it consumes and keeps paying the bound — the + // ceiling is weaker than the configured rate, but it is still a ceiling. + for (let index = 0; index < 11; index += 1) await open(`low-${index}`); + + // Every one of them opened a connection; none was turned away. + expect(fakeSockets).toHaveLength(11); + expect(diagnostics).not.toContainEqual(expect.objectContaining({ + event: 'ws_new_connection_paced', + outcome: 'refused', + })); + + // And pacing is still SHAPING, not switched off: the first arrival past + // the burst really did wait out the bound. Asserting only that sockets + // opened cannot tell this apart from disabled pacing. + const waits = diagnostics + .filter((event): event is typeof event & { waitedMs: number } => + event.event === 'ws_new_connection_paced' + && (event as { outcome?: string }).outcome === 'admitted' + && typeof (event as { waitedMs?: unknown }).waitedMs === 'number') + .map(event => event.waitedMs); + expect(waits).toHaveLength(1); + expect(waits[0]).toBeGreaterThan(1_000); + } finally { + delete process.env.CLODEX_WS_MAX_NEW_CONNECTIONS_PER_MIN; + resetResponsesWebSocketConnectionsForTests(); + } + // Deliberately at the SHIPPED timeouts, which is where the failure was + // measured, so one request really does wait out the ~4.8s bound. + }, 20_000); + + it('opens no connection for a request cancelled while it was queued', async () => { + const diagnostics: ResponsesWebSocketDiagnosticEvent[] = []; + const aborted = Object.assign(new Error('The operation was aborted'), { name: 'AbortError' }); + const wsFetch = createResponsesWebSocketFetch(WS_URL, undefined, { + accountId: 'acct-pacing-abort', + pacer: { admit: vi.fn(async () => { throw aborted; }) }, + onDiagnostic: event => diagnostics.push(event), + }); + + await expect(wsFetch('https://x', { + method: 'POST', headers: {}, body: JSON.stringify(sessionPayload([])), + })).rejects.toBe(aborted); + + expect(fakeSockets).toHaveLength(0); + expect(diagnostics).toContainEqual(expect.objectContaining({ + event: 'ws_new_connection_paced', + outcome: 'aborted', + })); + // The request never reached a head decision, so none is reported. + expect(diagnostics.some(event => event.event === 'ws_head_decision')).toBe(false); + }); +}); diff --git a/tests/ws-upgrade-pacer.test.ts b/tests/ws-upgrade-pacer.test.ts new file mode 100644 index 00000000..aa5d1a16 --- /dev/null +++ b/tests/ws-upgrade-pacer.test.ts @@ -0,0 +1,976 @@ +import { describe, it, expect, vi } from 'vitest'; +import { + DEFAULT_WS_NEW_CONNECTIONS_PER_MIN, + MAX_WS_NEW_CONNECTIONS_PER_MIN, + WsUpgradePacer, + WS_NEW_CONNECTIONS_PER_MIN_ENV, + WS_NEW_CONNECTION_BURST, + WS_NEW_CONNECTION_MAX_WAIT_CEILING_MS, + resetWsUpgradePacerForTests, + canRefuseAtRate, + pacedRetryAfterSeconds, + refusalScheduleMs, + resolvedPacingBudget, + wsNewConnectionMaxWaitMs, + wsNewConnectionsPerMinute, + type UpgradeAdmission, +} from '../src/oauth/ws-upgrade-pacer.js'; +import { installParentNoticeSink } from '../src/parent-notice.js'; +import { + UPSTREAM_IDLE_TIMEOUT_ENV, + UPSTREAM_MAX_RETRIES_ENV, + UPSTREAM_TOTAL_TIMEOUT_ENV, + upstreamRequestBudget, +} from '../src/upstream-retry.js'; + +/** The AI SDK's fallback ladder: 2s, 4s, 8s, … with no jitter. */ +function sdkBackoffMs(maxRetries: number): number { + return 2_000 * (2 ** maxRetries - 1); +} + +/** + * Run `body` against a known-clean upstream-budget environment. + * + * The budget is resolved from `process.env` in production, so a stray ambient + * value would silently change what these assertions mean. + */ +function withUpstreamEnv(values: Record, body: () => T): T { + const keys = [UPSTREAM_IDLE_TIMEOUT_ENV, UPSTREAM_TOTAL_TIMEOUT_ENV, UPSTREAM_MAX_RETRIES_ENV]; + const saved = Object.fromEntries(keys.map(key => [key, process.env[key]])); + try { + for (const key of keys) { + const value = values[key]; + if (value === undefined) delete process.env[key]; + else process.env[key] = value; + } + return body(); + } finally { + for (const key of keys) { + const value = saved[key]; + if (value === undefined) delete process.env[key]; + else process.env[key] = value; + } + } +} + +/** + * Deterministic clock. Nothing in this file sleeps: time only moves when a test + * advances it, so a rate expressed per minute is exercised in microseconds. + */ +class TestClock { + time = 0; + private timers: Array<{ at: number; fire: () => void }> = []; + + now = (): number => this.time; + + schedule = (ms: number, fire: () => void): (() => void) => { + const timer = { at: this.time + ms, fire }; + this.timers.push(timer); + return () => { this.timers = this.timers.filter(candidate => candidate !== timer); }; + }; + + /** Timers still armed. A leak shows up here. */ + get pending(): number { + return this.timers.length; + } + + async advance(ms: number): Promise { + const target = this.time + ms; + for (;;) { + const due = [...this.timers].sort((left, right) => left.at - right.at)[0]; + if (!due || due.at > target) break; + this.timers = this.timers.filter(candidate => candidate !== due); + this.time = due.at; + due.fire(); + await flush(); + } + this.time = target; + await flush(); + } +} + +/** Drain the microtask queue so awaiting callers resume before we assert. */ +function flush(): Promise { + return new Promise(resolve => { setImmediate(resolve); }); +} + +interface Tracked { + state: () => 'pending' | 'admitted' | 'refused' | 'rejected'; + admission: () => UpgradeAdmission | undefined; + /** Clock reading when this admission settled. */ + settledAt: () => number | undefined; + error: () => unknown; +} + +/** Track an admission's outcome without ever awaiting a pending one. */ +function track(clock: TestClock, promise: Promise): Tracked { + let state: 'pending' | 'admitted' | 'refused' | 'rejected' = 'pending'; + let admission: UpgradeAdmission | undefined; + let settledAt: number | undefined; + let error: unknown; + promise.then( + value => { state = value.kind; admission = value; settledAt = clock.time; }, + reason => { state = 'rejected'; error = reason; settledAt = clock.time; }, + ); + return { + state: () => state, + admission: () => admission, + settledAt: () => settledAt, + error: () => error, + }; +} + +/** + * A bound of exactly `ms`, expressed through the options production uses. + * `maxRetries` must be non-zero: with retries off the pacer never refuses, + * which is its own test below. + */ +function boundedBy(ms: number): { idleTimeoutMs: number; maxRetries: number } { + return { idleTimeoutMs: 4 * ms + 2_000, maxRetries: 1 }; +} + +describe('wsNewConnectionMaxWaitMs', () => { + // Whatever retry budget clodex ships, the ladder must fit the deadline. The + // matrix is the point: a bound derived for one retry count and used with + // another is how this becomes a merge-order bug. + it.each([0, 1, 2, 3, 5])('keeps a %i-retry ladder inside the deadline it shares', maxRetries => { + const idleTimeoutMs = 120_000; + const bound = wsNewConnectionMaxWaitMs(idleTimeoutMs, maxRetries); + const backoffMs = sdkBackoffMs(maxRetries); + + expect(bound).toBeGreaterThan(0); + expect((maxRetries + 1) * bound + backoffMs).toBeLessThan(idleTimeoutMs); + // Pacing takes at most half of what the backoff ladder leaves, so the + // request keeps an equal share for the provider's own first byte. + expect((maxRetries + 1) * bound).toBeLessThanOrEqual((idleTimeoutMs - backoffMs) / 2); + }); + + it('is total on degenerate input rather than failing open into an unbounded wait', () => { + // NaN would make every `waitMs > maxWaitMs` comparison false. + expect(wsNewConnectionMaxWaitMs(Number.NaN, 2)).toBe(0); + expect(wsNewConnectionMaxWaitMs(120_000, Number.NaN)).toBe(0); + expect(wsNewConnectionMaxWaitMs(120_000, -1)).toBe(0); + expect(wsNewConnectionMaxWaitMs(0, 2)).toBe(0); + expect(wsNewConnectionMaxWaitMs(120_000, 1_000)).toBe(0); + }); + + it('never exceeds the ceiling even when the deadline is generous', () => { + expect(wsNewConnectionMaxWaitMs(3_600_000, 0)).toBe(WS_NEW_CONNECTION_MAX_WAIT_CEILING_MS); + }); + + it('stops queueing entirely when the ladder alone would exhaust the deadline', () => { + // 10s deadline against a 62s ladder: there is no budget to wait in, so + // nothing is queued and the overflow is refused instead. + expect(wsNewConnectionMaxWaitMs(10_000, 5)).toBe(0); + }); +}); + +/** + * The bound against the budget the paced request actually spends. + * + * Both terms of `(maxRetries + 1) x bound + totalBackoff < idleTimeout` are + * user-configurable and they interact, so checking the arithmetic at one + * deadline proves nothing about the rest of the range. These drive real + * environments through the real budget resolver, which is the composition no + * single change's CI exercises. + * + * Two mutations this is built to catch, both of which ship a default-on hard + * abort: a flat 15s bound reads 6 x 15s + 62s = 152s against the default 120s + * deadline, and a bound derived from a hardcoded 120s is wrong at every + * deadline the user can actually configure. + */ +describe('the wait bound against the resolved request budget', () => { + const IDLE_TIMEOUTS = [ + undefined, // shipped default + '9000', // below the floor: clamps up to 10s + '10000', // the floor + '14001', // barely wider than its own backoff ladder + '20000', + '30000', + '120000', + '300000', + '3600000', // the ceiling + '7200000', // above the ceiling: clamps down + 'nonsense', // malformed: falls back to the default + ]; + const RETRIES = [undefined, '0', '1', '2', '5', '10', '99', '-1']; + + it.each(IDLE_TIMEOUTS.flatMap(idle => RETRIES.map(retries => [idle, retries] as const)))( + 'keeps the whole retry ladder inside the deadline (idle=%s, retries=%s)', + (idle, retries) => { + const budget = withUpstreamEnv( + { + [UPSTREAM_IDLE_TIMEOUT_ENV]: idle, + [UPSTREAM_MAX_RETRIES_ENV]: retries, + // Pinned high so the idle timeout is never lowered to meet it; the + // pair's interaction has its own case below. + [UPSTREAM_TOTAL_TIMEOUT_ENV]: '21600000', + }, + () => upstreamRequestBudget({ warn: () => {} }), + ); + const bound = wsNewConnectionMaxWaitMs(budget.idleTimeoutMs, budget.maxRetries); + + // The inequality, stated exactly as the module documents it. + expect((budget.maxRetries + 1) * bound + sdkBackoffMs(budget.maxRetries)) + .toBeLessThan(budget.idleTimeoutMs); + expect(bound).toBeGreaterThanOrEqual(0); + expect(bound).toBeLessThanOrEqual(WS_NEW_CONNECTION_MAX_WAIT_CEILING_MS); + + // And the same inequality against what the request ACTUALLY spends + // between attempts. The SDK substitutes a supplied hint for its own rung + // rather than taking the larger, so the real per-gap delay is the pacer's + // capped hint wherever that exceeds the rung. Budgeting only the ladder + // is what let a 30s hint overrun a 10s deadline. + // Derived here, NOT by calling the function under test. The first version + // of this line called `pacedRetryAfterSeconds(Number.POSITIVE_INFINITY, + // bound)` — and the helper maps every non-finite deficit to one second, + // so the oracle read 1000ms for EVERY bound and the matrix silently + // tested one degenerate value. A mutation that removed the cap above a 2s + // bound survived all 131 tests under it. + const hintCapMs = Math.max(1_000, Math.floor(bound / 1_000) * 1_000); + let worstGapsMs = 0; + for (let retry = 0; retry < budget.maxRetries; retry += 1) { + worstGapsMs += Math.max(hintCapMs, 2_000 * 2 ** retry); + } + expect((budget.maxRetries + 1) * bound + worstGapsMs) + .toBeLessThan(budget.idleTimeoutMs); + }, + ); + + it.each([ + // [bound, deficit, expected hint seconds] + [666, 30_000, 1], // sub-second bound: the floor, not the deficit + [2_000, 30_000, 2], + [4_833, 30_000, 4], // the shipped default bound + [15_000, 30_000, 15], // the ceiling + [15_000, 3_000, 3], // a deficit UNDER the cap is passed through + [4_833, 1_200, 2], // rounded up, still under the cap + ])('caps a refusal hint at bound=%ims (deficit %ims) to %is', (bound, deficit, expected) => { + // Direct cases at the two bounds production actually uses. The matrix above + // could not see these: it derived its own expectation, so a cap that only + // applied below 2s passed it. + expect(pacedRetryAfterSeconds(deficit, bound)).toBe(expected); + // Never longer than the bound, except where whole seconds cannot express + // it, and never zero. + expect(pacedRetryAfterSeconds(deficit, bound) * 1_000) + .toBeLessThanOrEqual(Math.max(1_000, bound)); + expect(pacedRetryAfterSeconds(deficit, bound)).toBeGreaterThanOrEqual(1); + }); + + it('is total on a degenerate deficit or bound', () => { + expect(pacedRetryAfterSeconds(Number.NaN, 15_000)).toBe(1); + expect(pacedRetryAfterSeconds(Number.POSITIVE_INFINITY, 15_000)).toBe(1); + expect(pacedRetryAfterSeconds(30_000, Number.NaN)).toBe(1); + expect(pacedRetryAfterSeconds(-5, 15_000)).toBe(1); + }); + + it('holds when a short total timeout drags the idle timeout down with it', () => { + // #171's pair rule: an explicit total below the idle lowers the idle. The + // pacer must size itself against the lowered value, not the requested one. + const budget = withUpstreamEnv( + { + [UPSTREAM_IDLE_TIMEOUT_ENV]: '600000', + [UPSTREAM_TOTAL_TIMEOUT_ENV]: '60000', + }, + () => upstreamRequestBudget({ warn: () => {} }), + ); + expect(budget.idleTimeoutMs).toBe(60_000); + + const bound = wsNewConnectionMaxWaitMs(budget.idleTimeoutMs, budget.maxRetries); + expect((budget.maxRetries + 1) * bound + sdkBackoffMs(budget.maxRetries)) + .toBeLessThan(60_000); + }); + + it('reads the budget in force rather than assuming one', () => { + // Unset, this is #171's derived default: five retries inside a 120s window. + expect(withUpstreamEnv({}, () => resolvedPacingBudget())) + .toEqual({ idleTimeoutMs: 120_000, maxRetries: 5 }); + expect(withUpstreamEnv({ [UPSTREAM_MAX_RETRIES_ENV]: '2' }, () => resolvedPacingBudget())) + .toEqual({ idleTimeoutMs: 120_000, maxRetries: 2 }); + expect(withUpstreamEnv({ [UPSTREAM_IDLE_TIMEOUT_ENV]: '30000' }, () => resolvedPacingBudget())) + // A 30s window cannot fund a fourth retry, so the ceiling caps it at three. + .toEqual({ idleTimeoutMs: 30_000, maxRetries: 3 }); + }); + + it('is what the pacer actually uses, at the shipped default', () => { + const pacer = withUpstreamEnv({}, () => new WsUpgradePacer()); + // Independent oracle: the arithmetic written out, not the function reused. + // (120000 - 62000) / 2 / 6 attempts. + expect(pacer.maxWaitMs).toBe(4_833); + expect(6 * pacer.maxWaitMs + 62_000).toBeLessThan(120_000); + }); + + it('shrinks its bound when the user shortens the deadline', () => { + // A pacer that derived its bound from a hardcoded 120s would read 13250ms + // here — over six times the deadline's actual share. + const pacer = withUpstreamEnv( + { [UPSTREAM_IDLE_TIMEOUT_ENV]: '30000' }, + () => new WsUpgradePacer(), + ); + expect(pacer.maxWaitMs).toBe(2_000); + expect(wsNewConnectionMaxWaitMs(120_000, 3)).toBe(13_250); + }); + + it('grows its bound only up to the ceiling when the user lengthens the deadline', () => { + const pacer = withUpstreamEnv( + { [UPSTREAM_IDLE_TIMEOUT_ENV]: '3600000' }, + () => new WsUpgradePacer(), + ); + expect(pacer.maxWaitMs).toBe(WS_NEW_CONNECTION_MAX_WAIT_CEILING_MS); + }); + + it('never asks the client to wait longer than the deadline funds', async () => { + // REGRESSION. The bound budgets the SDK's 2s/4s/8s ladder, but a refusal + // carries its own `retry-after`, and `getRetryDelayInMs` (ai@7.0.22, + // util/retry-with-exponential-backoff.ts) SUBSTITUTES that hint for the + // rung whenever it is under 60s. So the hint, not the ladder, is what the + // request spends between attempts — and an uncapped hint taken from the + // token deficit silently replaced the term the bound was derived against. + // + // 6/minute puts a token 10s away, so the third arrival needs 20s against a + // 15s bound. Before the cap it asked for 20s. The bound is deliberately the + // 15s ceiling here: a cap that only bit below 2s passed every other test. + const clock = new TestClock(); + const pacer = new WsUpgradePacer({ + ratePerMinute: 6, + burst: 1, + ...boundedBy(15_000), + now: clock.now, + schedule: clock.schedule, + }); + expect(pacer.maxWaitMs).toBe(15_000); + + const admissions = Array.from({ length: 4 }, () => track(clock, pacer.admit())); + await flush(); + + const refusals = admissions + .map(entry => entry.admission()) + .filter((value): value is Extract => + value?.kind === 'refused'); + expect(refusals.length).toBeGreaterThan(0); + + for (const refusal of refusals) { + // The deficit is still reported honestly; only the hint is capped. + expect(refusal.requiredWaitMs).toBeGreaterThan(pacer.maxWaitMs); + expect(refusal.retryAfterSeconds).toBe(15); + } + }); + + it('never asks for zero seconds when the bound is under a second', async () => { + // A sub-second bound is where flooring to whole seconds would produce + // `retry-after: 0` — retry immediately, burning the retry budget in + // milliseconds. 10s deadline funds a 666ms bound. + const clock = new TestClock(); + const pacer = new WsUpgradePacer({ + ratePerMinute: 60, + burst: 1, + idleTimeoutMs: 10_000, + maxRetries: 2, + now: clock.now, + schedule: clock.schedule, + }); + expect(pacer.maxWaitMs).toBe(666); + + const admissions = Array.from({ length: 6 }, () => track(clock, pacer.admit())); + await flush(); + const refusals = admissions + .map(entry => entry.admission()) + .filter((value): value is Extract => + value?.kind === 'refused'); + + expect(refusals.length).toBeGreaterThan(0); + for (const refusal of refusals) { + expect(refusal.retryAfterSeconds).toBe(1); + } + }); + + it.each([1, 2])( + 'admits late instead of failing fast at %i new connections per minute', + async ratePerMinute => { + // REGRESSION for the second failure the cap introduced. At 1/minute the + // first token is 60s away, but six attempts four seconds apart are all + // spent inside 20s — so every refused request exhausted its retries + // before a token could exist and died as a rate-limit error, which is the + // exact failure this feature exists to reduce. Measured against the real + // SDK: 10 of 10 terminal at both rates. + // + // Both rates are documented as supported (README: 1-600). + const clock = new TestClock(); + const pacer = new WsUpgradePacer({ + ratePerMinute, + now: clock.now, + schedule: clock.schedule, + }); + expect(pacer.maxWaitMs).toBeGreaterThan(0); + + const admissions = Array.from({ length: 20 }, () => track(clock, pacer.admit())); + await flush(); + await clock.advance(pacer.maxWaitMs); + + // Nothing is refused, because a refusal here could not be retried into + // an admission before the request's deadline. + expect(admissions.some(entry => entry.state() === 'refused')).toBe(false); + expect(admissions.every(entry => entry.state() === 'admitted')).toBe(true); + expect(clock.pending).toBe(0); + + // ...but this is NOT the same as pacing switched off. The burst is still + // shaped: the first arrival past it waits out the whole bound, and only + // the ones past the debt floor go straight through. Without this, zeroing + // the refill rate here would be indistinguishable from disabled pacing. + const delayed = admissions.filter(entry => (entry.settledAt() ?? 0) > 0); + expect(delayed).toHaveLength(1); + expect(delayed[0]!.settledAt()).toBe(pacer.maxWaitMs); + expect(admissions.filter(entry => entry.settledAt() === 0)).toHaveLength(19); + }, + ); + + it.each([ + // [bound, maxRetries, rate, canRefuse] — the shipped budget AND a + // user-shortened one. Every committed case used to be the shipped 4833ms + // bound with five retries, so a rule that ignored the budget entirely + // (`ratePerMinute >= 3`) passed the whole suite while recreating the + // terminal-failure defect under CLODEX_UPSTREAM_IDLE_TIMEOUT_MS=10000. + [4_833, 5, 60, true], // shipped defaults + [4_833, 5, 3, true], // lowest rate the shipped budget can still serve + [4_833, 5, 2, false], + [4_833, 5, 1, false], + [666, 2, 3, false], // 10s budget: schedule 2s vs a 20s refill + [666, 2, 29, false], // still short of the threshold + [666, 2, 30, true], // exactly at it: schedule 2s vs a 2s refill + [666, 2, 60, true], + [15_000, 1, 4, true], // 15s schedule vs a 15s refill + [15_000, 1, 3, false], // 15s schedule vs a 20s refill + ])( + 'decides refusability from the whole budget: bound=%i retries=%i rate=%i -> %s', + (bound, maxRetries, rate, expected) => { + expect(canRefuseAtRate(bound, maxRetries, rate)).toBe(expected); + }, + ); + + it.each([ + // The same thing behaviourally, on a NON-default budget, because a truth + // table over the helper cannot prove the constructor consults it. + [3, false], + [30, true], + ])('honours a shortened deadline when deciding to refuse (rate %i)', async (rate, refuses) => { + const clock = new TestClock(); + const pacer = new WsUpgradePacer({ + ratePerMinute: rate, + burst: 1, + idleTimeoutMs: 10_000, + maxRetries: 2, + now: clock.now, + schedule: clock.schedule, + }); + expect(pacer.maxWaitMs).toBe(666); + + const admissions = Array.from({ length: 6 }, () => track(clock, pacer.admit())); + await flush(); + await clock.advance(pacer.maxWaitMs); + + expect(admissions.some(entry => entry.state() === 'refused')).toBe(refuses); + }); + + it('still refuses at the shipped rate, where a retry can outlast a refill', () => { + // The guard above must not disable refusals at the default. One token per + // second is well inside a 5-retry schedule ~4s apart. + expect(canRefuseAtRate(4_833, 5, DEFAULT_WS_NEW_CONNECTIONS_PER_MIN)).toBe(true); + expect(refusalScheduleMs(4_833, 5)).toBe(20_000); + // ...and must disable them where it cannot. + expect(canRefuseAtRate(4_833, 5, 1)).toBe(false); + expect(canRefuseAtRate(4_833, 5, 2)).toBe(false); + // Retries off means nothing would retry a refusal at any rate. + expect(canRefuseAtRate(4_833, 0, 60)).toBe(false); + }); + + it('does not let a rate notice suppress the notice that pacing turned itself off', () => { + // The two notices shared one dedupe set keyed by an ARBITRARY environment + // string, so a rate value spelled like the disabled key silenced the one + // notice that must never go missing. Namespaced keys keep them apart. + resetWsUpgradePacerForTests(); + const notices: string[] = []; + const release = installParentNoticeSink(line => { notices.push(line); }); + try { + // Crafted to collide with the disabled notice's `${retries}:${idle}` key. + expect(wsNewConnectionsPerMinute({ + [WS_NEW_CONNECTIONS_PER_MIN_ENV]: 'disabled:3:14001', + })).toBe(DEFAULT_WS_NEW_CONNECTIONS_PER_MIN); + expect(notices.some(line => line.includes('ignoring'))).toBe(true); + + const pacer = withUpstreamEnv( + { [UPSTREAM_IDLE_TIMEOUT_ENV]: '14001' }, + () => new WsUpgradePacer(), + ); + expect(pacer.maxWaitMs).toBe(0); + } finally { + release(); + } + expect(notices.some(line => line.includes('not pacing new OpenAI connections'))).toBe(true); + }); + + it('refuses only past a large simultaneous fan-out at the shipped defaults', async () => { + // Reachability, measured rather than argued. This is where a user first + // meets a refusal, and it is NOT out of reach: 15 simultaneous agents is + // inside the many-agent workload this feature targets. #171's five-retry + // default tightens the bound and moves the threshold down from 26 to 15, + // so the rebase made refusals easier to reach, not harder. + const clock = new TestClock(); + const pacer = withUpstreamEnv( + {}, + () => new WsUpgradePacer({ now: clock.now, schedule: clock.schedule }), + ); + const admissions = Array.from({ length: 30 }, () => track(clock, pacer.admit())); + await flush(); + + expect(admissions.findIndex(entry => entry.state() === 'refused')).toBe(14); + expect(admissions.slice(0, 14).some(entry => entry.state() === 'refused')).toBe(false); + }); + + it('shapes 25 connections and then stops, with retries turned off', async () => { + // With retries off the ladder costs nothing, so the bound is the flat + // ceiling and the shaped burst is the doc's `burst + bound x refill` = 25. + // Pinned because the bound now differs between this mode and the default. + const clock = new TestClock(); + const pacer = withUpstreamEnv( + { [UPSTREAM_MAX_RETRIES_ENV]: '0' }, + () => new WsUpgradePacer({ now: clock.now, schedule: clock.schedule }), + ); + expect(pacer.maxWaitMs).toBe(WS_NEW_CONNECTION_MAX_WAIT_CEILING_MS); + + const admissions = Array.from({ length: 40 }, () => track(clock, pacer.admit())); + await flush(); + await clock.advance(pacer.maxWaitMs); + + expect(admissions.every(entry => entry.state() === 'admitted')).toBe(true); + expect(admissions.filter(entry => (entry.settledAt() ?? 0) > 0)).toHaveLength(15); + }); + + it('stops pacing rather than holding a request its deadline cannot fund', async () => { + // 14001ms funds a three-retry ladder costing 14000ms. There is 1ms left, so + // there is no room to queue: pacing off is the only safe reading. + const pacer = withUpstreamEnv( + { [UPSTREAM_IDLE_TIMEOUT_ENV]: '14001' }, + () => new WsUpgradePacer(), + ); + expect(pacer.maxWaitMs).toBe(0); + + // Degrading means admitting without delay, never refusing everything past + // the burst and never queueing past the deadline. + const clock = new TestClock(); + const disabled = withUpstreamEnv( + { [UPSTREAM_IDLE_TIMEOUT_ENV]: '14001' }, + () => new WsUpgradePacer({ now: clock.now, schedule: clock.schedule }), + ); + const admissions = await Promise.all( + Array.from({ length: 50 }, () => disabled.admit()), + ); + expect(admissions.every(entry => entry.kind === 'admitted' && entry.waitedMs === 0)).toBe(true); + expect(clock.pending).toBe(0); + }); +}); + +describe('WsUpgradePacer', () => { + it('admits a burst immediately and then holds arrivals to the configured rate', async () => { + const clock = new TestClock(); + // 60/minute is one connection per second, so each held arrival is a second. + const pacer = new WsUpgradePacer({ + ratePerMinute: 60, + burst: 10, + ...boundedBy(15_000), + now: clock.now, + schedule: clock.schedule, + }); + + const admissions = Array.from({ length: 13 }, () => track(clock, pacer.admit())); + await flush(); + + expect(admissions.slice(0, 10).map(entry => entry.state())).toEqual(Array(10).fill('admitted')); + expect(admissions.slice(0, 10).every(entry => entry.admission()?.kind === 'admitted' + && entry.admission().waitedMs === 0)).toBe(true); + expect(admissions.slice(10).map(entry => entry.state())).toEqual(['pending', 'pending', 'pending']); + + await clock.advance(1_000); + expect(admissions[10]!.admission()).toEqual({ kind: 'admitted', waitedMs: 1_000 }); + expect(admissions[11]!.state()).toBe('pending'); + + await clock.advance(1_000); + expect(admissions[11]!.admission()).toEqual({ kind: 'admitted', waitedMs: 2_000 }); + + await clock.advance(1_000); + expect(admissions[12]!.admission()).toEqual({ kind: 'admitted', waitedMs: 3_000 }); + expect(clock.pending).toBe(0); + }); + + it('holds SUSTAINED output to the rate when offered more than it', async () => { + // The regression this test exists for: an earlier design admitted anyway + // once the wait bound expired, which made sustained output equal sustained + // input delayed by the bound — 82/minute in, 82/minute out. 82/minute is + // the rate the logged rejections came from. + // + // Measured over the FINAL minute, after the burst and the queue transient + // have washed out, so this reads the steady-state rate and not the + // reservoir. The earlier design scored 82 here. + const clock = new TestClock(); + const pacer = new WsUpgradePacer({ + ratePerMinute: 60, + burst: 10, + ...boundedBy(15_000), + now: clock.now, + schedule: clock.schedule, + }); + + const offeredPerMinute = 82; + const windowMs = 240_000; + const intervalMs = Math.round(60_000 / offeredPerMinute); + const offered: Tracked[] = []; + while (clock.time < windowMs) { + offered.push(track(clock, pacer.admit())); + await clock.advance(intervalMs); + } + + const openedIn = (fromMs: number, toMs: number) => offered.filter(entry => { + const at = entry.settledAt(); + return entry.state() === 'admitted' && at !== undefined && at > fromMs && at <= toMs; + }).length; + + // ~328 requests offered across four minutes. + expect(offered.length).toBeGreaterThanOrEqual(320); + expect(openedIn(180_000, 240_000)).toBeLessThanOrEqual(61); + expect(openedIn(180_000, 240_000)).toBeGreaterThan(50); + // The overflow is refused, not silently admitted late. + expect(offered.filter(entry => entry.state() === 'refused').length).toBeGreaterThan(0); + // Whole-run ceiling, stated as the invariant: burst + one bound's worth of + // queue + the refill over the run. Not the offered count. + const admitted = offered.filter(entry => entry.state() === 'admitted').length; + expect(admitted).toBeLessThanOrEqual(10 + 15 + (windowMs / 60_000) * 60); + expect(admitted).toBeLessThan(offered.length); + }); + + it('shapes the opening burst but stops taxing once retries are turned off', async () => { + // With CLODEX_UPSTREAM_MAX_RETRIES=0 the SDK rethrows before it consults + // shouldRetry, so a refusal would be an immediate hard failure. It must + // admit — but delaying EVERY request by the bound past the debt floor + // shapes nothing (output would equal input, merely late) and only taxes the + // user, so shaping stops at the floor instead. + const clock = new TestClock(); + const pacer = new WsUpgradePacer({ + ratePerMinute: 60, + burst: 10, + idleTimeoutMs: 120_000, + maxRetries: 0, + now: clock.now, + schedule: clock.schedule, + }); + const maxDebt = pacer.maxWaitMs / 1_000; + + const admissions = Array.from({ length: 60 }, () => track(clock, pacer.admit())); + await flush(); + expect(admissions.some(entry => entry.state() === 'refused')).toBe(false); + + await clock.advance(pacer.maxWaitMs); + expect(admissions.every(entry => entry.state() === 'admitted')).toBe(true); + + // Shaping covers the burst plus one bound's worth of credit; within that + // window `maxDebt` requests are actually held back. + const delayed = admissions.filter(entry => (entry.settledAt() ?? 0) > 0).length; + expect(delayed).toBe(maxDebt); + // Everything past the floor is admitted with NO wait, rather than each + // paying the bound to achieve no shaping at all. + expect(admissions.filter(entry => entry.settledAt() === 0)).toHaveLength(60 - maxDebt); + }); + + it('does not pace at all when the deadline leaves no room to queue', async () => { + const clock = new TestClock(); + // A 10s deadline against a 62s backoff ladder: no budget to wait in. + const pacer = new WsUpgradePacer({ + ratePerMinute: 60, + burst: 1, + idleTimeoutMs: 10_000, + maxRetries: 5, + now: clock.now, + schedule: clock.schedule, + }); + expect(pacer.maxWaitMs).toBe(0); + + const admissions = Array.from({ length: 20 }, () => track(clock, pacer.admit())); + await flush(); + // Refusing everything past the burst would be worse than not pacing. + expect(admissions.every(entry => entry.admission()?.kind === 'admitted')).toBe(true); + expect(clock.pending).toBe(0); + }); + + it('refuses with a retryable backoff hint once the queue would exceed the bound', async () => { + const clock = new TestClock(); + const pacer = new WsUpgradePacer({ + ratePerMinute: 60, + burst: 1, + ...boundedBy(2_000), + now: clock.now, + schedule: clock.schedule, + }); + + const admissions = Array.from({ length: 5 }, () => track(clock, pacer.admit())); + await flush(); + + expect(admissions[0]!.admission()).toEqual({ kind: 'admitted', waitedMs: 0 }); + // The deficit is 3s but the bound is 2s, and the SDK spends this hint + // INSTEAD of its own backoff rung — so the hint is capped at the bound + // while `requiredWaitMs` still reports the deficit honestly. + expect(admissions[3]!.admission()).toEqual({ + kind: 'refused', requiredWaitMs: 3_000, retryAfterSeconds: 2, + }); + expect(admissions[4]!.admission()).toEqual({ + kind: 'refused', requiredWaitMs: 3_000, retryAfterSeconds: 2, + }); + + await clock.advance(2_000); + // The two that fitted inside the bound were queued, not refused. + expect(admissions[1]!.admission()).toEqual({ kind: 'admitted', waitedMs: 1_000 }); + expect(admissions[2]!.admission()).toEqual({ kind: 'admitted', waitedMs: 2_000 }); + expect(clock.pending).toBe(0); + }); + + it('charges a refusal nothing, so a retried request is not pushed further back', async () => { + // This is the livelock guard. A refused request opens no connection and the + // client retries it; if the refusal took a token, every retry would deepen + // the deficit that caused it and the ladder could never recover. + const clock = new TestClock(); + const pacer = new WsUpgradePacer({ + ratePerMinute: 60, + burst: 1, + ...boundedBy(2_000), + now: clock.now, + schedule: clock.schedule, + }); + + Array.from({ length: 3 }, () => track(clock, pacer.admit())); + const refused = Array.from({ length: 3 }, () => track(clock, pacer.admit())); + await flush(); + expect(refused.map(entry => entry.state())).toEqual(['refused', 'refused', 'refused']); + + await clock.advance(3_000); + const retried = track(clock, pacer.admit()); + await flush(); + // Three refusals debited nothing, so the bucket is back to full credit. + // Had they each taken a token, this would be refused too. + expect(retried.admission()).toEqual({ kind: 'admitted', waitedMs: 0 }); + }); + + it('lets a refused request through on a later attempt as the bucket refills', async () => { + const clock = new TestClock(); + const pacer = new WsUpgradePacer({ + ratePerMinute: 60, + burst: 1, + ...boundedBy(2_000), + now: clock.now, + schedule: clock.schedule, + }); + + Array.from({ length: 3 }, () => track(clock, pacer.admit())); + const firstAttempt = track(clock, pacer.admit()); + await flush(); + expect(firstAttempt.state()).toBe('refused'); + + // The AI SDK's first backoff step. + await clock.advance(2_000); + const secondAttempt = track(clock, pacer.admit()); + await flush(); + await clock.advance(2_000); + expect(secondAttempt.state()).toBe('admitted'); + }); + + it('releases a queued request as soon as its caller aborts, and returns its token', async () => { + const clock = new TestClock(); + const pacer = new WsUpgradePacer({ + ratePerMinute: 60, + burst: 1, + ...boundedBy(15_000), + now: clock.now, + schedule: clock.schedule, + }); + + track(clock, pacer.admit()); + const controller = new AbortController(); + const cancelled = track(clock, pacer.admit(controller.signal)); + await flush(); + expect(cancelled.state()).toBe('pending'); + expect(clock.pending).toBe(1); + + controller.abort(); + await flush(); + + // Promptly: the clock has not moved at all. + expect(clock.time).toBe(0); + expect(cancelled.state()).toBe('rejected'); + expect((cancelled.error() as Error).name).toBe('AbortError'); + // And its timer is gone, not left to fire into a dead request. + expect(clock.pending).toBe(0); + + // The abandoned reservation went back to the bucket, so the next arrival + // takes the slot it vacated: 1s, not the 2s it would inherit if the token + // had been consumed by a connection that was never opened. The control for + // this number is the "waits out the full queue" case below. + const next = track(clock, pacer.admit()); + await flush(); + await clock.advance(1_000); + expect(next.admission()).toEqual({ kind: 'admitted', waitedMs: 1_000 }); + }); + + it('makes a queued request wait out the full queue when nobody aborts', async () => { + const clock = new TestClock(); + const pacer = new WsUpgradePacer({ + ratePerMinute: 60, + burst: 1, + ...boundedBy(15_000), + now: clock.now, + schedule: clock.schedule, + }); + + track(clock, pacer.admit()); + track(clock, pacer.admit()); + const follower = track(clock, pacer.admit()); + await flush(); + + await clock.advance(1_000); + expect(follower.state()).toBe('pending'); + await clock.advance(1_000); + expect(follower.admission()).toEqual({ kind: 'admitted', waitedMs: 2_000 }); + }); + + it('spends no token on a request that was already cancelled', async () => { + const clock = new TestClock(); + const pacer = new WsUpgradePacer({ + ratePerMinute: 60, + burst: 1, + ...boundedBy(15_000), + now: clock.now, + schedule: clock.schedule, + }); + + const controller = new AbortController(); + controller.abort(); + await expect(pacer.admit(controller.signal)).rejects.toThrow(); + + const next = track(clock, pacer.admit()); + await flush(); + // Had the dead request taken the only token, this would have been queued. + expect(next.admission()).toEqual({ kind: 'admitted', waitedMs: 0 }); + expect(clock.pending).toBe(0); + }); + + it('surfaces the caller\'s own abort reason unchanged', async () => { + const clock = new TestClock(); + const pacer = new WsUpgradePacer({ + ratePerMinute: 60, + burst: 1, + ...boundedBy(15_000), + now: clock.now, + schedule: clock.schedule, + }); + + track(clock, pacer.admit()); + const controller = new AbortController(); + const queued = pacer.admit(controller.signal); + const settled = track(clock, queued); + await flush(); + + const reason = new Error('no data received from provider for 120s'); + controller.abort(reason); + await flush(); + expect(settled.error()).toBe(reason); + await expect(queued).rejects.toBe(reason); + }); + + it('never waits or refuses when pacing is turned off', async () => { + const clock = new TestClock(); + const pacer = new WsUpgradePacer({ + ratePerMinute: 0, + burst: 1, + now: clock.now, + schedule: clock.schedule, + }); + + const admissions = Array.from({ length: 50 }, () => track(clock, pacer.admit())); + await flush(); + expect(admissions.every(entry => entry.admission()?.kind === 'admitted')).toBe(true); + expect(clock.pending).toBe(0); + }); +}); + +describe('wsNewConnectionsPerMinute', () => { + it('defaults when unset or empty', () => { + expect(wsNewConnectionsPerMinute({})).toBe(DEFAULT_WS_NEW_CONNECTIONS_PER_MIN); + expect(wsNewConnectionsPerMinute({ [WS_NEW_CONNECTIONS_PER_MIN_ENV]: ' ' })) + .toBe(DEFAULT_WS_NEW_CONNECTIONS_PER_MIN); + }); + + it('accepts a valid rate and treats zero as off', () => { + expect(wsNewConnectionsPerMinute({ [WS_NEW_CONNECTIONS_PER_MIN_ENV]: ' 90 ' })).toBe(90); + expect(wsNewConnectionsPerMinute({ [WS_NEW_CONNECTIONS_PER_MIN_ENV]: '0' })).toBe(0); + // Any spelling `Number` reads as an integer is accepted, exactly as the + // existing CLODEX_UPSTREAM_MAX_RETRIES parser does. + expect(wsNewConnectionsPerMinute({ [WS_NEW_CONNECTIONS_PER_MIN_ENV]: '1e2' })).toBe(100); + }); + + it('treats the documented upper bound as inclusive', () => { + resetWsUpgradePacerForTests(); + const warn = vi.fn(); + // 600 is documented as allowed, so it must not clamp and must not warn. + // Literals, not the imported constants: an oracle that imports the value + // it checks cannot catch the value changing. + expect(wsNewConnectionsPerMinute({ [WS_NEW_CONNECTIONS_PER_MIN_ENV]: '600' }, warn)).toBe(600); + expect(warn).not.toHaveBeenCalled(); + expect(wsNewConnectionsPerMinute({ [WS_NEW_CONNECTIONS_PER_MIN_ENV]: '601' }, warn)).toBe(600); + expect(warn).toHaveBeenCalledTimes(1); + }); + + it('pins the shipped defaults against their literals', () => { + expect(DEFAULT_WS_NEW_CONNECTIONS_PER_MIN).toBe(60); + expect(MAX_WS_NEW_CONNECTIONS_PER_MIN).toBe(600); + expect(WS_NEW_CONNECTION_BURST).toBe(10); + expect(WS_NEW_CONNECTION_MAX_WAIT_CEILING_MS).toBe(15_000); + }); + + it('hands out exactly ten free connections before it starts queueing', async () => { + // Behavioural oracle for the burst: no test-only accessor, and it fails if + // the constructor stops using the documented default. + const clock = new TestClock(); + const pacer = new WsUpgradePacer({ now: clock.now, schedule: clock.schedule }); + const admissions = Array.from({ length: 12 }, () => track(clock, pacer.admit())); + await flush(); + const immediate = admissions.filter(entry => entry.state() === 'admitted' + && entry.admission()?.kind === 'admitted' && entry.settledAt() === 0).length; + expect(immediate).toBe(10); + expect(admissions[10]!.state()).toBe('pending'); + }); + + it('clamps an out-of-range rate and warns exactly once', () => { + resetWsUpgradePacerForTests(); + const warn = vi.fn(); + const env = { [WS_NEW_CONNECTIONS_PER_MIN_ENV]: '5000' }; + expect(wsNewConnectionsPerMinute(env, warn)).toBe(MAX_WS_NEW_CONNECTIONS_PER_MIN); + expect(wsNewConnectionsPerMinute(env, warn)).toBe(MAX_WS_NEW_CONNECTIONS_PER_MIN); + expect(warn).toHaveBeenCalledTimes(1); + expect(warn.mock.calls[0]![0]).toContain('clamping CLODEX_WS_MAX_NEW_CONNECTIONS_PER_MIN=5000 to 600'); + }); + + it.each(['abc', '-1', '12.5', 'Infinity', 'NaN', '1,000'])( + 'ignores the malformed value %s and warns once', + raw => { + resetWsUpgradePacerForTests(); + const warn = vi.fn(); + const env = { [WS_NEW_CONNECTIONS_PER_MIN_ENV]: raw }; + expect(wsNewConnectionsPerMinute(env, warn)).toBe(DEFAULT_WS_NEW_CONNECTIONS_PER_MIN); + expect(wsNewConnectionsPerMinute(env, warn)).toBe(DEFAULT_WS_NEW_CONNECTIONS_PER_MIN); + expect(warn).toHaveBeenCalledTimes(1); + expect(warn.mock.calls[0]![0]).toContain(`ignoring CLODEX_WS_MAX_NEW_CONNECTIONS_PER_MIN=${raw}`); + }, + ); + + it('never lets a failing notice channel break the caller', () => { + resetWsUpgradePacerForTests(); + const warn = vi.fn(() => { throw new Error('stderr is gone'); }); + expect(wsNewConnectionsPerMinute({ [WS_NEW_CONNECTIONS_PER_MIN_ENV]: 'nope' }, warn)) + .toBe(DEFAULT_WS_NEW_CONNECTIONS_PER_MIN); + }); +});