diff --git a/packages/agent/CHANGELOG.md b/packages/agent/CHANGELOG.md index e19ded1f3..acb3cbc0e 100644 --- a/packages/agent/CHANGELOG.md +++ b/packages/agent/CHANGELOG.md @@ -1,6 +1,9 @@ # Changelog ## [Unreleased] +### Added + +- Added `preferWebsockets` option to hint that websocket transport should be preferred when supported by the provider implementation ## [11.10.0] - 2026-02-10 ### Added diff --git a/packages/agent/src/agent.ts b/packages/agent/src/agent.ts index 533c2d929..799acb06b 100644 --- a/packages/agent/src/agent.ts +++ b/packages/agent/src/agent.ts @@ -74,6 +74,9 @@ export interface AgentOptions { */ kimiApiFormat?: "openai" | "anthropic"; + /** Hint that websocket transport should be preferred when supported by the provider implementation. */ + preferWebsockets?: boolean; + /** * Custom stream function (for proxy backends, etc.). Default uses streamSimple. */ @@ -168,6 +171,7 @@ export class Agent { #runningPrompt?: Promise; #resolveRunningPrompt?: () => void; #kimiApiFormat?: "openai" | "anthropic"; + #preferWebsockets?: boolean; /** Buffered Cursor tool results with text length at time of call (for correct ordering) */ #cursorToolResultBuffer: CursorToolResultEntry[] = []; @@ -192,6 +196,7 @@ export class Agent { this.#cursorExecHandlers = opts.cursorExecHandlers; this.#cursorOnToolResult = opts.cursorOnToolResult; this.#kimiApiFormat = opts.kimiApiFormat; + this.#preferWebsockets = opts.preferWebsockets; } /** @@ -591,6 +596,7 @@ export class Agent { thinkingBudgets: this.#thinkingBudgets, maxRetryDelayMs: this.#maxRetryDelayMs, kimiApiFormat: this.#kimiApiFormat, + preferWebsockets: this.#preferWebsockets, toolChoice: options?.toolChoice, convertToLlm: this.#convertToLlm, transformContext: this.#transformContext, diff --git a/packages/ai/CHANGELOG.md b/packages/ai/CHANGELOG.md index f0182d10a..ef5b4672f 100644 --- a/packages/ai/CHANGELOG.md +++ b/packages/ai/CHANGELOG.md @@ -1,8 +1,21 @@ # Changelog ## [Unreleased] + +### Added + +- Added `preferWebsockets` option to enable WebSocket transport for OpenAI Codex responses when supported +- Added `prewarmOpenAICodexResponses()` function to establish and reuse WebSocket connections across multiple requests +- Added `getOpenAICodexTransportDetails()` function to inspect transport layer details including WebSocket status and fallback information +- Added `getProviderDetails()` function to retrieve formatted provider configuration and transport information +- Added automatic fallback from WebSocket to SSE when connection fails, with transparent retry logic +- Added session state management to reuse WebSocket connections and enable request appending across turns +- Added support for x-codex-turn-state header to maintain conversation state across SSE requests + ### Changed +- Changed OpenAI Codex model configuration to prefer WebSocket transport by default with `preferWebsockets: true` +- Changed header handling to use appropriate OpenAI-Beta header values for WebSocket vs SSE transports - Perplexity OAuth token refresh now uses JWT expiry extraction instead of Socket.IO RPC, improving reliability when server is unreachable - Removed Socket.IO client implementation for Perplexity token refresh; tokens are now validated using embedded JWT expiry claims diff --git a/packages/ai/scripts/generate-models.ts b/packages/ai/scripts/generate-models.ts index 581bf0e52..6d3a2de8b 100644 --- a/packages/ai/scripts/generate-models.ts +++ b/packages/ai/scripts/generate-models.ts @@ -1285,6 +1285,7 @@ async function generateModels() { provider: "openai-codex", baseUrl: CODEX_BASE_URL, reasoning: true, + preferWebsockets: true, input: ["text"], cost: { input: 1.75, output: 14, cacheRead: 0.175, cacheWrite: 0 }, contextWindow: 128000, diff --git a/packages/ai/src/index.ts b/packages/ai/src/index.ts index 77fea2d4c..a4f8d4c4c 100644 --- a/packages/ai/src/index.ts +++ b/packages/ai/src/index.ts @@ -1,6 +1,7 @@ export type { Static, TSchema } from "@sinclair/typebox"; export { Type } from "@sinclair/typebox"; export * from "./models"; +export * from "./provider-details"; export * from "./providers/anthropic"; export * from "./providers/azure-openai-responses"; export * from "./providers/cursor"; diff --git a/packages/ai/src/models.json b/packages/ai/src/models.json index 6ef33b305..949b1ba0a 100644 --- a/packages/ai/src/models.json +++ b/packages/ai/src/models.json @@ -12965,6 +12965,7 @@ "provider": "openai-codex", "baseUrl": "https://chatgpt.com/backend-api", "reasoning": true, + "preferWebsockets": true, "input": [ "text" ], diff --git a/packages/ai/src/provider-details.ts b/packages/ai/src/provider-details.ts new file mode 100644 index 000000000..624e471b6 --- /dev/null +++ b/packages/ai/src/provider-details.ts @@ -0,0 +1,82 @@ +import { getOpenAICodexTransportDetails, type OpenAICodexTransportDetails } from "./providers/openai-codex-responses"; +import type { Api, Model, Provider } from "./types"; + +export interface ProviderDetailField { + label: string; + value: string; +} + +export interface ProviderDetails { + provider: Provider; + api: Api; + fields: ProviderDetailField[]; +} + +export interface ProviderDetailsContext { + model: Model; + sessionId?: string; + authMode?: string; + preferWebsockets?: boolean; +} + +export function getProviderDetails(context: ProviderDetailsContext): ProviderDetails { + const endpoint = formatEndpoint(context.model.baseUrl); + const fields: ProviderDetailField[] = [ + { label: "Model", value: context.model.id }, + { label: "API", value: context.model.api }, + { label: "Auth", value: context.authMode ?? "auto" }, + { label: "Endpoint", value: endpoint }, + ]; + + if (context.model.api === "openai-codex-responses") { + const codexDetails = getOpenAICodexTransportDetails(context.model as Model<"openai-codex-responses">, { + sessionId: context.sessionId, + baseUrl: context.model.baseUrl, + preferWebsockets: context.preferWebsockets, + }); + fields.push({ label: "Transport", value: formatCodexTransport(codexDetails) }); + fields.push({ label: "WebSocket", value: formatCodexWebSocket(codexDetails) }); + fields.push({ label: "Reuse", value: formatCodexReuse(codexDetails, context.sessionId) }); + } + + return { + provider: context.model.provider, + api: context.model.api, + fields, + }; +} + +function formatEndpoint(baseUrl: string): string { + try { + const parsed = new URL(baseUrl); + const path = parsed.pathname.replace(/\/$/, ""); + return `${parsed.origin}${path || "/"}`; + } catch { + return baseUrl; + } +} + +function formatCodexTransport(details: OpenAICodexTransportDetails): string { + if (details.lastTransport === "websocket") return "websocket"; + if (details.lastTransport === "sse" && (details.websocketDisabled || details.fallbackCount > 0)) { + return "sse (fallback)"; + } + if (details.lastTransport === "sse") return "sse"; + return details.websocketPreferred ? "websocket preferred" : "sse"; +} + +function formatCodexWebSocket(details: OpenAICodexTransportDetails): string { + if (!details.websocketPreferred) return "off"; + if (details.websocketDisabled) return "disabled after fallback"; + if (details.websocketConnected) return "connected"; + if (details.prewarmed) return "prewarmed"; + return details.hasSessionState ? "enabled" : "waiting for first request"; +} + +function formatCodexReuse( + details: OpenAICodexTransportDetails, + sessionId: string | undefined, +): string { + if (!sessionId) return "no session key"; + return details.canAppend ? "append enabled" : "full request"; +} diff --git a/packages/ai/src/providers/openai-codex-responses.ts b/packages/ai/src/providers/openai-codex-responses.ts index f5b611920..1c8c9ae7b 100644 --- a/packages/ai/src/providers/openai-codex-responses.ts +++ b/packages/ai/src/providers/openai-codex-responses.ts @@ -37,7 +37,12 @@ import { OPENAI_HEADERS, URL_PATHS, } from "./openai-codex/constants"; -import { type CodexRequestOptions, type RequestBody, transformRequestBody } from "./openai-codex/request-transformer"; +import { + type CodexRequestOptions, + type InputItem, + type RequestBody, + transformRequestBody, +} from "./openai-codex/request-transformer"; import { parseCodexError } from "./openai-codex/response-handler"; import { transformMessages } from "./transform-messages"; @@ -76,6 +81,30 @@ const CODEX_MAX_RETRIES = 5; const CODEX_RETRYABLE_STATUS = new Set([408, 429, 500, 502, 503, 504]); const CODEX_RETRY_DELAY_MS = 500; +function isCodexWebSocketEnvEnabled(): boolean { + return $env.PI_CODEX_WEBSOCKET === "1" || $env.PI_CODEX_WEBSOCKET === "true"; +} +const CODEX_WEBSOCKET_V2_ENABLED = $env.PI_CODEX_WEBSOCKET_V2 === "1" || $env.PI_CODEX_WEBSOCKET_V2 === "true"; + +type CodexWebSocketSessionState = { + disableWebsocket: boolean; + lastRequest?: RequestBody; + lastResponseId?: string; + canAppend: boolean; + turnState?: string; + modelsEtag?: string; + reasoningIncluded?: boolean; + connection?: CodexWebSocketConnection; + lastTransport?: "sse" | "websocket"; + fallbackCount: number; + lastFallbackAt?: number; + prewarmed: boolean; +}; +const codexWebSocketSessions = new Map(); +const codexWebSocketPublicToPrivate = new Map(); +const X_CODEX_TURN_STATE_HEADER = "x-codex-turn-state"; +const X_MODELS_ETAG_HEADER = "x-models-etag"; +const X_REASONING_INCLUDED_HEADER = "x-reasoning-included"; function normalizeResponsesToolCallId(id: string): { callId: string; itemId: string } { const [callId, itemId] = id.split("|"); if (callId && itemId) { @@ -185,56 +214,95 @@ export const streamOpenAICodexResponses: StreamFunction<"openai-codex-responses" options?.onPayload?.(transformedBody); const reasoningEffort = transformedBody.reasoning?.effort ?? null; - const headers = createCodexHeaders( - { ...(model.headers ?? {}), ...(options?.headers ?? {}) }, - accountId, - apiKey, - options?.sessionId, - ); - logCodexDebug("codex request", { - url, - model: params.model, - reasoningEffort, - headers: redactHeaders(headers), - }); - - const response = await fetchWithRetry( - url, - { - method: "POST", - headers, - body: JSON.stringify(transformedBody), - }, - options?.signal, - ); - - logCodexDebug("codex response", { - url: response.url, - status: response.status, - statusText: response.statusText, - contentType: response.headers.get("content-type") || null, - cfRay: response.headers.get("cf-ray") || null, - }); - - if (!response.ok) { - const info = await parseCodexError(response); - const error = new Error(info.friendlyMessage || info.message); - (error as { headers?: Headers }).headers = response.headers; - throw error; + const requestHeaders = { ...(model.headers ?? {}), ...(options?.headers ?? {}) }; + const sessionKey = getCodexWebSocketSessionKey(options?.sessionId, model, accountId, baseUrl); + const publicSessionKey = getCodexPublicSessionKey(options?.sessionId, model, baseUrl); + if (sessionKey && publicSessionKey) { + codexWebSocketPublicToPrivate.set(publicSessionKey, sessionKey); } + const websocketState = sessionKey ? getCodexWebSocketSessionState(sessionKey) : undefined; + let usingWebsocket = false; + let requestBodyForState = cloneRequestBody(transformedBody); + let eventStream: AsyncGenerator>; - if (!response.body) { - throw new Error("No response body"); + if (websocketState && shouldUseCodexWebSocket(model, websocketState, options?.preferWebsockets)) { + const websocketHeaders = createCodexHeaders( + requestHeaders, + accountId, + apiKey, + options?.sessionId, + "websocket", + websocketState, + ); + const websocketRequest = buildCodexWebSocketRequest( + transformedBody, + websocketState, + CODEX_WEBSOCKET_V2_ENABLED, + ); + requestBodyForState = cloneRequestBody(transformedBody); + logCodexDebug("codex websocket request", { + url: toWebSocketUrl(url), + model: params.model, + reasoningEffort, + headers: redactHeaders(websocketHeaders), + requestType: websocketRequest.type, + }); + try { + eventStream = await openCodexWebSocketEventStream( + toWebSocketUrl(url), + websocketHeaders, + websocketRequest, + websocketState, + options?.signal, + ); + usingWebsocket = true; + } catch (error) { + if (websocketState) { + websocketState.disableWebsocket = true; + websocketState.fallbackCount += 1; + websocketState.lastFallbackAt = Date.now(); + websocketState.canAppend = false; + websocketState.lastRequest = undefined; + websocketState.lastResponseId = undefined; + websocketState.connection?.close("fallback"); + websocketState.connection = undefined; + } + logCodexDebug("codex websocket fallback", { + error: error instanceof Error ? error.message : String(error), + }); + eventStream = await openCodexSseEventStream( + url, + requestHeaders, + accountId, + apiKey, + options?.sessionId, + transformedBody, + websocketState, + options?.signal, + ); + } + } else { + eventStream = await openCodexSseEventStream( + url, + requestHeaders, + accountId, + apiKey, + options?.sessionId, + transformedBody, + websocketState, + options?.signal, + ); + } + if (websocketState) { + websocketState.lastTransport = usingWebsocket ? "websocket" : "sse"; } stream.push({ type: "start", partial: output }); - let currentItem: ResponseReasoningItem | ResponseOutputMessage | ResponseFunctionToolCall | null = null; let currentBlock: ThinkingContent | TextContent | (ToolCall & { partialJson: string }) | null = null; const blocks = output.content; const blockIndex = () => blocks.length - 1; - - for await (const rawEvent of readSseJson>(response.body!, options?.signal)) { + for await (const rawEvent of eventStream) { const eventType = typeof rawEvent.type === "string" ? rawEvent.type : ""; if (!eventType) continue; @@ -394,10 +462,18 @@ export const streamOpenAICodexResponses: StreamFunction<"openai-codex-responses" }; stream.push({ type: "toolcall_end", contentIndex: blockIndex(), toolCall, partial: output }); } + } else if (eventType === "response.created") { + if (usingWebsocket && websocketState) { + const createdResponse = (rawEvent as { response?: { id?: string } }).response; + if (typeof createdResponse?.id === "string" && createdResponse.id.length > 0) { + websocketState.lastResponseId = createdResponse.id; + } + } } else if (eventType === "response.completed" || eventType === "response.done") { const response = ( rawEvent as { response?: { + id?: string; usage?: { input_tokens?: number; output_tokens?: number; @@ -419,6 +495,13 @@ export const streamOpenAICodexResponses: StreamFunction<"openai-codex-responses" cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 }, }; } + if (usingWebsocket && websocketState) { + websocketState.lastRequest = cloneRequestBody(requestBodyForState); + if (typeof response?.id === "string" && response.id.length > 0) { + websocketState.lastResponseId = response.id; + } + websocketState.canAppend = eventType === "response.done"; + } calculateCost(model, output.usage); output.stopReason = mapStopReason(response?.status); if (output.content.some(b => b.type === "toolCall") && output.stopReason === "stop") { @@ -459,20 +542,456 @@ export const streamOpenAICodexResponses: StreamFunction<"openai-codex-responses" return stream; }; +export async function prewarmOpenAICodexResponses( + model: Model<"openai-codex-responses">, + options?: Pick, +): Promise { + const apiKey = options?.apiKey || getEnvApiKey(model.provider) || ""; + if (!apiKey) return; + const accountId = getAccountId(apiKey); + const baseUrl = model.baseUrl || CODEX_BASE_URL; + const baseWithSlash = baseUrl.endsWith("/") ? baseUrl : `${baseUrl}/`; + const url = rewriteUrlForCodex(new URL(URL_PATHS.RESPONSES.slice(1), baseWithSlash).toString()); + const sessionKey = getCodexWebSocketSessionKey(options?.sessionId, model, accountId, baseUrl); + const publicSessionKey = getCodexPublicSessionKey(options?.sessionId, model, baseUrl); + if (publicSessionKey && sessionKey) { + codexWebSocketPublicToPrivate.set(publicSessionKey, sessionKey); + } + if (!sessionKey) return; + const state = getCodexWebSocketSessionState(sessionKey); + if (!shouldUseCodexWebSocket(model, state, options?.preferWebsockets)) return; + const headers = createCodexHeaders( + { ...(model.headers ?? {}), ...(options?.headers ?? {}) }, + accountId, + apiKey, + options?.sessionId, + "websocket", + state, + ); + await getOrCreateCodexWebSocketConnection(state, toWebSocketUrl(url), headers, options?.signal); + state.prewarmed = true; +} + +function cloneRequestBody(body: RequestBody): RequestBody { + return JSON.parse(JSON.stringify(body)) as RequestBody; +} + +function getCodexWebSocketSessionKey( + sessionId: string | undefined, + model: Model<"openai-codex-responses">, + accountId: string, + baseUrl: string, +): string | undefined { + if (!sessionId || sessionId.length === 0) return undefined; + return `${accountId}:${baseUrl}:${model.id}:${sessionId}`; +} + +function getCodexPublicSessionKey( + sessionId: string | undefined, + model: Model<"openai-codex-responses">, + baseUrl: string, +): string | undefined { + if (!sessionId || sessionId.length === 0) return undefined; + return `${baseUrl}:${model.id}:${sessionId}`; +} + +function getCodexWebSocketSessionState(sessionKey: string): CodexWebSocketSessionState { + const existing = codexWebSocketSessions.get(sessionKey); + if (existing) return existing; + const created: CodexWebSocketSessionState = { disableWebsocket: false, canAppend: false, fallbackCount: 0, prewarmed: false }; + codexWebSocketSessions.set(sessionKey, created); + return created; +} + +function shouldUseCodexWebSocket( + model: Model<"openai-codex-responses">, + state: CodexWebSocketSessionState | undefined, + preferWebsockets?: boolean, +): boolean { + if (!state || state.disableWebsocket) return false; + return isCodexWebSocketEnvEnabled() || preferWebsockets === true || model.preferWebsockets === true; +} + +export interface OpenAICodexTransportDetails { + websocketPreferred: boolean; + lastTransport?: "sse" | "websocket"; + websocketDisabled: boolean; + websocketConnected: boolean; + fallbackCount: number; + canAppend: boolean; + prewarmed: boolean; + hasSessionState: boolean; + lastFallbackAt?: number; +} + +export function getOpenAICodexTransportDetails( + model: Model<"openai-codex-responses">, + options?: { sessionId?: string; baseUrl?: string; preferWebsockets?: boolean }, +): OpenAICodexTransportDetails { + const baseUrl = options?.baseUrl || model.baseUrl || CODEX_BASE_URL; + const websocketPreferred = isCodexWebSocketEnvEnabled() || options?.preferWebsockets === true || model.preferWebsockets === true; + const publicSessionKey = getCodexPublicSessionKey(options?.sessionId, model, baseUrl); + const privateSessionKey = publicSessionKey ? codexWebSocketPublicToPrivate.get(publicSessionKey) : undefined; + const state = privateSessionKey ? codexWebSocketSessions.get(privateSessionKey) : undefined; + + return { + websocketPreferred, + lastTransport: state?.lastTransport, + websocketDisabled: state?.disableWebsocket ?? false, + websocketConnected: state?.connection?.isOpen() ?? false, + fallbackCount: state?.fallbackCount ?? 0, + canAppend: state?.canAppend ?? false, + prewarmed: state?.prewarmed ?? false, + hasSessionState: state !== undefined, + lastFallbackAt: state?.lastFallbackAt, + }; +} + +function buildAppendInput(previous: RequestBody | undefined, current: RequestBody): InputItem[] | null { + if (!previous) return null; + if (!Array.isArray(previous.input) || !Array.isArray(current.input)) return null; + if (current.input.length <= previous.input.length) return null; + const previousWithoutInput = { ...previous, input: undefined }; + const currentWithoutInput = { ...current, input: undefined }; + if (JSON.stringify(previousWithoutInput) !== JSON.stringify(currentWithoutInput)) { + return null; + } + for (let index = 0; index < previous.input.length; index += 1) { + if (JSON.stringify(previous.input[index]) !== JSON.stringify(current.input[index])) { + return null; + } + } + return current.input.slice(previous.input.length) as InputItem[]; +} + +function buildCodexWebSocketRequest( + requestBody: RequestBody, + state: CodexWebSocketSessionState | undefined, + v2Enabled: boolean, +): Record { + const appendInput = state?.canAppend ? buildAppendInput(state.lastRequest, requestBody) : null; + if (appendInput && appendInput.length > 0) { + if (v2Enabled && state?.lastResponseId) { + return { + type: "response.create", + ...requestBody, + previous_response_id: state.lastResponseId, + input: appendInput, + }; + } + return { + type: "response.append", + input: appendInput, + }; + } + return { + type: "response.create", + ...requestBody, + }; +} + +function toWebSocketUrl(url: string): string { + const parsed = new URL(url); + if (parsed.protocol === "https:") { + parsed.protocol = "wss:"; + } else if (parsed.protocol === "http:") { + parsed.protocol = "ws:"; + } + return parsed.toString(); +} + +function headersToRecord(headers: Headers): Record { + const result: Record = {}; + for (const [key, value] of headers.entries()) { + result[key] = value; + } + return result; +} + +class CodexWebSocketConnection { + #url: string; + #headers: Record; + #socket: WebSocket | null = null; + #queue: Array | Error | null> = []; + #waiters: Array<() => void> = []; + #connectPromise?: Promise; + #activeRequest = false; + + constructor(url: string, headers: Record) { + this.#url = url; + this.#headers = headers; + } + + isOpen(): boolean { + return this.#socket?.readyState === WebSocket.OPEN; + } + + close(reason = "done"): void { + if (this.#socket && (this.#socket.readyState === WebSocket.OPEN || this.#socket.readyState === WebSocket.CONNECTING)) { + this.#socket.close(1000, reason); + } + this.#socket = null; + } + + async connect(signal?: AbortSignal): Promise { + if (this.isOpen()) return; + if (this.#connectPromise) { + 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 }); + this.#socket = socket; + let settled = false; + let timeout: NodeJS.Timeout | undefined; + const onAbort = () => { + socket.close(1000, "aborted"); + if (!settled) { + settled = true; + reject(new Error("Request was aborted")); + } + }; + if (signal) { + if (signal.aborted) { + onAbort(); + } else { + signal.addEventListener("abort", onAbort, { once: true }); + } + } + + const clearPending = () => { + if (timeout) clearTimeout(timeout); + if (signal) signal.removeEventListener("abort", onAbort); + }; + + timeout = setTimeout(() => { + socket.close(1000, "connect-timeout"); + if (!settled) { + settled = true; + reject(new Error("WebSocket connection timeout")); + } + }, 10000); + + socket.addEventListener("open", () => { + if (!settled) { + settled = true; + clearPending(); + resolve(); + } + }); + + socket.addEventListener("error", event => { + const error = new Error(`WebSocket error: ${String(event.type)}`); + if (!settled) { + settled = true; + clearPending(); + reject(error); + return; + } + this.#push(error); + }); + + socket.addEventListener("close", event => { + this.#socket = null; + if (!settled) { + settled = true; + clearPending(); + reject(new Error(`WebSocket closed before open (${event.code})`)); + return; + } + this.#push(new Error(`WebSocket closed (${event.code})`)); + this.#push(null); + }); + + socket.addEventListener("message", event => { + if (typeof event.data !== "string") return; + try { + const parsed = JSON.parse(event.data) 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(error instanceof Error ? error : new Error(String(error))); + } + }); + + try { + await promise; + } finally { + this.#connectPromise = undefined; + } + } + + async *streamRequest(request: Record, signal?: AbortSignal): AsyncGenerator> { + if (!this.#socket || this.#socket.readyState !== WebSocket.OPEN) { + throw new Error("websocket connection is unavailable"); + } + if (this.#activeRequest) { + throw new Error("websocket request already in progress"); + } + + this.#activeRequest = true; + const onAbort = () => { + this.close("aborted"); + this.#push(new Error("Request was aborted")); + }; + if (signal) { + if (signal.aborted) { + onAbort(); + } else { + signal.addEventListener("abort", onAbort, { once: true }); + } + } + + try { + this.#socket.send(JSON.stringify(request)); + while (true) { + const next = await this.#nextMessage(); + if (next instanceof Error) { + throw next; + } + if (next === null) { + throw new Error("websocket closed before response completion"); + } + yield next; + const eventType = typeof next.type === "string" ? next.type : ""; + if ( + eventType === "response.completed" || + eventType === "response.done" || + eventType === "response.failed" || + eventType === "error" + ) { + break; + } + } + } finally { + this.#activeRequest = false; + if (signal) { + signal.removeEventListener("abort", onAbort); + } + } + } + + #push(item: Record | Error | null): void { + this.#queue.push(item); + const waiter = this.#waiters.shift(); + if (waiter) waiter(); + } + + async #nextMessage(): Promise | Error | null> { + while (this.#queue.length === 0) { + const { promise, resolve } = Promise.withResolvers(); + this.#waiters.push(resolve); + await promise; + } + return this.#queue.shift() ?? null; + } +} + +async function getOrCreateCodexWebSocketConnection( + state: CodexWebSocketSessionState, + url: string, + headers: Headers, + signal?: AbortSignal, +): Promise { + if (state.connection?.isOpen()) { + return state.connection; + } + state.connection?.close("reconnect"); + state.connection = new CodexWebSocketConnection(url, headersToRecord(headers)); + await state.connection.connect(signal); + return state.connection; +} +async function openCodexSseEventStream( + url: string, + requestHeaders: Record | undefined, + accountId: string, + apiKey: string, + sessionId: string | undefined, + body: RequestBody, + state: CodexWebSocketSessionState | undefined, + signal?: AbortSignal, +): Promise>> { + const headers = createCodexHeaders(requestHeaders, accountId, apiKey, sessionId, "sse", state); + logCodexDebug("codex request", { + url, + model: body.model, + headers: redactHeaders(headers), + }); + const response = await fetchWithRetry( + url, + { + method: "POST", + headers, + body: JSON.stringify(body), + }, + signal, + ); + logCodexDebug("codex response", { + url: response.url, + status: response.status, + statusText: response.statusText, + contentType: response.headers.get("content-type") || null, + cfRay: response.headers.get("cf-ray") || null, + }); + if (state) { + state.turnState = response.headers.get(X_CODEX_TURN_STATE_HEADER) ?? state.turnState; + state.modelsEtag = response.headers.get(X_MODELS_ETAG_HEADER) ?? state.modelsEtag; + const reasoningIncluded = response.headers.get(X_REASONING_INCLUDED_HEADER); + if (reasoningIncluded === "true" || reasoningIncluded === "false") { + state.reasoningIncluded = reasoningIncluded === "true"; + } + } + if (!response.ok) { + const info = await parseCodexError(response); + const error = new Error(info.friendlyMessage || info.message); + (error as { headers?: Headers }).headers = response.headers; + throw error; + } + if (!response.body) { + throw new Error("No response body"); + } + return readSseJson>(response.body, signal); +} +async function openCodexWebSocketEventStream( + url: string, + headers: Headers, + request: Record, + state: CodexWebSocketSessionState, + signal?: AbortSignal, +): Promise>> { + const connection = await getOrCreateCodexWebSocketConnection(state, url, headers, signal); + return connection.streamRequest(request, signal); +} function createCodexHeaders( initHeaders: Record | undefined, accountId: string, accessToken: string, promptCacheKey?: string, + transport: "sse" | "websocket" = "sse", + state?: CodexWebSocketSessionState, ): Headers { const headers = new Headers(initHeaders ?? {}); headers.delete("x-api-key"); headers.set("Authorization", `Bearer ${accessToken}`); headers.set(OPENAI_HEADERS.ACCOUNT_ID, accountId); - headers.set(OPENAI_HEADERS.BETA, OPENAI_HEADER_VALUES.BETA_RESPONSES); + headers.set( + OPENAI_HEADERS.BETA, + transport === "websocket" ? OPENAI_HEADER_VALUES.BETA_RESPONSES_WEBSOCKETS : OPENAI_HEADER_VALUES.BETA_RESPONSES, + ); headers.set(OPENAI_HEADERS.ORIGINATOR, OPENAI_HEADER_VALUES.ORIGINATOR_CODEX); headers.set("User-Agent", `pi/${packageJson.version} (${os.platform()} ${os.release()}; ${os.arch()})`); - if (promptCacheKey) { headers.set(OPENAI_HEADERS.CONVERSATION_ID, promptCacheKey); headers.set(OPENAI_HEADERS.SESSION_ID, promptCacheKey); @@ -480,8 +999,21 @@ function createCodexHeaders( headers.delete(OPENAI_HEADERS.CONVERSATION_ID); headers.delete(OPENAI_HEADERS.SESSION_ID); } - - headers.set("accept", "text/event-stream"); + if (state?.turnState) { + headers.set(X_CODEX_TURN_STATE_HEADER, state.turnState); + } else { + headers.delete(X_CODEX_TURN_STATE_HEADER); + } + if (state?.modelsEtag) { + headers.set(X_MODELS_ETAG_HEADER, state.modelsEtag); + } else { + headers.delete(X_MODELS_ETAG_HEADER); + } + if (transport === "sse") { + headers.set("accept", "text/event-stream"); + } else { + headers.delete("accept"); + } headers.set("content-type", "application/json"); return headers; } diff --git a/packages/ai/src/providers/openai-codex/constants.ts b/packages/ai/src/providers/openai-codex/constants.ts index cda31dd0c..2a07cbb6e 100644 --- a/packages/ai/src/providers/openai-codex/constants.ts +++ b/packages/ai/src/providers/openai-codex/constants.ts @@ -14,6 +14,7 @@ export const OPENAI_HEADERS = { export const OPENAI_HEADER_VALUES = { BETA_RESPONSES: "responses=experimental", + BETA_RESPONSES_WEBSOCKETS: "responses_websockets=2026-02-04", ORIGINATOR_CODEX: "pi", } as const; diff --git a/packages/ai/src/stream.ts b/packages/ai/src/stream.ts index 80f98b108..ff7dc12e6 100644 --- a/packages/ai/src/stream.ts +++ b/packages/ai/src/stream.ts @@ -470,6 +470,7 @@ function mapOptionsForApi( ...base, reasoningEffort: supportsXhigh(model) ? options?.reasoning : clampReasoning(options?.reasoning), toolChoice: mapOpenAiToolChoice(options?.toolChoice), + preferWebsockets: options?.preferWebsockets, } as OptionsForApi; case "google-generative-ai": { diff --git a/packages/ai/src/types.ts b/packages/ai/src/types.ts index 6762fa744..064ecf081 100644 --- a/packages/ai/src/types.ts +++ b/packages/ai/src/types.ts @@ -164,6 +164,8 @@ export interface SimpleStreamOptions extends StreamOptions { toolChoice?: ToolChoice; /** API format for Kimi Code provider: "openai" or "anthropic" (default: "anthropic") */ kimiApiFormat?: "openai" | "anthropic"; + /** Hint that websocket transport should be preferred when supported by the provider implementation. */ + preferWebsockets?: boolean; } // Generic StreamFunction with typed options @@ -390,6 +392,8 @@ export interface Model { contextWindow: number; maxTokens: number; headers?: Record; + /** Hint that websocket transport should be preferred when supported by the provider implementation. */ + preferWebsockets?: boolean; /** Compatibility overrides for openai-completions API. If not set, auto-detected from baseUrl. */ compat?: TApi extends "openai-completions" ? OpenAICompat : never; } diff --git a/packages/ai/test/openai-codex-stream.test.ts b/packages/ai/test/openai-codex-stream.test.ts index 1816498bd..08348b8f4 100644 --- a/packages/ai/test/openai-codex-stream.test.ts +++ b/packages/ai/test/openai-codex-stream.test.ts @@ -1,14 +1,16 @@ import { afterEach, describe, expect, it, vi } from "bun:test"; -import { streamOpenAICodexResponses } from "@oh-my-pi/pi-ai/providers/openai-codex-responses"; +import { getOpenAICodexTransportDetails, prewarmOpenAICodexResponses, streamOpenAICodexResponses } from "@oh-my-pi/pi-ai/providers/openai-codex-responses"; import type { Context, Model } from "@oh-my-pi/pi-ai/types"; import { TempDir } from "@oh-my-pi/pi-utils"; import { getAgentDir, setAgentDir } from "@oh-my-pi/pi-utils/dirs"; const originalFetch = global.fetch; const originalAgentDir = getAgentDir(); +const originalWebSocket = global.WebSocket; afterEach(() => { global.fetch = originalFetch; + global.WebSocket = originalWebSocket; setAgentDir(originalAgentDir); vi.restoreAllMocks(); }); @@ -408,4 +410,290 @@ describe("openai-codex streaming", () => { const streamResult = streamOpenAICodexResponses(model, context, { apiKey: token }); await streamResult.result(); }); + + it("falls back to SSE when websocket connect fails", async () => { + const tempDir = TempDir.createSync("@pi-codex-stream-"); + setAgentDir(tempDir.path()); + const payload = Buffer.from( + JSON.stringify({ "https://api.openai.com/auth": { chatgpt_account_id: "acc_test" } }), + "utf8", + ).toBase64(); + const token = `aaa.${payload}.bbb`; + const sse = `${[ + `data: ${JSON.stringify({ type: "response.content_part.added", part: { type: "output_text", text: "" } })}`, + `data: ${JSON.stringify({ type: "response.output_text.delta", delta: "Hello" })}`, + `data: ${JSON.stringify({ type: "response.output_item.done", item: { type: "message", id: "msg_1", role: "assistant", status: "completed", content: [{ type: "output_text", text: "Hello" }] } })}`, + `data: ${JSON.stringify({ type: "response.completed", response: { status: "completed", usage: { input_tokens: 5, output_tokens: 3, total_tokens: 8, input_tokens_details: { cached_tokens: 0 } } } })}`, + ].join("\n\n")}\n\n`; + + const fetchMock = vi.fn(async (input: string | URL) => { + const url = typeof input === "string" ? input : input.toString(); + if (url === "https://chatgpt.com/backend-api/codex/responses") { + return new Response(sse, { + status: 200, + headers: { "content-type": "text/event-stream" }, + }); + } + 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; + setTimeout(() => { + expect(this.options?.headers?.["OpenAI-Beta"] ?? this.options?.headers?.["openai-beta"]).toBe("responses_websockets=2026-02-04"); + this.#emit("error", new Event("error")); + this.#emit("close", new Event("close")); + this.readyState = FailingWebSocket.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; + const model: Model<"openai-codex-responses"> = { + id: "gpt-5.3-codex-spark", + name: "GPT-5.3 Codex Spark", + api: "openai-codex-responses", + provider: "openai-codex", + baseUrl: "https://chatgpt.com/backend-api", + reasoning: true, + preferWebsockets: true, + input: ["text"], + cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 }, + contextWindow: 128000, + maxTokens: 128000, + }; + const context: Context = { + systemPrompt: "You are a helpful assistant.", + messages: [{ role: "user", content: "Say hello", timestamp: Date.now() }], + }; + const streamResult = streamOpenAICodexResponses(model, context, { apiKey: token, sessionId: "ws-session" }); + const result = await streamResult.result(); + expect(result.role).toBe("assistant"); + expect(fetchMock).toHaveBeenCalled(); + const fallbackDetails = getOpenAICodexTransportDetails(model, { sessionId: "ws-session" }); + expect(fallbackDetails.lastTransport).toBe("sse"); + expect(fallbackDetails.websocketDisabled).toBe(true); + expect(fallbackDetails.fallbackCount).toBe(1); + }); + + it("reuses a prewarmed websocket connection across turns", async () => { + const tempDir = TempDir.createSync("@pi-codex-stream-"); + setAgentDir(tempDir.path()); + + const payload = Buffer.from( + JSON.stringify({ "https://api.openai.com/auth": { chatgpt_account_id: "acc_test" } }), + "utf8", + ).toBase64(); + const token = `aaa.${payload}.bbb`; + + const fetchMock = vi.fn(async () => { + throw new Error("SSE fallback should not be called"); + }); + 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 }) { + 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); + } + + 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); + } + } + } + + global.WebSocket = ReusableWebSocket as unknown as typeof WebSocket; + + const model: Model<"openai-codex-responses"> = { + id: "gpt-5.3-codex-spark", + name: "GPT-5.3 Codex Spark", + api: "openai-codex-responses", + provider: "openai-codex", + baseUrl: "https://chatgpt.com/backend-api", + reasoning: true, + preferWebsockets: true, + input: ["text"], + cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 }, + contextWindow: 128000, + maxTokens: 128000, + }; + + await prewarmOpenAICodexResponses(model, { apiKey: token, sessionId: "ws-reuse-session" }); + + const firstContext: Context = { + systemPrompt: "You are a helpful assistant.", + messages: [{ role: "user", content: "First", timestamp: Date.now() }], + }; + const secondContext: Context = { + systemPrompt: "You are a helpful assistant.", + messages: [ + { role: "user", content: "First", timestamp: Date.now() }, + { role: "user", content: "Second", timestamp: Date.now() }, + ], + }; + + await streamOpenAICodexResponses(model, firstContext, { apiKey: token, sessionId: "ws-reuse-session" }).result(); + await streamOpenAICodexResponses(model, secondContext, { apiKey: token, sessionId: "ws-reuse-session" }).result(); + + expect(constructorCount).toBe(1); + expect(sendCount).toBe(2); + expect(fetchMock).not.toHaveBeenCalled(); + const transportDetails = getOpenAICodexTransportDetails(model, { sessionId: "ws-reuse-session" }); + expect(transportDetails.lastTransport).toBe("websocket"); + expect(transportDetails.websocketConnected).toBe(true); + expect(transportDetails.prewarmed).toBe(true); + expect(transportDetails.canAppend).toBe(true); + }); + + it("replays x-codex-turn-state on subsequent SSE requests", async () => { + const tempDir = TempDir.createSync("@pi-codex-stream-"); + setAgentDir(tempDir.path()); + + const payload = Buffer.from( + JSON.stringify({ "https://api.openai.com/auth": { chatgpt_account_id: "acc_test" } }), + "utf8", + ).toBase64(); + const token = `aaa.${payload}.bbb`; + + const requestTurnStates: Array = []; + let callCount = 0; + const fetchMock = vi.fn(async (_input: string | URL, init?: RequestInit) => { + const headers = init?.headers instanceof Headers ? init.headers : new Headers(init?.headers); + requestTurnStates.push(headers.get("x-codex-turn-state")); + const sse = `${[ + `data: ${JSON.stringify({ type: "response.output_item.added", item: { type: "message", id: `msg_${callCount}`, role: "assistant", status: "in_progress", content: [] } })}`, + `data: ${JSON.stringify({ type: "response.content_part.added", part: { type: "output_text", text: "" } })}`, + `data: ${JSON.stringify({ type: "response.output_text.delta", delta: "Hello" })}`, + `data: ${JSON.stringify({ type: "response.output_item.done", item: { type: "message", id: `msg_${callCount}`, role: "assistant", status: "completed", content: [{ type: "output_text", text: "Hello" }] } })}`, + `data: ${JSON.stringify({ type: "response.completed", response: { status: "completed", usage: { input_tokens: 5, output_tokens: 3, total_tokens: 8, input_tokens_details: { cached_tokens: 0 } } } })}`, + ].join("\n\n")}\n\n`; + const responseHeaders = new Headers({ "content-type": "text/event-stream" }); + if (callCount === 0) { + responseHeaders.set("x-codex-turn-state", "turn-state-1"); + } + callCount += 1; + return new Response(sse, { status: 200, headers: responseHeaders }); + }); + global.fetch = fetchMock as unknown as typeof fetch; + + 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 context: Context = { + systemPrompt: "You are a helpful assistant.", + messages: [{ role: "user", content: "Say hello", timestamp: Date.now() }], + }; + + await streamOpenAICodexResponses(model, context, { apiKey: token, sessionId: "turn-state-session" }).result(); + await streamOpenAICodexResponses(model, context, { apiKey: token, sessionId: "turn-state-session" }).result(); + + expect(requestTurnStates[0]).toBeNull(); + expect(requestTurnStates[1]).toBe("turn-state-1"); + }); }); diff --git a/packages/coding-agent/CHANGELOG.md b/packages/coding-agent/CHANGELOG.md index d3c6c7d55..4e710570d 100644 --- a/packages/coding-agent/CHANGELOG.md +++ b/packages/coding-agent/CHANGELOG.md @@ -1,14 +1,17 @@ # Changelog ## [Unreleased] - ### Added +- Added `providers.openaiWebsockets` setting to prefer websocket transport for OpenAI Codex models +- Added provider details display in session info showing authentication mode, transport, and connection settings +- Added automatic prewarm of OpenAI Codex websocket connections on session creation for improved performance - Added real-time authentication validation in OAuth provider selector with visual status indicators (checking, valid, invalid) - Added `validateAuth` and `requestRender` options to OAuthSelectorComponent for custom authentication validation and UI refresh callbacks ### Changed +- Enhanced session info output to display active provider configuration and authentication details - Replaced `process.cwd()` with `getProjectDir()` throughout codebase for improved project directory detection and handling - Made `SessionManager.list()` async to support asynchronous session discovery operations - Preserved internal whitespace and indentation in bash command normalization to support heredocs and indentation-sensitive scripts diff --git a/packages/coding-agent/src/config/settings-schema.ts b/packages/coding-agent/src/config/settings-schema.ts index 9ce9ce075..cd5033b0b 100644 --- a/packages/coding-agent/src/config/settings-schema.ts +++ b/packages/coding-agent/src/config/settings-schema.ts @@ -592,6 +592,15 @@ export const SETTINGS_SCHEMA = { submenu: true, }, }, + "providers.openaiWebsockets": { + type: "boolean", + default: false, + ui: { + tab: "services", + label: "OpenAI websockets", + description: "Prefer websocket transport for OpenAI Codex models", + }, + }, // ───────────────────────────────────────────────────────────────────────── // Exa settings diff --git a/packages/coding-agent/src/modes/controllers/command-controller.ts b/packages/coding-agent/src/modes/controllers/command-controller.ts index b1b76f540..b1962c191 100644 --- a/packages/coding-agent/src/modes/controllers/command-controller.ts +++ b/packages/coding-agent/src/modes/controllers/command-controller.ts @@ -1,7 +1,7 @@ import * as fs from "node:fs/promises"; import * as os from "node:os"; import * as path from "node:path"; -import type { UsageLimit, UsageReport } from "@oh-my-pi/pi-ai"; +import { getEnvApiKey, getProviderDetails, type ProviderDetails, type UsageLimit, type UsageReport } from "@oh-my-pi/pi-ai"; import { copyToClipboard } from "@oh-my-pi/pi-natives"; import { Loader, Markdown, padding, Spacer, Text, visibleWidth } from "@oh-my-pi/pi-tui"; import { Snowflake } from "@oh-my-pi/pi-utils"; @@ -18,6 +18,7 @@ import { PythonExecutionComponent } from "../../modes/components/python-executio import { getMarkdownTheme, getSymbolTheme, theme } from "../../modes/theme/theme"; import type { InteractiveModeContext } from "../../modes/types"; import { createCompactionSummaryMessage } from "../../session/messages"; +import type { AuthStorage } from "../../session/auth-storage"; import { outputMeta } from "../../tools/output-meta"; import { resolveToCwd } from "../../tools/path-utils"; import { getChangelogPath, parseChangelog } from "../../utils/changelog"; @@ -207,6 +208,21 @@ export class CommandController { let info = `${theme.bold("Session Info")}\n\n`; info += `${theme.fg("dim", "File:")} ${stats.sessionFile ?? "In-memory"}\n`; info += `${theme.fg("dim", "ID:")} ${stats.sessionId}\n\n`; + info += `\n${theme.bold("Provider")}\n`; + const model = this.ctx.session.model; + if (!model) { + info += `${theme.fg("dim", "No model selected")}\n`; + } else { + const authMode = resolveProviderAuthMode(this.ctx.session.modelRegistry.authStorage, model.provider); + const providerDetails = getProviderDetails({ + model, + sessionId: stats.sessionId, + authMode, + preferWebsockets: this.ctx.settings.get("providers.openaiWebsockets") ?? false, + }); + info += renderProviderSection(providerDetails, theme); + } + info += `\n`; info += `${theme.bold("Messages")}\n`; info += `${theme.fg("dim", "User:")} ${stats.userMessages}\n`; info += `${theme.fg("dim", "Assistant:")} ${stats.assistantMessages}\n`; @@ -740,6 +756,32 @@ function formatDurationShort(ms: number): string { return `${totalSeconds}s`; } +function resolveProviderAuthMode(authStorage: AuthStorage, provider: string): string { + if (authStorage.hasOAuth(provider)) { + return "oauth"; + } + if (authStorage.has(provider)) { + return "api key"; + } + if (getEnvApiKey(provider)) { + return "env api key"; + } + if (authStorage.hasAuth(provider)) { + return "runtime/fallback"; + } + return "unknown"; +} + +export function renderProviderSection(details: ProviderDetails, uiTheme: Pick): string { + const lines: string[] = []; + lines.push(`${uiTheme.fg("dim", "Name:")} ${details.provider}`); + for (const field of details.fields) { + lines.push(`${uiTheme.fg("dim", `${field.label}:`)} ${field.value}`); + } + return `${lines.join("\n")}\n`; +} + + function resolveFraction(limit: UsageLimit): number | undefined { const amount = limit.amount; if (amount.usedFraction !== undefined) return amount.usedFraction; diff --git a/packages/coding-agent/src/sdk.ts b/packages/coding-agent/src/sdk.ts index 4d7e8685f..bc19ed42c 100644 --- a/packages/coding-agent/src/sdk.ts +++ b/packages/coding-agent/src/sdk.ts @@ -1,5 +1,6 @@ import { Agent, type AgentEvent, type AgentMessage, type AgentTool, type ThinkingLevel } from "@oh-my-pi/pi-agent-core"; import { type Message, type Model, supportsXhigh } from "@oh-my-pi/pi-ai"; +import { prewarmOpenAICodexResponses } from "@oh-my-pi/pi-ai/providers/openai-codex-responses"; import type { Component } from "@oh-my-pi/pi-tui"; import { $env, logger, postmortem } from "@oh-my-pi/pi-utils"; import { getAgentDbPath, getAgentDir, getProjectDir } from "@oh-my-pi/pi-utils/dirs"; @@ -1036,6 +1037,7 @@ export async function createAgentSession(options: CreateAgentSessionOptions = {} thinkingBudgets: settings.getGroup("thinkingBudgets"), temperature: settings.get("temperature") >= 0 ? settings.get("temperature") : undefined, kimiApiFormat: settings.get("providers.kimiApiFormat") ?? "anthropic", + preferWebsockets: settings.get("providers.openaiWebsockets") ?? false, getToolContext: tc => toolContextStore.getContext(tc), getApiKey: async provider => { // Use the provider argument from the in-flight request; @@ -1087,6 +1089,26 @@ export async function createAgentSession(options: CreateAgentSessionOptions = {} debugStartup("sdk:createAgentSession"); time("createAgentSession"); + if (model?.api === "openai-codex-responses") { + try { + debugStartup("sdk:prewarmCodexWebsocket:start"); + await prewarmOpenAICodexResponses(model, { + apiKey: await modelRegistry.getApiKey(model, sessionId), + sessionId, + preferWebsockets: settings.get("providers.openaiWebsockets") ?? false, + }); + debugStartup("sdk:prewarmCodexWebsocket:done"); + time("prewarmCodexWebsocket"); + } catch (error) { + logger.debug("Codex websocket prewarm failed", { + error: error instanceof Error ? error.message : String(error), + provider: model.provider, + model: model.id, + }); + } + } + + // Warm up LSP servers (connects to detected servers) let lspServers: CreateAgentSessionResult["lspServers"]; if (enableLsp && settings.get("lsp.diagnosticsOnWrite")) { diff --git a/packages/coding-agent/test/session-provider-section.test.ts b/packages/coding-agent/test/session-provider-section.test.ts new file mode 100644 index 000000000..4ad95f5e9 --- /dev/null +++ b/packages/coding-agent/test/session-provider-section.test.ts @@ -0,0 +1,34 @@ +import { describe, expect, it } from "bun:test"; +import { getProviderDetails, type Model } from "@oh-my-pi/pi-ai"; +import { renderProviderSection } from "@oh-my-pi/pi-coding-agent/modes/controllers/command-controller"; +describe("session provider section", () => { + it("renders codex provider details with transport fields", () => { + const model: Model<"openai-codex-responses"> = { + id: "gpt-5.3-codex-spark", + name: "GPT-5.3 Codex Spark", + api: "openai-codex-responses", + provider: "openai-codex", + baseUrl: "https://chatgpt.com/backend-api", + reasoning: true, + preferWebsockets: true, + input: ["text"], + cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 }, + contextWindow: 128000, + maxTokens: 128000, + }; + + const details = getProviderDetails({ + model, + sessionId: "session-1", + authMode: "oauth", + }); + const output = renderProviderSection(details, { fg: (_color: string, text: string) => text }); + + expect(output).toContain("Name:"); + expect(output).toContain("openai-codex"); + expect(output).toContain("Transport:"); + expect(output).toContain("WebSocket:"); + expect(output).toContain("Reuse:"); + expect(output).toContain("Auth:"); + }); +});