diff --git a/changelog.d/fixes/8052-codex-ws-compression.md b/changelog.d/fixes/8052-codex-ws-compression.md new file mode 100644 index 0000000000..9a7a679d02 --- /dev/null +++ b/changelog.d/fixes/8052-codex-ws-compression.md @@ -0,0 +1,2 @@ +- **fix(sse):** run the prompt-compression pipeline per turn in the Codex Responses WebSocket bridge instead of skipping it entirely — a reused WebSocket connection now re-runs `prepare()` (auth/policy/memory/reasoning-routing/compression) for every logical `response.create` turn, not just the first (#8052) + diff --git a/scripts/dev/responses-ws-proxy.mjs b/scripts/dev/responses-ws-proxy.mjs index 8825283aac..1b585e4618 100644 --- a/scripts/dev/responses-ws-proxy.mjs +++ b/scripts/dev/responses-ws-proxy.mjs @@ -579,6 +579,52 @@ class ResponsesWsSession { await this.forwardClientMessage(message); } + // #8052: shared by ensureUpstream() (first turn — also owns socket creation) and + // forwardClientMessage() (subsequent turns on a reused connection). Calls the internal + // "prepare" action — auth/policy/memory/reasoning-routing/compression — and refreshes + // preparedContext, but never touches this.upstream/this.upstreamReady; the caller decides + // whether a new upstream socket is needed. + async runPrepare(message, responseBody) { + const prepared = await callInternal(this.fetchImpl, this.baseUrl, this.bridgeSecret, "prepare", { + requestUrl: this.requestUrl, + headers: getAuthHeaders(this.requestUrl, this.requestHeaders), + message, + response: responseBody, + }); + + if (!prepared.ok) { + const message2 = + prepared.json?.error?.message || + prepared.json?.message || + prepared.text || + "Codex WS prepare failed"; + const code = prepared.json?.error?.code || "codex_ws_prepare_failed"; + const error = new Error(message2); + error.code = code; + error.status = prepared.status; + throw error; + } + + this.preparedContext = { + upstreamUrl: toStringOrNull(prepared.json?.upstreamUrl), + connectionId: toStringOrNull(prepared.json?.connectionId), + account: toStringOrNull(prepared.json?.account), + provider: toStringOrNull(prepared.json?.provider) || "codex", + model: toStringOrNull(prepared.json?.model) || toStringOrNull(responseBody.model), + requestedModel: toStringOrNull(responseBody.model), + reasoningRouting: + prepared.json?.reasoningRouting && + typeof prepared.json.reasoningRouting === "object" && + !Array.isArray(prepared.json.reasoningRouting) + ? prepared.json.reasoningRouting + : null, + serviceTier: + toStringOrNull(responseBody.service_tier) || toStringOrNull(responseBody.serviceTier), + }; + + return prepared; + } + async ensureUpstream(firstMessage) { if (this.upstreamReady) return this.upstreamReady; @@ -590,48 +636,7 @@ class ResponsesWsSession { this.firstResponseBody ||= responseBody; this.currentRequestBody = responseBody; - const prepared = await callInternal( - this.fetchImpl, - this.baseUrl, - this.bridgeSecret, - "prepare", - { - requestUrl: this.requestUrl, - headers: getAuthHeaders(this.requestUrl, this.requestHeaders), - message: firstMessage, - response: responseBody, - } - ); - - if (!prepared.ok) { - const message = - prepared.json?.error?.message || - prepared.json?.message || - prepared.text || - "Codex WS prepare failed"; - const code = prepared.json?.error?.code || "codex_ws_prepare_failed"; - const error = new Error(message); - error.code = code; - error.status = prepared.status; - throw error; - } - - this.preparedContext = { - upstreamUrl: toStringOrNull(prepared.json?.upstreamUrl), - connectionId: toStringOrNull(prepared.json?.connectionId), - account: toStringOrNull(prepared.json?.account), - provider: toStringOrNull(prepared.json?.provider) || "codex", - model: toStringOrNull(prepared.json?.model) || toStringOrNull(responseBody.model), - requestedModel: toStringOrNull(responseBody.model), - reasoningRouting: - prepared.json?.reasoningRouting && - typeof prepared.json.reasoningRouting === "object" && - !Array.isArray(prepared.json.reasoningRouting) - ? prepared.json.reasoningRouting - : null, - serviceTier: - toStringOrNull(responseBody.service_tier) || toStringOrNull(responseBody.serviceTier), - }; + const prepared = await this.runPrepare(firstMessage, responseBody); const wsOptions = { // #5591: chrome_149 is not a wreq-js 2.3.1 profile (max chrome_147); the @@ -704,7 +709,16 @@ class ResponsesWsSession { // turn's own request body so persistHistory() attaches the right // clientRequest instead of always the first turn's. const nextTurnBody = getResponseCreatePayload(message); - if (nextTurnBody !== null) this.currentRequestBody = nextTurnBody; + if (nextTurnBody !== null) { + this.currentRequestBody = nextTurnBody; + // #8052: a reused connection must re-run "prepare" (auth/policy/memory/ + // reasoning-routing/compression) for every logical turn, not just the first — + // otherwise every turn after the first bypasses the whole pipeline. This reuses + // the already-established upstream transport; it must NOT recreate the socket. + const prepared = await this.runPrepare(message, nextTurnBody); + this.upstream.send(jsonStringifySafe(withPreparedResponseCreate(message, prepared.json.response))); + return; + } this.upstream.send(jsonStringifySafe(message)); } catch (error) { const code = error?.code || "upstream_websocket_connect_failed"; diff --git a/src/app/api/internal/codex-responses-ws/compression.ts b/src/app/api/internal/codex-responses-ws/compression.ts new file mode 100644 index 0000000000..f4cd889c13 --- /dev/null +++ b/src/app/api/internal/codex-responses-ws/compression.ts @@ -0,0 +1,132 @@ +/** + * Codex Responses-over-WebSocket prompt-compression parity (#8052). + * + * The HTTP/SSE chat pipeline (open-sse/handlers/chatCore.ts) runs every request through the + * modular compression pipeline (settings → strategy selection → applyCompressionAsync → + * compression_analytics/compression_engine_breakdown writes) before dispatching upstream. The + * Codex Responses WebSocket bridge (`prepare()` in ./route.ts) never called any of that — it + * authenticated, injected memory, applied reasoning-routing, then went straight to + * `executor.transformRequest()`. + * + * This module wires the same core pipeline (settings resolution, mode selection via + * `selectCompressionStrategy`, `applyCompressionAsync`, analytics persistence) into the WS + * bridge's per-turn `prepare()` call, using `adaptBodyForCompression`'s existing Responses-API + * (`input[]`) adapter so the same engines that already understand chat `messages[]` bodies work + * unmodified here too. Deliberately scoped to the core pipeline (settings → mode → engines → + * analytics) — combo/output-style/live-zone/adaptive-budget refinements stay chatCore-only for + * now; the WS bridge previously had *zero* compression coverage, so this closes the primary gap. + */ + +import { logger } from "@omniroute/open-sse/utils/logger.ts"; +import { estimateTokens } from "@omniroute/open-sse/services/contextManager.ts"; +import { adaptBodyForCompression } from "@omniroute/open-sse/services/compression/bodyAdapter.ts"; +import type { + CompressionConfig, + CompressionResult, +} from "@omniroute/open-sse/services/compression/types.ts"; +import { resolveCompressionSettings } from "@omniroute/open-sse/handlers/chatCore/compressionSettings.ts"; +import { + writeCompressionAnalytics, + writeCompressionSkip, +} from "@omniroute/open-sse/handlers/chatCore/compressionAnalyticsWrite.ts"; + +const log = logger("RESPONSES_WS_COMPRESSION"); + +type JsonRecord = Record; + +export type ResponsesWsCompressionContext = { + provider: string; + model: string; + /** Distinct per logical turn — feeds the compression_analytics.request_id column. */ + requestId: string; +}; + +/** + * Runs the core compression pipeline against a Codex Responses `response` body and returns the + * (possibly rewritten) body. Best-effort: any resolution/engine failure logs and returns the + * original body untouched — a broken compression path must never break the WS turn. + */ +export async function applyResponsesWsCompression( + responseBody: JsonRecord, + ctx: ResponsesWsCompressionContext +): Promise { + try { + const { settings, enabled } = await resolveCompressionSettings(log); + if (!enabled || !settings) return responseBody; + + const adapter = adaptBodyForCompression(responseBody); + if (!adapter.adapted || !Array.isArray(adapter.body.messages) || adapter.body.messages.length === 0) { + return responseBody; + } + + const { selectCompressionStrategy, applyCompressionAsync } = await import( + "@omniroute/open-sse/services/compression/strategySelector.ts" + ); + + const estimatedTokens = estimateTokens(adapter.body.messages); + const cachingContext = { + provider: ctx.provider, + targetFormat: "openai-responses", + model: ctx.model, + connectionCacheOverride: null, + }; + const mode = selectCompressionStrategy( + settings, + null, + estimatedTokens, + adapter.body, + cachingContext, + {}, + null + ); + if (mode === "off") return responseBody; + + const result = await applyCompressionAsync(adapter.body, mode, { + model: ctx.model, + providerTransport: + ctx.provider === "anthropic" || ctx.provider === "claude" ? "direct" : "aggregator", + config: settings as CompressionConfig, + cachingContext, + }); + + return await persistAndRestore(result, adapter, responseBody, mode, settings, ctx); + } catch (err) { + log.warn( + `[codex-responses-ws] compression skipped: ${err instanceof Error ? err.message : String(err)}` + ); + return responseBody; + } +} + +async function persistAndRestore( + result: CompressionResult, + adapter: ReturnType, + responseBody: JsonRecord, + mode: string, + settings: CompressionConfig, + ctx: ResponsesWsCompressionContext +): Promise { + if (!result.stats) return responseBody; + + const writeOpts = { + stats: result.stats, + provider: ctx.provider, + effectiveModel: ctx.model, + effectiveServiceTier: undefined, + comboName: null, + mode, + compressionComboId: settings.compressionComboId, + skillRequestId: ctx.requestId, + cavemanOutputModeApplied: false, + cavemanOutputModeIntensity: null, + log, + }; + + if (!result.compressed) { + void writeCompressionSkip(writeOpts, "no_savings"); + return responseBody; + } + + void writeCompressionAnalytics(writeOpts); + return adapter.restore(result.body) as JsonRecord; +} diff --git a/src/app/api/internal/codex-responses-ws/route.ts b/src/app/api/internal/codex-responses-ws/route.ts index 772ccd9d5e..a865791a04 100644 --- a/src/app/api/internal/codex-responses-ws/route.ts +++ b/src/app/api/internal/codex-responses-ws/route.ts @@ -1,5 +1,5 @@ import { NextResponse } from "next/server"; -import { createHash, timingSafeEqual } from "node:crypto"; +import { createHash, randomUUID, timingSafeEqual } from "node:crypto"; import { z } from "zod"; import { CodexExecutor } from "@omniroute/open-sse/executors/codex.ts"; import { getApiKeyMetadata } from "@/lib/db/apiKeys"; @@ -32,6 +32,7 @@ import { import { resolveRequestRoutingTags } from "@/domain/tagRouter"; import { validateApiKeyRoutingTarget } from "@/shared/utils/apiKeyPolicy"; import { persistResponsesWsCallHistory } from "./history"; +import { applyResponsesWsCompression } from "./compression"; const CODEX_RESPONSES_WS_URL = "wss://chatgpt.com/backend-api/codex/responses"; const executor = new CodexExecutor(); @@ -512,6 +513,14 @@ async function prepare(body: JsonRecord) { responseBodyWithMemory = applyReasoningRuleDirective(withDirective) 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(), + }); const transformed = (await executor.transformRequest( model, responseBodyWithMemory, diff --git a/tests/unit/responses-ws-proxy-compression-parity.test.ts b/tests/unit/responses-ws-proxy-compression-parity.test.ts new file mode 100644 index 0000000000..013f202994 --- /dev/null +++ b/tests/unit/responses-ws-proxy-compression-parity.test.ts @@ -0,0 +1,218 @@ +// Regression test for issue #8052: the Codex Responses-over-WebSocket bridge bypassed the +// whole prompt-compression pipeline (and its analytics writes). +// +// Two distinct gaps made this happen: +// 1. `prepare()` in src/app/api/internal/codex-responses-ws/route.ts never called anything +// from open-sse/services/compression/* — unlike the HTTP/SSE path (chatCore.ts), which +// runs every request through applyCompressionAsync + writes compression_analytics. +// 2. scripts/dev/responses-ws-proxy.mjs memoized the upstream connection in +// `ensureUpstream()` — it only called the internal "prepare" action on the FIRST +// `response.create` of a WS session. `forwardClientMessage()` forwarded every +// subsequent turn on a reused connection straight to the upstream socket, so even if +// compression were wired into `prepare()`, only the first logical turn of a multi-turn +// WS session would ever see it. +// +// This test proves gap #2 by execution: it opens ONE WebSocket via createResponsesWsProxy(), +// sends two `response.create` messages sequentially on the SAME connection (mirroring the +// issue's "Codex clients may reuse one WebSocket connection for multiple logical turns" repro +// step), and asserts each logical turn triggers its own internal "prepare" call — the only +// place a per-turn compression pass can run. +import test from "node:test"; +import assert from "node:assert/strict"; +import http from "node:http"; + +const { createResponsesWsProxy } = await import("../../scripts/dev/responses-ws-proxy.mjs"); + +function listen(server: http.Server): Promise { + return new Promise((resolve) => { + server.listen(0, "127.0.0.1", () => { + const address = server.address(); + resolve((address as { port: number }).port); + }); + }); +} + +function close(server: http.Server): Promise { + return new Promise((resolve) => { + server.close(() => resolve()); + }); +} + +function readRequestBody(req: http.IncomingMessage): Promise { + return new Promise((resolve, reject) => { + const chunks: Buffer[] = []; + req.on("data", (chunk) => chunks.push(chunk)); + req.on("end", () => resolve(Buffer.concat(chunks).toString("utf8"))); + req.on("error", reject); + }); +} + +function waitFor( + predicate: () => T | undefined | null | false, + { timeoutMs = 3000, intervalMs = 10 }: { timeoutMs?: number; intervalMs?: number } = {} +): Promise { + const startedAt = Date.now(); + return new Promise((resolve, reject) => { + const timer = setInterval(() => { + try { + const value = predicate(); + if (value) { + clearInterval(timer); + resolve(value); + return; + } + if (Date.now() - startedAt >= timeoutMs) { + clearInterval(timer); + reject(new Error("Timed out waiting for condition")); + } + } catch (error) { + clearInterval(timer); + reject(error); + } + }, intervalMs); + }); +} + +test("#8052: every logical response.create turn on a reused WS should hit the internal prepare/compression path", async () => { + const internalRequests: Array> = []; + const downstreamMessages: Array> = []; + const upstreamSends: Array> = []; + + 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_1", + provider: "codex", + account: "codex@example.com", + model: "gpt-5.4-mini", + response: { ...body.response, model: "gpt-5.4-mini", stream: undefined }, + }) + ); + return; + } + + if (body.action === "log") { + res.writeHead(200, { "content-type": "application/json" }); + res.end(JSON.stringify({ ok: true, logged: true })); + return; + } + } + + res.writeHead(404, { "content-type": "application/json" }); + res.end(JSON.stringify({ error: "not_found" })); + }); + + let turn = 0; + const fakeUpstream = { + send(data: string) { + const parsed = JSON.parse(data); + upstreamSends.push(parsed); + if (parsed.type !== "response.create") return; + turn += 1; + const currentTurn = turn; + setTimeout(() => { + fakeUpstream.onmessage?.({ + data: JSON.stringify({ + type: "response.completed", + response: { + id: `resp_${currentTurn}`, + model: "gpt-5.4-mini", + status: "completed", + usage: { + input_tokens: 10 * currentTurn, + output_tokens: 20 * currentTurn, + total_tokens: 30 * currentTurn, + }, + }, + }), + }); + }, 10); + }, + close() {}, + onmessage: null as ((event: { data: string }) => void) | 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`); + ws.addEventListener("message", (event) => { + downstreamMessages.push(JSON.parse(String(event.data))); + }); + + try { + await new Promise((resolve) => ws.addEventListener("open", resolve, { once: true })); + + // Turn 1 on this single, reused WebSocket connection. + ws.send( + JSON.stringify({ + type: "response.create", + model: "gpt-5.4-mini", + input: [{ role: "user", content: "Reply with exactly: pong1" }], + stream: true, + }) + ); + + await waitFor( + () => downstreamMessages.filter((entry) => entry.type === "response.completed").length === 1 + ); + + // Turn 2 on the SAME WebSocket connection (client reuse), per the issue's repro: + // "Codex clients may reuse one WebSocket connection for multiple logical turns." + ws.send( + JSON.stringify({ + type: "response.create", + model: "gpt-5.4-mini", + input: [{ role: "user", content: "Reply with exactly: pong2" }], + stream: true, + }) + ); + + await waitFor( + () => downstreamMessages.filter((entry) => entry.type === "response.completed").length === 2 + ); + + const prepareRequests = internalRequests.filter((entry) => entry.action === "prepare"); + + assert.equal( + prepareRequests.length, + 2, + `expected 2 internal "prepare" calls (one per logical response.create turn), got ${prepareRequests.length} — ` + + "the second turn on a reused WebSocket connection bypasses prepare()/compression entirely (#8052)" + ); + } finally { + ws.close(); + await close(server); + } +});