Something went wrong. Try again.
Something went wrong. Try again.
21 kB · 510 lines
TypeScript
at main
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511import fs from "node:fs"import os from "node:os"import path from "node:path"import { listRequestMetrics, closeTokenUsageDatabases } from "./token-usage"import assert from "node:assert/strict"import { once } from "node:events"import { createServer } from "node:http"import test, { type TestContext } from "node:test"import { WebSocketServer, type WebSocket } from "ws"import { CodexWebSocketPool } from "./codex-websocket"import { sseFrames } from "./sse"
const input = [{ role: "user", content: "hello" }]const output = [ { type: "reasoning", id: "rs_1", encrypted_content: "opaque" }, { type: "message", id: "msg_1", role: "assistant", content: [{ type: "output_text", text: "answer", annotations: [] }], },]const complete = (socket: WebSocket, id = "resp_1") => { socket.send(JSON.stringify({ type: "response.output_text.delta", delta: "answer" })) socket.send( JSON.stringify({ type: "response.completed", response: { id, output, usage: { input_tokens: 16, output_tokens: 1 } }, }), )}const quota = { type: "error", status: 429, error: { code: "usage_limit_reached", message: "exhausted", resets_in_seconds: 60 },}
async function fixture(t: TestContext, respond: (socket: WebSocket, body: any, n: number) => void, options = {}) { const server = new WebSocketServer({ host: "127.0.0.1", port: 0 }) await once(server, "listening") const address = server.address() as { port: number } const pool = new CodexWebSocketPool({ url: `ws://127.0.0.1:${address.port}`, ...options }) const frames: any[] = [] const sockets: WebSocket[] = [] server.on("connection", (socket) => { sockets.push(socket) socket.on("message", (raw) => { const body = JSON.parse(raw.toString()) frames.push(body) respond(socket, body, frames.length) }) }) t.after(async () => { pool.close() for (const socket of sockets) socket.terminate() await new Promise<void>((resolve) => server.close(() => resolve())) }) const request = (body: any = {}, overrides: any = {}) => pool.request({ scope: "api-key:responses:conversation", url: "https://unused", headers: { authorization: "Bearer account-a" }, body: { model: "gpt-test", store: false, input, ...body }, chat: false, fallback: async () => { throw new Error("unexpected HTTP fallback") }, ...overrides, }) const run = async (body: any = {}, overrides: any = {}) => { const response = await request(body, overrides) assert.equal(response.status, 200, await (!response.ok ? response.text() : Promise.resolve(""))) const events: any[] = [] for await (const event of sseFrames(response.body!)) events.push(event) return events } return { pool, server, sockets, frames, request, run }}
test("reuses socket, preserves raw reasoning, and sends only append-only deltas", async (t) => { const f = await fixture(t, (socket) => complete(socket)) await f.run() const next = { role: "user", content: "next" } await f.run({ input: [...input, ...output, next], client_metadata: { turn_id: "fresh" } }) assert.equal(f.sockets.length, 1) assert.equal(f.frames[1].previous_response_id, "resp_1") assert.deepEqual(f.frames[1].input, [next]) const changed = [...input, ...output, next, ...output, next] await f.run({ input: changed, tools: [{ type: "function", name: "new_tool" }] }) assert.equal(f.frames[2].previous_response_id, undefined) assert.deepEqual(f.frames[2].input, changed)})
test("chat projection retains upstream reasoning without requiring it in chat messages", async (t) => { const f = await fixture(t, (socket) => complete(socket)) await f.run({}, { chat: true }) const next = { role: "user", content: "again" } await f.run( { input: [...input, { role: "assistant", content: [{ type: "output_text", text: "answer" }] }, next] }, { chat: true }, ) assert.equal(f.frames[1].previous_response_id, "resp_1") assert.deepEqual(f.frames[1].input, [next])})
test("explicit continuation, then a closed socket, replays the known full history", async (t) => { const f = await fixture(t, (socket) => complete(socket)) await f.run() f.sockets[0].close() await once(f.sockets[0], "close") await new Promise((resolve) => setImmediate(resolve)) const next = { role: "user", content: "again" } await f.run({ previous_response_id: "resp_1", input: [next] }) assert.equal(f.sockets.length, 2) assert.equal(f.frames[1].previous_response_id, undefined) assert.deepEqual(f.frames[1].input, [...input, ...output, next])})
test("open sockets keep continuation before the 30-minute idle prune", async (t) => { const f = await fixture(t, (socket) => complete(socket)) await f.run() const clock = Date.now let elapsed = 20 * 60_000 Date.now = () => clock() + elapsed t.after(() => { Date.now = clock }) const next = { role: "user", content: "again" } await f.run({ input: [...input, ...output, next] }) assert.equal(f.sockets.length, 1) assert.equal(f.frames[1].previous_response_id, "resp_1") assert.deepEqual(f.frames[1].input, [next])
elapsed += 31 * 60_000 const later = { role: "user", content: "later" } const history = [...input, ...output, next, ...output, later] await f.run({ input: history }) assert.equal(f.sockets.length, 2) assert.equal(f.frames[2].previous_response_id, undefined) assert.deepEqual(f.frames[2].input, history)})
test("unknown or differently scoped/accounted continuations never reach upstream", async (t) => { const f = await fixture(t, (socket) => complete(socket)) await f.run() const scoped = await f.request( { previous_response_id: "resp_1", input: [] }, { scope: "other-key:responses:conversation" }, ) assert.equal(scoped.status, 400) assert.match(await scoped.text(), /previous_response_not_found/) const stale = once(f.sockets[0], "close") const accounted = await f.request( { previous_response_id: "resp_1", input: [] }, { headers: { authorization: "Bearer account-b" } }, ) assert.equal(accounted.status, 400) assert.match(await accounted.text(), /previous_response_not_found/) await stale assert.equal(f.sockets.length, 1, "the idle socket for the other account is dropped, not reused") assert.equal(f.frames.length, 1)})
test("missing server continuation retries full history once before exposing frames", async (t) => { const f = await fixture(t, (socket, body) => { if (body.previous_response_id) { socket.send(JSON.stringify({ type: "response.created", response: { id: "unused" } })) socket.send(JSON.stringify({ type: "error", status: 400, error: { code: "previous_response_not_found" } })) } else complete(socket) }) await f.run() const events = await f.run({ input: [...input, ...output, ...input] }) assert.equal(f.frames.length, 3) assert.equal(f.frames[2].previous_response_id, undefined) assert.deepEqual(f.frames[2].input, [...input, ...output, ...input]) assert.ok(!JSON.stringify(events).includes("unused"))})
test("quota before output is HTTP 429; a partial response is never retried", async (t) => { const f = await fixture(t, (socket, _body, n) => { socket.send(JSON.stringify({ type: "response.created", response: { id: "resp_1" } })) if (n === 2) socket.send(JSON.stringify({ type: "response.output_text.delta", delta: "partial" })) socket.send(JSON.stringify(quota)) }) const first = await f.request() assert.equal(first.status, 429) assert.deepEqual(await first.json(), { error: quota.error }) const second = await f.run() assert.ok(JSON.stringify(second).includes("partial")) assert.equal(second.at(-1)?.type, "error") assert.equal(f.frames.length, 2)})
test("overlapping requests on one conversation keep separate sockets and both complete", async (t) => { const f = await fixture(t, (socket, _body, n) => { socket.send(JSON.stringify({ type: "response.output_text.delta", delta: `reply-${n}` })) }) const first = f.run({ client_metadata: { turn_id: "first" } }) for (let i = 0; i < 200 && f.frames.length < 1; i++) await new Promise((resolve) => setImmediate(resolve)) assert.equal(f.frames.length, 1) const second = f.run({ client_metadata: { turn_id: "second" } }) for (let i = 0; i < 200 && f.frames.length < 2; i++) await new Promise((resolve) => setImmediate(resolve)) assert.equal(f.frames.length, 2) assert.equal(f.sockets.length, 2, "a busy socket in the same conversation must not be reused") assert.equal(f.sockets[0].readyState, 1, "the first consumer's socket is left open") f.sockets[0].send(JSON.stringify({ type: "response.completed", response: { id: "resp-first", output } })) f.sockets[1].send(JSON.stringify({ type: "response.completed", response: { id: "resp-second", output } })) const [firstEvents, secondEvents] = await Promise.all([first, second]) assert.ok(JSON.stringify(firstEvents).includes("reply-1")) assert.ok(JSON.stringify(secondEvents).includes("reply-2"))})
test("cancelling one overlapping request leaves its peer streaming", async (t) => { const f = await fixture(t, (socket, _body, n) => { socket.send(JSON.stringify({ type: "response.output_text.delta", delta: `reply-${n}` })) }) const abort = new AbortController() const first = f.run({}, { signal: abort.signal }) for (let i = 0; i < 200 && f.frames.length < 1; i++) await new Promise((resolve) => setImmediate(resolve)) const second = f.run({ client_metadata: { turn_id: "peer" } }) for (let i = 0; i < 200 && f.frames.length < 2; i++) await new Promise((resolve) => setImmediate(resolve)) const firstClosed = once(f.sockets[0], "close") abort.abort() await assert.rejects(first, /abort/i) await firstClosed assert.equal(f.sockets.length, 2) assert.equal(f.sockets[0].readyState, 3, "only the aborted consumer's socket closes") assert.equal(f.sockets[1].readyState, 1, "the peer's socket keeps streaming") f.sockets[1].send(JSON.stringify({ type: "response.completed", response: { id: "resp-peer", output } })) assert.ok(JSON.stringify(await second).includes("reply-2"))})
test("a late close on a retired socket does not disturb a peer still streaming", async (t) => { const f = await fixture(t, (socket, _body, n) => { socket.send(JSON.stringify({ type: "response.output_text.delta", delta: `reply-${n}` })) }) const first = f.request() for (let i = 0; i < 200 && f.frames.length < 1; i++) await new Promise((resolve) => setImmediate(resolve)) const second = await f.request() const reader = second.body!.getReader() assert.match(new TextDecoder().decode((await reader.read()).value), /reply-2/) const closed = once(f.sockets[0], "close") f.sockets[0].close(1000, "normal") await closed assert.equal(f.sockets.length, 2) f.sockets[1].send(JSON.stringify({ type: "response.completed", response: { id: "resp-peer", output } })) const rest: string[] = [] for (;;) { const { value, done } = await reader.read() if (done) break rest.push(new TextDecoder().decode(value)) } assert.ok(rest.join("").includes("response.completed")) await assert.rejects( first.then(async (response) => response.text()), /closed before response.completed/, )})
test("a remote close at 1000 fails an in-flight response but is harmless once completed", async (t) => { const f = await fixture(t, (socket) => socket.send(JSON.stringify({ type: "response.output_text.delta", delta: "partial" })), ) const response = await f.request() const reader = response.body!.getReader() assert.match(new TextDecoder().decode((await reader.read()).value), /partial/) const closed = once(f.sockets[0], "close") f.sockets[0].close(1000, "normal") await assert.rejects(reader.read(), /closed before response.completed/) await closed
const done = await fixture(t, (socket) => { complete(socket) socket.close(1000, "normal") }) const events = await done.run() assert.ok(JSON.stringify(events).includes("response.completed")) assert.ok(!JSON.stringify(events).includes("error"))})
test("downstream cancellation closes the active upstream response", async (t) => { const f = await fixture(t, (socket) => socket.send(JSON.stringify({ type: "response.output_text.delta", delta: "early" })), ) const response = await f.request() const reader = response.body!.getReader() assert.match(new TextDecoder().decode((await reader.read()).value), /early/) const closed = once(f.sockets[0], "close") await reader.cancel() await closed})
test("abort releases the connection; pre-output timeout and transport failure fall back to HTTP", async (t) => { const f = await fixture( t, (socket, _body, n) => { if (n === 3) socket.terminate() }, { timeoutMs: 30 }, ) const abort = new AbortController() const pending = f.request({}, { signal: abort.signal }) setTimeout(() => abort.abort(), 10) await assert.rejects(pending, /abort/i) let fallbacks = 0 const fallback = async () => { fallbacks++ return new Response("fallback") } assert.equal(await (await f.request({}, { fallback })).text(), "fallback") assert.equal(await (await f.request({}, { fallback })).text(), "fallback") assert.equal(fallbacks, 2)})
test("upgrade-required rejection falls back to HTTP without generation replay", async (t) => { const server = createServer() server.on("upgrade", (_req, socket) => socket.end("HTTP/1.1 426 Upgrade Required\r\nContent-Length: 0\r\n\r\n")) server.listen(0, "127.0.0.1") await once(server, "listening") const pool = new CodexWebSocketPool({ url: `ws://127.0.0.1:${(server.address() as any).port}` }) t.after(() => { pool.close() server.close() }) let fallbacks = 0 for (let i = 0; i < 2; i++) { const response = await pool.request({ scope: "test", url: "unused", headers: {}, body: { input }, chat: false, fallback: async () => { fallbacks++ return new Response("fallback") }, }) assert.equal(await response.text(), "fallback") } assert.equal(fallbacks, 2)})
test("WebSocket rate-limit events refresh the HTTP usage headers each request", async (t) => { const f = await fixture(t, (socket, _body, n) => { socket.send( JSON.stringify({ type: "codex.rate_limits", rate_limits: { primary: { used_percent: n, window_minutes: 300, reset_at: 1800000000 } }, }), ) complete(socket) }) for (const n of [1, 2]) { const response = await f.request() assert.equal(response.headers.get("x-codex-primary-used-percent"), String(n)) assert.equal(response.headers.get("x-codex-primary-window-minutes"), "300") assert.equal(response.headers.get("x-codex-primary-reset-at"), "1800000000") await response.text() }})
test("output-item events supply continuation when response.completed omits its output", async (t) => { const tool = { type: "function_call", id: "fc_1", status: "completed", call_id: "call_1", name: "test", arguments: "{}", } const f = await fixture(t, (socket) => { socket.send(JSON.stringify({ type: "response.output_item.done", output_index: 0, item: tool })) socket.send(JSON.stringify({ type: "response.completed", response: { id: "resp_tool", output: [] } })) }) await f.run({}, { chat: true }) const result = { type: "function_call_output", call_id: "call_1", output: "done" } await f.run( { input: [...input, { type: "function_call", call_id: "call_1", name: "test", arguments: "{}" }, result] }, { chat: true }, ) assert.equal(f.frames[1].previous_response_id, "resp_tool") assert.deepEqual(f.frames[1].input, [result])})
test("the connection cap evicts idle sockets but never interrupts an active response", async (t) => { const idle = await fixture(t, (socket) => complete(socket), { maxConnections: 1 }) await idle.run() const closed = once(idle.sockets[0], "close") await idle.run({}, { scope: "new-conversation" }) await closed assert.equal(idle.sockets.length, 2)
const active = await fixture( t, (socket, _body, n) => { if (n === 1) socket.send(JSON.stringify({ type: "response.output_text.delta", delta: "working" })) else complete(socket) }, { maxConnections: 1 }, ) const response = await active.request() const reader = response.body!.getReader() assert.match(new TextDecoder().decode((await reader.read()).value), /working/) const rejected = await active.request({}, { scope: "other-conversation" }) assert.equal(rejected.status, 503) assert.equal(active.sockets[0].readyState, 1) await reader.cancel()})
test("overlapping sockets on one conversation count toward the connection cap", async (t) => { const f = await fixture( t, (socket, _body, n) => { socket.send(JSON.stringify({ type: "response.output_text.delta", delta: `reply-${n}` })) }, { maxConnections: 2 }, ) const first = await f.request({ client_metadata: { turn_id: "one" } }) const second = await f.request({ client_metadata: { turn_id: "two" } }) assert.equal(f.sockets.length, 2) const rejected = await f.request({}, { scope: "another-conversation" }) assert.equal(rejected.status, 503, "both active overlapping sockets occupy the cap") assert.equal(f.sockets.length, 2) await first.body!.cancel() await second.body!.cancel()})
test("unknown continuations at the cap do not evict a valid idle continuation", async (t) => { const f = await fixture(t, (socket) => complete(socket), { maxConnections: 1 }) await f.run() for (let attempt = 0; attempt < 2; attempt++) { const response = await f.request({ previous_response_id: "bogus", input: [] }) assert.equal(response.status, 400) assert.match(await response.text(), /previous_response_not_found/) } assert.equal(f.frames.length, 1, "invalid continuations never reach upstream") assert.equal(f.sockets.length, 1, "invalid continuations neither allocate nor evict") assert.equal(f.sockets[0].readyState, 1) const next = { role: "user", content: "next" } const events = await f.run({ previous_response_id: "resp_1", input: [next] }) assert.ok(JSON.stringify(events).includes("response.completed")) assert.equal(f.frames.length, 2) assert.equal(f.sockets.length, 1, "the valid continuation still reuses the preserved socket")})
test("tool continuations retain their turn metadata and first routing token; new user turns do not", async (t) => { const f = await fixture(t, (socket, _body, n) => { socket.send(JSON.stringify({ type: "response.metadata", headers: { "X-Codex-Turn-State": `state-${n}` } })) complete(socket) }) await f.run({ client_metadata: { turn_id: "first", "x-codex-turn-metadata": "first-metadata" } }) const result = { type: "function_call_output", call_id: "call_1", output: "done" } await f.run({ input: [...input, ...output, result], client_metadata: { turn_id: "second" } }) assert.equal(f.frames[1].client_metadata.turn_id, "first") assert.equal(f.frames[1].client_metadata["x-codex-turn-metadata"], "first-metadata") assert.equal(f.frames[1].client_metadata["x-codex-turn-state"], "state-1") await f.run({ input: [...input, ...output, result, ...output, result], client_metadata: { turn_id: "third" } }) assert.equal(f.frames[2].client_metadata["x-codex-turn-state"], "state-1") await f.run({ input: [...input, ...output, result, ...output, result, ...output, ...input], client_metadata: { turn_id: "new-user" }, }) assert.equal(f.frames[3].client_metadata.turn_id, "new-user") assert.equal(f.frames[3].client_metadata["x-codex-turn-state"], undefined)})
test("metrics distinguish missing-continuation replay, connection reuse, and reconnects", async (t) => { const stateDirectory = fs.mkdtempSync(path.join(os.tmpdir(), "ws-metrics-")) t.after(() => { closeTokenUsageDatabases() fs.rmSync(stateDirectory, { recursive: true, force: true }) }) const f = await fixture(t, (socket, body) => { if (body.previous_response_id) socket.send(JSON.stringify({ type: "error", status: 400, error: { code: "previous_response_not_found" } })) else complete(socket) }) let attempt = 0 const overrides = { trace: { stateDirectory, requestId: "logical", nextAttempt: () => ++attempt, apiKeyId: "hashed", protocol: "responses", downstreamTransport: "ws", downstreamHeaders: {}, }, } await f.run({}, overrides) await f.run({ input: [...input, ...output, ...input] }, overrides) f.sockets[0].close() await once(f.sockets[0], "close") await new Promise((resolve) => setImmediate(resolve)) await f.run({ previous_response_id: "resp_1", input }, overrides) const rows = listRequestMetrics(stateDirectory).records.sort((a, b) => a.attempt - b.attempt) assert.equal(rows.length, 4) assert.equal(new Set(rows.map((row) => row.id)).size, 4) assert.deepEqual( rows.map((row) => row.previous_response_id), [null, "resp_1", null, null], ) assert.equal(rows[1].status, "error") assert.equal(rows[1].input_tokens, null) assert.equal(rows[1].full_input_hash, rows[2].full_input_hash) assert.equal(rows[0].connection_id, rows[2].connection_id) assert.equal(rows[3].previous_connection_id, rows[2].connection_id) assert.notEqual(rows[3].connection_id, rows[2].connection_id) assert.equal(rows[3].reconnect_count, 1)})