diff --git a/pnpm-lock.yaml b/pnpm-lock.yaml index b12ecd1c77..395cb40b73 100644 --- a/pnpm-lock.yaml +++ b/pnpm-lock.yaml @@ -620,6 +620,9 @@ importers: workerpool: specifier: ^9.2.0 version: 9.2.0 + ws: + specifier: ^8.18.3 + version: 8.21.3 yaml: specifier: ^2.8.0 version: 2.9.0 @@ -666,6 +669,9 @@ importers: '@types/vscode': specifier: 1.100.0 version: 1.100.0 + '@types/ws': + specifier: ^8.18.1 + version: 8.18.2 '@vitest/coverage-v8': specifier: 4.1.9 version: 4.1.9(vitest@4.1.11) @@ -3573,6 +3579,9 @@ packages: '@types/vscode@1.100.0': resolution: {integrity: sha512-4uNyvzHoraXEeCamR3+fzcBlh7Afs4Ifjs4epINyUX/jvdk0uzLnwiDY35UKDKnkCHP5Nu3dljl2H8lR6s+rQw==} + '@types/ws@8.18.2': + resolution: {integrity: sha512-67MQl+fpWKVTT1NYdnmo3U4sc/xPo/zQBncVnI74qmQa0z/b+1g6iYqNmGCPbxO+zz2aklb08a0oHfegiVd0/w==} + '@types/yargs-parser@21.0.3': resolution: {integrity: sha512-I4q9QU9MQv4oEOz4tAHJtNz1cwuLxn2F3xcc2iV5WdqLPpUnj30aUuxt1mAxYTG+oe8CZMV/+6rU4S4gRDzqtQ==} @@ -11853,6 +11862,10 @@ snapshots: '@types/vscode@1.100.0': {} + '@types/ws@8.18.2': + dependencies: + '@types/node': 22.20.1 + '@types/yargs-parser@21.0.3': {} '@types/yargs@17.0.33': diff --git a/src/api/providers/CodexWebSocketTransport.ts b/src/api/providers/CodexWebSocketTransport.ts new file mode 100644 index 0000000000..e2dbf95935 --- /dev/null +++ b/src/api/providers/CodexWebSocketTransport.ts @@ -0,0 +1,118 @@ +import type WebSocket from "ws" + +import type { CodexWebSocketConnectionManager } from "./codex-websocket/managers/CodexWebSocketConnectionManager" +import type { CodexWebSocketRequestManager } from "./codex-websocket/managers/CodexWebSocketRequestManager" +import type { CodexWebSocketResponseManager } from "./codex-websocket/managers/CodexWebSocketResponseManager" +import type { PreparedCodexRequest } from "./codex-websocket/models/PreparedCodexRequest" +import type { CodexResponseEvent, CodexWebSocketOptions } from "./codex-websocket/models/protocol" +import type { CodexWebSocketContinuationRepository } from "./codex-websocket/repositories/CodexWebSocketContinuationRepository" +import type { CodexWebSocketRequestScope } from "./codex-websocket/scopes/CodexWebSocketRequestScope" +import { asJsonObject, parseResponseEvent } from "./codex-websocket/utils/protocol" + +export { CodexWebSocketUnavailableError } from "./codex-websocket/errors/CodexWebSocketUnavailableError" + +/** Task-local Responses transport. Connection and history management are independent of provider event handling. */ +export class CodexWebSocketTransport { + private busy = false + private requestScope?: CodexWebSocketRequestScope + + constructor( + private readonly connection: CodexWebSocketConnectionManager, + private readonly continuation: CodexWebSocketContinuationRepository, + private readonly createRequestScope: ( + options: CodexWebSocketOptions, + onAbort: () => void, + ) => CodexWebSocketRequestScope, + ) {} + + resetContinuation(): void { + this.continuation.reset() + } + async dispose(): Promise { + const scope = this.requestScope + this.requestScope = undefined + this.connection.dispose() + await scope?.dispose() + } + + async *stream(body: unknown, options: CodexWebSocketOptions): AsyncGenerator { + if (this.busy) throw new Error("Concurrent requests on a Codex WebSocket are not supported") + options.signal.throwIfAborted() + const scope = this.createRequestScope(options, () => { + void this.dispose() + }) + let request: CodexWebSocketRequestManager | undefined + let response: CodexWebSocketResponseManager | undefined + this.busy = true + this.requestScope = scope + try { + await scope.init(this.connection, this.continuation, body) + request = scope.requestManager + response = scope.responseManager + const prepared = response.preparedRequest + yield* this.readResponse(prepared, options.headers, request, response) + } catch (error) { + if (request?.signal.aborted) throw request.signal.reason + options.signal.throwIfAborted() + throw error + } finally { + await scope.dispose() + if (this.requestScope === scope) this.requestScope = undefined + this.busy = false + if (!response?.completed || request?.signal.aborted) await this.dispose() + else this.connection.release() + } + } + + private async *readResponse( + prepared: PreparedCodexRequest, + headers: Record, + request: CodexWebSocketRequestManager, + response: CodexWebSocketResponseManager, + ): AsyncGenerator { + request.refreshTimeout() + this.send(request.socket, prepared, headers, prepared.fullContextReason === undefined) + for await (const [data] of request.events) { + request.refreshTimeout() + const event = parseResponseEvent(String(data)) + if (response.accept(event) === "retry") { + this.send(request.socket, prepared, headers, false, "server cache miss") + continue + } + if (response.completed) request.clearTimeout() + yield event + if (response.completed) return + } + throw new Error("Codex WebSocket closed before response completed; retry the request") + } + + private send( + socket: WebSocket, + prepared: PreparedCodexRequest, + headers: Record, + incremental: boolean, + reason = prepared.fullContextReason, + ): void { + const { request, input, previousResponseId, offset } = prepared + const nextInput = incremental ? input.slice(offset) : input + const clientMetadata = { + ...(request.client_metadata ? asJsonObject(request.client_metadata) : {}), + ...(headers["x-openai-internal-codex-responses-lite"] === "true" + ? { ws_request_header_x_openai_internal_codex_responses_lite: "true" } + : {}), + } + console.info( + `[Codex WebSocket] Sending ${incremental ? "incremental" : "full"} context` + + ` (${nextInput.length}/${input.length} input items${incremental ? "" : `; reason: ${reason}`})`, + ) + socket.send( + JSON.stringify({ + ...request, + type: "response.create", + client_metadata: clientMetadata, + previous_response_id: incremental ? previousResponseId : undefined, + input: nextInput, + }), + ) + } +} diff --git a/src/api/providers/__tests__/CodexWebSocketTransport.spec.ts b/src/api/providers/__tests__/CodexWebSocketTransport.spec.ts new file mode 100644 index 0000000000..96cbda1fd2 --- /dev/null +++ b/src/api/providers/__tests__/CodexWebSocketTransport.spec.ts @@ -0,0 +1,470 @@ +import { once } from "node:events" +import { WebSocketServer, type WebSocket } from "ws" +import nock from "nock" +import { CodexWebSocketTransport, CodexWebSocketUnavailableError } from "../CodexWebSocketTransport" +import { CodexWebSocketTransportScope } from "../codex-websocket/scopes/CodexWebSocketTransportScope" +import { collectStream } from "../../../test-utils/stream" + +type Request = Record + +describe("CodexWebSocketTransport", () => { + let server: WebSocketServer + let transport: CodexWebSocketTransport + let transportScope: CodexWebSocketTransportScope + let requests: Request[] + let connections: number + let reply: (request: Request, socket: WebSocket) => void + const options = () => ({ + headers: { Authorization: "Bearer test-token" }, + signal: new AbortController().signal, + timeoutMs: 2_000, + }) + const body = (input: unknown[] = [{ role: "user", content: [{ type: "input_text", text: "Hello" }] }]) => ({ + model: "test-model", + stream: true, + store: false, + input, + }) + const complete = (socket: WebSocket, output: unknown[] = []) => + socket.send( + JSON.stringify({ + type: "response.completed", + response: { id: `resp_${requests.length}`, output, usage: { input_tokens: 1, output_tokens: 2 } }, + }), + ) + + beforeEach(async () => { + nock.enableNetConnect(/^127\.0\.0\.1:/) + requests = [] + connections = 0 + reply = (_request, socket) => complete(socket) + server = new WebSocketServer({ host: "127.0.0.1", port: 0 }) + await once(server, "listening") + const address = server.address() + if (typeof address === "string" || address === null) throw new Error("Missing test server address") + transportScope = new CodexWebSocketTransportScope(`ws://127.0.0.1:${address.port}/responses`) + transportScope.init() + transport = transportScope.transport + server.on("connection", (socket) => { + connections++ + socket.on("message", (data) => { + const request: Request = JSON.parse(String(data)) + requests.push(request) + reply(request, socket) + }) + }) + }) + + afterEach(async () => { + vi.restoreAllMocks() + await transportScope?.dispose() + nock.disableNetConnect() + for (const socket of server.clients) socket.terminate() + await new Promise((resolve, reject) => server.close((error) => (error ? reject(error) : resolve()))) + }) + + it("sends the Codex v2 handshake and Responses Lite metadata", async () => { + let headers: Record = {} + server.on("headers", (_headers, request) => { + headers = request.headers + }) + const events = await collectStream( + transport.stream(body(), { + ...options(), + headers: { Authorization: "Bearer test-token", "x-openai-internal-codex-responses-lite": "true" }, + }), + ) + expect(headers).toMatchObject({ + authorization: "Bearer test-token", + "openai-beta": "responses_websockets=2026-02-06", + }) + expect(requests[0]).toMatchObject({ + type: "response.create", + stream: true, + store: false, + client_metadata: { + ws_request_header_x_openai_internal_codex_responses_lite: "true", + }, + }) + expect(events[0]).toMatchObject({ type: "response.completed", response: { usage: { output_tokens: 2 } } }) + }) + + it("preserves caller metadata while adding Responses Lite transport metadata", async () => { + await collectStream( + transport.stream( + { ...body(), client_metadata: { trace_id: "trace-1", feature: "test" } }, + { + ...options(), + headers: { Authorization: "Bearer test-token", "x-openai-internal-codex-responses-lite": "true" }, + }, + ), + ) + expect(requests[0].client_metadata).toEqual({ + trace_id: "trace-1", + feature: "test", + ws_request_header_x_openai_internal_codex_responses_lite: "true", + }) + }) + + it("reuses a socket and sends only new tool output after reconstructed reasoning and tool calls", async () => { + const output = [ + { + type: "reasoning", + id: "rs_1", + summary: [{ type: "summary_text", text: "Thinking" }], + encrypted_content: "encrypted", + }, + { + type: "message", + id: "msg_1", + status: "completed", + role: "assistant", + content: [{ type: "output_text", text: "Reading", annotations: [] }], + }, + { + type: "function_call", + id: "fc_1", + status: "completed", + call_id: "call_1", + name: "read_file", + arguments: '{"path":"a","start":1}', + }, + ] + reply = (_request, socket) => complete(socket, requests.length === 1 ? output : []) + await collectStream(transport.stream(body(), options())) + const toolOutput = { type: "function_call_output", call_id: "call_1", output: "file content" } + await collectStream( + transport.stream( + body([ + ...body().input, + { type: "reasoning", id: "rs_1", encrypted_content: "encrypted" }, + { role: "assistant", content: [{ type: "output_text", text: "Reading" }] }, + { + type: "function_call", + call_id: "call_1", + name: "read_file", + arguments: '{"start":1,"path":"a"}', + }, + toolOutput, + ]), + options(), + ), + ) + expect(connections).toBe(1) + expect(requests[1]).toMatchObject({ previous_response_id: "resp_1", input: [toolOutput] }) + }) + + it.each(["history", "settings", "explicit reset"])("replays full context after %s changes", async (change) => { + const log = vi.spyOn(console, "info") + await collectStream(transport.stream(body(), options())) + if (change === "explicit reset") transport.resetContinuation() + const next = body(change === "history" ? [{ role: "user", content: "Compacted history" }] : body().input) + if (change === "settings") next.model = "another-model" + await collectStream(transport.stream(next, options())) + expect(requests[1].previous_response_id).toBeUndefined() + expect(requests[1].input).toEqual(next.input) + const reason = + change === "history" + ? "history changed" + : change === "settings" + ? "request settings changed" + : "no cached response" + expect(log).toHaveBeenLastCalledWith(expect.stringContaining(`reason: ${reason}`)) + }) + + it("recovers a missing previous response once with the full context", async () => { + const log = vi.spyOn(console, "info") + reply = (_request, socket) => { + if (requests.length === 2) + socket.send(JSON.stringify({ type: "error", error: { code: "previous_response_not_found" } })) + else complete(socket) + } + await collectStream(transport.stream(body(), options())) + const next = body([...body().input, { role: "user", content: "Next" }]) + await collectStream(transport.stream(next, options())) + expect(requests[1].previous_response_id).toBe("resp_1") + expect(requests[2]).toMatchObject({ input: next.input }) + expect(requests[2].previous_response_id).toBeUndefined() + expect(log).toHaveBeenLastCalledWith(expect.stringContaining("reason: server cache miss")) + }) + + it.each(["function_call", "message"])("continues after unencrypted reasoning followed by %s", async (type) => { + const call = { type: "function_call", call_id: "call_1", name: "list_files", arguments: "{}" } + const output = + type === "function_call" ? call : { role: "assistant", content: [{ type: "output_text", text: "Ready" }] } + reply = (_request, socket) => + complete( + socket, + requests.length === 1 + ? [{ type: "reasoning", id: "rs_1", summary: [{ type: "summary_text", text: "Thinking" }] }, output] + : [], + ) + await collectStream(transport.stream(body(), options())) + const next = + type === "function_call" + ? { type: "function_call_output", call_id: "call_1", output: "files" } + : { role: "user", content: "Continue" } + const history = [...body().input, output, next] + await collectStream(transport.stream(body(history), options())) + expect(requests[1]).toMatchObject({ previous_response_id: "resp_1", input: [next] }) + const followup = { role: "user", content: "Next step" } + await collectStream(transport.stream(body([...history, followup]), options())) + expect(requests[2]).toMatchObject({ previous_response_id: "resp_2", input: [followup] }) + expect(connections).toBe(1) + }) + + it("replays full context when encrypted reasoning is omitted", async () => { + const call = { type: "function_call", call_id: "call_1", name: "list_files", arguments: "{}" } + reply = (_request, socket) => + complete( + socket, + requests.length === 1 ? [{ type: "reasoning", id: "rs_1", encrypted_content: "encrypted" }, call] : [], + ) + await collectStream(transport.stream(body(), options())) + const next = body([...body().input, call, { type: "function_call_output", call_id: "call_1", output: "files" }]) + await collectStream(transport.stream(next, options())) + expect(requests[1].previous_response_id).toBeUndefined() + expect(requests[1].input).toEqual(next.input) + }) + + it("identifies changed message content without logging the content", async () => { + const log = vi.spyOn(console, "info") + const message = { role: "assistant", content: [{ type: "output_text", text: "private text" }] } + reply = (_request, socket) => complete(socket, requests.length === 1 ? [message] : []) + await collectStream(transport.stream(body(), options())) + await collectStream( + transport.stream( + body([ + ...body().input, + { ...message, content: [{ type: "output_text", text: "changed private text" }] }, + ]), + options(), + ), + ) + expect(log).toHaveBeenLastCalledWith(expect.stringContaining("message -> message; fields: content")) + expect(log.mock.calls.flat().join(" ")).not.toContain("private text") + }) + + it("reconnects with full context when credentials change", async () => { + await collectStream(transport.stream(body(), options())) + await collectStream( + transport.stream(body(), { ...options(), headers: { Authorization: "Bearer refreshed-token" } }), + ) + expect(connections).toBe(2) + expect(requests[1].previous_response_id).toBeUndefined() + }) + + it("rotates a connection before the server's one-hour limit", async () => { + await collectStream(transport.stream(body(), options())) + const later = Date.now() + 55 * 60_000 + vi.spyOn(Date, "now").mockReturnValue(later) + await collectStream(transport.stream(body(), options())) + expect(connections).toBe(2) + expect(requests[1].previous_response_id).toBeUndefined() + }) + + it("never labels a disconnect after sending as safe to replay over HTTP", async () => { + reply = (_request, socket) => socket.close() + const error = await collectStream(transport.stream(body(), options())).catch((error: unknown) => error) + expect(error).toBeInstanceOf(Error) + expect(error).not.toBeInstanceOf(CodexWebSocketUnavailableError) + expect(requests).toHaveLength(1) + }) + + it("permits HTTP fallback after a rejected upgrade and stops retrying the socket", async () => { + server.options.verifyClient = () => false + await expect(collectStream(transport.stream(body(), options()))).rejects.toBeInstanceOf( + CodexWebSocketUnavailableError, + ) + await expect(collectStream(transport.stream(body(), options()))).rejects.toBeInstanceOf( + CodexWebSocketUnavailableError, + ) + expect(connections).toBe(0) + expect(requests).toHaveLength(0) + }) + + it("attempts WebSocket again after credentials change following a failed upgrade", async () => { + server.options.verifyClient = () => false + await expect(collectStream(transport.stream(body(), options()))).rejects.toBeInstanceOf( + CodexWebSocketUnavailableError, + ) + server.options.verifyClient = () => true + await collectStream( + transport.stream(body(), { ...options(), headers: { Authorization: "Bearer refreshed-token" } }), + ) + expect(connections).toBe(1) + expect(requests[0].previous_response_id).toBeUndefined() + }) + + it("retries a transient upgrade failure after the cooldown", async () => { + server.options.verifyClient = () => false + await expect(collectStream(transport.stream(body(), options()))).rejects.toBeInstanceOf( + CodexWebSocketUnavailableError, + ) + server.options.verifyClient = () => true + vi.spyOn(Date, "now").mockReturnValue(Date.now() + 60_000) + await collectStream(transport.stream(body(), options())) + expect(connections).toBe(1) + expect(requests[0].previous_response_id).toBeUndefined() + }) + + it("reconnects with full history after the server closes an idle socket", async () => { + await collectStream(transport.stream(body(), options())) + const socket = [...server.clients][0] + const closed = once(socket, "close") + socket.close() + await closed + await collectStream(transport.stream(body(), options())) + expect(connections).toBe(2) + expect(requests[1].previous_response_id).toBeUndefined() + }) + + it("never replays a cache miss after any response event has arrived", async () => { + await collectStream(transport.stream(body(), options())) + reply = (_request, socket) => { + socket.send(JSON.stringify({ type: "response.output_text.delta", delta: "Accepted" })) + socket.send(JSON.stringify({ type: "error", error: { code: "previous_response_not_found" } })) + } + await expect(collectStream(transport.stream(body(), options()))).rejects.toThrow("previous_response_not_found") + expect(requests).toHaveLength(2) + }) + + it.each([ + { type: "response.failed", response: { error: { code: "server_error", message: "Service unavailable" } } }, + { type: "response.incomplete", response: { incomplete_details: { reason: "max_output_tokens" } } }, + ])("preserves server failure details for $type", async (event) => { + reply = (_request, socket) => socket.send(JSON.stringify(event)) + await expect(collectStream(transport.stream(body(), options()))).rejects.toThrow( + event.type === "response.failed" ? "Service unavailable" : "max_output_tokens", + ) + }) + + it.each(["response.failed", "response.incomplete"])("reports %s as a failed request", async (type) => { + reply = (_request, socket) => socket.send(JSON.stringify({ type, response: { id: "failed" } })) + await expect(collectStream(transport.stream(body(), options()))).rejects.toThrow(type) + }) + + it("does not loop when full-context recovery also fails", async () => { + await collectStream(transport.stream(body(), options())) + reply = (_request, socket) => + socket.send(JSON.stringify({ type: "error", error: { code: "previous_response_not_found" } })) + await expect(collectStream(transport.stream(body(), options()))).rejects.toThrow("previous_response_not_found") + expect(requests).toHaveLength(3) + }) + + it("closes the socket on Stop and starts fresh on the next request", async () => { + const controller = new AbortController() + reply = (_request, socket) => { + socket.send(JSON.stringify({ type: "response.output_text.delta", delta: "Partial" })) + } + const stream = transport.stream(body(), { ...options(), signal: controller.signal }) + expect((await stream.next()).value).toMatchObject({ delta: "Partial" }) + controller.abort() + await expect(stream.next()).rejects.toThrow() + reply = (_request, socket) => complete(socket) + await collectStream(transport.stream(body(), options())) + expect(connections).toBe(2) + expect(requests[1].previous_response_id).toBeUndefined() + }) + + it("disposes an active stream and reconnects with full context on the next request", async () => { + reply = (_request, socket) => + socket.send(JSON.stringify({ type: "response.output_text.delta", delta: "Partial" })) + const stream = transport.stream(body(), options()) + await stream.next() + await transportScope.dispose() + expect(() => transportScope.transport).toThrow("not initialized") + await expect(stream.next()).rejects.toThrow("closed before response completed") + transportScope.init() + transport = transportScope.transport + reply = (_request, socket) => complete(socket) + await collectStream(transport.stream(body(), options())) + expect(connections).toBe(2) + expect(requests[1].previous_response_id).toBeUndefined() + }) + + it("disposes a completed idle connection and reconnects without stale continuation", async () => { + await collectStream(transport.stream(body(), options())) + await transport.dispose() + await transport.dispose() + await collectStream(transport.stream(body(), options())) + expect(connections).toBe(2) + expect(requests[1].previous_response_id).toBeUndefined() + }) + + it("disposes request deadlines while the consumer is paused on a partial response", async () => { + vi.useFakeTimers({ toFake: ["setTimeout", "clearTimeout"] }) + reply = (_request, socket) => + socket.send(JSON.stringify({ type: "response.output_text.delta", delta: "Partial" })) + const stream = transport.stream(body(), options()) + try { + await stream.next() + expect(vi.getTimerCount()).toBe(1) + await transport.dispose() + expect(vi.getTimerCount()).toBe(0) + await expect(stream.next()).rejects.toThrow("closed before response completed") + } finally { + await stream.return(undefined) + vi.useRealTimers() + } + }) + + it("bounds a silent response with a timeout", async () => { + reply = () => {} + await expect(collectStream(transport.stream(body(), { ...options(), timeoutMs: 100 }))).rejects.toThrow() + }) + + it.each(["null", "[]", "{}", '{"type":42}', "invalid JSON"])( + "rejects malformed response %s and clears the connection", + async (payload) => { + reply = (_request, socket) => socket.send(payload) + const error = await collectStream(transport.stream(body(), options())).catch((error: unknown) => error) + expect(error).toBeInstanceOf(Error) + expect(error).not.toBeInstanceOf(CodexWebSocketUnavailableError) + reply = (_request, socket) => complete(socket) + await collectStream(transport.stream(body(), options())) + expect(connections).toBe(2) + expect(requests[1].previous_response_id).toBeUndefined() + }, + ) + + it("closes a completed idle connection after two minutes", async () => { + vi.useFakeTimers({ toFake: ["setTimeout", "clearTimeout"] }) + try { + await collectStream(transport.stream(body(), options())) + const closed = once([...server.clients][0], "close") + await vi.advanceTimersByTimeAsync(120_000) + await closed + await collectStream(transport.stream(body(), options())) + expect(connections).toBe(2) + expect(requests[1].previous_response_id).toBeUndefined() + } finally { + vi.useRealTimers() + } + }) + + it("clears the request deadline before yielding completion to a paused consumer", async () => { + vi.useFakeTimers({ toFake: ["setTimeout", "clearTimeout"] }) + const stream = transport.stream(body(), { ...options(), timeoutMs: 2_000 }) + try { + expect((await stream.next()).value).toMatchObject({ type: "response.completed" }) + await vi.advanceTimersByTimeAsync(10_000) + expect((await stream.next()).done).toBe(true) + await collectStream(transport.stream(body(), options())) + expect(connections).toBe(1) + expect(requests[1].previous_response_id).toBe("resp_1") + } finally { + await stream.return(undefined) + vi.useRealTimers() + } + }) + + it("rejects overlapping requests without disturbing the active stream", async () => { + reply = (_request, socket) => + socket.send(JSON.stringify({ type: "response.output_text.delta", delta: "Partial" })) + const stream = transport.stream(body(), options()) + await stream.next() + await expect(collectStream(transport.stream(body(), options()))).rejects.toThrow("Concurrent requests") + await stream.return(undefined) + }) +}) diff --git a/src/api/providers/codex-websocket/data/local/CodexWebSocketResponseLocalDataSource.ts b/src/api/providers/codex-websocket/data/local/CodexWebSocketResponseLocalDataSource.ts new file mode 100644 index 0000000000..e7801a17a5 --- /dev/null +++ b/src/api/providers/codex-websocket/data/local/CodexWebSocketResponseLocalDataSource.ts @@ -0,0 +1,18 @@ +import type { CachedCodexResponse } from "../../models/CachedCodexResponse" + +/** Task-local, in-memory continuation cache. Stores snapshots only, never response content. */ +export class CodexWebSocketResponseLocalDataSource { + private response?: CachedCodexResponse + + read(): CachedCodexResponse | undefined { + return this.response + } + + write(response: CachedCodexResponse | undefined): void { + this.response = response + } + + clear(): void { + this.response = undefined + } +} diff --git a/src/api/providers/codex-websocket/data/local/__tests__/CodexWebSocketResponseLocalDataSource.spec.ts b/src/api/providers/codex-websocket/data/local/__tests__/CodexWebSocketResponseLocalDataSource.spec.ts new file mode 100644 index 0000000000..3bfdb7e696 --- /dev/null +++ b/src/api/providers/codex-websocket/data/local/__tests__/CodexWebSocketResponseLocalDataSource.spec.ts @@ -0,0 +1,42 @@ +import { CodexWebSocketResponseLocalDataSource } from "../CodexWebSocketResponseLocalDataSource" +import type { CachedCodexResponse } from "../../../models/CachedCodexResponse" + +describe("CodexWebSocketResponseLocalDataSource", () => { + const response = (): CachedCodexResponse => ({ id: "resp_1", settings: "settings-hash", input: [] }) + + it("starts empty and retains the latest response snapshot", () => { + const store = new CodexWebSocketResponseLocalDataSource() + expect(store.read()).toBeUndefined() + const cached = response() + store.write(cached) + expect(store.read()).toBe(cached) + const next = { ...response(), id: "resp_2" } + store.write(next) + expect(store.read()).toBe(next) + }) + + it("clears cached state idempotently", () => { + const store = new CodexWebSocketResponseLocalDataSource() + store.write(response()) + store.clear() + store.clear() + expect(store.read()).toBeUndefined() + }) + + it("replaces a cached response with an absent response", () => { + const store = new CodexWebSocketResponseLocalDataSource() + store.write(response()) + store.write(undefined) + expect(store.read()).toBeUndefined() + }) + + it("isolates state between task-local stores", () => { + const first = new CodexWebSocketResponseLocalDataSource() + const second = new CodexWebSocketResponseLocalDataSource() + first.write(response()) + expect(second.read()).toBeUndefined() + second.write({ ...response(), id: "resp_2" }) + first.clear() + expect(second.read()?.id).toBe("resp_2") + }) +}) diff --git a/src/api/providers/codex-websocket/data/remote/CodexWebSocketSocketRemoteDataSource.ts b/src/api/providers/codex-websocket/data/remote/CodexWebSocketSocketRemoteDataSource.ts new file mode 100644 index 0000000000..d1ecf368be --- /dev/null +++ b/src/api/providers/codex-websocket/data/remote/CodexWebSocketSocketRemoteDataSource.ts @@ -0,0 +1,83 @@ +import { once } from "node:events" +import WebSocket from "ws" + +import type { CodexWebSocketOptions } from "../../models/protocol" + +const HANDSHAKE_TIMEOUT_MS = 10_000 + +/** Owns one socket and its subscriptions, from initialization through terminal close. */ +export class CodexWebSocketSocketRemoteDataSource { + private _socket?: WebSocket + private disposeSocket?: () => void + + constructor( + private readonly url: string, + private readonly onError: () => void, + private readonly onClose: () => void, + ) {} + + get socket(): WebSocket { + if (!this._socket) throw new Error("Codex WebSocket connection scope is not initialized") + return this._socket + } + + async init(options: CodexWebSocketOptions): Promise { + if (this._socket) throw new Error("Codex WebSocket connection scope is already initialized") + options.signal.throwIfAborted() + const socket = this.createSocket(options) + this._socket = socket + this.disposeSocket = this.registerHandlers(socket) + try { + await this.waitForOpen(socket, options.signal) + } catch (error) { + if (this._socket === socket) this.dispose() + throw error + } + } + + dispose(): void { + const disposeSocket = this.disposeSocket + this._socket = undefined + this.disposeSocket = undefined + disposeSocket?.() + } + + private createSocket(options: CodexWebSocketOptions): WebSocket { + return new WebSocket(this.url, { + headers: { ...options.headers, "OpenAI-Beta": "responses_websockets=2026-02-06" }, + handshakeTimeout: Math.min(options.timeoutMs, HANDSHAKE_TIMEOUT_MS), + }) + } + + private registerHandlers(socket: WebSocket): () => void { + const onError = () => { + if (this._socket === socket) this.onError() + } + const onClose = () => { + if (this._socket === socket) this.onClose() + } + socket.on("error", onError) + socket.on("close", onClose) + return () => { + socket.off("close", onClose) + this.terminateSocket(socket, onError) + } + } + + private terminateSocket(socket: WebSocket, onError: () => void): void { + if (socket.readyState === WebSocket.CLOSED) { + socket.off("error", onError) + return + } + // An upgrade terminated in flight can emit an asynchronous error. Keep its inert + // error handler until terminal close, then remove it as well. + socket.once("close", () => socket.off("error", onError)) + socket.terminate() + } + + private async waitForOpen(socket: WebSocket, signal: AbortSignal): Promise { + await once(socket, "open", { signal }) + signal.throwIfAborted() + if (this._socket !== socket) throw new Error("Codex WebSocket connection scope was disposed") + } +} diff --git a/src/api/providers/codex-websocket/errors/CodexWebSocketUnavailableError.ts b/src/api/providers/codex-websocket/errors/CodexWebSocketUnavailableError.ts new file mode 100644 index 0000000000..5c1a05b4f4 --- /dev/null +++ b/src/api/providers/codex-websocket/errors/CodexWebSocketUnavailableError.ts @@ -0,0 +1,2 @@ +/** Only a failed upgrade permits automatic replay over HTTP. */ +export class CodexWebSocketUnavailableError extends Error {} diff --git a/src/api/providers/codex-websocket/managers/CodexWebSocketConnectionManager.ts b/src/api/providers/codex-websocket/managers/CodexWebSocketConnectionManager.ts new file mode 100644 index 0000000000..15228e8945 --- /dev/null +++ b/src/api/providers/codex-websocket/managers/CodexWebSocketConnectionManager.ts @@ -0,0 +1,119 @@ +import WebSocket from "ws" + +import { CodexWebSocketConnectionScope } from "../scopes/CodexWebSocketConnectionScope" +import { CodexWebSocketUnavailableError } from "../errors/CodexWebSocketUnavailableError" +import type { CodexWebSocketConnectionStateHolder } from "../state-holders/CodexWebSocketConnectionStateHolder" +import type { CodexWebSocketOptions } from "../models/protocol" +import { fingerprint } from "../utils/protocol" + +const MAX_CONNECTION_AGE_MS = 55 * 60_000 +const IDLE_CONNECTION_TIMEOUT_MS = 120_000 +const UPGRADE_RETRY_DELAY_MS = 60_000 + +type ConnectionAttempt = ReturnType + +/** Owns a single authenticated connection and bounds its lifetime between requests. */ +export class CodexWebSocketConnectionManager { + constructor( + private readonly url: string, + private readonly resetContinuation: () => void, + private readonly stateHolder: CodexWebSocketConnectionStateHolder, + ) {} + + private get state() { + return this.stateHolder.state + } + + async acquire(options: CodexWebSocketOptions): Promise { + options.signal.throwIfAborted() + const key = fingerprint(options.headers) + this.assertAcquisitionAllowed(key) + const socket = this.tryReuseConnection(key) + if (socket) return socket + return this.openConnection(key, options) + } + + release(): void { + const state = this.state + if (state.status !== "active" && state.status !== "idle") return + if (state.status === "idle") clearTimeout(state.idleTimer) + const idleTimer = setTimeout(() => { + if (this.stateHolder.isCurrent(idleState)) this.dispose() + }, IDLE_CONNECTION_TIMEOUT_MS) + const idleState = this.stateHolder.idle(state, idleTimer) + idleTimer.unref() + } + + dispose(): void { + // Request cleanup must not erase the safe-upgrade-failure cooldown. + const state = this.stateHolder.disconnect() + if (state.status === "idle") clearTimeout(state.idleTimer) + this.resetContinuation() + if ("scope" in state) state.scope.dispose() + } + + private assertAcquisitionAllowed(key: string): void { + const state = this.state + if (state.status === "connecting") { + throw new Error("Concurrent Codex WebSocket connection attempts are not supported") + } + if (state.status === "unavailable" && state.key === key && Date.now() < state.retryAt) { + throw new CodexWebSocketUnavailableError("Codex WebSocket unavailable; using HTTP") + } + } + + private tryReuseConnection(key: string): WebSocket | undefined { + const state = this.state + if (state.status !== "active" && state.status !== "idle") return undefined + if (state.key !== key || Date.now() - state.connectedAt >= MAX_CONNECTION_AGE_MS) return undefined + const socket = state.scope.socketRemoteDataSource.socket + if (socket.readyState !== WebSocket.OPEN) return undefined + if (state.status === "idle") clearTimeout(state.idleTimer) + this.stateHolder.activate(state) + return socket + } + + private async openConnection(key: string, options: CodexWebSocketOptions): Promise { + this.dispose() + const scope = this.createConnectionScope() + const attempt = this.stateHolder.beginConnection(key, scope) + try { + await scope.init(options) + options.signal.throwIfAborted() + if (!this.stateHolder.isCurrent(attempt)) throw new Error("Codex WebSocket connection attempt was disposed") + } catch (error) { + this.handleUpgradeFailure(attempt, options.signal, error) + } + this.stateHolder.connect(attempt, Date.now()) + console.info("[Codex WebSocket] Connected") + return scope.socketRemoteDataSource.socket + } + + private createConnectionScope(): CodexWebSocketConnectionScope { + const scope = new CodexWebSocketConnectionScope( + this.url, + () => { + if (this.stateHolder.ownsScope(scope)) this.resetContinuation() + }, + () => this.handleScopeClose(scope), + ) + return scope + } + + private handleScopeClose(scope: CodexWebSocketConnectionScope): void { + if (!this.stateHolder.ownsScope(scope)) return + // A connecting attempt publishes its own outcome after init settles. + if (this.state.status === "connecting") scope.dispose() + else this.dispose() + } + + private handleUpgradeFailure(attempt: ConnectionAttempt, signal: AbortSignal, error: unknown): never { + const current = this.stateHolder.isCurrent(attempt) + if (current) this.dispose() + signal.throwIfAborted() + if (!current) throw error + this.stateHolder.markUnavailable(attempt.key, Date.now() + UPGRADE_RETRY_DELAY_MS) + console.warn("[Codex WebSocket] Upgrade failed; falling back to HTTP") + throw new CodexWebSocketUnavailableError("Codex WebSocket upgrade failed", { cause: error }) + } +} diff --git a/src/api/providers/codex-websocket/managers/CodexWebSocketRequestManager.ts b/src/api/providers/codex-websocket/managers/CodexWebSocketRequestManager.ts new file mode 100644 index 0000000000..f854109bbd --- /dev/null +++ b/src/api/providers/codex-websocket/managers/CodexWebSocketRequestManager.ts @@ -0,0 +1,96 @@ +import { on } from "node:events" +import type WebSocket from "ws" + +import type { CodexWebSocketConnectionManager } from "./CodexWebSocketConnectionManager" +import type { CodexWebSocketRequestStateModel } from "../models/CodexWebSocketRequestStateModel" +import type { CodexWebSocketRequestStateHolder } from "../state-holders/CodexWebSocketRequestStateHolder" +import type { CodexWebSocketOptions } from "../models/protocol" + +type RequestState = CodexWebSocketRequestStateModel, NodeJS.Timeout> + +/** Owns cancellation, the inactivity timer, and response subscriptions for a single request. */ +export class CodexWebSocketRequestManager { + constructor( + private readonly options: CodexWebSocketOptions, + private readonly onAbort: () => void, + private readonly stateHolder: CodexWebSocketRequestStateHolder, + ) {} + + private get state(): RequestState { + return this.stateHolder.state + } + + get signal(): AbortSignal { + return "signal" in this.state ? this.state.signal : this.options.signal + } + + get socket(): WebSocket { + if (this.state.status !== "active") throw new Error("Codex WebSocket request scope is not initialized") + return this.state.socket + } + + get events(): AsyncIterableIterator { + if (this.state.status !== "active") throw new Error("Codex WebSocket request scope is not initialized") + return this.state.events + } + + async init(connection: CodexWebSocketConnectionManager): Promise { + if (this.state.status === "active" || this.state.status === "initializing") { + throw new Error("Codex WebSocket request scope is already initialized") + } + this.options.signal.throwIfAborted() + const controller = new AbortController() + const signal = AbortSignal.any([this.options.signal, controller.signal]) + const initializing = this.stateHolder.beginInitialization(controller, signal) + initializing.signal.addEventListener("abort", this.onAbort, { once: true }) + try { + const socket = await connection.acquire({ ...this.options, signal: initializing.signal }) + initializing.signal.throwIfAborted() + if (!this.stateHolder.isCurrent(initializing)) throw new Error("Codex WebSocket request scope was disposed") + // Register before sending so a fast reply cannot be lost. + const events = on(socket, "message", { signal: initializing.signal, close: ["close"] }) + this.stateHolder.activate(initializing, socket, events) + } catch (error) { + if (this.stateHolder.isCurrent(initializing)) { + this.stateHolder.dispose() + await this.disposeResources(initializing) + } + initializing.signal.throwIfAborted() + throw error + } + } + + refreshTimeout(): void { + const state = this.state + if (state.status !== "active") throw new Error("Codex WebSocket request scope is not initialized") + clearTimeout(state.timeout) + const timeout = setTimeout(() => { + if (this.stateHolder.isCurrent(timedState)) + state.controller.abort(new Error("Codex WebSocket stream timed out")) + }, this.options.timeoutMs) + const timedState = this.stateHolder.setDeadline(state, timeout) + } + + clearTimeout(): void { + const state = this.state + if (state.status !== "active") return + clearTimeout(state.timeout) + this.stateHolder.clearDeadline(state) + } + + async dispose(): Promise { + const state = this.stateHolder.dispose() + if (state.status !== "initializing" && state.status !== "active") return + if (state.status === "initializing") + state.controller.abort(new Error("Codex WebSocket request scope was disposed")) + await this.disposeResources(state) + } + + private async disposeResources(state: Extract): Promise { + state.signal.removeEventListener("abort", this.onAbort) + if (state.status === "active") { + clearTimeout(state.timeout) + await state.events.return?.() + } + } +} diff --git a/src/api/providers/codex-websocket/managers/CodexWebSocketResponseManager.ts b/src/api/providers/codex-websocket/managers/CodexWebSocketResponseManager.ts new file mode 100644 index 0000000000..a7c5b8b880 --- /dev/null +++ b/src/api/providers/codex-websocket/managers/CodexWebSocketResponseManager.ts @@ -0,0 +1,64 @@ +import { CodexWebSocketContinuationRepository } from "../repositories/CodexWebSocketContinuationRepository" +import type { PreparedCodexRequest } from "../models/PreparedCodexRequest" +import type { CodexWebSocketResponseStateHolder } from "../state-holders/CodexWebSocketResponseStateHolder" +import type { CodexResponseEvent } from "../models/protocol" +import { asJsonObject, responseError } from "../utils/protocol" + +/** Tracks response completion and the one safe cache-miss recovery, independently of socket IO. */ +export class CodexWebSocketResponseManager { + constructor( + private readonly continuation: CodexWebSocketContinuationRepository, + private readonly stateHolder: CodexWebSocketResponseStateHolder, + ) {} + + init(body: unknown): void { + if (this.stateHolder.state.status !== "uninitialized") { + throw new Error("Codex WebSocket response manager is already initialized") + } + this.stateHolder.initialize(this.continuation.prepare(body)) + } + + get preparedRequest(): PreparedCodexRequest { + return this.stateHolder.preparedRequest + } + + get completed(): boolean { + return this.stateHolder.state.status === "completed" + } + + accept(event: CodexResponseEvent): "emit" | "retry" { + this.stateHolder.assertAcceptingEvents() + switch (event.type) { + case "error": + return this.handleError(event) + case "response.failed": + case "response.incomplete": + throw responseError(event) + case "response.output_item.done": + this.stateHolder.appendOutputItem(event.item) + break + case "response.completed": + case "response.done": + this.completeResponse(event.response) + break + default: + this.stateHolder.startStreaming() + } + return "emit" + } + + private completeResponse(response: unknown): void { + const prepared = this.preparedRequest + this.continuation.record(prepared, asJsonObject(response), [...this.stateHolder.streamedOutput]) + this.stateHolder.complete() + } + + private handleError(event: CodexResponseEvent): "retry" { + const error = asJsonObject(event.error) + const canRecover = error.code === "previous_response_not_found" && this.stateHolder.canRecoverCacheMiss + if (!canRecover) throw responseError(event) + this.stateHolder.beginRecovery() + this.continuation.reset() + return "retry" + } +} diff --git a/src/api/providers/codex-websocket/managers/__tests__/CodexWebSocketConnectionManager.spec.ts b/src/api/providers/codex-websocket/managers/__tests__/CodexWebSocketConnectionManager.spec.ts new file mode 100644 index 0000000000..eb6b653bb2 --- /dev/null +++ b/src/api/providers/codex-websocket/managers/__tests__/CodexWebSocketConnectionManager.spec.ts @@ -0,0 +1,206 @@ +import WebSocket from "ws" + +import { CodexWebSocketConnectionManager } from "../CodexWebSocketConnectionManager" +import { CodexWebSocketUnavailableError } from "../../errors/CodexWebSocketUnavailableError" +import { CodexWebSocketConnectionStateHolder } from "../../state-holders/CodexWebSocketConnectionStateHolder" +import type { CodexWebSocketOptions } from "../../models/protocol" + +interface TestScope { + socketRemoteDataSource: { socket: { readyState: number } } + init: ReturnType Promise>> + dispose: ReturnType void>> + open(): void + fail(error: Error): void + onError(): void + onClose(): void +} + +const { scopes } = vi.hoisted(() => ({ scopes: [] as TestScope[] })) + +vi.mock("../../scopes/CodexWebSocketConnectionScope", () => { + class TestConnectionScope implements TestScope { + readonly socketRemoteDataSource = { socket: { readyState: WebSocket.CONNECTING as number } } + private resolveInit?: () => void + private rejectInit?: (error: Error) => void + + constructor( + _url: string, + readonly onError: () => void, + readonly onClose: () => void, + ) { + scopes.push(this) + } + + init = vi.fn( + (_options: CodexWebSocketOptions) => + new Promise((resolve, reject) => { + this.resolveInit = resolve + this.rejectInit = reject + }), + ) + dispose = vi.fn(() => { + this.socketRemoteDataSource.socket.readyState = WebSocket.CLOSED + }) + open(): void { + this.socketRemoteDataSource.socket.readyState = WebSocket.OPEN + this.resolveInit?.() + } + fail(error: Error): void { + this.rejectInit?.(error) + } + } + return { CodexWebSocketConnectionScope: TestConnectionScope } +}) + +describe("CodexWebSocketConnectionManager state transitions", () => { + let manager: CodexWebSocketConnectionManager + let reset = vi.fn<() => void>() + const options = (): CodexWebSocketOptions => ({ + headers: { Authorization: "Bearer test-token" }, + signal: new AbortController().signal, + timeoutMs: 2_000, + }) + const lastScope = (): TestScope => { + const scope = scopes.at(-1) + if (!scope) throw new Error("No test connection scope was created") + return scope + } + + beforeEach(() => { + vi.useFakeTimers() + scopes.length = 0 + reset = vi.fn() + manager = new CodexWebSocketConnectionManager( + "ws://test/responses", + reset, + new CodexWebSocketConnectionStateHolder(), + ) + }) + + afterEach(() => { + manager.dispose() + vi.useRealTimers() + vi.restoreAllMocks() + }) + + it("moves from connecting to active to idle and back without retaining the idle timer", async () => { + const acquired = manager.acquire(options()) + const scope = lastScope() + manager.release() + expect(vi.getTimerCount()).toBe(0) + scope.open() + await acquired + manager.release() + expect(vi.getTimerCount()).toBe(1) + await expect(manager.acquire(options())).resolves.toBe(scope.socketRemoteDataSource.socket) + expect(vi.getTimerCount()).toBe(0) + await vi.advanceTimersByTimeAsync(120_000) + expect(scope.dispose).not.toHaveBeenCalled() + expect(scopes).toHaveLength(1) + }) + + it("rejects overlapping upgrades without disturbing the current attempt", async () => { + const acquired = manager.acquire(options()) + const scope = lastScope() + await expect(manager.acquire(options())).rejects.toThrow("Concurrent") + expect(scopes).toHaveLength(1) + expect(scope.dispose).not.toHaveBeenCalled() + scope.open() + await expect(acquired).resolves.toBe(scope.socketRemoteDataSource.socket) + }) + + it("replaces a socket that is no longer open even before its close callback runs", async () => { + const acquired = manager.acquire(options()) + const first = lastScope() + first.open() + await acquired + first.socketRemoteDataSource.socket.readyState = WebSocket.CLOSING + const next = manager.acquire(options()) + const current = lastScope() + current.open() + await expect(next).resolves.toBe(current.socketRemoteDataSource.socket) + expect(first.dispose).toHaveBeenCalledOnce() + expect(scopes).toHaveLength(2) + }) + + it("preserves upgrade cooldown through cleanup without retaining a socket or timer", async () => { + const failed = expect(manager.acquire(options())).rejects.toBeInstanceOf(CodexWebSocketUnavailableError) + const scope = lastScope() + scope.fail(new Error("Upgrade rejected")) + await failed + manager.dispose() + manager.release() + await expect(manager.acquire(options())).rejects.toBeInstanceOf(CodexWebSocketUnavailableError) + expect(scope.dispose).toHaveBeenCalledOnce() + expect(vi.getTimerCount()).toBe(0) + expect(scopes).toHaveLength(1) + await vi.advanceTimersByTimeAsync(60_000) + const retried = manager.acquire(options()) + lastScope().open() + await retried + expect(scopes).toHaveLength(2) + }) + + it("preserves cancellation rather than reporting HTTP fallback during cooldown", async () => { + const failed = expect(manager.acquire(options())).rejects.toBeInstanceOf(CodexWebSocketUnavailableError) + lastScope().fail(new Error("Upgrade rejected")) + await failed + const controller = new AbortController() + const reason = new Error("Stopped during cooldown") + controller.abort(reason) + await expect(manager.acquire({ ...options(), signal: controller.signal })).rejects.toBe(reason) + await expect(manager.acquire(options())).rejects.toBeInstanceOf(CodexWebSocketUnavailableError) + expect(scopes).toHaveLength(1) + }) + + it.each(["success", "failure"])( + "ignores late upgrade %s after disposal and replacement", + async (outcome: string) => { + const first = manager.acquire(options()) + const rejected = expect(first).rejects.toThrow() + const staleScope = lastScope() + manager.dispose() + const second = manager.acquire(options()) + const currentScope = lastScope() + currentScope.open() + await second + if (outcome === "success") staleScope.open() + else staleScope.fail(new Error("Late upgrade failure")) + await rejected + const resets = reset.mock.calls.length + staleScope.onError() + staleScope.onClose() + expect(reset).toHaveBeenCalledTimes(resets) + await expect(manager.acquire(options())).resolves.toBe(currentScope.socketRemoteDataSource.socket) + expect(currentScope.dispose).not.toHaveBeenCalled() + expect(scopes).toHaveLength(2) + }, + ) + + it("does not turn disposal of an in-flight upgrade into an HTTP fallback cooldown", async () => { + const failed = expect(manager.acquire(options())).rejects.not.toBeInstanceOf(CodexWebSocketUnavailableError) + const staleScope = lastScope() + manager.dispose() + staleScope.fail(new Error("Disposed upgrade")) + await failed + const next = manager.acquire(options()) + lastScope().open() + await next + expect(scopes).toHaveLength(2) + }) + + it("ignores a stale idle-timer callback after reacquiring the connection", async () => { + const acquired = manager.acquire(options()) + const scope = lastScope() + scope.open() + await acquired + const setTimeoutSpy = vi.spyOn(globalThis, "setTimeout") + manager.release() + const callback = setTimeoutSpy.mock.calls.at(-1)?.[0] + if (typeof callback !== "function") throw new Error("No idle timeout callback was registered") + await manager.acquire(options()) + callback() + expect(scope.dispose).not.toHaveBeenCalled() + await expect(manager.acquire(options())).resolves.toBe(scope.socketRemoteDataSource.socket) + }) +}) diff --git a/src/api/providers/codex-websocket/managers/__tests__/CodexWebSocketRequestManager.spec.ts b/src/api/providers/codex-websocket/managers/__tests__/CodexWebSocketRequestManager.spec.ts new file mode 100644 index 0000000000..1311bdfbab --- /dev/null +++ b/src/api/providers/codex-websocket/managers/__tests__/CodexWebSocketRequestManager.spec.ts @@ -0,0 +1,173 @@ +import WebSocket from "ws" + +import { CodexWebSocketConnectionManager } from "../CodexWebSocketConnectionManager" +import { CodexWebSocketRequestManager } from "../CodexWebSocketRequestManager" +import { CodexWebSocketUnavailableError } from "../../errors/CodexWebSocketUnavailableError" +import { CodexWebSocketConnectionStateHolder } from "../../state-holders/CodexWebSocketConnectionStateHolder" +import { CodexWebSocketRequestStateHolder } from "../../state-holders/CodexWebSocketRequestStateHolder" + +vi.mock("ws", async () => { + const { EventEmitter } = await import("node:events") + return { + default: class TestSocket extends EventEmitter { + constructor(_url: string) { + super() + } + }, + } +}) + +describe("CodexWebSocketRequestManager state transitions", () => { + let manager: CodexWebSocketRequestManager + let connection: CodexWebSocketConnectionManager + let socket: WebSocket + let stateHolder: CodexWebSocketRequestStateHolder + let onAbort = vi.fn<() => void>() + + beforeEach(() => { + vi.useFakeTimers() + onAbort = vi.fn() + socket = new WebSocket("ws://test/responses") + connection = new CodexWebSocketConnectionManager( + "ws://test/responses", + vi.fn(), + new CodexWebSocketConnectionStateHolder(), + ) + vi.spyOn(connection, "acquire").mockResolvedValue(socket) + stateHolder = new CodexWebSocketRequestStateHolder() + manager = new CodexWebSocketRequestManager( + { headers: {}, signal: new AbortController().signal, timeoutMs: 2_000 }, + onAbort, + stateHolder, + ) + }) + + afterEach(async () => { + await manager.dispose() + vi.useRealTimers() + vi.restoreAllMocks() + }) + + it("makes socket, subscriptions, and deadline unavailable together after disposal", async () => { + await manager.init(connection) + manager.refreshTimeout() + expect(manager.socket).toBe(socket) + expect(socket.listenerCount("message")).toBe(1) + expect(vi.getTimerCount()).toBe(1) + await manager.dispose() + await manager.dispose() + expect(() => manager.socket).toThrow("not initialized") + expect(() => manager.events).toThrow("not initialized") + expect(() => manager.refreshTimeout()).toThrow("not initialized") + expect(socket.eventNames()).toEqual([]) + expect(vi.getTimerCount()).toBe(0) + }) + + it("reads the caller signal and clears deadlines safely before initialization and after disposal", async () => { + expect(manager.signal).toBeInstanceOf(AbortSignal) + expect(manager.signal.aborted).toBe(false) + manager.clearTimeout() + await manager.init(connection) + await manager.dispose() + manager.clearTimeout() + expect(vi.getTimerCount()).toBe(0) + }) + + it("rejects repeated initialization without disturbing an active request", async () => { + await manager.init(connection) + const signal = manager.signal + await expect(manager.init(connection)).rejects.toThrow("already initialized") + expect(manager.signal).toBe(signal) + expect(manager.socket).toBe(socket) + expect(socket.listenerCount("message")).toBe(1) + expect(connection.acquire).toHaveBeenCalledOnce() + }) + + it("rejects overlapping initialization while acquisition is still pending", async () => { + let resolve: ((value: WebSocket) => void) | undefined + vi.mocked(connection.acquire).mockImplementationOnce( + () => + new Promise((complete) => { + resolve = complete + }), + ) + const initialized = manager.init(connection) + await expect(manager.init(connection)).rejects.toThrow("already initialized") + if (!resolve) throw new Error("Acquisition was not started") + resolve(socket) + await initialized + expect(manager.socket).toBe(socket) + expect(connection.acquire).toHaveBeenCalledOnce() + }) + + it("does not convert a failed upgrade into cancellation during cleanup", async () => { + const error = new CodexWebSocketUnavailableError("Upgrade rejected") + vi.mocked(connection.acquire).mockRejectedValueOnce(error) + await expect(manager.init(connection)).rejects.toBe(error) + expect(manager.signal.aborted).toBe(false) + expect(onAbort).not.toHaveBeenCalled() + await manager.init(connection) + expect(manager.socket).toBe(socket) + }) + + it("ignores stale request initialization after disposal and reinitialization", async () => { + let resolveFirst: ((value: WebSocket) => void) | undefined + vi.mocked(connection.acquire).mockImplementationOnce( + () => new Promise((resolve) => (resolveFirst = resolve)), + ) + const first = expect(manager.init(connection)).rejects.toThrow("scope was disposed") + await manager.dispose() + await manager.init(connection) + const currentSignal = manager.signal + const staleSocket = new WebSocket("ws://stale/responses") + if (!resolveFirst) throw new Error("First acquisition was not started") + resolveFirst(staleSocket) + await first + expect(manager.socket).toBe(socket) + expect(manager.signal).toBe(currentSignal) + expect(currentSignal.aborted).toBe(false) + expect(socket.listenerCount("message")).toBe(1) + expect(staleSocket.listenerCount("message")).toBe(0) + }) + + it("rejects acquisition completing after the injected holder has discarded the attempt", async () => { + const initialized = expect(manager.init(connection)).rejects.toThrow("scope was disposed") + const abandoned = stateHolder.dispose() + await initialized + expect(socket.eventNames()).toEqual([]) + expect(() => manager.socket).toThrow("not initialized") + // Discarding state is not IO cleanup: the owner must release its abandoned listener. + if (abandoned.status !== "initializing") throw new Error("Expected an initializing request") + abandoned.signal.removeEventListener("abort", onAbort) + expect(abandoned.signal.aborted).toBe(false) + }) + + it.each(["refresh", "clear", "dispose"])( + "ignores a stale deadline callback after %s", + async (operation: string) => { + await manager.init(connection) + const setTimeoutSpy = vi.spyOn(globalThis, "setTimeout") + manager.refreshTimeout() + const callback = setTimeoutSpy.mock.calls.at(-1)?.[0] + if (typeof callback !== "function") throw new Error("No deadline callback was registered") + if (operation === "refresh") manager.refreshTimeout() + else if (operation === "clear") manager.clearTimeout() + else await manager.dispose() + callback() + expect(manager.signal.aborted).toBe(false) + expect(onAbort).not.toHaveBeenCalled() + }, + ) + + it("preserves timeout cancellation as the outcome of a disposed request", async () => { + await manager.init(connection) + manager.refreshTimeout() + await vi.advanceTimersByTimeAsync(2_000) + expect(onAbort).toHaveBeenCalledOnce() + const signal = manager.signal + expect(signal.reason).toEqual(new Error("Codex WebSocket stream timed out")) + await manager.dispose() + expect(manager.signal).toBe(signal) + expect(socket.eventNames()).toEqual([]) + }) +}) diff --git a/src/api/providers/codex-websocket/managers/__tests__/CodexWebSocketResponseManager.spec.ts b/src/api/providers/codex-websocket/managers/__tests__/CodexWebSocketResponseManager.spec.ts new file mode 100644 index 0000000000..e37d1f9820 --- /dev/null +++ b/src/api/providers/codex-websocket/managers/__tests__/CodexWebSocketResponseManager.spec.ts @@ -0,0 +1,81 @@ +import { CodexWebSocketResponseManager } from "../CodexWebSocketResponseManager" +import { CodexWebSocketResponseLocalDataSource } from "../../data/local/CodexWebSocketResponseLocalDataSource" +import { CodexWebSocketContinuationRepository } from "../../repositories/CodexWebSocketContinuationRepository" +import { CodexWebSocketResponseStateHolder } from "../../state-holders/CodexWebSocketResponseStateHolder" +import type { CodexResponseEvent } from "../../models/protocol" + +describe("CodexWebSocketResponseManager state transitions", () => { + const body = { model: "test-model", input: [{ role: "user", content: "Hello" }] } + const cacheMiss: CodexResponseEvent = { type: "error", error: { code: "previous_response_not_found" } } + let repository: CodexWebSocketContinuationRepository + let manager: CodexWebSocketResponseManager + + beforeEach(() => { + repository = new CodexWebSocketContinuationRepository(new CodexWebSocketResponseLocalDataSource()) + manager = new CodexWebSocketResponseManager(repository, new CodexWebSocketResponseStateHolder()) + }) + + it("requires explicit initialization and rejects repeated initialization", () => { + const prepare = vi.spyOn(repository, "prepare") + expect(prepare).not.toHaveBeenCalled() + expect(() => manager.preparedRequest).toThrow("not initialized") + expect(() => manager.accept({ type: "response.created" })).toThrow("not initialized") + manager.init(body) + expect(prepare).toHaveBeenCalledOnce() + expect(manager.completed).toBe(false) + expect(() => manager.init(body)).toThrow("already initialized") + }) + + it.each(["response.completed", "response.done"])("records streamed items in order on %s", (type: string) => { + manager.init(body) + const prepared = manager.preparedRequest + const record = vi.spyOn(repository, "record") + const first = { role: "assistant", content: "First" } + const second = { role: "assistant", content: "Second" } + manager.accept({ type: "response.output_item.done", item: first }) + manager.accept({ type: "response.output_text.delta", delta: "Text" }) + manager.accept({ type: "response.output_item.done", item: second }) + expect(manager.completed).toBe(false) + const response = { id: "resp_1" } + expect(manager.accept({ type, response })).toBe("emit") + expect(record).toHaveBeenCalledWith(prepared, response, [first, second]) + expect(manager.completed).toBe(true) + expect(() => manager.accept({ type: "response.created" })).toThrow("already completed") + expect(record).toHaveBeenCalledOnce() + }) + + it("completes an empty response without entering streaming", () => { + manager.init(body) + const record = vi.spyOn(repository, "record") + manager.accept({ type: "response.completed", response: { id: "resp_1", output: [] } }) + expect(record).toHaveBeenCalledWith(manager.preparedRequest, { id: "resp_1", output: [] }, []) + expect(manager.completed).toBe(true) + }) + + it("allows cache-miss recovery once while awaiting the first event", () => { + repository.record(repository.prepare(body), { id: "resp_cached", output: [] }, []) + manager.init(body) + const reset = vi.spyOn(repository, "reset") + expect(manager.accept(cacheMiss)).toBe("retry") + expect(reset).toHaveBeenCalledOnce() + expect(() => manager.accept(cacheMiss)).toThrow("previous_response_not_found") + expect(reset).toHaveBeenCalledOnce() + }) + + it("rejects cache-miss recovery after streaming starts", () => { + repository.record(repository.prepare(body), { id: "resp_cached", output: [] }, []) + manager.init(body) + manager.accept({ type: "response.created" }) + const reset = vi.spyOn(repository, "reset") + expect(() => manager.accept(cacheMiss)).toThrow("previous_response_not_found") + expect(reset).not.toHaveBeenCalled() + }) + + it.each(["response.failed", "response.incomplete"])("does not record completion for %s", (type: string) => { + manager.init(body) + const record = vi.spyOn(repository, "record") + expect(() => manager.accept({ type, response: { error: { message: "Failed" } } })).toThrow("Failed") + expect(record).not.toHaveBeenCalled() + expect(manager.completed).toBe(false) + }) +}) diff --git a/src/api/providers/codex-websocket/models/CachedCodexResponse.ts b/src/api/providers/codex-websocket/models/CachedCodexResponse.ts new file mode 100644 index 0000000000..268090c9c5 --- /dev/null +++ b/src/api/providers/codex-websocket/models/CachedCodexResponse.ts @@ -0,0 +1,7 @@ +import type { CodexWebSocketItemSnapshotModel } from "./CodexWebSocketItemSnapshotModel" + +export interface CachedCodexResponse { + id: string + settings: string + input: CodexWebSocketItemSnapshotModel[] +} diff --git a/src/api/providers/codex-websocket/models/CodexWebSocketConnectionStateModel.ts b/src/api/providers/codex-websocket/models/CodexWebSocketConnectionStateModel.ts new file mode 100644 index 0000000000..89a656f418 --- /dev/null +++ b/src/api/providers/codex-websocket/models/CodexWebSocketConnectionStateModel.ts @@ -0,0 +1,18 @@ +/** Valid runtime states. Resource types are injected so the model does not depend on scopes or IO. */ +export type CodexWebSocketConnectionStateModel = + | { readonly status: "disconnected" } + | { readonly status: "connecting"; readonly key: string; readonly scope: TScope } + | { + readonly status: "active" + readonly key: string + readonly scope: TScope + readonly connectedAt: number + } + | { + readonly status: "idle" + readonly key: string + readonly scope: TScope + readonly connectedAt: number + readonly idleTimer: TTimer + } + | { readonly status: "unavailable"; readonly key: string; readonly retryAt: number } diff --git a/src/api/providers/codex-websocket/models/CodexWebSocketItemSnapshotModel.ts b/src/api/providers/codex-websocket/models/CodexWebSocketItemSnapshotModel.ts new file mode 100644 index 0000000000..c2f06e9e46 --- /dev/null +++ b/src/api/providers/codex-websocket/models/CodexWebSocketItemSnapshotModel.ts @@ -0,0 +1,59 @@ +import type { JsonObject } from "./protocol" +import { fingerprint, asJsonObject } from "../utils/protocol" + +/** Retain hashes only; diagnostic logs must never expose messages, arguments or encrypted reasoning. */ +export class CodexWebSocketItemSnapshotModel { + private constructor( + readonly hash: string, + readonly type: string, + private readonly fields: Record, + ) {} + + static create(value: unknown): CodexWebSocketItemSnapshotModel { + const item = CodexWebSocketItemSnapshotModel.normalize(value) + return new CodexWebSocketItemSnapshotModel( + fingerprint(item), + String(item.type ?? "message"), + Object.fromEntries(Object.entries(item).map(([key, field]) => [key, fingerprint(field)])), + ) + } + + changedFields(other: CodexWebSocketItemSnapshotModel): string[] { + return Object.keys({ ...this.fields, ...other.fields }).filter((key) => this.fields[key] !== other.fields[key]) + } + + /** Compare wire history with Zoo's reconstructed history, ignoring response-only metadata. */ + private static normalize(value: unknown): JsonObject { + const item = asJsonObject(value) + switch (item.type) { + case "message": + case undefined: + return { + role: item.role, + content: Array.isArray(item.content) + ? item.content.map((part: unknown) => { + const content = asJsonObject(part) + return content.type === "output_text" + ? { type: content.type, text: content.text } + : content + }) + : item.content, + } + case "function_call": { + let args = item.arguments + if (typeof args === "string") { + try { + args = JSON.parse(args) + } catch { + // Malformed arguments must still compare exactly. + } + } + return { type: item.type, call_id: item.call_id, name: item.name, arguments: args } + } + case "reasoning": + return { type: item.type, id: item.id, encrypted_content: item.encrypted_content } + default: + return item + } + } +} diff --git a/src/api/providers/codex-websocket/models/CodexWebSocketRequestStateModel.ts b/src/api/providers/codex-websocket/models/CodexWebSocketRequestStateModel.ts new file mode 100644 index 0000000000..b96666c021 --- /dev/null +++ b/src/api/providers/codex-websocket/models/CodexWebSocketRequestStateModel.ts @@ -0,0 +1,13 @@ +/** Request resources exist together; disposed requests retain only their cancellation outcome. */ +export type CodexWebSocketRequestStateModel = + | { readonly status: "uninitialized" } + | { readonly status: "initializing"; readonly controller: AbortController; readonly signal: AbortSignal } + | { + readonly status: "active" + readonly controller: AbortController + readonly signal: AbortSignal + readonly socket: TSocket + readonly events: TEvents + readonly timeout?: TTimer + } + | { readonly status: "disposed"; readonly signal: AbortSignal } diff --git a/src/api/providers/codex-websocket/models/CodexWebSocketResponseStateModel.ts b/src/api/providers/codex-websocket/models/CodexWebSocketResponseStateModel.ts new file mode 100644 index 0000000000..ac2b1d8af9 --- /dev/null +++ b/src/api/providers/codex-websocket/models/CodexWebSocketResponseStateModel.ts @@ -0,0 +1,9 @@ +import type { PreparedCodexRequest } from "./PreparedCodexRequest" + +/** Cache-miss recovery is possible only while awaiting the first response event. */ +export type CodexWebSocketResponseStateModel = + | { readonly status: "uninitialized" } + | { readonly status: "awaiting"; readonly prepared: PreparedCodexRequest } + | { readonly status: "recovering"; readonly prepared: PreparedCodexRequest } + | { readonly status: "streaming"; readonly prepared: PreparedCodexRequest; readonly output: readonly unknown[] } + | { readonly status: "completed"; readonly prepared: PreparedCodexRequest } diff --git a/src/api/providers/codex-websocket/models/PreparedCodexRequest.ts b/src/api/providers/codex-websocket/models/PreparedCodexRequest.ts new file mode 100644 index 0000000000..09058e2dfd --- /dev/null +++ b/src/api/providers/codex-websocket/models/PreparedCodexRequest.ts @@ -0,0 +1,12 @@ +import type { CodexWebSocketItemSnapshotModel } from "./CodexWebSocketItemSnapshotModel" +import type { JsonObject } from "./protocol" + +export interface PreparedCodexRequest { + request: JsonObject + input: unknown[] + snapshots: CodexWebSocketItemSnapshotModel[] + settings: string + previousResponseId?: string + offset: number + fullContextReason?: string +} diff --git a/src/api/providers/codex-websocket/models/__tests__/CodexWebSocketItemSnapshotModel.spec.ts b/src/api/providers/codex-websocket/models/__tests__/CodexWebSocketItemSnapshotModel.spec.ts new file mode 100644 index 0000000000..24b15e28de --- /dev/null +++ b/src/api/providers/codex-websocket/models/__tests__/CodexWebSocketItemSnapshotModel.spec.ts @@ -0,0 +1,21 @@ +import { CodexWebSocketItemSnapshotModel } from "../CodexWebSocketItemSnapshotModel" + +describe("CodexWebSocketItemSnapshotModel", () => { + const call = (args: unknown) => ({ type: "function_call", call_id: "call_1", name: "read_file", arguments: args }) + + it("compares structured function arguments without requiring JSON strings", () => { + const structured = CodexWebSocketItemSnapshotModel.create(call({ path: "file.ts", start: 1 })) + const serialized = CodexWebSocketItemSnapshotModel.create(call('{"start":1,"path":"file.ts"}')) + expect(structured.hash).toBe(serialized.hash) + expect(structured.changedFields(serialized)).toEqual([]) + }) + + it("preserves malformed function arguments for exact comparison", () => { + const first = CodexWebSocketItemSnapshotModel.create(call("invalid JSON")) + const same = CodexWebSocketItemSnapshotModel.create(call("invalid JSON")) + const changed = CodexWebSocketItemSnapshotModel.create(call("different invalid JSON")) + expect(first.hash).toBe(same.hash) + expect(first.hash).not.toBe(changed.hash) + expect(first.changedFields(changed)).toEqual(["arguments"]) + }) +}) diff --git a/src/api/providers/codex-websocket/models/protocol.ts b/src/api/providers/codex-websocket/models/protocol.ts new file mode 100644 index 0000000000..2bd8bc58b5 --- /dev/null +++ b/src/api/providers/codex-websocket/models/protocol.ts @@ -0,0 +1,8 @@ +export type JsonObject = Record +export type CodexResponseEvent = JsonObject & { type: string } + +export interface CodexWebSocketOptions { + headers: Record + signal: AbortSignal + timeoutMs: number +} diff --git a/src/api/providers/codex-websocket/repositories/CodexWebSocketContinuationRepository.ts b/src/api/providers/codex-websocket/repositories/CodexWebSocketContinuationRepository.ts new file mode 100644 index 0000000000..769a2740ed --- /dev/null +++ b/src/api/providers/codex-websocket/repositories/CodexWebSocketContinuationRepository.ts @@ -0,0 +1,69 @@ +import type { CodexWebSocketResponseLocalDataSource } from "../data/local/CodexWebSocketResponseLocalDataSource" +import { CodexWebSocketItemSnapshotModel } from "../models/CodexWebSocketItemSnapshotModel" +import type { PreparedCodexRequest } from "../models/PreparedCodexRequest" +import type { JsonObject } from "../models/protocol" +import { fingerprint, asJsonObject } from "../utils/protocol" + +/** Compares the server output with the history Zoo can reconstruct, retaining only hashes. */ +export class CodexWebSocketContinuationRepository { + constructor(private readonly store: CodexWebSocketResponseLocalDataSource) {} + + reset(): void { + this.store.clear() + } + + prepare(body: unknown): PreparedCodexRequest { + const request = asJsonObject(body) + if (!Array.isArray(request.input)) throw new Error("Codex WebSocket input must be an array") + const { input: _input, ...settings } = request + const settingsKey = fingerprint(settings) + const snapshots = request.input.map((item) => CodexWebSocketItemSnapshotModel.create(item)) + const cached = this.store.read() + return { + request, + input: request.input, + snapshots, + settings: settingsKey, + previousResponseId: cached?.id, + offset: cached?.input.length ?? 0, + fullContextReason: this.getFullContextReason(settingsKey, snapshots), + } + } + + record(prepared: PreparedCodexRequest, response: JsonObject, streamedOutput: unknown[]): void { + const output = Array.isArray(response.output) ? response.output : streamedOutput + // Zoo replays only encrypted reasoning. Plain reasoning remains in the server cache, + // but must not occupy a slot in the local history prefix. + const replayable = output.filter((value) => { + const item = asJsonObject(value) + return item.type !== "reasoning" || Boolean(item.encrypted_content) + }) + this.store.write( + typeof response.id === "string" + ? { + id: response.id, + settings: prepared.settings, + input: [ + ...prepared.snapshots, + ...replayable.map((item) => CodexWebSocketItemSnapshotModel.create(item)), + ], + } + : undefined, + ) + } + + private getFullContextReason(settings: string, input: CodexWebSocketItemSnapshotModel[]): string | undefined { + const previous = this.store.read() + if (!previous) return "no cached response" + if (previous.settings !== settings) return "request settings changed" + if (previous.input.length > input.length) return "history shortened" + const index = previous.input.findIndex((item, index) => item.hash !== input[index].hash) + if (index !== -1) { + const before = previous.input[index] + const after = input[index] + const fields = before.changedFields(after) + return `history changed at item ${index}: ${before.type} -> ${after.type}; fields: ${fields.join(", ")}` + } + return undefined + } +} diff --git a/src/api/providers/codex-websocket/repositories/__tests__/CodexWebSocketContinuationRepository.spec.ts b/src/api/providers/codex-websocket/repositories/__tests__/CodexWebSocketContinuationRepository.spec.ts new file mode 100644 index 0000000000..8bb9ca7b4f --- /dev/null +++ b/src/api/providers/codex-websocket/repositories/__tests__/CodexWebSocketContinuationRepository.spec.ts @@ -0,0 +1,91 @@ +import { CodexWebSocketContinuationRepository } from "../CodexWebSocketContinuationRepository" +import { CodexWebSocketResponseLocalDataSource } from "../../data/local/CodexWebSocketResponseLocalDataSource" + +describe("CodexWebSocketContinuationRepository", () => { + const body = (input: unknown[] = [{ role: "user", content: "private prompt" }]) => ({ + model: "test-model", + stream: true, + input, + }) + let store: CodexWebSocketResponseLocalDataSource + let repository: CodexWebSocketContinuationRepository + + beforeEach(() => { + store = new CodexWebSocketResponseLocalDataSource() + repository = new CodexWebSocketContinuationRepository(store) + }) + + it("only receives dependencies in its constructor", () => { + const read = vi.spyOn(store, "read") + const write = vi.spyOn(store, "write") + const clear = vi.spyOn(store, "clear") + new CodexWebSocketContinuationRepository(store) + expect(read).not.toHaveBeenCalled() + expect(write).not.toHaveBeenCalled() + expect(clear).not.toHaveBeenCalled() + }) + + it("uses the injected store to continue from a previously recorded response", () => { + const prepared = repository.prepare(body()) + expect(prepared.fullContextReason).toBe("no cached response") + repository.record(prepared, { id: "resp_1", output: [] }, []) + const nextInput = [...body().input, { role: "user", content: "Next" }] + const nextRepository = new CodexWebSocketContinuationRepository(store) + expect(nextRepository.prepare(body(nextInput))).toMatchObject({ + previousResponseId: "resp_1", + offset: 1, + fullContextReason: undefined, + }) + }) + + it("stores only hashes and replayable reasoning snapshots", () => { + const reasoning = { type: "reasoning", id: "rs_1", encrypted_content: "private encrypted reasoning" } + const message = { role: "assistant", content: [{ type: "output_text", text: "private answer" }] } + const output = [ + { type: "reasoning", summary: [{ type: "summary_text", text: "private unencrypted reasoning" }] }, + reasoning, + message, + ] + repository.record(repository.prepare(body()), { id: "resp_1", output }, []) + const cached = store.read() + expect(cached?.input).toHaveLength(3) + expect(cached?.input.map((item) => item.type)).toEqual(["message", "reasoning", "message"]) + expect(JSON.stringify(cached)).not.toContain("private") + expect(repository.prepare(body([...body().input, reasoning, message]))).toMatchObject({ + previousResponseId: "resp_1", + offset: 3, + fullContextReason: undefined, + }) + }) + + it("uses streamed output when the completed response omits output", () => { + const output = [{ role: "assistant", content: "Answer" }] + repository.record(repository.prepare(body()), { id: "resp_1" }, output) + expect(store.read()?.input).toHaveLength(2) + expect(repository.prepare(body([...body().input, ...output])).fullContextReason).toBeUndefined() + }) + + it("clears stored continuation when a response has no usable ID", () => { + const prepared = repository.prepare(body()) + repository.record(prepared, { id: "resp_1", output: [] }, []) + repository.record(prepared, { output: [] }, []) + expect(store.read()).toBeUndefined() + expect(repository.prepare(body()).fullContextReason).toBe("no cached response") + }) + + it("resets continuation by clearing the local store", () => { + repository.record(repository.prepare(body()), { id: "resp_1", output: [] }, []) + repository.reset() + expect(store.read()).toBeUndefined() + expect(repository.prepare(body()).previousResponseId).toBeUndefined() + }) + + it("requires full context when the reconstructed history is shortened", () => { + const input = [...body().input, { role: "user", content: "Follow-up" }] + repository.record(repository.prepare(body(input)), { id: "resp_1", output: [] }, []) + expect(repository.prepare(body())).toMatchObject({ + previousResponseId: "resp_1", + fullContextReason: "history shortened", + }) + }) +}) diff --git a/src/api/providers/codex-websocket/scopes/CodexWebSocketConnectionScope.ts b/src/api/providers/codex-websocket/scopes/CodexWebSocketConnectionScope.ts new file mode 100644 index 0000000000..6fd071219e --- /dev/null +++ b/src/api/providers/codex-websocket/scopes/CodexWebSocketConnectionScope.ts @@ -0,0 +1,45 @@ +import type * as vscode from "vscode" + +import { CodexWebSocketSocketRemoteDataSource } from "../data/remote/CodexWebSocketSocketRemoteDataSource" +import type { CodexWebSocketOptions } from "../models/protocol" + +/** Assembles and owns the services for one authenticated socket. */ +export class CodexWebSocketConnectionScope implements vscode.Disposable { + private _socketRemoteDataSource?: CodexWebSocketSocketRemoteDataSource + private initialized = false + + constructor( + private readonly url: string, + private readonly onError: () => void, + private readonly onClose: () => void, + ) {} + + get socketRemoteDataSource(): CodexWebSocketSocketRemoteDataSource { + if (!this.initialized || !this._socketRemoteDataSource) { + throw new Error("Codex WebSocket connection scope is not initialized") + } + return this._socketRemoteDataSource + } + + async init(options: CodexWebSocketOptions): Promise { + if (this._socketRemoteDataSource) throw new Error("Codex WebSocket connection scope is already initialized") + const socketRemoteDataSource = new CodexWebSocketSocketRemoteDataSource(this.url, this.onError, this.onClose) + this._socketRemoteDataSource = socketRemoteDataSource + try { + await socketRemoteDataSource.init(options) + if (this._socketRemoteDataSource !== socketRemoteDataSource) + throw new Error("Codex WebSocket connection scope was disposed") + this.initialized = true + } catch (error) { + if (this._socketRemoteDataSource === socketRemoteDataSource) this.dispose() + throw error + } + } + + dispose(): void { + const socketRemoteDataSource = this._socketRemoteDataSource + this._socketRemoteDataSource = undefined + this.initialized = false + socketRemoteDataSource?.dispose() + } +} diff --git a/src/api/providers/codex-websocket/scopes/CodexWebSocketRequestScope.ts b/src/api/providers/codex-websocket/scopes/CodexWebSocketRequestScope.ts new file mode 100644 index 0000000000..b2f8ae29b9 --- /dev/null +++ b/src/api/providers/codex-websocket/scopes/CodexWebSocketRequestScope.ts @@ -0,0 +1,71 @@ +import type * as vscode from "vscode" + +import type { CodexWebSocketConnectionManager } from "../managers/CodexWebSocketConnectionManager" +import { CodexWebSocketRequestManager } from "../managers/CodexWebSocketRequestManager" +import { CodexWebSocketResponseManager } from "../managers/CodexWebSocketResponseManager" +import type { CodexWebSocketOptions } from "../models/protocol" +import type { CodexWebSocketContinuationRepository } from "../repositories/CodexWebSocketContinuationRepository" +import { CodexWebSocketRequestStateHolder } from "../state-holders/CodexWebSocketRequestStateHolder" +import { CodexWebSocketResponseStateHolder } from "../state-holders/CodexWebSocketResponseStateHolder" + +/** Assembles and owns request-local services without owning the reusable connection. */ +export class CodexWebSocketRequestScope implements vscode.Disposable { + private _requestManager?: CodexWebSocketRequestManager + private _responseManager?: CodexWebSocketResponseManager + private initializing = false + private initialized = false + + constructor( + private readonly options: CodexWebSocketOptions, + private readonly onAbort: () => void, + ) {} + + get requestManager(): CodexWebSocketRequestManager { + if (!this.initialized || !this._requestManager) + throw new Error("Codex WebSocket request scope is not initialized") + return this._requestManager + } + + get responseManager(): CodexWebSocketResponseManager { + if (!this.initialized || !this._responseManager) { + throw new Error("Codex WebSocket request scope is not initialized") + } + return this._responseManager + } + + async init( + connection: CodexWebSocketConnectionManager, + continuation: CodexWebSocketContinuationRepository, + body: unknown, + ): Promise { + if (this._requestManager || this.initializing) { + throw new Error("Codex WebSocket request scope is already initialized") + } + const requestStateHolder = new CodexWebSocketRequestStateHolder() + const request = new CodexWebSocketRequestManager(this.options, this.onAbort, requestStateHolder) + this._requestManager = request + this.initializing = true + try { + await request.init(connection) + if (this._requestManager !== request) throw new Error("Codex WebSocket request scope was disposed") + const responseStateHolder = new CodexWebSocketResponseStateHolder() + const responseManager = new CodexWebSocketResponseManager(continuation, responseStateHolder) + responseManager.init(body) + this._responseManager = responseManager + this.initialized = true + } catch (error) { + if (this._requestManager === request) await this.dispose() + throw error + } finally { + this.initializing = false + } + } + + async dispose(): Promise { + const request = this._requestManager + this._requestManager = undefined + this._responseManager = undefined + this.initialized = false + await request?.dispose() + } +} diff --git a/src/api/providers/codex-websocket/scopes/CodexWebSocketTransportScope.ts b/src/api/providers/codex-websocket/scopes/CodexWebSocketTransportScope.ts new file mode 100644 index 0000000000..5744c600e4 --- /dev/null +++ b/src/api/providers/codex-websocket/scopes/CodexWebSocketTransportScope.ts @@ -0,0 +1,46 @@ +import type * as vscode from "vscode" + +import { CodexWebSocketTransport } from "../../CodexWebSocketTransport" +import { CodexWebSocketResponseLocalDataSource } from "../data/local/CodexWebSocketResponseLocalDataSource" +import { CodexWebSocketConnectionManager } from "../managers/CodexWebSocketConnectionManager" +import { CodexWebSocketContinuationRepository } from "../repositories/CodexWebSocketContinuationRepository" +import { CodexWebSocketConnectionStateHolder } from "../state-holders/CodexWebSocketConnectionStateHolder" +import { CodexWebSocketRequestScope } from "./CodexWebSocketRequestScope" + +const CODEX_WEBSOCKET_URL = "wss://chatgpt.com/backend-api/codex/responses" + +/** Assembles and owns the task-local WebSocket transport graph. */ +export class CodexWebSocketTransportScope implements vscode.Disposable { + private _transport?: CodexWebSocketTransport + + constructor(private readonly url = CODEX_WEBSOCKET_URL) {} + + get transport(): CodexWebSocketTransport { + if (!this._transport) throw new Error("Codex WebSocket transport scope is not initialized") + return this._transport + } + + /** Assembles services without opening a socket; authentication is supplied per request. */ + init(): void { + if (this._transport) throw new Error("Codex WebSocket transport scope is already initialized") + const responseLocalDataSource = new CodexWebSocketResponseLocalDataSource() + const continuationRepository = new CodexWebSocketContinuationRepository(responseLocalDataSource) + const connectionStateHolder = new CodexWebSocketConnectionStateHolder() + const connectionManager = new CodexWebSocketConnectionManager( + this.url, + () => continuationRepository.reset(), + connectionStateHolder, + ) + this._transport = new CodexWebSocketTransport( + connectionManager, + continuationRepository, + (options, onAbort) => new CodexWebSocketRequestScope(options, onAbort), + ) + } + + async dispose(): Promise { + const transport = this._transport + this._transport = undefined + await transport?.dispose() + } +} diff --git a/src/api/providers/codex-websocket/scopes/__tests__/CodexWebSocketScopes.spec.ts b/src/api/providers/codex-websocket/scopes/__tests__/CodexWebSocketScopes.spec.ts new file mode 100644 index 0000000000..1320d283d4 --- /dev/null +++ b/src/api/providers/codex-websocket/scopes/__tests__/CodexWebSocketScopes.spec.ts @@ -0,0 +1,435 @@ +import WebSocket from "ws" + +import { CodexWebSocketTransport } from "../../../CodexWebSocketTransport" +import { CodexWebSocketResponseLocalDataSource } from "../../data/local/CodexWebSocketResponseLocalDataSource" +import { CodexWebSocketContinuationRepository } from "../../repositories/CodexWebSocketContinuationRepository" +import { CodexWebSocketConnectionManager } from "../../managers/CodexWebSocketConnectionManager" +import { CodexWebSocketConnectionScope } from "../CodexWebSocketConnectionScope" +import { CodexWebSocketRequestScope } from "../CodexWebSocketRequestScope" +import { CodexWebSocketTransportScope } from "../CodexWebSocketTransportScope" +import { CodexWebSocketRequestManager } from "../../managers/CodexWebSocketRequestManager" +import { CodexWebSocketResponseManager } from "../../managers/CodexWebSocketResponseManager" +import { CodexWebSocketSocketRemoteDataSource } from "../../data/remote/CodexWebSocketSocketRemoteDataSource" +import { CodexWebSocketUnavailableError } from "../../errors/CodexWebSocketUnavailableError" +import type { CodexWebSocketOptions } from "../../models/protocol" +import { CodexWebSocketConnectionStateHolder } from "../../state-holders/CodexWebSocketConnectionStateHolder" +import { CodexWebSocketRequestStateHolder } from "../../state-holders/CodexWebSocketRequestStateHolder" + +const { sockets } = vi.hoisted(() => ({ sockets: [] as WebSocket[] })) + +vi.mock("ws", async () => { + const { EventEmitter } = await import("node:events") + class MockWebSocket extends EventEmitter { + static readonly CONNECTING = 0 + static readonly OPEN = 1 + static readonly CLOSING = 2 + static readonly CLOSED = 3 + readyState = MockWebSocket.CONNECTING + + constructor() { + super() + // The mock implements the socket surface used by these scopes, not all of ws. + sockets.push(this as unknown as WebSocket) + } + + override emit(event: string | symbol, ...args: unknown[]): boolean { + if (event === "open") this.readyState = MockWebSocket.OPEN + if (event === "close") this.readyState = MockWebSocket.CLOSED + return super.emit(event, ...args) + } + + terminate = vi.fn(() => { + const connecting = this.readyState === MockWebSocket.CONNECTING + this.readyState = MockWebSocket.CLOSING + queueMicrotask(() => { + if (connecting) this.emit("error", new Error("WebSocket closed before the connection was established")) + this.emit("close") + }) + }) + } + return { default: MockWebSocket } +}) + +describe("Codex WebSocket scopes", () => { + const resources: { dispose(): void | Promise }[] = [] + let controller: AbortController + let options: CodexWebSocketOptions + let reset = vi.fn<() => void>() + let connection: CodexWebSocketConnectionManager + let repository: CodexWebSocketContinuationRepository + const body = { model: "test-model", input: [] } + + const own = }>(resource: T): T => { + resources.push(resource) + return resource + } + const lastSocket = (): WebSocket => { + const socket = sockets.at(-1) + if (!socket) throw new Error("No test socket was created") + return socket + } + const expectNoListeners = (socket: WebSocket) => { + expect(socket.eventNames()).toEqual([]) + } + + beforeEach(() => { + sockets.length = 0 + controller = new AbortController() + options = { headers: { Authorization: "Bearer test-token" }, signal: controller.signal, timeoutMs: 2_000 } + reset = vi.fn() + connection = own( + new CodexWebSocketConnectionManager( + "ws://test/responses", + reset, + new CodexWebSocketConnectionStateHolder(), + ), + ) + repository = new CodexWebSocketContinuationRepository(new CodexWebSocketResponseLocalDataSource()) + }) + + afterEach(async () => { + for (const resource of resources.splice(0).reverse()) await resource.dispose() + vi.useRealTimers() + vi.restoreAllMocks() + }) + + it("does not create sockets, subscriptions, or timers in constructors", async () => { + vi.useFakeTimers() + const addListener = vi.spyOn(AbortSignal.prototype, "addEventListener") + const onAbort = vi.fn() + const scope = own(new CodexWebSocketConnectionScope("ws://test/responses", vi.fn(), vi.fn())) + const requestScope = own(new CodexWebSocketRequestScope(options, onAbort)) + const request = own(new CodexWebSocketRequestManager(options, onAbort, new CodexWebSocketRequestStateHolder())) + own(new CodexWebSocketSocketRemoteDataSource("ws://test/responses", vi.fn(), vi.fn())) + expect(sockets).toHaveLength(0) + expect(addListener).not.toHaveBeenCalled() + expect(vi.getTimerCount()).toBe(0) + expect(() => scope.socketRemoteDataSource).toThrow("not initialized") + expect(() => requestScope.requestManager).toThrow("not initialized") + expect(() => request.socket).toThrow("not initialized") + expect(() => request.events).toThrow("not initialized") + expect(() => request.refreshTimeout()).toThrow("not initialized") + scope.dispose() + await requestScope.dispose() + await request.dispose() + controller.abort() + expect(onAbort).not.toHaveBeenCalled() + }) + + it("exposes only owned services and lifecycle methods on scopes", () => { + expect(Object.getOwnPropertyNames(CodexWebSocketConnectionScope.prototype).sort()).toEqual( + ["constructor", "socketRemoteDataSource", "init", "dispose"].sort(), + ) + expect(Object.getOwnPropertyNames(CodexWebSocketRequestScope.prototype).sort()).toEqual( + ["constructor", "requestManager", "responseManager", "init", "dispose"].sort(), + ) + expect(Object.getOwnPropertyNames(CodexWebSocketTransportScope.prototype).sort()).toEqual( + ["constructor", "transport", "init", "dispose"].sort(), + ) + }) + + it("assembles the transport only on root-scope initialization without creating runtime resources", async () => { + vi.useFakeTimers() + const addListener = vi.spyOn(AbortSignal.prototype, "addEventListener") + const scope = own(new CodexWebSocketTransportScope("ws://test/responses")) + expect(() => scope.transport).toThrow("not initialized") + await scope.dispose() + scope.init() + const transport = scope.transport + expect(transport).toBeInstanceOf(CodexWebSocketTransport) + expect(() => scope.init()).toThrow("already initialized") + expect(sockets).toHaveLength(0) + expect(addListener).not.toHaveBeenCalled() + expect(vi.getTimerCount()).toBe(0) + const dispose = vi.spyOn(transport, "dispose") + await scope.dispose() + await scope.dispose() + expect(dispose).toHaveBeenCalledOnce() + expect(() => scope.transport).toThrow("not initialized") + scope.init() + expect(scope.transport).not.toBe(transport) + }) + + it("creates the connection on init and removes its subscriptions on disposal", async () => { + const onError = vi.fn() + const onClose = vi.fn() + const scope = own(new CodexWebSocketConnectionScope("ws://test/responses", onError, onClose)) + const initialized = scope.init(options) + const socket = lastSocket() + socket.emit("open") + await initialized + expect(scope.socketRemoteDataSource).toBeInstanceOf(CodexWebSocketSocketRemoteDataSource) + expect(scope.socketRemoteDataSource.socket).toBe(socket) + await expect(scope.init(options)).rejects.toThrow("already initialized") + socket.emit("error", new Error("Idle error")) + expect(onError).toHaveBeenCalledOnce() + scope.dispose() + scope.dispose() + expect(socket.terminate).toHaveBeenCalledOnce() + expect(() => scope.socketRemoteDataSource).toThrow("not initialized") + await Promise.resolve() + expect(onClose).not.toHaveBeenCalled() + expectNoListeners(socket) + }) + + it("guards direct socket access and initialization without duplicating handlers", async () => { + const source = own(new CodexWebSocketSocketRemoteDataSource("ws://test/responses", vi.fn(), vi.fn())) + expect(() => source.socket).toThrow("not initialized") + source.dispose() + const initialized = source.init(options) + const socket = lastSocket() + socket.emit("open") + await initialized + expect(source.socket).toBe(socket) + await expect(source.init(options)).rejects.toThrow("already initialized") + expect(socket.listenerCount("error")).toBe(1) + expect(socket.listenerCount("close")).toBe(1) + }) + + it("ignores stale socket-close dispatch after a replacement has initialized", async () => { + const onClose = vi.fn() + const source = own(new CodexWebSocketSocketRemoteDataSource("ws://test/responses", vi.fn(), onClose)) + const firstInit = source.init(options) + const first = lastSocket() + first.emit("open") + await firstInit + const staleClose = first.listeners("close")[0] + if (!staleClose) throw new Error("No socket-close handler was registered") + source.dispose() + const nextInit = source.init(options) + const next = lastSocket() + next.emit("open") + await nextInit + staleClose() + expect(onClose).not.toHaveBeenCalled() + expect(source.socket).toBe(next) + }) + + it("preserves a replacement socket when a previous handshake finishes after disposal", async () => { + const source = own(new CodexWebSocketSocketRemoteDataSource("ws://test/responses", vi.fn(), vi.fn())) + const firstInit = expect(source.init(options)).rejects.toThrow("scope was disposed") + const first = lastSocket() + first.emit("open") + source.dispose() + const nextInit = source.init(options) + const next = lastSocket() + next.emit("open") + await firstInit + await nextInit + expect(source.socket).toBe(next) + expectNoListeners(first) + }) + + it("does not let late service initialization dispose a replacement connection scope", async () => { + vi.spyOn(CodexWebSocketSocketRemoteDataSource.prototype, "init").mockResolvedValue() + const scope = own(new CodexWebSocketConnectionScope("ws://test/responses", vi.fn(), vi.fn())) + const firstInit = expect(scope.init(options)).rejects.toThrow("connection scope was disposed") + scope.dispose() + await scope.init(options) + const replacement = scope.socketRemoteDataSource + await firstInit + expect(scope.socketRemoteDataSource).toBe(replacement) + expect(sockets).toHaveLength(0) + }) + + it("safely disposes an in-flight upgrade without retaining listeners", async () => { + const onError = vi.fn() + const onClose = vi.fn() + const scope = own(new CodexWebSocketConnectionScope("ws://test/responses", onError, onClose)) + const initialized = expect(scope.init(options)).rejects.toThrow("before the connection was established") + const socket = lastSocket() + scope.dispose() + await initialized + expect(onError).not.toHaveBeenCalled() + expect(onClose).not.toHaveBeenCalled() + expectNoListeners(socket) + }) + + it("cleans up cancellation after socket open but before initialization resolves", async () => { + const scope = own(new CodexWebSocketConnectionScope("ws://test/responses", vi.fn(), vi.fn())) + const initialized = scope.init(options) + const reason = new Error("Stopped after socket open") + const rejected = expect(initialized).rejects.toBe(reason) + const socket = lastSocket() + socket.emit("open") + controller.abort(reason) + await rejected + expect(socket.terminate).toHaveBeenCalledOnce() + expect(() => scope.socketRemoteDataSource).toThrow("not initialized") + expectNoListeners(socket) + }) + + it("removes subscriptions from an already-closed socket", async () => { + const onClose = vi.fn() + const scope = own(new CodexWebSocketConnectionScope("ws://test/responses", vi.fn(), onClose)) + const initialized = scope.init(options) + const socket = lastSocket() + socket.emit("open") + await initialized + socket.emit("close") + expect(onClose).toHaveBeenCalledOnce() + scope.dispose() + expect(socket.terminate).not.toHaveBeenCalled() + expectNoListeners(socket) + }) + + it("can initialize again after disposal without stale callbacks affecting the new socket", async () => { + const onError = vi.fn() + const onClose = vi.fn() + const scope = own(new CodexWebSocketConnectionScope("ws://test/responses", onError, onClose)) + const firstInit = scope.init(options) + const first = lastSocket() + first.emit("open") + await firstInit + scope.dispose() + const nextInit = scope.init(options) + const next = lastSocket() + first.emit("error", new Error("Stale error")) + next.emit("open") + await nextInit + expectNoListeners(first) + expect(onError).not.toHaveBeenCalled() + expect(onClose).not.toHaveBeenCalled() + expect(scope.socketRemoteDataSource.socket).toBe(next) + }) + + it("does not open a connection for an already-aborted request", async () => { + controller.abort(new Error("Stopped")) + const scope = own(new CodexWebSocketConnectionScope("ws://test/responses", vi.fn(), vi.fn())) + const requestScope = own(new CodexWebSocketRequestScope(options, vi.fn())) + await expect(scope.init(options)).rejects.toThrow("Stopped") + await expect(requestScope.init(connection, repository, body)).rejects.toThrow("Stopped") + expect(sockets).toHaveLength(0) + }) + + it("owns request subscriptions and deadlines without disposing the reusable connection", async () => { + vi.useFakeTimers() + const onAbort = vi.fn() + const requestScope = own(new CodexWebSocketRequestScope(options, onAbort)) + const initialized = requestScope.init(connection, repository, body) + expect(() => requestScope.requestManager).toThrow("not initialized") + const socket = lastSocket() + socket.emit("open") + await initialized + const request = requestScope.requestManager + expect(request).toBeInstanceOf(CodexWebSocketRequestManager) + expect(requestScope.responseManager).toBeInstanceOf(CodexWebSocketResponseManager) + expect(requestScope.responseManager.preparedRequest.request).toBe(body) + expect(request.socket).toBe(socket) + expect(socket.listenerCount("message")).toBe(1) + await expect(requestScope.init(connection, repository, body)).rejects.toThrow("already initialized") + request.refreshTimeout() + request.refreshTimeout() + expect(vi.getTimerCount()).toBe(1) + await requestScope.dispose() + await requestScope.dispose() + expect(() => requestScope.requestManager).toThrow("not initialized") + expect(() => request.socket).toThrow("not initialized") + expect(() => request.events).toThrow("not initialized") + expect(vi.getTimerCount()).toBe(0) + expect(socket.listenerCount("message")).toBe(0) + expect(socket.listenerCount("error")).toBe(1) + expect(socket.listenerCount("close")).toBe(1) + expect(socket.terminate).not.toHaveBeenCalled() + controller.abort() + expect(onAbort).not.toHaveBeenCalled() + }) + + it("can reinitialize a disposed request on the same connection without duplicate listeners", async () => { + const requestScope = own(new CodexWebSocketRequestScope(options, () => connection.dispose())) + const initialized = requestScope.init(connection, repository, body) + const socket = lastSocket() + socket.emit("open") + await initialized + const firstRequest = requestScope.requestManager + await requestScope.dispose() + await requestScope.init(connection, repository, body) + const request = requestScope.requestManager + expect(request).not.toBe(firstRequest) + expect(sockets).toHaveLength(1) + expect(socket.listenerCount("message")).toBe(1) + const event = request.events.next() + const reason = new Error("Stopped") + controller.abort(reason) + await expect(event).rejects.toThrow() + expect(request.signal.reason).toBe(reason) + await requestScope.dispose() + expectNoListeners(socket) + }) + + it("cleans up a rejected upgrade and preserves the safe HTTP fallback error", async () => { + const onAbort = vi.fn() + const requestScope = own(new CodexWebSocketRequestScope(options, onAbort)) + const initialized = expect(requestScope.init(connection, repository, body)).rejects.toBeInstanceOf( + CodexWebSocketUnavailableError, + ) + const socket = lastSocket() + socket.emit("error", new Error("Upgrade rejected")) + await initialized + expect(() => requestScope.requestManager).toThrow("not initialized") + expect(onAbort).not.toHaveBeenCalled() + expectNoListeners(socket) + controller.abort() + expect(onAbort).not.toHaveBeenCalled() + }) + + it("cleans up request subscriptions when response-manager initialization fails", async () => { + const onAbort = vi.fn() + const scope = own(new CodexWebSocketRequestScope(options, onAbort)) + const initialized = expect(scope.init(connection, repository, { input: "invalid" })).rejects.toThrow( + "input must be an array", + ) + const socket = lastSocket() + socket.emit("open") + await initialized + expect(() => scope.requestManager).toThrow("not initialized") + expect(() => scope.responseManager).toThrow("not initialized") + expect(socket.listenerCount("message")).toBe(0) + expect(socket.listenerCount("error")).toBe(1) + expect(socket.listenerCount("close")).toBe(1) + controller.abort() + expect(onAbort).not.toHaveBeenCalled() + }) + + it("cancels and cleans up a request disposed during initialization", async () => { + const requestScope = own(new CodexWebSocketRequestScope(options, () => connection.dispose())) + const initialized = expect(requestScope.init(connection, repository, body)).rejects.toThrow() + const socket = lastSocket() + await requestScope.dispose() + await initialized + expectNoListeners(socket) + expect(() => requestScope.requestManager).toThrow("not initialized") + }) + + it("does not become initialized if disposed after acquiring a socket but before init resolves", async () => { + const socket = new WebSocket("ws://test/responses") + socket.emit("open") + vi.spyOn(connection, "acquire").mockResolvedValue(socket) + const requestScope = own(new CodexWebSocketRequestScope(options, vi.fn())) + const initialized = expect(requestScope.init(connection, repository, body)).rejects.toThrow( + "scope was disposed", + ) + await Promise.resolve() + await requestScope.dispose() + await initialized + expectNoListeners(socket) + expect(() => requestScope.requestManager).toThrow("not initialized") + }) + + it("clears an idle connection timer on disposal and does not schedule one without a socket", async () => { + vi.useFakeTimers() + connection.release() + expect(vi.getTimerCount()).toBe(0) + const acquired = connection.acquire(options) + const socket = lastSocket() + socket.emit("open") + await acquired + connection.release() + connection.release() + expect(vi.getTimerCount()).toBe(1) + connection.dispose() + expect(vi.getTimerCount()).toBe(0) + await Promise.resolve() + expectNoListeners(socket) + }) +}) diff --git a/src/api/providers/codex-websocket/state-holders/CodexWebSocketConnectionStateHolder.ts b/src/api/providers/codex-websocket/state-holders/CodexWebSocketConnectionStateHolder.ts new file mode 100644 index 0000000000..bb18f75419 --- /dev/null +++ b/src/api/providers/codex-websocket/state-holders/CodexWebSocketConnectionStateHolder.ts @@ -0,0 +1,64 @@ +import type { CodexWebSocketConnectionStateModel } from "../models/CodexWebSocketConnectionStateModel" +import type { CodexWebSocketConnectionScope } from "../scopes/CodexWebSocketConnectionScope" +import { StateHolder } from "./StateHolder" + +type ConnectionState = CodexWebSocketConnectionStateModel +type ConnectingState = Extract +type LiveState = Extract + +export class CodexWebSocketConnectionStateHolder extends StateHolder { + constructor() { + super({ status: "disconnected" }) + } + + beginConnection(key: string, scope: CodexWebSocketConnectionScope): ConnectingState { + if (this.state.status !== "disconnected" && this.state.status !== "unavailable") { + throw new Error("Codex WebSocket connection is already active") + } + const state = { status: "connecting", key, scope } as const + this.replaceState(state) + return state + } + + connect(attempt: ConnectingState, connectedAt: number): void { + this.ensureCurrent(attempt) + this.replaceState({ status: "active", key: attempt.key, scope: attempt.scope, connectedAt }) + } + + activate(state: LiveState): void { + this.ensureCurrent(state) + this.replaceState({ status: "active", key: state.key, scope: state.scope, connectedAt: state.connectedAt }) + } + + idle(state: LiveState, idleTimer: NodeJS.Timeout): Extract { + this.ensureCurrent(state) + const next = { + status: "idle", + key: state.key, + scope: state.scope, + connectedAt: state.connectedAt, + idleTimer, + } as const + this.replaceState(next) + return next + } + + markUnavailable(key: string, retryAt: number): void { + if (this.state.status !== "disconnected") throw new Error("Codex WebSocket connection must be disconnected") + this.replaceState({ status: "unavailable", key, retryAt }) + } + + disconnect(): ConnectionState { + const previous = this.state + if (previous.status !== "unavailable") this.replaceState({ status: "disconnected" }) + return previous + } + + ownsScope(scope: CodexWebSocketConnectionScope): boolean { + return "scope" in this.state && this.state.scope === scope + } + + private ensureCurrent(state: ConnectionState): void { + if (!this.isCurrent(state)) throw new Error("Codex WebSocket connection attempt was disposed") + } +} diff --git a/src/api/providers/codex-websocket/state-holders/CodexWebSocketRequestStateHolder.ts b/src/api/providers/codex-websocket/state-holders/CodexWebSocketRequestStateHolder.ts new file mode 100644 index 0000000000..ee328c24e0 --- /dev/null +++ b/src/api/providers/codex-websocket/state-holders/CodexWebSocketRequestStateHolder.ts @@ -0,0 +1,59 @@ +import type WebSocket from "ws" + +import type { CodexWebSocketRequestStateModel } from "../models/CodexWebSocketRequestStateModel" +import { StateHolder } from "./StateHolder" + +type RequestState = CodexWebSocketRequestStateModel, NodeJS.Timeout> +type InitializingState = Extract +type ActiveState = Extract + +export class CodexWebSocketRequestStateHolder extends StateHolder { + constructor() { + super({ status: "uninitialized" }) + } + + beginInitialization(controller: AbortController, signal: AbortSignal): InitializingState { + if (this.state.status === "active" || this.state.status === "initializing") { + throw new Error("Codex WebSocket request scope is already initialized") + } + const state = { status: "initializing", controller, signal } as const + this.replaceState(state) + return state + } + + activate(initializing: InitializingState, socket: WebSocket, events: AsyncIterableIterator): void { + this.ensureCurrent(initializing) + this.replaceState({ + status: "active", + controller: initializing.controller, + signal: initializing.signal, + socket, + events, + }) + } + + setDeadline(state: ActiveState, timeout: NodeJS.Timeout): ActiveState { + this.ensureCurrent(state) + const timed = { ...state, timeout } + this.replaceState(timed) + return timed + } + + clearDeadline(state: ActiveState): void { + this.ensureCurrent(state) + const { timeout: _timeout, ...active } = state + this.replaceState(active) + } + + dispose(): RequestState { + const previous = this.state + if (previous.status === "initializing" || previous.status === "active") { + this.replaceState({ status: "disposed", signal: previous.signal }) + } + return previous + } + + private ensureCurrent(state: RequestState): void { + if (!this.isCurrent(state)) throw new Error("Codex WebSocket request scope was disposed") + } +} diff --git a/src/api/providers/codex-websocket/state-holders/CodexWebSocketResponseStateHolder.ts b/src/api/providers/codex-websocket/state-holders/CodexWebSocketResponseStateHolder.ts new file mode 100644 index 0000000000..89f7adfdde --- /dev/null +++ b/src/api/providers/codex-websocket/state-holders/CodexWebSocketResponseStateHolder.ts @@ -0,0 +1,60 @@ +import type { CodexWebSocketResponseStateModel } from "../models/CodexWebSocketResponseStateModel" +import type { PreparedCodexRequest } from "../models/PreparedCodexRequest" +import { StateHolder } from "./StateHolder" + +export class CodexWebSocketResponseStateHolder extends StateHolder { + constructor() { + super({ status: "uninitialized" }) + } + + initialize(prepared: PreparedCodexRequest): void { + if (this.state.status !== "uninitialized") + throw new Error("Codex WebSocket response manager is already initialized") + this.replaceState({ status: "awaiting", prepared }) + } + + get preparedRequest(): PreparedCodexRequest { + if (this.state.status === "uninitialized") + throw new Error("Codex WebSocket response manager is not initialized") + return this.state.prepared + } + + get streamedOutput(): readonly unknown[] { + return this.state.status === "streaming" ? this.state.output : [] + } + + assertAcceptingEvents(): void { + if (this.state.status === "uninitialized") + throw new Error("Codex WebSocket response manager is not initialized") + if (this.state.status === "completed") throw new Error("Codex WebSocket response is already completed") + } + + appendOutputItem(item: unknown): void { + this.assertAcceptingEvents() + this.replaceState({ + status: "streaming", + prepared: this.preparedRequest, + output: [...this.streamedOutput, item], + }) + } + + startStreaming(): void { + this.assertAcceptingEvents() + if (this.state.status === "streaming") return + this.replaceState({ status: "streaming", prepared: this.preparedRequest, output: [] }) + } + + complete(): void { + this.assertAcceptingEvents() + this.replaceState({ status: "completed", prepared: this.preparedRequest }) + } + + get canRecoverCacheMiss(): boolean { + return this.state.status === "awaiting" && this.state.prepared.fullContextReason === undefined + } + + beginRecovery(): void { + if (!this.canRecoverCacheMiss) throw new Error("Codex WebSocket response cannot recover a cache miss") + this.replaceState({ status: "recovering", prepared: this.preparedRequest }) + } +} diff --git a/src/api/providers/codex-websocket/state-holders/StateHolder.ts b/src/api/providers/codex-websocket/state-holders/StateHolder.ts new file mode 100644 index 0000000000..0a948e142a --- /dev/null +++ b/src/api/providers/codex-websocket/state-holders/StateHolder.ts @@ -0,0 +1,16 @@ +/** Owns immutable state snapshots; subclasses expose domain-specific transitions. */ +export abstract class StateHolder { + protected constructor(private current: TState) {} + + get state(): TState { + return this.current + } + + isCurrent(snapshot: TState): boolean { + return this.current === snapshot + } + + protected replaceState(state: TState): void { + this.current = state + } +} diff --git a/src/api/providers/codex-websocket/state-holders/__tests__/CodexWebSocketStateHolders.spec.ts b/src/api/providers/codex-websocket/state-holders/__tests__/CodexWebSocketStateHolders.spec.ts new file mode 100644 index 0000000000..c4b8c7a66e --- /dev/null +++ b/src/api/providers/codex-websocket/state-holders/__tests__/CodexWebSocketStateHolders.spec.ts @@ -0,0 +1,159 @@ +import WebSocket from "ws" + +import { asyncStreamFrom } from "../../../../../test-utils/stream" +import { CodexWebSocketResponseLocalDataSource } from "../../data/local/CodexWebSocketResponseLocalDataSource" +import { CodexWebSocketContinuationRepository } from "../../repositories/CodexWebSocketContinuationRepository" +import { CodexWebSocketConnectionScope } from "../../scopes/CodexWebSocketConnectionScope" +import { CodexWebSocketConnectionStateHolder } from "../CodexWebSocketConnectionStateHolder" +import { CodexWebSocketRequestStateHolder } from "../CodexWebSocketRequestStateHolder" +import { CodexWebSocketResponseStateHolder } from "../CodexWebSocketResponseStateHolder" +import { StateHolder } from "../StateHolder" + +vi.mock("ws", async () => { + const { EventEmitter } = await import("node:events") + return { default: class TestSocket extends EventEmitter {} } +}) + +describe("Codex WebSocket state holders", () => { + beforeEach(() => vi.useFakeTimers()) + afterEach(() => vi.useRealTimers()) + const scope = () => new CodexWebSocketConnectionScope("ws://test/responses", vi.fn(), vi.fn()) + const prepared = () => + new CodexWebSocketContinuationRepository(new CodexWebSocketResponseLocalDataSource()).prepare({ + model: "test-model", + input: [], + }) + + it("uses the generic holder and replaces snapshots rather than mutating prior state", () => { + const holder = new CodexWebSocketConnectionStateHolder() + expect(holder).toBeInstanceOf(StateHolder) + const initial = holder.state + const attempt = holder.beginConnection("credentials-hash", scope()) + expect(holder.isCurrent(attempt)).toBe(true) + expect(holder.isCurrent(initial)).toBe(false) + holder.connect(attempt, 123) + expect(initial).toEqual({ status: "disconnected" }) + expect(attempt.status).toBe("connecting") + expect(holder.state).toMatchObject({ status: "active", key: "credentials-hash", connectedAt: 123 }) + }) + + it("associates idle resources only with the idle snapshot", () => { + const holder = new CodexWebSocketConnectionStateHolder() + const attempt = holder.beginConnection("key", scope()) + holder.connect(attempt, 123) + const active = holder.state + if (active.status !== "active") throw new Error("Expected an active connection") + const timer = setTimeout(() => {}, 1_000) + const idle = holder.idle(active, timer) + holder.activate(idle) + expect(active).not.toHaveProperty("idleTimer") + expect(idle.idleTimer).toBe(timer) + expect(holder.state).not.toHaveProperty("idleTimer") + expect(vi.getTimerCount()).toBe(1) + clearTimeout(timer) + }) + + it("rejects transitions from a replaced connection snapshot", () => { + const holder = new CodexWebSocketConnectionStateHolder() + const stale = holder.beginConnection("old", scope()) + holder.disconnect() + const current = holder.beginConnection("new", scope()) + expect(() => holder.connect(stale, 123)).toThrow("disposed") + expect(holder.state).toBe(current) + }) + + it("rejects overlapping connection state and cooldown publication while connecting", () => { + const holder = new CodexWebSocketConnectionStateHolder() + const attempt = holder.beginConnection("key", scope()) + expect(() => holder.beginConnection("other", scope())).toThrow("already active") + expect(() => holder.markUnavailable("key", 60_000)).toThrow("must be disconnected") + expect(holder.state).toBe(attempt) + }) + + it("preserves credential cooldown without retaining connection resources", () => { + const holder = new CodexWebSocketConnectionStateHolder() + holder.markUnavailable("key", 60_000) + const unavailable = holder.state + holder.disconnect() + expect(holder.state).toBe(unavailable) + expect(holder.state).toEqual({ status: "unavailable", key: "key", retryAt: 60_000 }) + expect(vi.getTimerCount()).toBe(0) + }) + + it("changes request lifecycle state without aborting signals or owning subscriptions", () => { + const holder = new CodexWebSocketRequestStateHolder() + const controller = new AbortController() + const initializing = holder.beginInitialization(controller, controller.signal) + holder.dispose() + expect(holder.state).toEqual({ status: "disposed", signal: controller.signal }) + expect(controller.signal.aborted).toBe(false) + expect(initializing.status).toBe("initializing") + expect(vi.getTimerCount()).toBe(0) + }) + + it("rejects repeated initialization and activation from an obsolete request snapshot", () => { + const holder = new CodexWebSocketRequestStateHolder() + const controller = new AbortController() + const first = holder.beginInitialization(controller, controller.signal) + expect(() => holder.beginInitialization(controller, controller.signal)).toThrow("already initialized") + holder.dispose() + const current = holder.beginInitialization(controller, controller.signal) + const socket = new WebSocket("ws://test") + const events = asyncStreamFrom([]) + expect(() => holder.activate(first, socket, events)).toThrow("scope was disposed") + expect(holder.state).toBe(current) + holder.activate(current, socket, events) + expect(() => holder.beginInitialization(controller, controller.signal)).toThrow("already initialized") + }) + + it("replaces request deadline snapshots without clearing timers itself", () => { + const holder = new CodexWebSocketRequestStateHolder() + const controller = new AbortController() + const initializing = holder.beginInitialization(controller, controller.signal) + holder.activate(initializing, new WebSocket("ws://test"), asyncStreamFrom([])) + const active = holder.state + if (active.status !== "active") throw new Error("Expected an active request") + const timer = setTimeout(() => {}, 1_000) + const timed = holder.setDeadline(active, timer) + holder.clearDeadline(timed) + expect(active).not.toHaveProperty("timeout") + expect(timed.timeout).toBe(timer) + expect(holder.state).not.toHaveProperty("timeout") + expect(vi.getTimerCount()).toBe(1) + clearTimeout(timer) + }) + + it("keeps prior response output snapshots unchanged and makes completion terminal", () => { + const holder = new CodexWebSocketResponseStateHolder() + holder.initialize(prepared()) + holder.appendOutputItem("first") + const first = holder.state + holder.appendOutputItem("second") + expect(first).toMatchObject({ status: "streaming", output: ["first"] }) + expect(holder.streamedOutput).toEqual(["first", "second"]) + holder.complete() + expect(holder.state.status).toBe("completed") + expect(holder.state).not.toHaveProperty("output") + expect(() => holder.appendOutputItem("late")).toThrow("already completed") + }) + + it("rejects repeated response initialization without replacing the prepared request", () => { + const holder = new CodexWebSocketResponseStateHolder() + holder.initialize(prepared()) + const initial = holder.state + expect(() => holder.initialize(prepared())).toThrow("already initialized") + expect(holder.state).toBe(initial) + }) + + it("permits recovery only in the initial awaiting state", () => { + const holder = new CodexWebSocketResponseStateHolder() + holder.initialize({ ...prepared(), fullContextReason: undefined, previousResponseId: "resp_1" }) + expect(holder.canRecoverCacheMiss).toBe(true) + holder.beginRecovery() + expect(holder.state.status).toBe("recovering") + expect(holder.canRecoverCacheMiss).toBe(false) + expect(() => holder.beginRecovery()).toThrow("cannot recover") + holder.startStreaming() + expect(holder.state.status).toBe("streaming") + }) +}) diff --git a/src/api/providers/codex-websocket/utils/__tests__/protocol.spec.ts b/src/api/providers/codex-websocket/utils/__tests__/protocol.spec.ts new file mode 100644 index 0000000000..e88696bd99 --- /dev/null +++ b/src/api/providers/codex-websocket/utils/__tests__/protocol.spec.ts @@ -0,0 +1,12 @@ +import { fingerprint } from "../protocol" + +describe("Codex WebSocket fingerprints", () => { + it("uses the null fingerprint when the input cannot be serialized", () => { + expect(fingerprint(undefined)).toBe(fingerprint(null)) + }) + + it("ignores object key ordering while distinguishing changed values", () => { + expect(fingerprint({ model: "test", stream: true })).toBe(fingerprint({ stream: true, model: "test" })) + expect(fingerprint({ model: "test" })).not.toBe(fingerprint({ model: "changed" })) + }) +}) diff --git a/src/api/providers/codex-websocket/utils/protocol.ts b/src/api/providers/codex-websocket/utils/protocol.ts new file mode 100644 index 0000000000..73a635ed08 --- /dev/null +++ b/src/api/providers/codex-websocket/utils/protocol.ts @@ -0,0 +1,32 @@ +import { createHash } from "node:crypto" +import stringify from "safe-stable-stringify" + +import type { JsonObject, CodexResponseEvent } from "../models/protocol" + +export function asJsonObject(value: unknown): JsonObject { + if (!value || typeof value !== "object" || Array.isArray(value)) { + throw new Error("Invalid Codex WebSocket payload") + } + return value as JsonObject +} + +export function parseResponseEvent(data: string): CodexResponseEvent { + const event = asJsonObject(JSON.parse(data)) + if (typeof event.type !== "string") throw new Error("Invalid Codex WebSocket event type") + return event as CodexResponseEvent +} + +export function fingerprint(value: unknown): string { + return createHash("sha256") + .update(stringify(value) ?? "null") + .digest("hex") +} + +export function responseError(event: CodexResponseEvent): Error { + const response = event.response ? asJsonObject(event.response) : undefined + const error = event.error ?? response?.error ?? response?.incomplete_details + const details = error ? asJsonObject(error) : {} + return new Error( + `Codex WebSocket (${String(event.status ?? event.type)}): ${String(details.message ?? details.code ?? details.reason ?? "request failed")}`, + ) +} diff --git a/src/package.json b/src/package.json index 0931ca779d..02f0b050ee 100644 --- a/src/package.json +++ b/src/package.json @@ -540,6 +540,7 @@ "vscode-material-icons": "^0.1.1", "web-tree-sitter": "^0.25.6", "workerpool": "^9.2.0", + "ws": "^8.18.3", "yaml": "^2.8.0", "zod": "3.25.76" }, @@ -557,6 +558,7 @@ "@types/ps-tree": "1.1.6", "@types/semver-compare": "1.0.3", "@types/vscode": "1.100.0", + "@types/ws": "^8.18.1", "@vitest/coverage-v8": "4.1.9", "@vscode/vsce": "3.9.2", "ai": "6.0.218",