feat(ai): add cached codex websocket transport

packages/ai/test/openai-codex-stream.test.ts

@@ -3,21 +3,26 @@ import { tmpdir } from "node:os";
 import { join } from "node:path";
 import { afterEach, describe, expect, it, vi } from "vitest";
 import {
+	getOpenAICodexWebSocketDebugStats,
+	resetOpenAICodexWebSocketDebugStats,
 	streamOpenAICodexResponses,
 	streamSimpleOpenAICodexResponses,
 } from "../src/providers/openai-codex-responses.js";
 import type { Context, Model } from "../src/types.js";
 
 const originalFetch = global.fetch;
+const originalWebSocket = globalThis.WebSocket;
 const originalAgentDir = process.env.PI_CODING_AGENT_DIR;
 
 afterEach(() => {
 	global.fetch = originalFetch;
+	globalThis.WebSocket = originalWebSocket;
 	if (originalAgentDir === undefined) {
 		delete process.env.PI_CODING_AGENT_DIR;
 	} else {
 		process.env.PI_CODING_AGENT_DIR = originalAgentDir;
 	}
+	resetOpenAICodexWebSocketDebugStats();
 	vi.restoreAllMocks();
 });
 
@@ -746,4 +751,150 @@ describe("openai-codex streaming", () => {
 		const streamResult = streamOpenAICodexResponses(model, context, { apiKey: token });
 		await streamResult.result();
 	});
+	it("sends only response input deltas in websocket-cached mode", async () => {
+		const token = mockToken();
+		const sentBodies: unknown[] = [];
+		const responses = [
+			{ responseId: "resp_1", messageId: "msg_1", text: "Hello" },
+			{ responseId: "resp_2", messageId: "msg_2", text: "Done" },
+		];
+
+		class MockWebSocket {
+			static OPEN = 1;
+			readyState = MockWebSocket.OPEN;
+			private listeners = new Map<string, Set<(event: unknown) => void>>();
+
+			constructor(_url: string, _protocols?: string | string[] | { headers?: Record<string, string> }) {
+				queueMicrotask(() => this.dispatch("open", {}));
+			}
+
+			addEventListener(type: string, listener: (event: unknown) => void): void {
+				let listeners = this.listeners.get(type);
+				if (!listeners) {
+					listeners = new Set();
+					this.listeners.set(type, listeners);
+				}
+				listeners.add(listener);
+			}
+
+			removeEventListener(type: string, listener: (event: unknown) => void): void {
+				this.listeners.get(type)?.delete(listener);
+			}
+
+			send(data: string): void {
+				sentBodies.push(JSON.parse(data));
+				const response = responses.shift();
+				if (!response) throw new Error("unexpected websocket request");
+				const events = [
+					{ type: "response.created", response: { id: response.responseId } },
+					{
+						type: "response.output_item.added",
+						item: {
+							type: "message",
+							id: response.messageId,
+							role: "assistant",
+							status: "in_progress",
+							content: [],
+						},
+					},
+					{ type: "response.content_part.added", part: { type: "output_text", text: "" } },
+					{ type: "response.output_text.delta", delta: response.text },
+					{
+						type: "response.output_item.done",
+						item: {
+							type: "message",
+							id: response.messageId,
+							role: "assistant",
+							status: "completed",
+							content: [{ type: "output_text", text: response.text }],
+						},
+					},
+					{
+						type: "response.completed",
+						response: {
+							id: response.responseId,
+							status: "completed",
+							usage: {
+								input_tokens: 5,
+								output_tokens: 3,
+								total_tokens: 8,
+								input_tokens_details: { cached_tokens: 0 },
+							},
+						},
+					},
+				];
+				queueMicrotask(() => {
+					for (const event of events) {
+						this.dispatch("message", { data: JSON.stringify(event) });
+					}
+				});
+			}
+
+			close(): void {
+				this.readyState = 3;
+			}
+
+			private dispatch(type: string, event: unknown): void {
+				for (const listener of this.listeners.get(type) ?? []) {
+					listener(event);
+				}
+			}
+		}
+
+		globalThis.WebSocket = MockWebSocket as unknown as typeof WebSocket;
+
+		const model: Model<"openai-codex-responses"> = {
+			id: "gpt-5.1-codex",
+			name: "GPT-5.1 Codex",
+			api: "openai-codex-responses",
+			provider: "openai-codex",
+			baseUrl: "https://chatgpt.com/backend-api",
+			reasoning: true,
+			input: ["text"],
+			cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 },
+			contextWindow: 400000,
+			maxTokens: 128000,
+		};
+		const firstContext: Context = {
+			systemPrompt: "You are a helpful assistant.",
+			messages: [{ role: "user", content: "Say hello", timestamp: 1 }],
+		};
+
+		const first = await streamOpenAICodexResponses(model, firstContext, {
+			apiKey: token,
+			sessionId: "session-1",
+			transport: "websocket-cached",
+		}).result();
+
+		const secondContext: Context = {
+			systemPrompt: "You are a helpful assistant.",
+			messages: [...firstContext.messages, first, { role: "user", content: "Now finish", timestamp: 2 }],
+		};
+		await streamOpenAICodexResponses(model, secondContext, {
+			apiKey: token,
+			sessionId: "session-1",
+			transport: "websocket-cached",
+		}).result();
+
+		expect(sentBodies).toHaveLength(2);
+		const firstBody = sentBodies[0] as { input: unknown[]; previous_response_id?: string; store?: boolean };
+		const secondBody = sentBodies[1] as { input: unknown[]; previous_response_id?: string; store?: boolean };
+		expect(firstBody.store).toBe(false);
+		expect(firstBody.previous_response_id).toBeUndefined();
+		expect(firstBody.input).toEqual([{ role: "user", content: [{ type: "input_text", text: "Say hello" }] }]);
+		expect(secondBody.store).toBe(false);
+		expect(secondBody.previous_response_id).toBe("resp_1");
+		expect(secondBody.input).toEqual([{ role: "user", content: [{ type: "input_text", text: "Now finish" }] }]);
+		expect(getOpenAICodexWebSocketDebugStats("session-1")).toMatchObject({
+			requests: 2,
+			connectionsCreated: 1,
+			connectionsReused: 1,
+			cachedContextRequests: 2,
+			storeTrueRequests: 0,
+			fullContextRequests: 1,
+			deltaRequests: 1,
+			lastDeltaInputItems: 1,
+			lastPreviousResponseId: "resp_1",
+		});
+	});
 });