diff --git a/packages/ai/CHANGELOG.md b/packages/ai/CHANGELOG.md index b1822dab4..fda71e30c 100644 --- a/packages/ai/CHANGELOG.md +++ b/packages/ai/CHANGELOG.md @@ -1,8 +1,10 @@ # Changelog ## [Unreleased] + ### Added +- Added `disableReasoning` to stream and OpenAI completion options to force reasoning off for models that support it, sending `reasoning: { enabled: false }` for OpenRouter-compatible requests - Added `thinkingDisplay` option to Anthropic options to control whether adaptive and explicit reasoning is returned as `summarized` or `omitted` - Added Anthropic model compatibility flags `supportsEagerToolInputStreaming` and `supportsLongCacheRetention` for API-capability-specific request behavior @@ -24,6 +26,7 @@ - Fixed Anthropic stream handling to parse raw SSE envelopes directly, ignore unrelated events, and repair malformed JSON in SSE payloads - Fixed Anthropic streaming to emit an explicit error when the SSE stream ends without a `message_stop` event - Fixed OpenAI Codex websocket continuations to send true `previous_response_id` deltas for `store: false` transcripts, expose request stats, and default text verbosity to `low` unless explicitly overridden. +- Fixed OpenAI Codex websocket append reuse after `response.completed` terminal events. ## [14.5.14] - 2026-05-01 ### Added diff --git a/packages/ai/src/providers/openai-codex-responses.ts b/packages/ai/src/providers/openai-codex-responses.ts index ade597dfd..509434e7d 100644 --- a/packages/ai/src/providers/openai-codex-responses.ts +++ b/packages/ai/src/providers/openai-codex-responses.ts @@ -1202,7 +1202,7 @@ function handleResponseCompleted( state.lastResponseId = response.id; state.lastResponseItems = stripInputItemIds(structuredCloneJSON(runtime.nativeOutputItems)); } - state.canAppend = rawEvent.type === "response.done"; + state.canAppend = rawEvent.type === "response.done" || rawEvent.type === "response.completed"; } calculateCost(model, output.usage); @@ -1843,12 +1843,10 @@ class CodexWebSocketConnection { await this.#connectPromise; return; } - const WebSocketWithHeaders = WebSocket as unknown as { - new (url: string, options?: { headers?: Record }): WebSocket; - }; const { promise, resolve, reject } = Promise.withResolvers(); this.#connectPromise = promise; - const socket = new WebSocketWithHeaders(this.#url, { headers: this.#headers }); + const socket = new WebSocket(this.#url, { headers: this.#headers }); + socket.binaryType = "nodebuffer"; this.#socket = socket; let settled = false; let timeout: NodeJS.Timeout | undefined; @@ -1878,15 +1876,15 @@ class CodexWebSocketConnection { } }, CODEX_WEBSOCKET_CONNECT_TIMEOUT_MS); - socket.addEventListener("open", event => { + socket.onopen = event => { if (!settled) { settled = true; clearPending(); this.#captureHandshakeHeaders(socket, event); resolve(); } - }); - socket.addEventListener("error", event => { + }; + socket.onerror = event => { const eventRecord = event as unknown as Record; const detail = (typeof eventRecord.message === "string" && eventRecord.message) || @@ -1900,8 +1898,8 @@ class CodexWebSocketConnection { return; } this.#push(error); - }); - socket.addEventListener("close", event => { + }; + socket.onclose = event => { this.#socket = null; if (!settled) { settled = true; @@ -1911,28 +1909,26 @@ class CodexWebSocketConnection { } this.#push(createCodexWebSocketTransportError(`websocket closed (${event.code})`)); this.#push(null); - }); - socket.addEventListener("message", event => { - void (async () => { - try { - const text = await decodeCodexWebSocketData(event.data); - if (!text) return; - const parsed = JSON.parse(text) as Record; - if (parsed.type === "error" && typeof parsed.error === "object" && parsed.error) { - const inner = parsed.error as Record; - if (typeof parsed.code !== "string" && typeof inner.code === "string") { - parsed.code = inner.code; - } - if (typeof parsed.message !== "string" && typeof inner.message === "string") { - parsed.message = inner.message; - } + }; + socket.onmessage = event => { + try { + const text = typeof event.data === "string" ? event.data : Buffer.from(event.data).toString("utf-8"); + if (!text) return; + const parsed = JSON.parse(text) as Record; + if (parsed.type === "error" && typeof parsed.error === "object" && parsed.error) { + const inner = parsed.error as Record; + if (typeof parsed.code !== "string" && typeof inner.code === "string") { + parsed.code = inner.code; + } + if (typeof parsed.message !== "string" && typeof inner.message === "string") { + parsed.message = inner.message; } - this.#push(parsed); - } catch (error) { - this.#push(createCodexWebSocketTransportError(String(error))); } - })(); - }); + this.#push(parsed); + } catch (error) { + this.#push(createCodexWebSocketTransportError(String(error))); + } + }; logger.time("codexWs:awaitTcpHandshake"); try { @@ -2280,23 +2276,6 @@ function resolveCodexResponsesUrl(baseUrl: string | undefined): string { return `${normalized}/codex/responses`; } -async function decodeCodexWebSocketData(data: unknown): Promise { - if (typeof data === "string") return data; - if (data instanceof ArrayBuffer) { - return new TextDecoder().decode(new Uint8Array(data)); - } - if (ArrayBuffer.isView(data)) { - const view = data; - return new TextDecoder().decode(new Uint8Array(view.buffer, view.byteOffset, view.byteLength)); - } - if (data && typeof data === "object" && "arrayBuffer" in data) { - const blobLike = data as { arrayBuffer: () => Promise }; - const arrayBuffer = await blobLike.arrayBuffer(); - return new TextDecoder().decode(new Uint8Array(arrayBuffer)); - } - return null; -} - function getAccountId(accessToken: string): string { const accountId = getCodexAccountId(accessToken); if (!accountId) { diff --git a/packages/ai/src/providers/openai-completions.ts b/packages/ai/src/providers/openai-completions.ts index 2895510a7..72ece308d 100644 --- a/packages/ai/src/providers/openai-completions.ts +++ b/packages/ai/src/providers/openai-completions.ts @@ -9,6 +9,7 @@ import type { ChatCompletionMessageParam, ChatCompletionToolMessageParam, } from "openai/resources/chat/completions"; +import type { Effort } from "../model-thinking"; import { calculateCost } from "../models"; import { getEnvApiKey } from "../stream"; import { @@ -17,6 +18,7 @@ import { type Message, type MessageAttribution, type Model, + type OpenAICompat, type ProviderSessionState, type ServiceTier, type StopReason, @@ -125,13 +127,21 @@ function hasToolHistory(messages: Message[]): boolean { export interface OpenAICompletionsOptions extends StreamOptions { toolChoice?: ToolChoice; reasoning?: "minimal" | "low" | "medium" | "high" | "xhigh"; + /** Force-disable reasoning for OpenRouter-format requests (sends `reasoning: { enabled: false }`). */ + disableReasoning?: boolean; serviceTier?: ServiceTier; } -type OpenAICompletionsSamplingParams = OpenAI.Chat.Completions.ChatCompletionCreateParamsStreaming & { +type OpenAICompletionsParams = OpenAI.Chat.Completions.ChatCompletionCreateParamsStreaming & { top_k?: number; min_p?: number; repetition_penalty?: number; + thinking?: { type: "enabled" | "disabled" }; + enable_thinking?: boolean; + chat_template_kwargs?: { enable_thinking: boolean }; + reasoning?: { effort?: string } | { enabled: false }; + provider?: OpenAICompat["openRouterRouting"]; + providerOptions?: { gateway?: { only?: string[]; order?: string[] } }; }; type AppliedToolStrictMode = "mixed" | "all_strict" | "none"; @@ -824,7 +834,7 @@ function buildParams( options: OpenAICompletionsOptions | undefined, resolvedBaseUrl?: string, toolStrictModeOverride?: ToolStrictModeOverride, -): { params: OpenAICompletionsSamplingParams; toolStrictMode: AppliedToolStrictMode } { +): { params: OpenAICompletionsParams; toolStrictMode: AppliedToolStrictMode } { const compat = getCompat(model, resolvedBaseUrl); const messages = convertMessages(model, context, compat); maybeAddOpenRouterAnthropicCacheControl(model, messages); @@ -837,7 +847,7 @@ function buildParams( const effectiveMaxTokens = options?.maxTokens ?? (isKimi ? model.maxTokens : undefined); const requestModelId = model.provider === "fireworks" ? toFireworksWireModelId(model.id) : model.id; - const params: OpenAICompletionsSamplingParams = { + const params: OpenAICompletionsParams = { model: requestModelId, messages, stream: true, @@ -845,7 +855,7 @@ function buildParams( let toolStrictMode: AppliedToolStrictMode = "none"; if (compat.supportsUsageInStreaming !== false) { - (params as { stream_options?: { include_usage: boolean } }).stream_options = { include_usage: true }; + params.stream_options = { include_usage: true }; } if (compat.supportsStore) { @@ -854,7 +864,7 @@ function buildParams( if (effectiveMaxTokens) { if (compat.maxTokensField === "max_tokens") { - (params as any).max_tokens = effectiveMaxTokens; + params.max_tokens = effectiveMaxTokens; } else { params.max_completion_tokens = effectiveMaxTokens; } @@ -897,27 +907,40 @@ function buildParams( if (supportsReasoningParams && compat.thinkingFormat === "zai" && model.reasoning) { // Z.ai uses binary thinking: { type: "enabled" | "disabled" } - // Must explicitly disable since z.ai defaults to thinking enabled - Reflect.set(params, "thinking", { type: options?.reasoning ? "enabled" : "disabled" }); + // Must explicitly disable since z.ai defaults to thinking enabled. + const enabled = options?.reasoning && !options?.disableReasoning; + params.thinking = { type: enabled ? "enabled" : "disabled" }; } else if (supportsReasoningParams && compat.thinkingFormat === "qwen" && model.reasoning) { // Qwen uses top-level enable_thinking: boolean - Reflect.set(params, "enable_thinking", !!options?.reasoning); + params.enable_thinking = !!options?.reasoning && !options?.disableReasoning; } else if (supportsReasoningParams && compat.thinkingFormat === "qwen-chat-template" && model.reasoning) { - Reflect.set(params, "chat_template_kwargs", { enable_thinking: !!options?.reasoning }); + params.chat_template_kwargs = { + enable_thinking: !!options?.reasoning && !options?.disableReasoning, + }; + } else if (supportsReasoningParams && compat.thinkingFormat === "openrouter" && model.reasoning) { + // OpenRouter normalizes reasoning across providers via a nested reasoning object. + // Without an explicit signal, OpenRouter defaults reasoning models to thinking, which + // silently consumes the entire output budget on small `max_tokens` requests (e.g. + // title generation). Honor `disableReasoning` to opt out cleanly. + const openRouterParams = params as typeof params & { + reasoning?: { effort?: string } | { enabled: false }; + }; + if (options?.disableReasoning) { + openRouterParams.reasoning = { enabled: false }; + } else if (options?.reasoning) { + openRouterParams.reasoning = { + effort: mapReasoningEffort(options.reasoning, compat.reasoningEffortMap), + }; + } } else if ( supportsReasoningParams && - compat.thinkingFormat === "openrouter" && options?.reasoning && - model.reasoning + !options?.disableReasoning && + model.reasoning && + compat.supportsReasoningEffort ) { - // OpenRouter normalizes reasoning across providers via a nested reasoning object. - const openRouterParams = params as typeof params & { reasoning?: { effort?: string } }; - openRouterParams.reasoning = { - effort: mapReasoningEffort(options.reasoning, compat.reasoningEffortMap), - }; - } else if (supportsReasoningParams && options?.reasoning && model.reasoning && compat.supportsReasoningEffort) { // OpenAI-style reasoning_effort - Reflect.set(params, "reasoning_effort", mapReasoningEffort(options.reasoning, compat.reasoningEffortMap)); + params.reasoning_effort = mapReasoningEffort(options.reasoning, compat.reasoningEffortMap) as Effort; } if (compat.disableReasoningOnForcedToolChoice && isForcedToolChoice(params.tool_choice)) { @@ -925,13 +948,13 @@ function buildParams( // Kimi 400 with `tool_choice 'specified' is incompatible with thinking // enabled`. Drop reasoning for this turn instead of dropping tool_choice; // the agent still gets the forced tool call, just without thinking. - delete (params as { reasoning_effort?: unknown }).reasoning_effort; - delete (params as { reasoning?: unknown }).reasoning; + delete params.reasoning_effort; + delete params.reasoning; } // OpenRouter provider routing preferences if (model.baseUrl.includes("openrouter.ai") && compat.openRouterRouting) { - Reflect.set(params, "provider", compat.openRouterRouting); + params.provider = compat.openRouterRouting; } // Vercel AI Gateway provider routing preferences @@ -941,7 +964,7 @@ function buildParams( const gatewayOptions: Record = {}; if (routing.only) gatewayOptions.only = routing.only; if (routing.order) gatewayOptions.order = routing.order; - Reflect.set(params, "providerOptions", { gateway: gatewayOptions }); + params.providerOptions = { gateway: gatewayOptions }; } } @@ -949,13 +972,6 @@ function buildParams( Object.assign(params, compat.extraBody); } - return buildParamsResult(params, toolStrictMode); -} - -function buildParamsResult( - params: OpenAICompletionsSamplingParams, - toolStrictMode: AppliedToolStrictMode, -): { params: OpenAICompletionsSamplingParams; toolStrictMode: AppliedToolStrictMode } { return { params, toolStrictMode }; } diff --git a/packages/ai/src/stream.ts b/packages/ai/src/stream.ts index f7887739d..a3a85df28 100644 --- a/packages/ai/src/stream.ts +++ b/packages/ai/src/stream.ts @@ -553,6 +553,7 @@ function mapOptionsForApi( return castApi<"openai-completions">({ ...base, reasoning: resolveOpenAiReasoningEffort(model, options), + disableReasoning: options?.disableReasoning, toolChoice: mapOpenAiToolChoice(options?.toolChoice), serviceTier: options?.serviceTier, }); diff --git a/packages/ai/src/types.ts b/packages/ai/src/types.ts index 95eb5e7ad..da1e8c3ed 100644 --- a/packages/ai/src/types.ts +++ b/packages/ai/src/types.ts @@ -246,6 +246,15 @@ export interface StreamOptions { // Unified options with reasoning passed to streamSimple() and completeSimple() export interface SimpleStreamOptions extends StreamOptions { reasoning?: Effort; + /** + * Force-disable reasoning for the request even when the model supports it. + * Takes precedence over `reasoning`. Useful for fast utility calls + * (e.g. title generation) where the model would otherwise burn the entire + * output budget on internal thinking. Currently honored by OpenRouter + * (sends `reasoning: { enabled: false }`); other providers already behave + * this way when `reasoning` is undefined. + */ + disableReasoning?: boolean; /** Custom token budgets for thinking levels (token-based providers only) */ thinkingBudgets?: ThinkingBudgets; /** Cursor exec handlers for local tool execution */ diff --git a/packages/ai/test/openai-codex-stream.test.ts b/packages/ai/test/openai-codex-stream.test.ts index 9442f0ae1..7d77f32bf 100644 --- a/packages/ai/test/openai-codex-stream.test.ts +++ b/packages/ai/test/openai-codex-stream.test.ts @@ -82,15 +82,108 @@ function encodeWebSocketMessage(value: Record): Uint8Array { return new TextEncoder().encode(JSON.stringify(value)); } -function createMessageEvent(data: unknown): MessageEvent { - return new MessageEvent("message", { data }); -} +type WsHeaders = Record; +type WsEventType = "open" | "message" | "error" | "close"; -function emitWebSocketEvent(listeners: Map void>>, type: string, event: Event): void { - const eventListeners = listeners.get(type); - if (!eventListeners) return; - for (const listener of eventListeners) { - listener(event); +const DEFAULT_USAGE = { + input_tokens: 5, + output_tokens: 3, + total_tokens: 8, + input_tokens_details: { cached_tokens: 0 }, +}; + +/** + * Drop-in mock for the global `WebSocket` used by the codex websocket transport. + * + * Production code wires lifecycle handlers via `onopen`/`onmessage`/`onerror`/`onclose` + * properties; tests drive the connection by calling `emit()`, `scheduleOpen()`, + * `sendJson()`, or the `emitCodexResponse()` convenience. + */ +class MockWebSocket { + static readonly CONNECTING = 0; + static readonly OPEN = 1; + static readonly CLOSING = 2; + static readonly CLOSED = 3; + + readyState: number = MockWebSocket.CONNECTING; + binaryType: "blob" | "arraybuffer" | "nodebuffer" = "blob"; + + onopen: ((event: Event) => void) | null = null; + onmessage: ((event: MessageEvent) => void) | null = null; + onerror: ((event: Event) => void) | null = null; + onclose: ((event: Event) => void) | null = null; + + constructor( + public readonly url: string, + public readonly options?: { headers?: WsHeaders }, + ) {} + + send(_data: string): void {} + + close(): void { + this.readyState = MockWebSocket.CLOSED; + } + + /** Dispatch an event to the matching `on{type}` handler. */ + emit(type: WsEventType, event: Event): void { + const handler = (this as unknown as Record)[`on${type}`]; + if (typeof handler === "function") (handler as (e: Event) => void).call(this, event); + } + + /** Asynchronously transition to OPEN and emit `open`. */ + scheduleOpen(): void { + setTimeout(() => { + this.readyState = MockWebSocket.OPEN; + this.emit("open", new Event("open")); + }, 0); + } + + /** Emit a message frame with arbitrary data. */ + sendMessage(data: unknown): void { + this.emit("message", { data } as unknown as MessageEvent); + } + + /** Emit a message frame with stringified-JSON data. */ + sendJson(payload: Record): void { + this.sendMessage(JSON.stringify(payload)); + } + + /** Emit the standard Codex completed-response sequence. */ + emitCodexResponse(opts: { + messageId: string; + responseId: string; + text: string; + terminalType?: "response.done" | "response.completed"; + includeCreated?: boolean; + }): void { + const { messageId, responseId, text, terminalType = "response.done", includeCreated = false } = opts; + if (includeCreated) { + this.sendJson({ type: "response.created", response: { id: responseId } }); + } + this.sendJson({ + type: "response.output_item.added", + item: { type: "message", id: messageId, role: "assistant", status: "in_progress", content: [] }, + }); + this.sendJson({ type: "response.content_part.added", part: { type: "output_text", text: "" } }); + this.sendJson({ type: "response.output_text.delta", delta: text }); + this.sendJson({ + type: "response.output_item.done", + item: { + type: "message", + id: messageId, + role: "assistant", + status: "completed", + content: [{ type: "output_text", text }], + }, + }); + this.sendJson({ + type: terminalType, + response: { + id: responseId, + status: "completed", + usage: DEFAULT_USAGE, + }, + }); } } @@ -130,32 +223,10 @@ describe("openai-codex streaming", () => { const tempDir = TempDir.createSync("@pi-codex-stream-"); setAgentDir(tempDir.path()); const token = createCodexTestToken(); - type WsListener = (event: Event) => void; - class BinaryPayloadWebSocket { - static readonly CONNECTING = 0; - static readonly OPEN = 1; - static readonly CLOSING = 2; - static readonly CLOSED = 3; - readyState = BinaryPayloadWebSocket.CONNECTING; - #listeners = new Map>(); - - constructor(_url: string, _options?: { headers?: Record }) { - setTimeout(() => { - this.readyState = BinaryPayloadWebSocket.OPEN; - this.#emit("open", new Event("open")); - }, 0); - } - - addEventListener(type: string, listener: unknown): void { - if (typeof listener !== "function") return; - const listeners = this.#listeners.get(type) ?? new Set(); - listeners.add(listener as WsListener); - this.#listeners.set(type, listeners); - } - - removeEventListener(type: string, listener: unknown): void { - if (typeof listener !== "function") return; - this.#listeners.get(type)?.delete(listener as WsListener); + class BinaryPayloadWebSocket extends MockWebSocket { + constructor(url: string, options?: { headers?: WsHeaders }) { + super(url, options); + this.scheduleOpen(); } send(): void { @@ -180,38 +251,14 @@ describe("openai-codex streaming", () => { }); const completed = encodeWebSocketMessage({ type: "response.done", - response: { - id: "resp_ws", - status: "completed", - usage: { - input_tokens: 5, - output_tokens: 3, - total_tokens: 8, - input_tokens_details: { cached_tokens: 0 }, - }, - }, + response: { id: "resp_ws", status: "completed", usage: DEFAULT_USAGE }, }); - this.#emit( - "message", - createMessageEvent(added.buffer.slice(added.byteOffset, added.byteOffset + added.byteLength)), - ); - this.#emit("message", createMessageEvent(contentPart)); - this.#emit("message", createMessageEvent(new DataView(delta.buffer, delta.byteOffset, delta.byteLength))); - this.#emit("message", createMessageEvent(new Blob([done]))); - this.#emit( - "message", - createMessageEvent( - completed.buffer.slice(completed.byteOffset, completed.byteOffset + completed.byteLength), - ), - ); - } - - close(): void { - this.readyState = BinaryPayloadWebSocket.CLOSED; - } - - #emit(type: string, event: Event): void { - emitWebSocketEvent(this.#listeners, type, event); + // Exercise every payload shape the production decoder must accept. + this.sendMessage(added.buffer.slice(added.byteOffset, added.byteOffset + added.byteLength)); + this.sendMessage(contentPart); + this.sendMessage(Buffer.from(delta)); + this.sendMessage(Buffer.from(done)); + this.sendMessage(completed.buffer.slice(completed.byteOffset, completed.byteOffset + completed.byteLength)); } } @@ -234,62 +281,27 @@ describe("openai-codex streaming", () => { setAgentDir(tempDir.path()); const token = createCodexTestToken(); let capturedHeaders: Record | undefined; - type WsListener = (event: Event) => void; - class HeaderCaptureWebSocket { - static readonly CONNECTING = 0; - static readonly OPEN = 1; - static readonly CLOSING = 2; - static readonly CLOSED = 3; - readyState = HeaderCaptureWebSocket.CONNECTING; - #listeners = new Map>(); - - constructor(_url: string, options?: { headers?: Record }) { + class HeaderCaptureWebSocket extends MockWebSocket { + constructor(url: string, options?: { headers?: WsHeaders }) { + super(url, options); capturedHeaders = options?.headers; - setTimeout(() => { - this.readyState = HeaderCaptureWebSocket.OPEN; - this.#emit("open", new Event("open")); - }, 0); - } - - addEventListener(type: string, listener: unknown): void { - if (typeof listener !== "function") return; - const listeners = this.#listeners.get(type) ?? new Set(); - listeners.add(listener as WsListener); - this.#listeners.set(type, listeners); - } - - removeEventListener(type: string, listener: unknown): void { - if (typeof listener !== "function") return; - this.#listeners.get(type)?.delete(listener as WsListener); + this.scheduleOpen(); } send(): void { - this.#emit( - "message", - createMessageEvent( - JSON.stringify({ - type: "response.done", - response: { - id: "resp_ws", - status: "completed", - usage: { - input_tokens: 1, - output_tokens: 1, - total_tokens: 2, - input_tokens_details: { cached_tokens: 0 }, - }, - }, - }), - ), - ); - } - - close(): void { - this.readyState = HeaderCaptureWebSocket.CLOSED; - } - - #emit(type: string, event: Event): void { - emitWebSocketEvent(this.#listeners, type, event); + this.sendJson({ + type: "response.done", + response: { + id: "resp_ws", + status: "completed", + usage: { + input_tokens: 1, + output_tokens: 1, + total_tokens: 2, + input_tokens_details: { cached_tokens: 0 }, + }, + }, + }); } } @@ -1027,52 +1039,18 @@ describe("openai-codex streaming", () => { return new Response("not found", { status: 404 }); }); global.fetch = fetchMock as unknown as typeof fetch; - type WsListener = (event: Event) => void; - class FailingWebSocket { - static readonly CONNECTING = 0; - static readonly OPEN = 1; - static readonly CLOSING = 2; - static readonly CLOSED = 3; - readyState = FailingWebSocket.CONNECTING; - #listeners = new Map>(); - url: string; - options?: { headers?: Record }; - - constructor(url: string, options?: { headers?: Record }) { - this.url = url; - this.options = options; + class FailingWebSocket extends MockWebSocket { + constructor(url: string, options?: { headers?: WsHeaders }) { + super(url, options); setTimeout(() => { expect(this.options?.headers?.["OpenAI-Beta"] ?? this.options?.headers?.["openai-beta"]).toStartWith( "responses_websockets=", ); - this.#emit("error", new Event("error")); - this.#emit("close", new Event("close")); - this.readyState = FailingWebSocket.CLOSED; + this.emit("error", new Event("error")); + this.emit("close", new Event("close")); + this.readyState = MockWebSocket.CLOSED; }, 0); } - addEventListener(type: string, listener: unknown): void { - if (typeof listener !== "function") return; - const listeners = this.#listeners.get(type) ?? new Set(); - listeners.add(listener as WsListener); - this.#listeners.set(type, listeners); - } - removeEventListener(type: string, listener: unknown): void { - if (typeof listener !== "function") return; - const listeners = this.#listeners.get(type); - listeners?.delete(listener as WsListener); - } - - send(): void {} - close(): void { - this.readyState = FailingWebSocket.CLOSED; - } - #emit(type: string, event: Event): void { - const listeners = this.#listeners.get(type); - if (!listeners) return; - for (const listener of listeners) { - listener(event); - } - } } global.WebSocket = FailingWebSocket as unknown as typeof WebSocket; @@ -1130,50 +1108,17 @@ describe("openai-codex streaming", () => { }); global.fetch = fetchMock as unknown as typeof fetch; - type WsListener = (event: Event) => void; let constructorCount = 0; - class FailingConnectWebSocket { - static readonly CONNECTING = 0; - static readonly OPEN = 1; - static readonly CLOSING = 2; - static readonly CLOSED = 3; - readyState = FailingConnectWebSocket.CONNECTING; - #listeners = new Map>(); - - constructor(_url: string, _options?: { headers?: Record }) { + class FailingConnectWebSocket extends MockWebSocket { + constructor(url: string, options?: { headers?: WsHeaders }) { + super(url, options); constructorCount += 1; setTimeout(() => { - this.#emit("error", new Event("error")); - this.#emit("close", new Event("close")); - this.readyState = FailingConnectWebSocket.CLOSED; + this.emit("error", new Event("error")); + this.emit("close", new Event("close")); + this.readyState = MockWebSocket.CLOSED; }, 0); } - - addEventListener(type: string, listener: unknown): void { - if (typeof listener !== "function") return; - const listeners = this.#listeners.get(type) ?? new Set(); - listeners.add(listener as WsListener); - this.#listeners.set(type, listeners); - } - - removeEventListener(type: string, listener: unknown): void { - if (typeof listener !== "function") return; - const listeners = this.#listeners.get(type); - listeners?.delete(listener as WsListener); - } - - send(): void {} - close(): void { - this.readyState = FailingConnectWebSocket.CLOSED; - } - - #emit(type: string, event: Event): void { - const listeners = this.#listeners.get(type); - if (!listeners) return; - for (const listener of listeners) { - listener(event); - } - } } global.WebSocket = FailingConnectWebSocket as unknown as typeof WebSocket; @@ -1238,92 +1183,20 @@ describe("openai-codex streaming", () => { }); global.fetch = fetchMock as unknown as typeof fetch; - type WsListener = (event: Event) => void; - class HandshakeWebSocket { - static readonly CONNECTING = 0; - static readonly OPEN = 1; - static readonly CLOSING = 2; - static readonly CLOSED = 3; - readyState = HandshakeWebSocket.CONNECTING; + class HandshakeWebSocket extends MockWebSocket { handshakeHeaders = { "x-codex-turn-state": "ws-turn-state-1", "x-models-etag": "models-etag-1", "x-reasoning-included": "true", }; - #listeners = new Map>(); - constructor(_url: string, _options?: { headers?: Record }) { - setTimeout(() => { - this.readyState = HandshakeWebSocket.OPEN; - this.#emit("open", new Event("open")); - }, 0); - } - - addEventListener(type: string, listener: unknown): void { - if (typeof listener !== "function") return; - const listeners = this.#listeners.get(type) ?? new Set(); - listeners.add(listener as WsListener); - this.#listeners.set(type, listeners); - } - - removeEventListener(type: string, listener: unknown): void { - if (typeof listener !== "function") return; - const listeners = this.#listeners.get(type); - listeners?.delete(listener as WsListener); + constructor(url: string, options?: { headers?: WsHeaders }) { + super(url, options); + this.scheduleOpen(); } send(): void { - this.#emit("message", { - data: JSON.stringify({ - type: "response.output_item.added", - item: { type: "message", id: "msg_ws", role: "assistant", status: "in_progress", content: [] }, - }), - } as unknown as Event); - this.#emit("message", { - data: JSON.stringify({ type: "response.content_part.added", part: { type: "output_text", text: "" } }), - } as unknown as Event); - this.#emit("message", { - data: JSON.stringify({ type: "response.output_text.delta", delta: "Hello WS" }), - } as unknown as Event); - this.#emit("message", { - data: JSON.stringify({ - type: "response.output_item.done", - item: { - type: "message", - id: "msg_ws", - role: "assistant", - status: "completed", - content: [{ type: "output_text", text: "Hello WS" }], - }, - }), - } as unknown as Event); - this.#emit("message", { - data: JSON.stringify({ - type: "response.done", - response: { - id: "resp_ws", - status: "completed", - usage: { - input_tokens: 5, - output_tokens: 3, - total_tokens: 8, - input_tokens_details: { cached_tokens: 0 }, - }, - }, - }), - } as unknown as Event); - } - - close(): void { - this.readyState = HandshakeWebSocket.CLOSED; - } - - #emit(type: string, event: Event): void { - const listeners = this.#listeners.get(type); - if (!listeners) return; - for (const listener of listeners) { - listener(event); - } + this.emitCodexResponse({ messageId: "msg_ws", responseId: "resp_ws", text: "Hello WS" }); } } @@ -1380,91 +1253,35 @@ describe("openai-codex streaming", () => { }); global.fetch = fetchMock as unknown as typeof fetch; - type WsListener = (event: Event) => void; - class ServiceTierWebSocket { - static readonly CONNECTING = 0; - static readonly OPEN = 1; - static readonly CLOSING = 2; - static readonly CLOSED = 3; - readyState = ServiceTierWebSocket.CONNECTING; - #listeners = new Map>(); - - constructor(_url: string, _options?: { headers?: Record }) { - setTimeout(() => { - this.readyState = ServiceTierWebSocket.OPEN; - this.#emit("open", new Event("open")); - }, 0); - } - - addEventListener(type: string, listener: unknown): void { - if (typeof listener !== "function") return; - const listeners = this.#listeners.get(type) ?? new Set(); - listeners.add(listener as WsListener); - this.#listeners.set(type, listeners); - } - - removeEventListener(type: string, listener: unknown): void { - if (typeof listener !== "function") return; - const listeners = this.#listeners.get(type); - listeners?.delete(listener as WsListener); + class ServiceTierWebSocket extends MockWebSocket { + constructor(url: string, options?: { headers?: WsHeaders }) { + super(url, options); + this.scheduleOpen(); } send(data: string): void { sentRequests.push(JSON.parse(data) as Record); - this.#emit("message", { - data: JSON.stringify({ - type: "response.output_item.added", - item: { type: "message", id: "msg_ws", role: "assistant", status: "in_progress", content: [] }, - }), - } as unknown as Event); - this.#emit("message", { - data: JSON.stringify({ type: "response.content_part.added", part: { type: "output_text", text: "" } }), - } as unknown as Event); - this.#emit("message", { - data: JSON.stringify({ type: "response.output_text.delta", delta: "Hello WS" }), - } as unknown as Event); - this.#emit("message", { - data: JSON.stringify({ - type: "response.output_item.done", - item: { - type: "message", - id: "msg_ws", - role: "assistant", - status: "completed", - content: [{ type: "output_text", text: "Hello WS" }], - }, - }), - } as unknown as Event); - this.#emit("message", { - data: JSON.stringify({ type: "response.created", response: { id: "resp_ws" } }), - } as unknown as Event); - this.#emit("message", { - data: JSON.stringify({ - type: "response.done", - response: { - id: "resp_ws", - status: "completed", - usage: { - input_tokens: 5, - output_tokens: 3, - total_tokens: 8, - input_tokens_details: { cached_tokens: 0 }, - }, - }, - }), - } as unknown as Event); - } - - close(): void { - this.readyState = ServiceTierWebSocket.CLOSED; - } - - #emit(type: string, event: Event): void { - const listeners = this.#listeners.get(type); - if (!listeners) return; - for (const listener of listeners) { - listener(event); - } + this.sendJson({ + type: "response.output_item.added", + item: { type: "message", id: "msg_ws", role: "assistant", status: "in_progress", content: [] }, + }); + this.sendJson({ type: "response.content_part.added", part: { type: "output_text", text: "" } }); + this.sendJson({ type: "response.output_text.delta", delta: "Hello WS" }); + this.sendJson({ + type: "response.output_item.done", + item: { + type: "message", + id: "msg_ws", + role: "assistant", + status: "completed", + content: [{ type: "output_text", text: "Hello WS" }], + }, + }); + this.sendJson({ type: "response.created", response: { id: "resp_ws" } }); + this.sendJson({ + type: "response.done", + response: { id: "resp_ws", status: "completed", usage: DEFAULT_USAGE }, + }); } } @@ -1514,96 +1331,22 @@ describe("openai-codex streaming", () => { }); global.fetch = fetchMock as unknown as typeof fetch; - type WsListener = (event: Event) => void; - class DeltaWebSocket { - static readonly CONNECTING = 0; - static readonly OPEN = 1; - static readonly CLOSING = 2; - static readonly CLOSED = 3; - readyState = DeltaWebSocket.CONNECTING; - #listeners = new Map>(); - - constructor(_url: string, _options?: { headers?: Record }) { - setTimeout(() => { - this.readyState = DeltaWebSocket.OPEN; - this.#emit("open", new Event("open")); - }, 0); - } - - addEventListener(type: string, listener: unknown): void { - if (typeof listener !== "function") return; - const listeners = this.#listeners.get(type) ?? new Set(); - listeners.add(listener as WsListener); - this.#listeners.set(type, listeners); - } - - removeEventListener(type: string, listener: unknown): void { - if (typeof listener !== "function") return; - const listeners = this.#listeners.get(type); - listeners?.delete(listener as WsListener); + class DeltaWebSocket extends MockWebSocket { + constructor(url: string, options?: { headers?: WsHeaders }) { + super(url, options); + this.scheduleOpen(); } send(data: string): void { - const request = JSON.parse(data) as Record; - sentRequests.push(request); + sentRequests.push(JSON.parse(data) as Record); const responseIndex = sentRequests.length; - const responseId = `resp_${responseIndex}`; - const messageId = `msg_${responseIndex}`; - const text = responseIndex === 1 ? "First answer" : "Second answer"; - this.#emit("message", { - data: JSON.stringify({ type: "response.created", response: { id: responseId } }), - } as unknown as Event); - this.#emit("message", { - data: JSON.stringify({ - type: "response.output_item.added", - item: { type: "message", id: messageId, role: "assistant", status: "in_progress", content: [] }, - }), - } as unknown as Event); - this.#emit("message", { - data: JSON.stringify({ type: "response.content_part.added", part: { type: "output_text", text: "" } }), - } as unknown as Event); - this.#emit("message", { - data: JSON.stringify({ type: "response.output_text.delta", delta: text }), - } as unknown as Event); - this.#emit("message", { - data: JSON.stringify({ - type: "response.output_item.done", - item: { - type: "message", - id: messageId, - role: "assistant", - status: "completed", - content: [{ type: "output_text", text }], - }, - }), - } as unknown as Event); - this.#emit("message", { - data: JSON.stringify({ - type: "response.done", - response: { - id: responseId, - status: "completed", - usage: { - input_tokens: 5, - output_tokens: 3, - total_tokens: 8, - input_tokens_details: { cached_tokens: 0 }, - }, - }, - }), - } as unknown as Event); - } - - close(): void { - this.readyState = DeltaWebSocket.CLOSED; - } - - #emit(type: string, event: Event): void { - const listeners = this.#listeners.get(type); - if (!listeners) return; - for (const listener of listeners) { - listener(event); - } + this.emitCodexResponse({ + messageId: `msg_${responseIndex}`, + responseId: `resp_${responseIndex}`, + text: responseIndex === 1 ? "First answer" : "Second answer", + terminalType: "response.completed", + includeCreated: true, + }); } } @@ -1734,90 +1477,17 @@ describe("openai-codex streaming", () => { }); global.fetch = fetchMock as unknown as typeof fetch; - type WsListener = (event: Event) => void; - class WebSocketV2HeaderProbe { - static readonly CONNECTING = 0; - static readonly OPEN = 1; - static readonly CLOSING = 2; - static readonly CLOSED = 3; - readyState = WebSocketV2HeaderProbe.CONNECTING; - #listeners = new Map>(); - - constructor(_url: string, options?: { headers?: Record }) { + class WebSocketV2HeaderProbe extends MockWebSocket { + constructor(url: string, options?: { headers?: WsHeaders }) { + super(url, options); expect(options?.headers?.["OpenAI-Beta"] ?? options?.headers?.["openai-beta"]).toBe( "responses_websockets=2026-02-06", ); - setTimeout(() => { - this.readyState = WebSocketV2HeaderProbe.OPEN; - this.#emit("open", new Event("open")); - }, 0); - } - - addEventListener(type: string, listener: unknown): void { - if (typeof listener !== "function") return; - const listeners = this.#listeners.get(type) ?? new Set(); - listeners.add(listener as WsListener); - this.#listeners.set(type, listeners); - } - - removeEventListener(type: string, listener: unknown): void { - if (typeof listener !== "function") return; - const listeners = this.#listeners.get(type); - listeners?.delete(listener as WsListener); + this.scheduleOpen(); } send(): void { - this.#emit("message", { - data: JSON.stringify({ - type: "response.output_item.added", - item: { type: "message", id: "msg_v2", role: "assistant", status: "in_progress", content: [] }, - }), - } as unknown as Event); - this.#emit("message", { - data: JSON.stringify({ type: "response.content_part.added", part: { type: "output_text", text: "" } }), - } as unknown as Event); - this.#emit("message", { - data: JSON.stringify({ type: "response.output_text.delta", delta: "Hello v2" }), - } as unknown as Event); - this.#emit("message", { - data: JSON.stringify({ - type: "response.output_item.done", - item: { - type: "message", - id: "msg_v2", - role: "assistant", - status: "completed", - content: [{ type: "output_text", text: "Hello v2" }], - }, - }), - } as unknown as Event); - this.#emit("message", { - data: JSON.stringify({ - type: "response.done", - response: { - id: "resp_v2", - status: "completed", - usage: { - input_tokens: 5, - output_tokens: 3, - total_tokens: 8, - input_tokens_details: { cached_tokens: 0 }, - }, - }, - }), - } as unknown as Event); - } - - close(): void { - this.readyState = WebSocketV2HeaderProbe.CLOSED; - } - - #emit(type: string, event: Event): void { - const listeners = this.#listeners.get(type); - if (!listeners) return; - for (const listener of listeners) { - listener(event); - } + this.emitCodexResponse({ messageId: "msg_v2", responseId: "resp_v2", text: "Hello v2" }); } } @@ -1873,51 +1543,16 @@ describe("openai-codex streaming", () => { }); global.fetch = fetchMock as unknown as typeof fetch; - type WsListener = (event: Event) => void; let sendCount = 0; - class IdleWebSocket { - static readonly CONNECTING = 0; - static readonly OPEN = 1; - static readonly CLOSING = 2; - static readonly CLOSED = 3; - readyState = IdleWebSocket.CONNECTING; - #listeners = new Map>(); - - constructor(_url: string, _options?: { headers?: Record }) { - setTimeout(() => { - this.readyState = IdleWebSocket.OPEN; - this.#emit("open", new Event("open")); - }, 0); - } - - addEventListener(type: string, listener: unknown): void { - if (typeof listener !== "function") return; - const listeners = this.#listeners.get(type) ?? new Set(); - listeners.add(listener as WsListener); - this.#listeners.set(type, listeners); - } - - removeEventListener(type: string, listener: unknown): void { - if (typeof listener !== "function") return; - const listeners = this.#listeners.get(type); - listeners?.delete(listener as WsListener); + class IdleWebSocket extends MockWebSocket { + constructor(url: string, options?: { headers?: WsHeaders }) { + super(url, options); + this.scheduleOpen(); } send(): void { sendCount += 1; } - - close(): void { - this.readyState = IdleWebSocket.CLOSED; - } - - #emit(type: string, event: Event): void { - const listeners = this.#listeners.get(type); - if (!listeners) return; - for (const listener of listeners) { - listener(event); - } - } } global.WebSocket = IdleWebSocket as unknown as typeof WebSocket; @@ -1979,104 +1614,29 @@ describe("openai-codex streaming", () => { }); global.fetch = fetchMock as unknown as typeof fetch; - type WsListener = (event: Event) => void; let constructorCount = 0; const requestTypes: string[] = []; - class FlakyCloseWebSocket { - static readonly CONNECTING = 0; - static readonly OPEN = 1; - static readonly CLOSING = 2; - static readonly CLOSED = 3; - readyState = FlakyCloseWebSocket.CONNECTING; - #listeners = new Map>(); - - constructor(_url: string, _options?: { headers?: Record }) { + class FlakyCloseWebSocket extends MockWebSocket { + constructor(url: string, options?: { headers?: WsHeaders }) { + super(url, options); constructorCount += 1; - setTimeout(() => { - this.readyState = FlakyCloseWebSocket.OPEN; - this.#emit("open", new Event("open")); - }, 0); - } - - addEventListener(type: string, listener: unknown): void { - if (typeof listener !== "function") return; - const listeners = this.#listeners.get(type) ?? new Set(); - listeners.add(listener as WsListener); - this.#listeners.set(type, listeners); - } - - removeEventListener(type: string, listener: unknown): void { - if (typeof listener !== "function") return; - const listeners = this.#listeners.get(type); - listeners?.delete(listener as WsListener); + this.scheduleOpen(); } send(data: string): void { const request = JSON.parse(data) as { type?: string }; requestTypes.push(typeof request.type === "string" ? request.type : ""); if (requestTypes.length === 1) { - this.readyState = FlakyCloseWebSocket.CLOSED; - this.#emit("close", { code: 1012 } as unknown as Event); + this.readyState = MockWebSocket.CLOSED; + this.emit("close", { code: 1012 } as unknown as Event); return; } - this.#emit("message", { - data: JSON.stringify({ - type: "response.output_item.added", - item: { - type: "message", - id: "msg_retry_close", - role: "assistant", - status: "in_progress", - content: [], - }, - }), - } as unknown as Event); - this.#emit("message", { - data: JSON.stringify({ type: "response.content_part.added", part: { type: "output_text", text: "" } }), - } as unknown as Event); - this.#emit("message", { - data: JSON.stringify({ type: "response.output_text.delta", delta: "Hello retry close" }), - } as unknown as Event); - this.#emit("message", { - data: JSON.stringify({ - type: "response.output_item.done", - item: { - type: "message", - id: "msg_retry_close", - role: "assistant", - status: "completed", - content: [{ type: "output_text", text: "Hello retry close" }], - }, - }), - } as unknown as Event); - this.#emit("message", { - data: JSON.stringify({ - type: "response.done", - response: { - id: "resp_retry_close", - status: "completed", - usage: { - input_tokens: 5, - output_tokens: 3, - total_tokens: 8, - input_tokens_details: { cached_tokens: 0 }, - }, - }, - }), - } as unknown as Event); - } - - close(): void { - this.readyState = FlakyCloseWebSocket.CLOSED; - } - - #emit(type: string, event: Event): void { - const listeners = this.#listeners.get(type); - if (!listeners) return; - for (const listener of listeners) { - listener(event); - } + this.emitCodexResponse({ + messageId: "msg_retry_close", + responseId: "resp_retry_close", + text: "Hello retry close", + }); } } @@ -2142,50 +1702,16 @@ describe("openai-codex streaming", () => { }); global.fetch = fetchMock as unknown as typeof fetch; - type WsListener = (event: Event) => void; - class UnavailableBeforeStreamWebSocket { - static readonly CONNECTING = 0; - static readonly OPEN = 1; - static readonly CLOSING = 2; - static readonly CLOSED = 3; - readyState = UnavailableBeforeStreamWebSocket.CONNECTING; - #listeners = new Map>(); - - constructor(_url: string, _options?: { headers?: Record }) { + class UnavailableBeforeStreamWebSocket extends MockWebSocket { + constructor(url: string, options?: { headers?: WsHeaders }) { + super(url, options); setTimeout(() => { - this.readyState = UnavailableBeforeStreamWebSocket.OPEN; - this.#emit("open", new Event("open")); - this.readyState = UnavailableBeforeStreamWebSocket.CLOSED; - this.#emit("close", { code: 1006 } as unknown as Event); + this.readyState = MockWebSocket.OPEN; + this.emit("open", new Event("open")); + this.readyState = MockWebSocket.CLOSED; + this.emit("close", { code: 1006 } as unknown as Event); }, 0); } - - addEventListener(type: string, listener: unknown): void { - if (typeof listener !== "function") return; - const listeners = this.#listeners.get(type) ?? new Set(); - listeners.add(listener as WsListener); - this.#listeners.set(type, listeners); - } - - removeEventListener(type: string, listener: unknown): void { - if (typeof listener !== "function") return; - const listeners = this.#listeners.get(type); - listeners?.delete(listener as WsListener); - } - - send(): void {} - - close(): void { - this.readyState = UnavailableBeforeStreamWebSocket.CLOSED; - } - - #emit(type: string, event: Event): void { - const listeners = this.#listeners.get(type); - if (!listeners) return; - for (const listener of listeners) { - listener(event); - } - } } global.WebSocket = UnavailableBeforeStreamWebSocket as unknown as typeof WebSocket; @@ -2241,41 +1767,19 @@ describe("openai-codex streaming", () => { }); global.fetch = fetchMock as unknown as typeof fetch; - type WsListener = (event: Event) => void; const sentTypesByConnection: string[][] = []; let constructorCount = 0; let abortSecondRequest: (() => void) | undefined; - class AbortResetWebSocket { - static readonly CONNECTING = 0; - static readonly OPEN = 1; - static readonly CLOSING = 2; - static readonly CLOSED = 3; - readyState = AbortResetWebSocket.CONNECTING; - #listeners = new Map>(); + class AbortResetWebSocket extends MockWebSocket { #connectionIndex: number; - constructor(_url: string, _options?: { headers?: Record }) { + constructor(url: string, options?: { headers?: WsHeaders }) { + super(url, options); this.#connectionIndex = constructorCount; constructorCount += 1; sentTypesByConnection[this.#connectionIndex] = []; - setTimeout(() => { - this.readyState = AbortResetWebSocket.OPEN; - this.#emit("open", new Event("open")); - }, 0); - } - - addEventListener(type: string, listener: unknown): void { - if (typeof listener !== "function") return; - const listeners = this.#listeners.get(type) ?? new Set(); - listeners.add(listener as WsListener); - this.#listeners.set(type, listeners); - } - - removeEventListener(type: string, listener: unknown): void { - if (typeof listener !== "function") return; - const listeners = this.#listeners.get(type); - listeners?.delete(listener as WsListener); + this.scheduleOpen(); } send(data: string): void { @@ -2285,25 +1789,16 @@ describe("openai-codex streaming", () => { const requestIndex = sentTypesByConnection[this.#connectionIndex]?.length ?? 0; if (this.#connectionIndex === 0 && requestIndex === 1) { - this.#emitCompleted("msg_1", "resp_1", "Hello one"); + this.emitCodexResponse({ messageId: "msg_1", responseId: "resp_1", text: "Hello one" }); return; } if (this.#connectionIndex === 0 && requestIndex === 2) { - this.#emit("message", { - data: JSON.stringify({ - type: "response.output_item.added", - item: { type: "message", id: "msg_2", role: "assistant", status: "in_progress", content: [] }, - }), - } as unknown as Event); - this.#emit("message", { - data: JSON.stringify({ - type: "response.content_part.added", - part: { type: "output_text", text: "" }, - }), - } as unknown as Event); - this.#emit("message", { - data: JSON.stringify({ type: "response.output_text.delta", delta: "Still streaming" }), - } as unknown as Event); + this.sendJson({ + type: "response.output_item.added", + item: { type: "message", id: "msg_2", role: "assistant", status: "in_progress", content: [] }, + }); + this.sendJson({ type: "response.content_part.added", part: { type: "output_text", text: "" } }); + this.sendJson({ type: "response.output_text.delta", delta: "Still streaming" }); setTimeout(() => { abortSecondRequest?.(); }, 0); @@ -2311,65 +1806,11 @@ describe("openai-codex streaming", () => { } if (this.#connectionIndex === 1 && requestIndex === 1) { expect(requestType).toBe("response.create"); - this.#emitCompleted("msg_3", "resp_3", "Hello three"); + this.emitCodexResponse({ messageId: "msg_3", responseId: "resp_3", text: "Hello three" }); return; } throw new Error(`Unexpected websocket send sequence: ${this.#connectionIndex}:${requestIndex}`); } - - close(): void { - this.readyState = AbortResetWebSocket.CLOSED; - } - - #emitCompleted(messageId: string, responseId: string, text: string): void { - this.#emit("message", { - data: JSON.stringify({ - type: "response.output_item.added", - item: { type: "message", id: messageId, role: "assistant", status: "in_progress", content: [] }, - }), - } as unknown as Event); - this.#emit("message", { - data: JSON.stringify({ type: "response.content_part.added", part: { type: "output_text", text: "" } }), - } as unknown as Event); - this.#emit("message", { - data: JSON.stringify({ type: "response.output_text.delta", delta: text }), - } as unknown as Event); - this.#emit("message", { - data: JSON.stringify({ - type: "response.output_item.done", - item: { - type: "message", - id: messageId, - role: "assistant", - status: "completed", - content: [{ type: "output_text", text }], - }, - }), - } as unknown as Event); - this.#emit("message", { - data: JSON.stringify({ - type: "response.done", - response: { - id: responseId, - status: "completed", - usage: { - input_tokens: 5, - output_tokens: 3, - total_tokens: 8, - input_tokens_details: { cached_tokens: 0 }, - }, - }, - }), - } as unknown as Event); - } - - #emit(type: string, event: Event): void { - const listeners = this.#listeners.get(type); - if (!listeners) return; - for (const listener of listeners) { - listener(event); - } - } } global.WebSocket = AbortResetWebSocket as unknown as typeof WebSocket; @@ -2452,37 +1893,14 @@ describe("openai-codex streaming", () => { }); global.fetch = fetchMock as unknown as typeof fetch; - type WsListener = (event: Event) => void; const sentTypes: string[] = []; let constructorCount = 0; - class ErrorResetWebSocket { - static readonly CONNECTING = 0; - static readonly OPEN = 1; - static readonly CLOSING = 2; - static readonly CLOSED = 3; - readyState = ErrorResetWebSocket.CONNECTING; - #listeners = new Map>(); - - constructor(_url: string, _options?: { headers?: Record }) { + class ErrorResetWebSocket extends MockWebSocket { + constructor(url: string, options?: { headers?: WsHeaders }) { + super(url, options); constructorCount += 1; - setTimeout(() => { - this.readyState = ErrorResetWebSocket.OPEN; - this.#emit("open", new Event("open")); - }, 0); - } - - addEventListener(type: string, listener: unknown): void { - if (typeof listener !== "function") return; - const listeners = this.#listeners.get(type) ?? new Set(); - listeners.add(listener as WsListener); - this.#listeners.set(type, listeners); - } - - removeEventListener(type: string, listener: unknown): void { - if (typeof listener !== "function") return; - const listeners = this.#listeners.get(type); - listeners?.delete(listener as WsListener); + this.scheduleOpen(); } send(data: string): void { @@ -2492,80 +1910,24 @@ describe("openai-codex streaming", () => { const requestIndex = sentTypes.length; if (requestIndex === 1) { - this.#emitCompleted("msg_1", "resp_1", "Hello one"); + this.emitCodexResponse({ messageId: "msg_1", responseId: "resp_1", text: "Hello one" }); return; } if (requestIndex === 2) { - this.#emit("message", { - data: JSON.stringify({ - type: "error", - code: "invalid_request_error", - message: "simulated request error", - }), - } as unknown as Event); + this.sendJson({ + type: "error", + code: "invalid_request_error", + message: "simulated request error", + }); return; } if (requestIndex === 3) { expect(requestType).toBe("response.create"); - this.#emitCompleted("msg_3", "resp_3", "Hello three"); + this.emitCodexResponse({ messageId: "msg_3", responseId: "resp_3", text: "Hello three" }); return; } throw new Error(`Unexpected websocket request index: ${requestIndex}`); } - - close(): void { - this.readyState = ErrorResetWebSocket.CLOSED; - } - - #emitCompleted(messageId: string, responseId: string, text: string): void { - this.#emit("message", { - data: JSON.stringify({ - type: "response.output_item.added", - item: { type: "message", id: messageId, role: "assistant", status: "in_progress", content: [] }, - }), - } as unknown as Event); - this.#emit("message", { - data: JSON.stringify({ type: "response.content_part.added", part: { type: "output_text", text: "" } }), - } as unknown as Event); - this.#emit("message", { - data: JSON.stringify({ type: "response.output_text.delta", delta: text }), - } as unknown as Event); - this.#emit("message", { - data: JSON.stringify({ - type: "response.output_item.done", - item: { - type: "message", - id: messageId, - role: "assistant", - status: "completed", - content: [{ type: "output_text", text }], - }, - }), - } as unknown as Event); - this.#emit("message", { - data: JSON.stringify({ - type: "response.done", - response: { - id: responseId, - status: "completed", - usage: { - input_tokens: 5, - output_tokens: 3, - total_tokens: 8, - input_tokens_details: { cached_tokens: 0 }, - }, - }, - }), - } as unknown as Event); - } - - #emit(type: string, event: Event): void { - const listeners = this.#listeners.get(type); - if (!listeners) return; - for (const listener of listeners) { - listener(event); - } - } } global.WebSocket = ErrorResetWebSocket as unknown as typeof WebSocket; @@ -2653,49 +2015,14 @@ describe("openai-codex streaming", () => { ); global.fetch = fetchMock as unknown as typeof fetch; - type WsListener = (event: Event) => void; - class MalformedMessageWebSocket { - static readonly CONNECTING = 0; - static readonly OPEN = 1; - static readonly CLOSING = 2; - static readonly CLOSED = 3; - readyState = MalformedMessageWebSocket.CONNECTING; - #listeners = new Map>(); - - constructor(_url: string, _options?: { headers?: Record }) { - setTimeout(() => { - this.readyState = MalformedMessageWebSocket.OPEN; - this.#emit("open", new Event("open")); - }, 0); - } - - addEventListener(type: string, listener: unknown): void { - if (typeof listener !== "function") return; - const listeners = this.#listeners.get(type) ?? new Set(); - listeners.add(listener as WsListener); - this.#listeners.set(type, listeners); - } - - removeEventListener(type: string, listener: unknown): void { - if (typeof listener !== "function") return; - const listeners = this.#listeners.get(type); - listeners?.delete(listener as WsListener); + class MalformedMessageWebSocket extends MockWebSocket { + constructor(url: string, options?: { headers?: WsHeaders }) { + super(url, options); + this.scheduleOpen(); } send(): void { - this.#emit("message", { data: "{" } as unknown as Event); - } - - close(): void { - this.readyState = MalformedMessageWebSocket.CLOSED; - } - - #emit(type: string, event: Event): void { - const listeners = this.#listeners.get(type); - if (!listeners) return; - for (const listener of listeners) { - listener(event); - } + this.sendMessage("{"); } } @@ -2755,68 +2082,27 @@ describe("openai-codex streaming", () => { ); global.fetch = fetchMock as unknown as typeof fetch; - type WsListener = (event: Event) => void; - class BufferedCloseWebSocket { - static readonly CONNECTING = 0; - static readonly OPEN = 1; - static readonly CLOSING = 2; - static readonly CLOSED = 3; - readyState = BufferedCloseWebSocket.CONNECTING; - #listeners = new Map>(); - - constructor(_url: string, _options?: { headers?: Record }) { - setTimeout(() => { - this.readyState = BufferedCloseWebSocket.OPEN; - this.#emit("open", new Event("open")); - }, 0); - } - - addEventListener(type: string, listener: unknown): void { - if (typeof listener !== "function") return; - const listeners = this.#listeners.get(type) ?? new Set(); - listeners.add(listener as WsListener); - this.#listeners.set(type, listeners); - } - - removeEventListener(type: string, listener: unknown): void { - if (typeof listener !== "function") return; - const listeners = this.#listeners.get(type); - listeners?.delete(listener as WsListener); + class BufferedCloseWebSocket extends MockWebSocket { + constructor(url: string, options?: { headers?: WsHeaders }) { + super(url, options); + this.scheduleOpen(); } send(): void { - this.#emit("message", { - data: JSON.stringify({ - type: "response.output_item.added", - item: { - type: "message", - id: "msg_ws_partial", - role: "assistant", - status: "in_progress", - content: [], - }, - }), - } as unknown as Event); - this.#emit("message", { - data: JSON.stringify({ type: "response.content_part.added", part: { type: "output_text", text: "" } }), - } as unknown as Event); - this.#emit("message", { - data: JSON.stringify({ type: "response.output_text.delta", delta: "Partial output" }), - } as unknown as Event); - this.readyState = BufferedCloseWebSocket.CLOSED; - this.#emit("close", { code: 1006 } as unknown as Event); - } - - close(): void { - this.readyState = BufferedCloseWebSocket.CLOSED; - } - - #emit(type: string, event: Event): void { - const listeners = this.#listeners.get(type); - if (!listeners) return; - for (const listener of listeners) { - listener(event); - } + this.sendJson({ + type: "response.output_item.added", + item: { + type: "message", + id: "msg_ws_partial", + role: "assistant", + status: "in_progress", + content: [], + }, + }); + this.sendJson({ type: "response.content_part.added", part: { type: "output_text", text: "" } }); + this.sendJson({ type: "response.output_text.delta", delta: "Partial output" }); + this.readyState = MockWebSocket.CLOSED; + this.emit("close", { code: 1006 } as unknown as Event); } } @@ -2879,39 +2165,17 @@ describe("openai-codex streaming", () => { }); global.fetch = fetchMock as unknown as typeof fetch; - type WsListener = (event: Event) => void; const requestTypes: string[] = []; - class DivergedAppendWebSocket { - static readonly CONNECTING = 0; - static readonly OPEN = 1; - static readonly CLOSING = 2; - static readonly CLOSED = 3; - readyState = DivergedAppendWebSocket.CONNECTING; + class DivergedAppendWebSocket extends MockWebSocket { handshakeHeaders = { "x-codex-turn-state": "ws-turn-state-1", "x-models-etag": "ws-models-etag-1", }; - #listeners = new Map>(); #sendCount = 0; - constructor(_url: string, _options?: { headers?: Record }) { - setTimeout(() => { - this.readyState = DivergedAppendWebSocket.OPEN; - this.#emit("open", new Event("open")); - }, 0); - } - - addEventListener(type: string, listener: unknown): void { - if (typeof listener !== "function") return; - const listeners = this.#listeners.get(type) ?? new Set(); - listeners.add(listener as WsListener); - this.#listeners.set(type, listeners); - } - - removeEventListener(type: string, listener: unknown): void { - if (typeof listener !== "function") return; - const listeners = this.#listeners.get(type); - listeners?.delete(listener as WsListener); + constructor(url: string, options?: { headers?: WsHeaders }) { + super(url, options); + this.scheduleOpen(); } send(data: string): void { @@ -2919,61 +2183,11 @@ describe("openai-codex streaming", () => { const request = JSON.parse(data) as { type?: string }; requestTypes.push(typeof request.type === "string" ? request.type : ""); const idSuffix = String(this.#sendCount); - this.#emitCompleted(`msg_${idSuffix}`, `resp_${idSuffix}`, `Hello WS ${idSuffix}`); - } - - close(): void { - this.readyState = DivergedAppendWebSocket.CLOSED; - } - - #emitCompleted(messageId: string, responseId: string, text: string): void { - this.#emit("message", { - data: JSON.stringify({ - type: "response.output_item.added", - item: { type: "message", id: messageId, role: "assistant", status: "in_progress", content: [] }, - }), - } as unknown as Event); - this.#emit("message", { - data: JSON.stringify({ type: "response.content_part.added", part: { type: "output_text", text: "" } }), - } as unknown as Event); - this.#emit("message", { - data: JSON.stringify({ type: "response.output_text.delta", delta: text }), - } as unknown as Event); - this.#emit("message", { - data: JSON.stringify({ - type: "response.output_item.done", - item: { - type: "message", - id: messageId, - role: "assistant", - status: "completed", - content: [{ type: "output_text", text }], - }, - }), - } as unknown as Event); - this.#emit("message", { - data: JSON.stringify({ - type: "response.done", - response: { - id: responseId, - status: "completed", - usage: { - input_tokens: 5, - output_tokens: 3, - total_tokens: 8, - input_tokens_details: { cached_tokens: 0 }, - }, - }, - }), - } as unknown as Event); - } - - #emit(type: string, event: Event): void { - const listeners = this.#listeners.get(type); - if (!listeners) return; - for (const listener of listeners) { - listener(event); - } + this.emitCodexResponse({ + messageId: `msg_${idSuffix}`, + responseId: `resp_${idSuffix}`, + text: `Hello WS ${idSuffix}`, + }); } } @@ -3043,103 +2257,24 @@ describe("openai-codex streaming", () => { }); global.fetch = fetchMock as unknown as typeof fetch; - type WsListener = (event: Event) => void; let constructorCount = 0; let sendCount = 0; - class ReusableWebSocket { - static readonly CONNECTING = 0; - static readonly OPEN = 1; - static readonly CLOSING = 2; - static readonly CLOSED = 3; - - readyState = ReusableWebSocket.CONNECTING; - #listeners = new Map>(); - - constructor( - public readonly url: string, - public readonly options?: { headers?: Record }, - ) { + class ReusableWebSocket extends MockWebSocket { + constructor(url: string, options?: { headers?: WsHeaders }) { + super(url, options); constructorCount += 1; - setTimeout(() => { - this.readyState = ReusableWebSocket.OPEN; - this.#emit("open", new Event("open")); - }, 0); - } - - addEventListener(type: string, listener: unknown): void { - if (typeof listener !== "function") return; - const listeners = this.#listeners.get(type) ?? new Set(); - listeners.add(listener as WsListener); - this.#listeners.set(type, listeners); - } - - removeEventListener(type: string, listener: unknown): void { - if (typeof listener !== "function") return; - const listeners = this.#listeners.get(type); - listeners?.delete(listener as WsListener); + this.scheduleOpen(); } send(data: string): void { sendCount += 1; const request = JSON.parse(data) as Record; expect(typeof request.type).toBe("string"); - this.#emit("message", { - data: JSON.stringify({ - type: "response.output_item.added", - item: { - type: "message", - id: `msg_${sendCount}`, - role: "assistant", - status: "in_progress", - content: [], - }, - }), - } as unknown as Event); - this.#emit("message", { - data: JSON.stringify({ type: "response.content_part.added", part: { type: "output_text", text: "" } }), - } as unknown as Event); - this.#emit("message", { - data: JSON.stringify({ type: "response.output_text.delta", delta: `Hello ${sendCount}` }), - } as unknown as Event); - this.#emit("message", { - data: JSON.stringify({ - type: "response.output_item.done", - item: { - type: "message", - id: `msg_${sendCount}`, - role: "assistant", - status: "completed", - content: [{ type: "output_text", text: `Hello ${sendCount}` }], - }, - }), - } as unknown as Event); - this.#emit("message", { - data: JSON.stringify({ - type: "response.done", - response: { - id: `resp_${sendCount}`, - status: "completed", - usage: { - input_tokens: 5, - output_tokens: 3, - total_tokens: 8, - input_tokens_details: { cached_tokens: 0 }, - }, - }, - }), - } as unknown as Event); - } - - close(): void { - this.readyState = ReusableWebSocket.CLOSED; - } - - #emit(type: string, event: Event): void { - const listeners = this.#listeners.get(type); - if (!listeners) return; - for (const listener of listeners) { - listener(event); - } + this.emitCodexResponse({ + messageId: `msg_${sendCount}`, + responseId: `resp_${sendCount}`, + text: `Hello ${sendCount}`, + }); } } diff --git a/packages/coding-agent/src/utils/title-generator.ts b/packages/coding-agent/src/utils/title-generator.ts index d25f2155b..d29c1c664 100644 --- a/packages/coding-agent/src/utils/title-generator.ts +++ b/packages/coding-agent/src/utils/title-generator.ts @@ -2,7 +2,7 @@ * Generate session titles using a smol, fast model. */ import * as path from "node:path"; -import type { ThinkingLevel } from "@oh-my-pi/pi-agent-core"; + import { type Api, completeSimple, type Model } from "@oh-my-pi/pi-ai"; import { logger, prompt } from "@oh-my-pi/pi-utils"; import type { ModelRegistry } from "../config/model-registry"; @@ -17,22 +17,14 @@ const TERMINAL_TITLE_CONTROL_CHARS = /[\u0000-\u001f\u007f-\u009f]/g; const MAX_INPUT_CHARS = 2000; -function getTitleModel( - registry: ModelRegistry, - settings: Settings, - currentModel?: Model, -): { model: Model; thinkingLevel?: ThinkingLevel } | undefined { +function getTitleModel(registry: ModelRegistry, settings: Settings, currentModel?: Model): Model | undefined { const availableModels = registry.getAvailable(); if (availableModels.length === 0) return undefined; - const titleModel = resolveRoleSelection(["commit", "smol"], settings, availableModels, registry); - if (titleModel) { - return { model: titleModel.model, thinkingLevel: titleModel.thinkingLevel }; - } + const titleModel = resolveRoleSelection(["commit", "smol"], settings, availableModels, registry)?.model; + if (titleModel) return titleModel; - if (currentModel) { - return { model: currentModel }; - } + if (currentModel) return currentModel; return undefined; } @@ -42,7 +34,7 @@ function getTitleModel( * * @param firstMessage The first user message * @param registry Model registry - * @param settings Settings used to resolve the smol role, including per-role thinking + * @param settings Settings used to resolve the smol role * @param sessionId Optional session id for sticky API key selection */ export async function generateSessionTitle( @@ -52,8 +44,8 @@ export async function generateSessionTitle( sessionId?: string, currentModel?: Model, ): Promise { - const candidate = getTitleModel(registry, settings, currentModel); - if (!candidate) { + const model = getTitleModel(registry, settings, currentModel); + if (!model) { logger.debug("title-generator: no title model found"); return null; } @@ -65,11 +57,11 @@ export async function generateSessionTitle( ${truncatedMessage} `; - const apiKey = await registry.getApiKey(candidate.model, sessionId); + const apiKey = await registry.getApiKey(model, sessionId); if (!apiKey) { logger.debug("title-generator: no API key for smol model", { - provider: candidate.model.provider, - id: candidate.model.id, + provider: model.provider, + id: model.id, }); return null; } @@ -78,7 +70,7 @@ ${truncatedMessage} // don't burn the entire output budget on internal thinking and return an empty // string. With reasoning disabled, 30 tokens of output is plenty. const request = { - model: `${candidate.model.provider}/${candidate.model.id}`, + model: `${model.provider}/${model.id}`, systemPrompt: TITLE_SYSTEM_PROMPT, userMessage, maxTokens: 30, @@ -87,7 +79,7 @@ ${truncatedMessage} try { const response = await completeSimple( - candidate.model, + model, { systemPrompt: request.systemPrompt, messages: [{ role: "user", content: request.userMessage, timestamp: Date.now() }], diff --git a/packages/coding-agent/test/role-thinking-helper-propagation.test.ts b/packages/coding-agent/test/role-thinking-helper-propagation.test.ts index 638170cc3..90f96dbf7 100644 --- a/packages/coding-agent/test/role-thinking-helper-propagation.test.ts +++ b/packages/coding-agent/test/role-thinking-helper-propagation.test.ts @@ -50,7 +50,7 @@ describe("role thinking helper propagation", () => { expect(completeSimpleMock.mock.calls[0]?.[2]).toMatchObject({ reasoning: Effort.Minimal }); }); - it("passes smol-role thinking to title generation", async () => { + it("disables reasoning for title generation even when smol role has thinking", async () => { const model = getModelOrThrow("claude-sonnet-4-5"); const settings = createSettings({ default: `${model.provider}/${model.id}:high`, @@ -67,6 +67,6 @@ describe("role thinking helper propagation", () => { const title = await generateSessionTitle("Investigate resolver", registry as never, settings); expect(title).toBe("Investigate resolver"); - expect(completeSimpleMock.mock.calls[0]?.[2]).toMatchObject({ reasoning: Effort.Low }); + expect(completeSimpleMock.mock.calls[0]?.[2]).toMatchObject({ disableReasoning: true }); }); });