diff --git a/changelog.d/fixes/12911-codex-ws-lease-fail-fast.md b/changelog.d/fixes/12911-codex-ws-lease-fail-fast.md new file mode 100644 index 0000000000..01856ae353 --- /dev/null +++ b/changelog.d/fixes/12911-codex-ws-lease-fail-fast.md @@ -0,0 +1 @@ +- **fix(sse):** the Codex Responses WebSocket bridge now fails fast to another eligible account instead of queueing behind a saturated one, releasing its per-account lease exactly once when the session ends ([#12911](https://github.com/diegosouzapw/OmniRoute/pull/12911) — thanks @initguru). This also fixes `accountSemaphore`'s `maxQueueSize: 0` handling, which previously behaved as an unbounded queue instead of failing over immediately — benefiting every caller that configures `queueDepth: 0` (for example combo routing), not just the Codex WS bridge. diff --git a/open-sse/services/accountSemaphore.ts b/open-sse/services/accountSemaphore.ts index 4d0896ea30..a26a1df44a 100644 --- a/open-sse/services/accountSemaphore.ts +++ b/open-sse/services/accountSemaphore.ts @@ -240,26 +240,31 @@ export function acquireMany( for (const key of keys) { const gate = ensureGate(key, enabled.get(key)!); clearCleanupTimer(gate); - if (maxQueueSize > 0 && gate.queue.length >= maxQueueSize) { - return Promise.reject( - createSemaphoreError( - "SEMAPHORE_QUEUE_FULL", - `Semaphore queue full (${maxQueueSize}) for ${key}` - ) - ); - } } - if ( - keys.every((key) => { - const gate = gates.get(key)!; - return gate.queue.length === 0 && gate.running < gate.maxConcurrency && !isBlocked(gate); - }) - ) { + const canAcquireImmediately = keys.every((key) => { + const gate = gates.get(key)!; + return gate.queue.length === 0 && gate.running < gate.maxConcurrency && !isBlocked(gate); + }); + if (canAcquireImmediately) { for (const key of keys) gates.get(key)!.running++; return Promise.resolve(createCompositeReleaseFn(keys)); } + if (maxQueueSize >= 0) { + for (const key of keys) { + const gate = gates.get(key)!; + if (gate.queue.length >= maxQueueSize) { + return Promise.reject( + createSemaphoreError( + "SEMAPHORE_QUEUE_FULL", + `Semaphore queue full (${maxQueueSize}) for ${key}` + ) + ); + } + } + } + return new Promise((resolve, reject) => { const request: AcquireRequest = { keys, diff --git a/scripts/dev/responses-ws-proxy.mjs b/scripts/dev/responses-ws-proxy.mjs index 9255d268ff..d9894481d5 100644 --- a/scripts/dev/responses-ws-proxy.mjs +++ b/scripts/dev/responses-ws-proxy.mjs @@ -430,6 +430,9 @@ class ResponsesWsSession { this.firstResponseBody = null; this.currentRequestBody = null; this.preparedContext = null; + this.leaseId = null; + this.leaseReleased = false; + this.leaseReleaseInFlight = false; // #7388: logging must be scoped per logical turn (one `response.create` // through its terminal event), not once for the lifetime of the WS // connection — a single boolean here silently dropped every turn after @@ -640,6 +643,23 @@ class ResponsesWsSession { toStringOrNull(responseBody.service_tier) || toStringOrNull(responseBody.serviceTier), }; + // A reused WS connection re-runs prepare per logical turn, and each prepare + // acquires a fresh per-account lease. Release the previous turn before + // adopting the new lease so one session cannot hoard account slots. + const previousLeaseId = this.leaseId; + const newLeaseId = toStringOrNull(prepared.json?.leaseId); + if (this.closed) { + this.leaseId = null; + this.releaseLeaseId(newLeaseId); + return prepared; + } + this.leaseId = newLeaseId; + if (previousLeaseId && previousLeaseId !== newLeaseId) { + this.releaseLeaseId(previousLeaseId); + } + this.leaseReleased = false; + this.leaseReleaseInFlight = false; + return prepared; } @@ -770,6 +790,35 @@ class ResponsesWsSession { } } + releaseLease() { + if (this.leaseReleased || this.leaseReleaseInFlight || !this.leaseId) return; + this.leaseReleaseInFlight = true; + const leaseId = this.leaseId; + void callInternal(this.fetchImpl, this.baseUrl, this.bridgeSecret, "release", { leaseId }) + .then((response) => { + if (!response.ok) throw new Error("lease release rejected"); + this.leaseReleased = true; + this.leaseId = null; + }) + .catch(() => { + this.leaseReleaseInFlight = false; + const retry = setTimeout(() => this.releaseLease(), 1000); + retry.unref?.(); + }); + } + + releaseLeaseId(leaseId) { + if (!leaseId) return; + void callInternal(this.fetchImpl, this.baseUrl, this.bridgeSecret, "release", { leaseId }) + .then((response) => { + if (!response.ok) throw new Error("lease release rejected"); + }) + .catch(() => { + const retry = setTimeout(() => this.releaseLeaseId(leaseId), 1000); + retry.unref?.(); + }); + } + async persistHistory({ status = 200, success = true, @@ -820,6 +869,7 @@ class ResponsesWsSession { close(code = 1000, reason = "normal_closure") { if (this.closed) return; this.closed = true; + this.releaseLease(); clearInterval(this.pingTimer); this.cleanupBuffers(); @@ -847,6 +897,7 @@ class ResponsesWsSession { dispose() { if (this.closed) return; this.closed = true; + this.releaseLease(); clearInterval(this.pingTimer); this.cleanupBuffers(); try { diff --git a/src/app/api/internal/codex-responses-ws/route.ts b/src/app/api/internal/codex-responses-ws/route.ts index 4b26af0316..ce3c5e428d 100644 --- a/src/app/api/internal/codex-responses-ws/route.ts +++ b/src/app/api/internal/codex-responses-ws/route.ts @@ -7,6 +7,7 @@ import { authorizeWebSocketHandshake, extractWsTokenFromRequest } from "@/lib/ws import { getModelInfo } from "@/sse/services/model"; import { resolveCcDiscoveryAliasStrip } from "@/lib/ccDiscoveryAliasResolve"; import { getProviderCredentialsWithQuotaPreflight } from "@/sse/services/auth"; +import { acquireCodexWsLease, releaseCodexWsLease } from "@/sse/services/codexWsLease"; import { enforceApiKeyPolicy } from "@/shared/utils/apiKeyPolicy"; import { checkAndRefreshToken } from "@/sse/services/tokenRefresh"; import { resolveCodexWsModelInfo } from "./modelResolution"; @@ -22,8 +23,8 @@ import { sanitizeErrorMessage } from "@omniroute/open-sse/utils/error.ts"; import { logger } from "@omniroute/open-sse/utils/logger.ts"; import { resolveProxy } from "@omniroute/open-sse/utils/networkProxy.ts"; import { withCodexFingerprintCredentials } from "@omniroute/open-sse/config/codexIdentity.ts"; -import { proxyConfigToUrl } from "@omniroute/open-sse/utils/proxyDispatcher.ts"; import { withReasoningRuleContext } from "@omniroute/open-sse/utils/reasoningRuleContext.ts"; +import { proxyConfigToUrl } from "@omniroute/open-sse/utils/proxyDispatcher.ts"; import { attachReasoningRuleDirective, applyReasoningRuleDirective, @@ -373,28 +374,58 @@ async function resolveCodexCredentials( model: string, allowedConnections: string[] | null ) { - const credentials = await getProviderCredentialsWithQuotaPreflight( - provider, - null, - allowedConnections, - model - ); - if (!credentials || "allRateLimited" in credentials) { - return { - error: jsonError( - 503, - "codex_credentials_unavailable", - "No available Codex OAuth connection for Responses WebSocket" - ), - }; + const excludedConnectionIds: string[] = []; + let credentials: Awaited> = null; + + // A saturated account is excluded and another eligible account is selected; + // never queue a Responses WS session behind an existing tool turn — a queued + // session times out client-side as 499/502. + for (let attempt = 0; attempt < 8; attempt += 1) { + credentials = await getProviderCredentialsWithQuotaPreflight( + provider, + null, + allowedConnections, + model, + { excludeConnectionIds: excludedConnectionIds } + ); + if (!credentials || "allRateLimited" in credentials || !credentials.connectionId) break; + + const leaseId = await acquireCodexWsLease(credentials.connectionId, credentials.maxConcurrent); + if (!leaseId) { + excludedConnectionIds.push(credentials.connectionId); + continue; + } + + let refreshed: Awaited>; + try { + refreshed = await checkAndRefreshToken(provider, credentials); + } catch (error) { + releaseCodexWsLease(leaseId); + return { + error: jsonError( + 502, + "codex_ws_prepare_failed", + sanitizeErrorMessage(error instanceof Error ? error.message : String(error)) + ), + }; + } + + if (!refreshed?.accessToken) { + releaseCodexWsLease(leaseId); + return { + error: jsonError(401, "codex_oauth_token_missing", "Codex OAuth access token is missing"), + }; + } + return { credentials: refreshed, leaseId }; } - const refreshed = await checkAndRefreshToken(provider, credentials); - if (!refreshed?.accessToken) { - return { - error: jsonError(401, "codex_oauth_token_missing", "Codex OAuth access token is missing"), - }; - } - return { credentials: refreshed }; + + return { + error: jsonError( + 503, + "codex_credentials_unavailable", + "No available Codex OAuth connection for Responses WebSocket" + ), + }; } async function resolveCodexRequestContext(body: JsonRecord) { @@ -484,19 +515,31 @@ async function resolveCodexUpstreamContext( if (credentialResult.error) return credentialResult; let reasoningDecision = context.decision; if (!reasoningDecision) { - reasoningDecision = await resolveReasoningRoutingRule({ - sourceModel: context.intent.model, - sourceModelAliases: context.sourceModels.aliases, - sourceEffort: context.intent.sourceEffort, - hasReasoningSignal: context.intent.hasReasoningSignal, - hasThinkingBudget: context.intent.hasThinkingBudget, - apiKeyId: context.metadata?.id ?? null, - connectionId: credentialResult.credentials.connectionId, - requestTags: context.routingTags.tags, - connectionOnly: true, - capabilityModel: `codex/${model}`, - }); + try { + reasoningDecision = await resolveReasoningRoutingRule({ + sourceModel: context.intent.model, + sourceModelAliases: context.sourceModels.aliases, + sourceEffort: context.intent.sourceEffort, + hasReasoningSignal: context.intent.hasReasoningSignal, + hasThinkingBudget: context.intent.hasThinkingBudget, + apiKeyId: context.metadata?.id ?? null, + connectionId: credentialResult.credentials.connectionId, + requestTags: context.routingTags.tags, + connectionOnly: true, + capabilityModel: `codex/${model}`, + }); + } catch (error) { + releaseCodexWsLease(credentialResult.leaseId); + return { + error: jsonError( + 502, + "codex_ws_prepare_failed", + sanitizeErrorMessage(error instanceof Error ? error.message : String(error)) + ), + }; + } if (reasoningDecision?.capability === "unsupported") { + releaseCodexWsLease(credentialResult.leaseId); return { error: jsonError( 400, @@ -511,6 +554,7 @@ async function resolveCodexUpstreamContext( provider, model, credentials: credentialResult.credentials, + leaseId: credentialResult.leaseId, reasoningDecision, }; } @@ -538,56 +582,87 @@ async function prepare(body: JsonRecord) { ); } } + const upstream = await resolveCodexUpstreamContext(context); if ("error" in upstream) return upstream.error; - const { responseBody, metadata, provider, model, credentials: refreshedCredentials } = upstream; - const reasoningDecision = upstream.reasoningDecision; - - let responseBodyWithMemory = await maybeInjectResponsesWsMemory(responseBody, metadata); - let reasoningRouting: JsonRecord | null = null; - let reasoningRuleDirective: unknown; - if (reasoningDecision) { - const withDirective = attachReasoningRuleDirective(responseBodyWithMemory, reasoningDecision); - reasoningRuleDirective = withDirective._omnirouteReasoningRule; - reasoningRouting = isRecord(withDirective._omnirouteReasoningRouteTrace) - ? withDirective._omnirouteReasoningRouteTrace - : null; - responseBodyWithMemory = applyReasoningRuleDirective( - withDirective, - "openai-responses" - ) as JsonRecord; - delete responseBodyWithMemory._omnirouteReasoningRouteTrace; - } - // #8052: the WS bridge previously skipped the whole prompt-compression pipeline that the - // HTTP/SSE path (chatCore.ts) runs on every request — wire the same core pipeline in here, - // per logical turn, before handing off to the executor. - responseBodyWithMemory = await applyResponsesWsCompression(responseBodyWithMemory, { + const { + responseBody, + metadata, provider, model, - requestId: randomUUID(), - }); - const credentialsWithFingerprint = withCodexFingerprintCredentials( - withReasoningRuleContext(refreshedCredentials, reasoningRuleDirective), - context.clientHeaders, - responseBodyWithMemory - ); - const transformed = (await executor.transformRequest( - model, - responseBodyWithMemory, - true, - credentialsWithFingerprint - )) as JsonRecord; - transformed.model = model; - delete transformed.stream; - delete transformed.stream_options; + credentials: refreshedCredentials, + leaseId, + } = upstream; + const reasoningDecision = upstream.reasoningDecision; - const headers = normalizeUpstreamHeaders(executor.buildHeaders(credentialsWithFingerprint, true)); + let responseBodyWithMemory: JsonRecord; + let reasoningRouting: JsonRecord | null = null; + let transformed: JsonRecord; + let credentialsWithFingerprint: typeof refreshedCredentials; + let reasoningRuleDirective: unknown; + try { + responseBodyWithMemory = await maybeInjectResponsesWsMemory(responseBody, metadata); + if (reasoningDecision) { + const withDirective = attachReasoningRuleDirective(responseBodyWithMemory, reasoningDecision); + reasoningRuleDirective = withDirective._omnirouteReasoningRule; + reasoningRouting = isRecord(withDirective._omnirouteReasoningRouteTrace) + ? withDirective._omnirouteReasoningRouteTrace + : null; + responseBodyWithMemory = applyReasoningRuleDirective( + withDirective, + "openai-responses" + ) as JsonRecord; + delete responseBodyWithMemory._omnirouteReasoningRouteTrace; + } + // #8052: the WS bridge previously skipped the whole prompt-compression pipeline that the + // HTTP/SSE path (chatCore.ts) runs on every request — wire the same core pipeline in here, + // per logical turn, before handing off to the executor. + responseBodyWithMemory = await applyResponsesWsCompression(responseBodyWithMemory, { + provider, + model, + requestId: randomUUID(), + }); + credentialsWithFingerprint = withCodexFingerprintCredentials( + withReasoningRuleContext(refreshedCredentials, reasoningRuleDirective), + context.clientHeaders, + responseBodyWithMemory + ); + transformed = (await executor.transformRequest( + model, + responseBodyWithMemory, + true, + credentialsWithFingerprint + )) as JsonRecord; + transformed.model = model; + delete transformed.stream; + delete transformed.stream_options; + } catch (error) { + releaseCodexWsLease(leaseId); + return jsonError( + 502, + "codex_ws_prepare_failed", + sanitizeErrorMessage(error instanceof Error ? error.message : String(error)) + ); + } - // #5611: apply the configured Global/provider proxy to the upstream Codex - // Responses WebSocket too. The downstream client→OmniRoute hop works, but the - // upstream wreq-js.websocket() connect previously ignored the Proxy Registry, - // so a no-direct-egress container failed with a DNS lookup error. - const proxy = await resolveCodexProxy(provider); + let headers: Record; + let proxy: string | undefined; + try { + headers = normalizeUpstreamHeaders(executor.buildHeaders(credentialsWithFingerprint, true)); + + // #5611: apply the configured Global/provider proxy to the upstream Codex + // Responses WebSocket too. The downstream client→OmniRoute hop works, but the + // upstream wreq-js.websocket() connect previously ignored the Proxy Registry, + // so a no-direct-egress container failed with a DNS lookup error. + proxy = await resolveCodexProxy(provider); + } catch (error) { + releaseCodexWsLease(leaseId); + return jsonError( + 502, + "codex_ws_prepare_failed", + sanitizeErrorMessage(error instanceof Error ? error.message : String(error)) + ); + } return NextResponse.json({ ok: true, @@ -599,6 +674,7 @@ async function prepare(body: JsonRecord) { browser: "chrome_142", os: "windows", connectionId: refreshedCredentials.connectionId, + leaseId, provider, account: refreshedCredentials.email || null, model, @@ -634,6 +710,12 @@ export async function POST(request: Request) { if (action === "prepare") { return prepare(body); } + if (action === "release") { + return NextResponse.json({ + ok: true, + released: releaseCodexWsLease(toStringOrNull(body.leaseId)), + }); + } if (action === "log") { try { return await persistResponsesWsCallHistory(body); diff --git a/src/sse/services/codexWsLease.ts b/src/sse/services/codexWsLease.ts new file mode 100644 index 0000000000..9ef1675fe3 --- /dev/null +++ b/src/sse/services/codexWsLease.ts @@ -0,0 +1,43 @@ +import { randomUUID } from "node:crypto"; +import { + acquire as acquireAccountSemaphore, + buildAccountSemaphoreKey, +} from "@omniroute/open-sse/services/accountSemaphore.ts"; + +const leases = new Map void>(); + +/** Acquire a non-queued, process-local account slot for one Responses WS session. */ +export async function acquireCodexWsLease( + connectionId: string, + configuredMaxConcurrent: number | null | undefined +): Promise { + const key = buildAccountSemaphoreKey({ provider: "codex", accountKey: connectionId }); + try { + const release = await acquireAccountSemaphore(key, { + maxConcurrency: + typeof configuredMaxConcurrent === "number" && configuredMaxConcurrent > 0 + ? configuredMaxConcurrent + : 1, + maxQueueSize: 0, + }); + const leaseId = randomUUID(); + leases.set(leaseId, release); + return leaseId; + } catch { + return null; + } +} + +/** Release a bridge lease once; unknown/already-released leases are harmless. */ +export function releaseCodexWsLease(leaseId: string | null | undefined): boolean { + if (!leaseId) return false; + const release = leases.get(leaseId); + if (!release) return false; + leases.delete(leaseId); + release(); + return true; +} + +export function clearCodexWsLeasesForTest(): void { + for (const leaseId of leases.keys()) releaseCodexWsLease(leaseId); +} diff --git a/tests/unit/accountSemaphore.test.ts b/tests/unit/accountSemaphore.test.ts index 145fea317f..92cd8ce5d2 100644 --- a/tests/unit/accountSemaphore.test.ts +++ b/tests/unit/accountSemaphore.test.ts @@ -115,6 +115,22 @@ describe("accountSemaphore acquireMany", () => { releaseGlobal(); (await queued)(); }); + + it("fails immediately when maxQueueSize is zero", async () => { + const release = await acquire("codex:account-a", { maxConcurrency: 1 }); + + await assert.rejects( + acquire("codex:account-a", { + maxConcurrency: 1, + maxQueueSize: 0, + timeoutMs: 200, + }), + (error: Error & { code?: string }) => error.code === "SEMAPHORE_QUEUE_FULL" + ); + assert.equal(getStats()["codex:account-a"]?.queued ?? 0, 0); + + release(); + }); }); describe("accountSemaphore", async () => { diff --git a/tests/unit/codex-ws-lease-contract.test.ts b/tests/unit/codex-ws-lease-contract.test.ts new file mode 100644 index 0000000000..99144bdb19 --- /dev/null +++ b/tests/unit/codex-ws-lease-contract.test.ts @@ -0,0 +1,156 @@ +/** + * Regression guard for the Codex Responses-WS per-account concurrency lease + * (non-queued per-account slot). + * + * Covers the route-level lease contract: + * 1. the internal `release` action releases an acquired lease exactly once, + * 2. a saturated account is excluded and prepare returns 503 + * `codex_credentials_unavailable` instead of queuing, + * 3. once the lease is released the same account is eligible again and + * prepare succeeds with a fresh leaseId. + */ +import assert from "node:assert/strict"; +import fs from "node:fs"; +import os from "node:os"; +import path from "node:path"; +import test from "node:test"; + +const TEST_DATA_DIR = fs.mkdtempSync(path.join(os.tmpdir(), "omniroute-codex-ws-lease-")); +process.env.DATA_DIR = TEST_DATA_DIR; +process.env.API_KEY_SECRET = process.env.API_KEY_SECRET || "codex-ws-lease-api-key-secret"; +process.env.OMNIROUTE_WS_BRIDGE_SECRET = "codex-ws-lease-bridge-secret"; +process.env.OMNIROUTE_CODEX_WS_ENABLED = "true"; + +const coreDb = await import("../../src/lib/db/core.ts"); +const apiKeysDb = await import("../../src/lib/db/apiKeys.ts"); +const providersDb = await import("../../src/lib/db/providers.ts"); +const costRules = await import("../../src/domain/costRules.ts"); +const rateLimiter = await import("../../src/shared/utils/rateLimiter.ts"); +const route = await import("../../src/app/api/internal/codex-responses-ws/route.ts"); +const codexWsLease = await import("../../src/sse/services/codexWsLease.ts"); + +rateLimiter.setRateLimiterTestMode(true); + +type BridgeBody = Record; + +function getFsErrorCode(error: unknown): string | undefined { + if (typeof error !== "object" || error === null || !("code" in error)) return undefined; + const { code } = error as { code?: unknown }; + return typeof code === "string" ? code : undefined; +} + +async function resetStorage() { + apiKeysDb.resetApiKeyState(); + costRules.resetCostData(); + coreDb.resetDbInstance(); + + for (let attempt = 0; attempt < 10; attempt++) { + try { + if (fs.existsSync(TEST_DATA_DIR)) { + fs.rmSync(TEST_DATA_DIR, { recursive: true, force: true }); + } + break; + } catch (error: unknown) { + const code = getFsErrorCode(error); + if ((code === "EBUSY" || code === "EPERM") && attempt < 9) { + await new Promise((resolve) => setTimeout(resolve, 50 * (attempt + 1))); + } else { + throw error; + } + } + } + + fs.mkdirSync(TEST_DATA_DIR, { recursive: true }); +} + +test.beforeEach(async () => { + codexWsLease.clearCodexWsLeasesForTest(); + await resetStorage(); +}); + +test.after(async () => { + codexWsLease.clearCodexWsLeasesForTest(); + apiKeysDb.resetApiKeyState(); + costRules.resetCostData(); + coreDb.resetDbInstance(); + fs.rmSync(TEST_DATA_DIR, { recursive: true, force: true }); +}); + +function buildBridgeRequest(body: BridgeBody): Request { + return new Request("http://localhost/api/internal/codex-responses-ws", { + method: "POST", + headers: { + "content-type": "application/json", + "x-omniroute-ws-bridge-secret": process.env.OMNIROUTE_WS_BRIDGE_SECRET as string, + }, + body: JSON.stringify(body), + }); +} + +function buildPrepareRequest(apiKey: string, model: string): Request { + return buildBridgeRequest({ + action: "prepare", + requestUrl: `/api/v1/responses?api_key=${encodeURIComponent(apiKey)}`, + response: { model }, + }); +} + +async function seedCodexConnection(name = "Codex WS lease test") { + const connection = (await providersDb.createProviderConnection({ + provider: "codex", + authType: "oauth", + name, + accessToken: "test-codex-access-token", + refreshToken: "test-codex-refresh-token", + expiresAt: Date.now() + 60 * 60 * 1000, + isActive: true, + })) as { id: string }; + return connection.id; +} + +test("internal release action releases an acquired lease exactly once", async () => { + const leaseId = await codexWsLease.acquireCodexWsLease("release-action-conn", 1); + assert.ok(leaseId, "expected an acquired lease"); + + const first = await route.POST(buildBridgeRequest({ action: "release", leaseId })); + const firstBody = (await first.json()) as { ok?: boolean; released?: boolean }; + assert.equal(first.status, 200); + assert.equal(firstBody.ok, true); + assert.equal(firstBody.released, true); + + const second = await route.POST(buildBridgeRequest({ action: "release", leaseId })); + const secondBody = (await second.json()) as { ok?: boolean; released?: boolean }; + assert.equal(secondBody.released, false, "a released lease must not release twice"); +}); + +test("saturated account is excluded (503) instead of queued; eligible again after release", async () => { + const key = await apiKeysDb.createApiKey("Lease Saturation Key", "machine-lease-sat"); + await apiKeysDb.updateApiKeyPermissions(key.id, { + allowedModels: ["gpt-5.5"], + }); + const connectionId = await seedCodexConnection(); + + const lease = await codexWsLease.acquireCodexWsLease(connectionId, 1); + assert.ok(lease, "expected to hold the account lease"); + + const saturated = await route.POST(buildPrepareRequest(key.key, "gpt-5.5")); + const saturatedBody = (await saturated.json()) as { error?: { code?: string } }; + assert.equal( + saturated.status, + 503, + `expected 503 on a saturated account, got ${saturated.status}: ${JSON.stringify(saturatedBody)}` + ); + assert.equal(saturatedBody.error?.code, "codex_credentials_unavailable"); + + assert.equal(codexWsLease.releaseCodexWsLease(lease), true); + + const ok = await route.POST(buildPrepareRequest(key.key, "gpt-5.5")); + const okBody = (await ok.json()) as { ok?: boolean; leaseId?: string | null }; + assert.equal( + ok.status, + 200, + `expected prepare to succeed after release: ${JSON.stringify(okBody)}` + ); + assert.equal(okBody.ok, true); + assert.ok(okBody.leaseId, "a successful prepare must return a leaseId"); +}); diff --git a/tests/unit/codex-ws-load-balancer.test.ts b/tests/unit/codex-ws-load-balancer.test.ts new file mode 100644 index 0000000000..6560a79c47 --- /dev/null +++ b/tests/unit/codex-ws-load-balancer.test.ts @@ -0,0 +1,26 @@ +import assert from "node:assert/strict"; +import fs from "node:fs"; +import path from "node:path"; +import test from "node:test"; + +const balancer = await import("../../src/sse/services/codexWsLease.ts"); + +test("Codex WS lease module has no quota-based account selector", () => { + const source = fs.readFileSync( + path.resolve(import.meta.dirname, "../../src/sse/services/codexWsLease.ts"), + "utf8" + ); + + assert.doesNotMatch(source, /selectCodex(?:Fair|Ws)Connection|remainingQuota|nextCodexFair/); +}); + +test("Codex WS leases enforce per-account concurrency and are released exactly once", async () => { + balancer.clearCodexWsLeasesForTest(); + + const lease = await balancer.acquireCodexWsLease("account-a", 1); + assert.ok(lease); + assert.equal(await balancer.acquireCodexWsLease("account-a", 1), null); + + assert.equal(balancer.releaseCodexWsLease(lease), true); + assert.equal(balancer.releaseCodexWsLease(lease), false); +}); diff --git a/tests/unit/responses-ws-proxy.test.mjs b/tests/unit/responses-ws-proxy.test.mjs index 95890e344b..e61c745fc9 100644 --- a/tests/unit/responses-ws-proxy.test.mjs +++ b/tests/unit/responses-ws-proxy.test.mjs @@ -76,6 +76,7 @@ test("responses ws proxy prepares and forwards OpenAI Responses websocket events upstreamUrl: "wss://chatgpt.com/backend-api/codex/responses", headers: { Authorization: "Bearer upstream-token" }, connectionId: "conn_1", + leaseId: "lease_1", provider: "codex", account: "codex@example.com", // #5611: prepare resolves the configured proxy and threads it through. @@ -416,3 +417,172 @@ test("responses ws proxy closes oversized client messages with 1009", async () = await close(server); }); + +test("responses ws proxy releases the prepared per-account lease on session close", async () => { + const internalRequests = []; + + const server = http.createServer(async (req, res) => { + const url = new URL(req.url || "/", `http://${req.headers.host}`); + if (url.pathname === "/api/internal/codex-responses-ws") { + const body = JSON.parse((await readRequestBody(req)) || "{}"); + internalRequests.push(body); + + if (body.action === "authenticate") { + res.writeHead(200, { "content-type": "application/json" }); + res.end(JSON.stringify({ ok: true, authenticated: true, authType: "api_key" })); + return; + } + + if (body.action === "prepare") { + res.writeHead(200, { "content-type": "application/json" }); + res.end( + JSON.stringify({ + ok: true, + upstreamUrl: "wss://chatgpt.com/backend-api/codex/responses", + headers: { Authorization: "Bearer upstream-token" }, + connectionId: "conn_lease", + leaseId: "lease_1", + response: { ...body.response, model: "gpt-5.5", stream: undefined }, + }) + ); + return; + } + + if (body.action === "release") { + res.writeHead(200, { "content-type": "application/json" }); + res.end(JSON.stringify({ ok: true, released: true })); + return; + } + } + + res.writeHead(404, { "content-type": "application/json" }); + res.end(JSON.stringify({ error: "not_found" })); + }); + + const fakeUpstream = { + send() {}, + close() {}, + onmessage: null, + onerror: null, + onclose: null, + }; + + const port = await listen(server); + const proxy = createResponsesWsProxy({ + baseUrl: `http://127.0.0.1:${port}`, + bridgeSecret: "bridge-secret", + pingIntervalMs: 1000, + idleTimeoutMs: 10000, + wsFactory: async () => fakeUpstream, + }); + + server.on("upgrade", async (req, socket, head) => { + const handled = await proxy.handleUpgrade(req, socket, head); + if (!handled && !socket.destroyed) { + socket.destroy(); + } + }); + + const ws = new WebSocket(`ws://127.0.0.1:${port}/api/v1/responses?api_key=local-token`); + await new Promise((resolve) => ws.addEventListener("open", resolve, { once: true })); + + ws.send( + JSON.stringify({ + type: "response.create", + model: "gpt-5.5", + input: [{ role: "user", content: "hello" }], + }) + ); + + await waitFor(() => internalRequests.find((entry) => entry.action === "prepare")); + ws.close(); + + const releaseRequest = await waitFor(() => + internalRequests.find((entry) => entry.action === "release") + ); + assert.equal(releaseRequest.leaseId, "lease_1"); + + await close(server); +}); + +test("responses ws proxy releases a replaced lease before adopting a reused-turn lease", async () => { + const internalRequests = []; + let prepareCount = 0; + + const server = http.createServer(async (req, res) => { + const url = new URL(req.url || "/", `http://${req.headers.host}`); + if (url.pathname === "/api/internal/codex-responses-ws") { + const body = JSON.parse((await readRequestBody(req)) || "{}"); + internalRequests.push(body); + if (body.action === "authenticate") { + res.writeHead(200, { "content-type": "application/json" }); + res.end(JSON.stringify({ ok: true, authenticated: true, authType: "api_key" })); + return; + } + if (body.action === "prepare") { + prepareCount += 1; + res.writeHead(200, { "content-type": "application/json" }); + res.end( + JSON.stringify({ + ok: true, + upstreamUrl: "wss://chatgpt.com/backend-api/codex/responses", + headers: { Authorization: "Bearer upstream-token" }, + connectionId: "conn_lease", + leaseId: `lease_${prepareCount}`, + response: { ...body.response, model: "gpt-5.5", stream: undefined }, + }) + ); + return; + } + if (body.action === "release") { + res.writeHead(200, { "content-type": "application/json" }); + res.end(JSON.stringify({ ok: true, released: true })); + return; + } + } + res.writeHead(404, { "content-type": "application/json" }); + res.end(JSON.stringify({ error: "not_found" })); + }); + + const fakeUpstream = { send() {}, close() {}, onmessage: null, onerror: null, onclose: null }; + const port = await listen(server); + const proxy = createResponsesWsProxy({ + baseUrl: `http://127.0.0.1:${port}`, + bridgeSecret: "bridge-secret", + pingIntervalMs: 1000, + idleTimeoutMs: 10000, + wsFactory: async () => fakeUpstream, + }); + server.on("upgrade", async (req, socket, head) => { + const handled = await proxy.handleUpgrade(req, socket, head); + if (!handled && !socket.destroyed) socket.destroy(); + }); + + const ws = new WebSocket(`ws://127.0.0.1:${port}/api/v1/responses?api_key=local-token`); + await new Promise((resolve) => ws.addEventListener("open", resolve, { once: true })); + ws.send( + JSON.stringify({ + type: "response.create", + model: "gpt-5.5", + input: [{ role: "user", content: "first" }], + }) + ); + await waitFor(() => internalRequests.filter((entry) => entry.action === "prepare").length === 1); + ws.send( + JSON.stringify({ + type: "response.create", + model: "gpt-5.5", + input: [{ role: "user", content: "second" }], + }) + ); + await waitFor(() => internalRequests.filter((entry) => entry.action === "prepare").length === 2); + await waitFor(() => + internalRequests.some((entry) => entry.action === "release" && entry.leaseId === "lease_1") + ); + + ws.close(); + await waitFor(() => + internalRequests.some((entry) => entry.action === "release" && entry.leaseId === "lease_2") + ); + await close(server); +});