feat: added WebSocket transport support for OpenAI Codex with SSE fallback
- Added WebSocket transport support for OpenAI Codex responses with automatic fallback to SSE on connection failure. - Added preferWebsockets option to Agent and Model configurations to hint that WebSocket transport should be preferred when supported by provider implementations. - Added prewarmOpenAICodexResponses() function to pre-establish WebSocket connections for improved performance. - Added getProviderDetails() function and getOpenAICodexTransportDetails() function to expose transport state and provider configuration information. - Added provider details display in session info showing active provider configuration and authentication details. - Added OpenAI websockets setting to enable WebSocket transport preference for OpenAI Codex models in coding agent configuration.
This commit is contained in:
can1357
@@ -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
|
||||
|
||||
@@ -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<void>;
|
||||
#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,
|
||||
|
||||
@@ -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
|
||||
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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";
|
||||
|
||||
@@ -12965,6 +12965,7 @@
|
||||
"provider": "openai-codex",
|
||||
"baseUrl": "https://chatgpt.com/backend-api",
|
||||
"reasoning": true,
|
||||
"preferWebsockets": true,
|
||||
"input": [
|
||||
"text"
|
||||
],
|
||||
|
||||
@@ -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<Api>;
|
||||
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";
|
||||
}
|
||||
@@ -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<string, CodexWebSocketSessionState>();
|
||||
const codexWebSocketPublicToPrivate = new Map<string, string>();
|
||||
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<Record<string, unknown>>;
|
||||
|
||||
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<Record<string, unknown>>(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<OpenAICodexResponsesOptions, "apiKey" | "headers" | "sessionId" | "signal" | "preferWebsockets">,
|
||||
): Promise<void> {
|
||||
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<string, unknown> {
|
||||
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<string, string> {
|
||||
const result: Record<string, string> = {};
|
||||
for (const [key, value] of headers.entries()) {
|
||||
result[key] = value;
|
||||
}
|
||||
return result;
|
||||
}
|
||||
|
||||
class CodexWebSocketConnection {
|
||||
#url: string;
|
||||
#headers: Record<string, string>;
|
||||
#socket: WebSocket | null = null;
|
||||
#queue: Array<Record<string, unknown> | Error | null> = [];
|
||||
#waiters: Array<() => void> = [];
|
||||
#connectPromise?: Promise<void>;
|
||||
#activeRequest = false;
|
||||
|
||||
constructor(url: string, headers: Record<string, string>) {
|
||||
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<void> {
|
||||
if (this.isOpen()) return;
|
||||
if (this.#connectPromise) {
|
||||
await this.#connectPromise;
|
||||
return;
|
||||
}
|
||||
const WebSocketWithHeaders = WebSocket as unknown as {
|
||||
new (url: string, options?: { headers?: Record<string, string> }): WebSocket;
|
||||
};
|
||||
const { promise, resolve, reject } = Promise.withResolvers<void>();
|
||||
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<string, unknown>;
|
||||
if (parsed.type === "error" && typeof parsed.error === "object" && parsed.error) {
|
||||
const inner = parsed.error as Record<string, unknown>;
|
||||
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<string, unknown>, signal?: AbortSignal): AsyncGenerator<Record<string, unknown>> {
|
||||
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<string, unknown> | Error | null): void {
|
||||
this.#queue.push(item);
|
||||
const waiter = this.#waiters.shift();
|
||||
if (waiter) waiter();
|
||||
}
|
||||
|
||||
async #nextMessage(): Promise<Record<string, unknown> | Error | null> {
|
||||
while (this.#queue.length === 0) {
|
||||
const { promise, resolve } = Promise.withResolvers<void>();
|
||||
this.#waiters.push(resolve);
|
||||
await promise;
|
||||
}
|
||||
return this.#queue.shift() ?? null;
|
||||
}
|
||||
}
|
||||
|
||||
async function getOrCreateCodexWebSocketConnection(
|
||||
state: CodexWebSocketSessionState,
|
||||
url: string,
|
||||
headers: Headers,
|
||||
signal?: AbortSignal,
|
||||
): Promise<CodexWebSocketConnection> {
|
||||
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<string, string> | undefined,
|
||||
accountId: string,
|
||||
apiKey: string,
|
||||
sessionId: string | undefined,
|
||||
body: RequestBody,
|
||||
state: CodexWebSocketSessionState | undefined,
|
||||
signal?: AbortSignal,
|
||||
): Promise<AsyncGenerator<Record<string, unknown>>> {
|
||||
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<Record<string, unknown>>(response.body, signal);
|
||||
}
|
||||
async function openCodexWebSocketEventStream(
|
||||
url: string,
|
||||
headers: Headers,
|
||||
request: Record<string, unknown>,
|
||||
state: CodexWebSocketSessionState,
|
||||
signal?: AbortSignal,
|
||||
): Promise<AsyncGenerator<Record<string, unknown>>> {
|
||||
const connection = await getOrCreateCodexWebSocketConnection(state, url, headers, signal);
|
||||
return connection.streamRequest(request, signal);
|
||||
}
|
||||
function createCodexHeaders(
|
||||
initHeaders: Record<string, string> | 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;
|
||||
}
|
||||
|
||||
@@ -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;
|
||||
|
||||
|
||||
@@ -470,6 +470,7 @@ function mapOptionsForApi<TApi extends Api>(
|
||||
...base,
|
||||
reasoningEffort: supportsXhigh(model) ? options?.reasoning : clampReasoning(options?.reasoning),
|
||||
toolChoice: mapOpenAiToolChoice(options?.toolChoice),
|
||||
preferWebsockets: options?.preferWebsockets,
|
||||
} as OptionsForApi<TApi>;
|
||||
|
||||
case "google-generative-ai": {
|
||||
|
||||
@@ -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<TApi extends Api = any> {
|
||||
contextWindow: number;
|
||||
maxTokens: number;
|
||||
headers?: Record<string, string>;
|
||||
/** 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;
|
||||
}
|
||||
|
||||
@@ -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<string, Set<WsListener>>();
|
||||
url: string;
|
||||
options?: { headers?: Record<string, string> };
|
||||
|
||||
constructor(url: string, options?: { headers?: Record<string, string> }) {
|
||||
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<WsListener>();
|
||||
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<string, Set<WsListener>>();
|
||||
|
||||
constructor(public readonly url: string, public readonly options?: { headers?: Record<string, string> }) {
|
||||
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<WsListener>();
|
||||
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<string, unknown>;
|
||||
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<string | null> = [];
|
||||
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");
|
||||
});
|
||||
});
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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<typeof theme, "fg">): 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;
|
||||
|
||||
@@ -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")) {
|
||||
|
||||
@@ -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:");
|
||||
});
|
||||
});
|
||||
Reference in New Issue
Block a user