Something went wrong. Try again.
Something went wrong. Try again.
12 kB · 314 lines
TypeScript
at main
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315import assert from "node:assert/strict"import fs from "node:fs"import os from "node:os"import path from "node:path"import test, { type TestContext } from "node:test"import { createCodexBridgeServer, handleCodexCompletion, __codexBridgeTest } from "./codex"import { createServer } from "./server"import { accountCredentialPath, configureKeyRouting, createAccount } from "./accounts"import { closeTokenUsageDatabases, getTokenUsage } from "./token-usage"
const encode = (event: unknown) => new TextEncoder().encode(`data: ${JSON.stringify(event)}\r\n\r\n`)const call = { type: "function_call", id: "item_py", call_id: "call_py", name: "python", arguments: "" }const completed = { type: "response.completed", response: { id: "resp_stream", usage: { input_tokens: 3, output_tokens: 2, total_tokens: 5 } },}
function auth(t: TestContext): string { const dir = fs.mkdtempSync(path.join(os.tmpdir(), "codex-stream-")) const authPath = path.join(dir, "auth.json") fs.writeFileSync(authPath, JSON.stringify({ tokens: { access_token: "header.payload.signature" } })) t.after(() => fs.rmSync(dir, { recursive: true, force: true })) return authPath}
async function translate(t: TestContext, events: unknown[]): Promise<string> { const result = await handleCodexCompletion({ model: "gpt-test", stream: true, messages: [] }, undefined, { authPath: auth(t), fetchImpl: (async () => new Response(events.map((e) => new TextDecoder().decode(encode(e))).join(""))) as typeof fetch, }) let wire = "" for await (const chunk of result.stream!) wire += chunk return wire}
for (const [unified, includeUsage] of [ [false, true], [true, true], [true, false],]) { test( `streams before completion (${unified ? "unified" : "standalone"}, usage=${includeUsage})`, { timeout: 10000 }, async (t) => { const authPath = auth(t) const previousState = process.env.BRIDGE_STATE_DIR const stateDir = path.dirname(authPath) process.env.BRIDGE_STATE_DIR = stateDir t.after(() => { closeTokenUsageDatabases() if (previousState === undefined) delete process.env.BRIDGE_STATE_DIR else process.env.BRIDGE_STATE_DIR = previousState }) let controller!: ReadableStreamDefaultController<Uint8Array> let cancelled = false const upstream = new ReadableStream<Uint8Array>({ start(c) { controller = c }, cancel() { cancelled = true }, }) const fetchImpl = (async () => new Response(upstream, { headers: { "content-type": "text/event-stream" } })) as typeof fetch const key = "cccccccc-cccc-4ccc-8ccc-cccccccccccc" if (unified) { const account = createAccount(stateDir, "codex") fs.copyFileSync(authPath, accountCredentialPath(stateDir, account.id, "codex")!) configureKeyRouting(stateDir, key, { usageScheme: "fill-first", accounts: { codex: [account.id] }, preferredAccounts: {}, }) } const app = unified ? createServer( { port: 0, host: "127.0.0.1", models: { "gpt-test": { id: "gpt-test", name: "test", provider: "codex", apiKey: "" } }, providers: { codex: { auth_path: authPath } }, }, new Map([[key, { key, allowedModels: ["*"], caps: [] }]]), { fetchImpl }, ) : createCodexBridgeServer({ authPath, fetchImpl }) t.after(() => app.close()) const address = await app.listen({ port: 0, host: "127.0.0.1" }) const response = await fetch(`${address}/v1/chat/completions`, { method: "POST", headers: { "content-type": "application/json", authorization: `Bearer ${key}` }, body: JSON.stringify({ model: "gpt-test", messages: [{ role: "user", content: "write" }], stream: true, stream_options: { include_usage: includeUsage }, }), signal: AbortSignal.timeout(8000), }) assert.equal(response.status, 200, response.ok ? undefined : await response.text()) const reader = response.body!.getReader() const decoder = new TextDecoder() let wire = "" async function until(needle: string) { while (!wire.includes(needle)) { const next = await reader.read() assert.equal(next.done, false) wire += decoder.decode(next.value, { stream: true }) } } // Upstream deliberately stays open: these must arrive without waiting for EOF. controller.enqueue(encode({ type: "response.output_item.added", output_index: 2, item: call })) await until("call_py") controller.enqueue(encode({ type: "response.reasoning_summary_text.delta", delta: "checking" })) await until("checking") controller.enqueue(encode({ type: "response.output_text.delta", delta: "hello" })) await until("hello") const first = '{"code":"print(' controller.enqueue( encode({ type: "response.function_call_arguments.delta", output_index: 2, item_id: call.id, delta: first }), ) await until("print(") assert.equal(wire.includes("[DONE]"), false) const last = "'界')\"}" const bytes = encode({ type: "response.function_call_arguments.delta", output_index: 2, item_id: call.id, delta: last, }) // Split UTF-8 sequences and SSE boundaries across arbitrary transport chunks. for (const byte of bytes) controller.enqueue(Uint8Array.of(byte)) controller.enqueue( encode({ type: "response.function_call_arguments.done", output_index: 2, arguments: first + last }), ) controller.enqueue( encode({ type: "response.output_item.done", output_index: 2, item: { ...call, arguments: first + last } }), ) controller.enqueue( encode({ ...completed, response: { ...completed.response, output: [{ type: "reasoning" }, { type: "message" }, { ...call, arguments: first + last }], }, }), ) while (true) { const next = await reader.read() if (next.done) break wire += decoder.decode(next.value, { stream: true }) } wire += decoder.decode() const payloads = __codexBridgeTest.parseSse(wire) const parts = payloads.flatMap((p) => p.choices?.[0]?.delta?.tool_calls ?? []) assert.equal(parts.filter((p) => p.function.name).length, 1) assert.deepEqual( parts.map((p) => p.index), [0, 0, 0], ) assert.equal(parts.map((p) => p.function.arguments).join(""), first + last) assert.equal(payloads.filter((p) => p.usage).length, includeUsage ? 1 : 0) assert.equal(payloads.filter((p) => p.choices?.[0]?.finish_reason === "tool_calls").length, 1) assert.equal(cancelled, true, "terminal event releases the still-open upstream") if (unified) { const usage = getTokenUsage(stateDir, key) assert.equal(usage.requests, 1) assert.equal(usage.totalTokens, 5, "internal accounting survives client usage opt-out") } }, )}
test("done-only tool fallback, multiple calls and usage opt-out", async (t) => { const wire = await translate(t, [ { type: "response.output_item.done", output_index: 1, item: { ...call, arguments: '{"a":1}' } }, { type: "response.output_item.added", output_index: 3, item: { ...call, id: "item_2", call_id: "call_2" } }, { type: "response.function_call_arguments.delta", output_index: 3, delta: '{"b":' }, { type: "response.function_call_arguments.done", output_index: 3, arguments: '{"b":2}' }, completed, ]) const payloads = __codexBridgeTest.parseSse(wire) const calls = payloads.flatMap((p) => p.choices?.[0]?.delta?.tool_calls ?? []) assert.equal( calls .filter((c) => c.index === 0) .map((c) => c.function.arguments) .join(""), '{"a":1}', ) assert.equal( calls .filter((c) => c.index === 1) .map((c) => c.function.arguments) .join(""), '{"b":2}', ) assert.equal( payloads.some((p) => p.usage), false, )})
test("failed and truncated streams cannot become successful completions", async (t) => { for (const ending of [[], [{ type: "response.failed", response: { error: { message: "provider failed" } } }]]) { const wire = await translate(t, [{ type: "response.output_text.delta", delta: "partial" }, ...ending]) const events = __codexBridgeTest.parseSse(wire) assert.ok(events.some((e) => e.error)) assert.equal( events.some((e) => e.choices?.[0]?.finish_reason), false, ) assert.equal(wire.match(/\[DONE\]/g)?.length, 1) }})
test("client disconnect aborts an idle upstream fetch", { timeout: 10000 }, async (t) => { let upstreamSignal: AbortSignal | null | undefined let aborted!: () => void const stopped = new Promise<void>((resolve) => (aborted = resolve)) const fetchImpl = (async (_url: unknown, init?: RequestInit) => { upstreamSignal = init?.signal const upstream = new ReadableStream<Uint8Array>({ start(controller) { upstreamSignal?.addEventListener( "abort", () => { controller.error(new Error("cancelled")) aborted() }, { once: true }, ) }, }) return new Response(upstream) }) as typeof fetch const app = createCodexBridgeServer({ authPath: auth(t), fetchImpl }) t.after(() => app.close()) const address = await app.listen({ host: "127.0.0.1", port: 0 }) const abort = new AbortController() const response = await fetch(`${address}/v1/chat/completions`, { method: "POST", headers: { "content-type": "application/json" }, body: JSON.stringify({ stream: true, messages: [] }), signal: abort.signal, }) await response.body!.getReader().read() abort.abort() await stopped assert.equal(upstreamSignal?.aborted, true)})
test("interleaved tool deltas retain independent indices and final arguments", async (t) => { const second = { ...call, id: "item_2", call_id: "call_2", name: "other" } const wire = await translate(t, [ { type: "response.output_item.added", output_index: 1, item: call }, { type: "response.output_item.added", output_index: 2, item: second }, { type: "response.function_call_arguments.delta", output_index: 1, delta: '{"a":' }, { type: "response.function_call_arguments.delta", output_index: 2, delta: '{"b":2}' }, { type: "response.function_call_arguments.delta", output_index: 1, delta: "1}" }, { ...completed, response: { ...completed.response, output: [{ type: "reasoning" }, { ...call, arguments: '{"a":1}' }, { ...second, arguments: '{"b":2}' }], }, }, ]) const parts = __codexBridgeTest.parseSse(wire).flatMap((p) => p.choices?.[0]?.delta?.tool_calls ?? []) assert.deepEqual( parts.filter((p) => p.id).map((p) => [p.index, p.id]), [ [0, "call_py"], [1, "call_2"], ], ) assert.equal( parts .filter((p) => p.index === 0) .map((p) => p.function.arguments) .join(""), '{"a":1}', ) assert.equal( parts .filter((p) => p.index === 1) .map((p) => p.function.arguments) .join(""), '{"b":2}', )})
test("revised arguments fail explicitly; incomplete responses retain their finish reason", async (t) => { const invalid = await translate(t, [ { type: "response.output_item.added", output_index: 0, item: call }, { type: "response.function_call_arguments.delta", output_index: 0, delta: '{"old":' }, { type: "response.function_call_arguments.done", output_index: 0, arguments: '{"new":1}' }, completed, ]) assert.match(invalid, /revised already-streamed/) assert.equal( __codexBridgeTest.parseSse(invalid).some((p) => p.choices?.[0]?.finish_reason), false, ) for (const [reason, finish] of [ ["max_output_tokens", "length"], ["content_filter", "content_filter"], ]) { const wire = await translate(t, [{ type: "response.incomplete", response: { incomplete_details: { reason } } }]) assert.equal(__codexBridgeTest.parseSse(wire).at(-1).choices[0].finish_reason, finish) }})