fix(codex): timeouts for websockets (#4979)

This commit is contained in:
Armin Ronacher
2026-05-27 12:11:29 +02:00
committed by GitHub
parent 26f1e00f71
commit 493efd422d
10 changed files with 551 additions and 48 deletions

View File

@@ -7,6 +7,7 @@
- Fixed OpenAI Codex Responses cache-affinity headers to send `session-id` instead of proxy-incompatible `session_id` ([#4967](https://github.com/earendil-works/pi/issues/4967)).
- Fixed `openai-codex/gpt-5.3-codex-spark` generated metadata to use its 128k context window ([#4969](https://github.com/earendil-works/pi/issues/4969)).
- Fixed OpenRouter/Poolside context overflow detection for `maximum allowed input length` errors ([#4943](https://github.com/earendil-works/pi/issues/4943)).
- Fixed OpenAI Codex Responses WebSocket streams to apply bounded connect and idle timeouts instead of waiting indefinitely when no events arrive ([#4945](https://github.com/earendil-works/pi/issues/4945)).
## [0.75.5] - 2026-05-23

View File

@@ -53,6 +53,7 @@ const JWT_CLAIM_PATH = "https://api.openai.com/auth" as const;
const DEFAULT_MAX_RETRIES = 0;
const BASE_DELAY_MS = 1000;
const DEFAULT_MAX_RETRY_DELAY_MS = 60_000;
const DEFAULT_WEBSOCKET_CONNECT_TIMEOUT_MS = 15_000;
const CODEX_TOOL_CALL_PROVIDERS = new Set(["openai", "openai-codex", "opencode"]);
const WEBSOCKET_MESSAGE_TOO_BIG_CLOSE_CODE = 1009;
@@ -163,6 +164,14 @@ function sleep(ms: number, signal?: AbortSignal): Promise<void> {
});
}
function normalizeTimeoutMs(value: number | undefined): number | undefined {
if (value === undefined) return undefined;
if (!Number.isFinite(value) || value < 0) {
throw new Error(`Invalid timeoutMs: ${String(value)}`);
}
return Math.floor(value);
}
// ============================================================================
// Main Stream Function
// ============================================================================
@@ -215,6 +224,8 @@ export const streamOpenAICodexResponses: StreamFunction<"openai-codex-responses"
websocketRequestId,
);
const bodyJson = JSON.stringify(body);
const idleTimeoutMs = normalizeTimeoutMs(options?.timeoutMs);
const websocketConnectTimeoutMs = normalizeTimeoutMs(options?.websocketConnectTimeoutMs);
const transport = options?.transport || "auto";
const websocketDisabledForSession = transport !== "sse" && isWebSocketSseFallbackActive(options?.sessionId);
if (websocketDisabledForSession) {
@@ -234,6 +245,8 @@ export const streamOpenAICodexResponses: StreamFunction<"openai-codex-responses"
() => {
websocketStarted = true;
},
idleTimeoutMs,
websocketConnectTimeoutMs,
options,
);
@@ -838,7 +851,12 @@ function scheduleSessionWebSocketExpiry(sessionId: string, entry: CachedWebSocke
}, SESSION_WEBSOCKET_CACHE_TTL_MS);
}
async function connectWebSocket(url: string, headers: Headers, signal?: AbortSignal): Promise<WebSocketLike> {
async function connectWebSocket(
url: string,
headers: Headers,
signal?: AbortSignal,
connectTimeoutMs = DEFAULT_WEBSOCKET_CONNECT_TIMEOUT_MS,
): Promise<WebSocketLike> {
const WebSocketCtor = await getWebSocketConstructor();
if (!WebSocketCtor) {
throw new Error("WebSocket transport is not available in this runtime");
@@ -849,6 +867,7 @@ async function connectWebSocket(url: string, headers: Headers, signal?: AbortSig
return new Promise<WebSocketLike>((resolve, reject) => {
let settled = false;
let timeout: ReturnType<typeof setTimeout> | undefined;
let socket: WebSocketLike;
try {
@@ -858,6 +877,25 @@ async function connectWebSocket(url: string, headers: Headers, signal?: AbortSig
return;
}
const cleanup = () => {
if (timeout) {
clearTimeout(timeout);
timeout = undefined;
}
socket.removeEventListener("open", onOpen);
socket.removeEventListener("error", onError);
socket.removeEventListener("close", onClose);
signal?.removeEventListener("abort", onAbort);
};
const fail = (error: Error, closeReason?: string) => {
if (settled) return;
settled = true;
cleanup();
if (closeReason) {
closeWebSocketSilently(socket, 1000, closeReason);
}
reject(error);
};
const onOpen: WebSocketListener = () => {
if (settled) return;
settled = true;
@@ -865,38 +903,28 @@ async function connectWebSocket(url: string, headers: Headers, signal?: AbortSig
resolve(socket);
};
const onError: WebSocketListener = (event) => {
const error = extractWebSocketError(event);
if (settled) return;
settled = true;
cleanup();
reject(error);
fail(extractWebSocketError(event));
};
const onClose: WebSocketListener = (event) => {
const error = extractWebSocketCloseError(event);
if (settled) return;
settled = true;
cleanup();
reject(error);
fail(extractWebSocketCloseError(event));
};
const onAbort = () => {
if (settled) return;
settled = true;
cleanup();
socket.close(1000, "aborted");
reject(new Error("Request was aborted"));
};
const cleanup = () => {
socket.removeEventListener("open", onOpen);
socket.removeEventListener("error", onError);
socket.removeEventListener("close", onClose);
signal?.removeEventListener("abort", onAbort);
fail(new Error("Request was aborted"), "aborted");
};
socket.addEventListener("open", onOpen);
socket.addEventListener("error", onError);
socket.addEventListener("close", onClose);
signal?.addEventListener("abort", onAbort);
if (connectTimeoutMs > 0) {
timeout = setTimeout(() => {
fail(new Error(`WebSocket connect timeout after ${connectTimeoutMs}ms`), "connect_timeout");
}, connectTimeoutMs);
}
if (signal?.aborted) {
onAbort();
}
});
}
@@ -905,6 +933,7 @@ async function acquireWebSocket(
headers: Headers,
sessionId: string | undefined,
signal?: AbortSignal,
connectTimeoutMs?: number,
): Promise<{
socket: WebSocketLike;
entry?: CachedWebSocketConnection;
@@ -912,17 +941,11 @@ async function acquireWebSocket(
release: (options?: { keep?: boolean }) => void;
}> {
if (!sessionId) {
const socket = await connectWebSocket(url, headers, signal);
const socket = await connectWebSocket(url, headers, signal, connectTimeoutMs);
return {
socket,
reused: false,
release: ({ keep } = {}) => {
if (keep === false) {
closeWebSocketSilently(socket);
return;
}
closeWebSocketSilently(socket);
},
release: () => closeWebSocketSilently(socket),
};
}
@@ -950,7 +973,7 @@ async function acquireWebSocket(
};
}
if (cached.busy) {
const socket = await connectWebSocket(url, headers, signal);
const socket = await connectWebSocket(url, headers, signal, connectTimeoutMs);
return {
socket,
reused: false,
@@ -965,7 +988,7 @@ async function acquireWebSocket(
}
}
const socket = await connectWebSocket(url, headers, signal);
const socket = await connectWebSocket(url, headers, signal, connectTimeoutMs);
const entry: CachedWebSocketConnection = { socket, busy: true };
websocketSessionCache.set(sessionId, entry);
return {
@@ -1044,7 +1067,11 @@ async function decodeWebSocketData(data: unknown): Promise<string | null> {
return null;
}
async function* parseWebSocket(socket: WebSocketLike, signal?: AbortSignal): AsyncGenerator<Record<string, unknown>> {
async function* parseWebSocket(
socket: WebSocketLike,
signal?: AbortSignal,
idleTimeoutMs?: number,
): AsyncGenerator<Record<string, unknown>> {
const queue: Record<string, unknown>[] = [];
let pending: (() => void) | null = null;
let done = false;
@@ -1124,8 +1151,23 @@ async function* parseWebSocket(socket: WebSocketLike, signal?: AbortSignal): Asy
continue;
}
if (done) break;
await new Promise<void>((resolve) => {
let timeout: ReturnType<typeof setTimeout> | undefined;
await new Promise<void>((resolve, reject) => {
pending = resolve;
if (idleTimeoutMs !== undefined && idleTimeoutMs > 0) {
timeout = setTimeout(() => {
const error = new Error(`WebSocket idle timeout after ${idleTimeoutMs}ms`);
failed = error;
done = true;
pending = null;
closeWebSocketSilently(socket, 1000, "idle_timeout");
reject(error);
}, idleTimeoutMs);
}
}).finally(() => {
if (timeout) {
clearTimeout(timeout);
}
});
}
@@ -1222,9 +1264,17 @@ async function processWebSocketStream(
stream: AssistantMessageEventStream,
model: Model<"openai-codex-responses">,
onStart: () => void,
idleTimeoutMs: number | undefined,
websocketConnectTimeoutMs: number | undefined,
options?: OpenAICodexResponsesOptions,
): Promise<void> {
const { socket, entry, reused, release } = await acquireWebSocket(url, headers, options?.sessionId, options?.signal);
const { socket, entry, reused, release } = await acquireWebSocket(
url,
headers,
options?.sessionId,
options?.signal,
websocketConnectTimeoutMs,
);
let keepConnection = true;
const useCachedContext = options?.transport === "websocket-cached" || options?.transport === "auto";
// ChatGPT Codex Responses rejects `store: true` ("Store must be set to false").
@@ -1253,7 +1303,7 @@ async function processWebSocketStream(
socket.send(JSON.stringify({ type: "response.create", ...requestBody }));
await processResponsesStream(
startWebSocketOutputOnFirstEvent(
mapCodexEvents(parseWebSocket(socket, options?.signal)),
mapCodexEvents(parseWebSocket(socket, options?.signal, idleTimeoutMs)),
output,
stream,
onStart,

View File

@@ -13,6 +13,7 @@ export function buildBaseOptions(_model: Model<Api>, options?: SimpleStreamOptio
onPayload: options?.onPayload,
onResponse: options?.onResponse,
timeoutMs: options?.timeoutMs,
websocketConnectTimeoutMs: options?.websocketConnectTimeoutMs,
maxRetries: options?.maxRetries,
maxRetryDelayMs: options?.maxRetryDelayMs,
metadata: options?.metadata,

View File

@@ -123,6 +123,12 @@ export interface StreamOptions {
* For example, OpenAI and Anthropic SDK clients default to 10 minutes.
*/
timeoutMs?: number;
/**
* WebSocket connect timeout in milliseconds for providers that support
* WebSocket transports. This covers the connection/open handshake only;
* stream idleness after connection uses timeoutMs.
*/
websocketConnectTimeoutMs?: number;
/**
* Maximum retry attempts for providers/SDKs that support client-side retries.
* For example, OpenAI and Anthropic SDK clients default to 2.

View File

@@ -933,6 +933,280 @@ describe("openai-codex streaming", () => {
});
});
it("falls back to SSE when websocket connect does not open before the connect timeout", async () => {
vi.useFakeTimers();
const token = mockToken();
const encoder = new TextEncoder();
const sse = buildSSEPayload({ status: "completed" });
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") {
throw new Error(`Unexpected URL: ${url}`);
}
return new Response(
new ReadableStream<Uint8Array>({
start(controller) {
controller.enqueue(encoder.encode(sse));
controller.close();
},
}),
{ status: 200, headers: { "content-type": "text/event-stream" } },
);
});
vi.stubGlobal("fetch", fetchMock);
class MockWebSocket {
private listeners = new Map<string, Set<(event: unknown) => void>>();
addEventListener(type: string, listener: (event: unknown) => void): void {
let listeners = this.listeners.get(type);
if (!listeners) {
listeners = new Set();
this.listeners.set(type, listeners);
}
listeners.add(listener);
}
removeEventListener(type: string, listener: (event: unknown) => void): void {
this.listeners.get(type)?.delete(listener);
}
send(): void {
throw new Error("send should not be called before websocket open");
}
close(): void {}
}
vi.stubGlobal("WebSocket", MockWebSocket);
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: 1 }],
};
const resultPromise = streamOpenAICodexResponses(model, context, {
apiKey: token,
sessionId: "ws-connect-timeout",
transport: "auto",
timeoutMs: 300_000,
websocketConnectTimeoutMs: 50,
}).result();
await vi.advanceTimersByTimeAsync(50);
const result = await resultPromise;
expect(result.content.find((content) => content.type === "text")?.text).toBe("Hello");
expect(fetchMock).toHaveBeenCalledTimes(1);
expect(getOpenAICodexWebSocketDebugStats("ws-connect-timeout")).toMatchObject({
websocketFailures: 1,
sseFallbacks: 1,
websocketFallbackActive: true,
lastWebSocketError: "WebSocket connect timeout after 50ms",
});
});
it("falls back to SSE when a websocket is idle before the first event", async () => {
vi.useFakeTimers();
const token = mockToken();
const sentBodies: unknown[] = [];
const encoder = new TextEncoder();
const sse = buildSSEPayload({ status: "completed" });
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") {
throw new Error(`Unexpected URL: ${url}`);
}
return new Response(
new ReadableStream<Uint8Array>({
start(controller) {
controller.enqueue(encoder.encode(sse));
controller.close();
},
}),
{ status: 200, headers: { "content-type": "text/event-stream" } },
);
});
vi.stubGlobal("fetch", fetchMock);
class MockWebSocket {
static OPEN = 1;
readyState = MockWebSocket.OPEN;
private listeners = new Map<string, Set<(event: unknown) => void>>();
constructor(_url: string, _protocols?: string | string[] | { headers?: Record<string, string> }) {
queueMicrotask(() => this.dispatch("open", {}));
}
addEventListener(type: string, listener: (event: unknown) => void): void {
let listeners = this.listeners.get(type);
if (!listeners) {
listeners = new Set();
this.listeners.set(type, listeners);
}
listeners.add(listener);
}
removeEventListener(type: string, listener: (event: unknown) => void): void {
this.listeners.get(type)?.delete(listener);
}
send(data: string): void {
sentBodies.push(JSON.parse(data));
}
close(): void {
this.readyState = 3;
}
private dispatch(type: string, event: unknown): void {
for (const listener of this.listeners.get(type) ?? []) {
listener(event);
}
}
}
vi.stubGlobal("WebSocket", MockWebSocket);
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: 1 }],
};
const resultPromise = streamOpenAICodexResponses(model, context, {
apiKey: token,
sessionId: "ws-idle-before-start",
transport: "auto",
timeoutMs: 50,
}).result();
await vi.advanceTimersByTimeAsync(0);
expect(sentBodies).toHaveLength(1);
await vi.advanceTimersByTimeAsync(50);
const result = await resultPromise;
expect(result.content.find((content) => content.type === "text")?.text).toBe("Hello");
expect(fetchMock).toHaveBeenCalledTimes(1);
expect(getOpenAICodexWebSocketDebugStats("ws-idle-before-start")).toMatchObject({
websocketFailures: 1,
sseFallbacks: 1,
websocketFallbackActive: true,
});
});
it("errors when a websocket is idle after the stream started", async () => {
vi.useFakeTimers();
const token = mockToken();
const fetchMock = vi.fn(async () => new Response("unexpected fetch", { status: 500 }));
vi.stubGlobal("fetch", fetchMock);
class MockWebSocket {
static OPEN = 1;
readyState = MockWebSocket.OPEN;
private listeners = new Map<string, Set<(event: unknown) => void>>();
constructor(_url: string, _protocols?: string | string[] | { headers?: Record<string, string> }) {
queueMicrotask(() => this.dispatch("open", {}));
}
addEventListener(type: string, listener: (event: unknown) => void): void {
let listeners = this.listeners.get(type);
if (!listeners) {
listeners = new Set();
this.listeners.set(type, listeners);
}
listeners.add(listener);
}
removeEventListener(type: string, listener: (event: unknown) => void): void {
this.listeners.get(type)?.delete(listener);
}
send(): void {
queueMicrotask(() => {
this.dispatch("message", {
data: JSON.stringify({
type: "response.output_item.added",
item: { type: "message", id: "msg_1", role: "assistant", status: "in_progress", content: [] },
}),
});
});
}
close(): void {
this.readyState = 3;
}
private dispatch(type: string, event: unknown): void {
for (const listener of this.listeners.get(type) ?? []) {
listener(event);
}
}
}
vi.stubGlobal("WebSocket", MockWebSocket);
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: 1 }],
};
const resultPromise = streamOpenAICodexResponses(model, context, {
apiKey: token,
transport: "auto",
timeoutMs: 50,
}).result();
await vi.advanceTimersByTimeAsync(0);
await vi.advanceTimersByTimeAsync(50);
const result = await resultPromise;
expect(result.stopReason).toBe("error");
expect(result.errorMessage).toBe("WebSocket idle timeout after 50ms");
expect(fetchMock).not.toHaveBeenCalled();
});
it("sends only response input deltas in websocket-cached mode", async () => {
const token = mockToken();
const sentBodies: unknown[] = [];

View File

@@ -7,6 +7,7 @@
- Fixed user message transcript rendering to preserve user-authored ordered-list markers ([#5013](https://github.com/earendil-works/pi/issues/5013)).
- Fixed self-update commands to bypass npm, pnpm, and Bun minimum release age gates for explicit `pi update` runs ([#4929](https://github.com/earendil-works/pi/issues/4929)).
- Fixed context token estimates to count user image attachments consistently with tool result images ([#4983](https://github.com/earendil-works/pi/issues/4983)).
- Fixed `httpIdleTimeoutMs` to apply to OpenAI Codex Responses WebSocket idle waits and added `websocketConnectTimeoutMs` for bounded WebSocket connect waits ([#4945](https://github.com/earendil-works/pi/issues/4945)).
- Fixed `RpcClient` to reject pending requests and consume stdin pipe errors when the child process exits unexpectedly ([#4764](https://github.com/earendil-works/pi/issues/4764)).
- Fixed managed npm extension updates to avoid package managers installing or resolving pi host packages as peer dependencies ([#4907](https://github.com/earendil-works/pi/issues/4907)).
- Fixed RPC mode raw stdout writes to retry transient backpressure errors and flush queued protocol output during shutdown ([#4897](https://github.com/earendil-works/pi/issues/4897)).

View File

@@ -129,7 +129,9 @@ Keep `retry.provider.maxRetries` at `0` unless provider-level retries are explic
|---------|------|---------|-------------|
| `steeringMode` | string | `"one-at-a-time"` | How steering messages are sent: `"all"` or `"one-at-a-time"` |
| `followUpMode` | string | `"one-at-a-time"` | How follow-up messages are sent: `"all"` or `"one-at-a-time"` |
| `transport` | string | `"sse"` | Preferred transport for providers that support multiple transports: `"sse"`, `"websocket"`, or `"auto"` |
| `transport` | string | `"auto"` | Preferred transport for providers that support multiple transports: `"sse"`, `"websocket"`, `"websocket-cached"`, or `"auto"` |
| `httpIdleTimeoutMs` | number | `300000` | HTTP header/body idle timeout in milliseconds, also used by providers with explicit stream idle timeouts. Set to `0` to disable. |
| `websocketConnectTimeoutMs` | number | `15000` | WebSocket connect/open handshake timeout in milliseconds for providers that support WebSocket transports. Set to `0` to disable. |
### Terminal & Images

View File

@@ -340,11 +340,18 @@ export async function createAgentSession(options: CreateAgentSessionOptions = {}
throw new Error(auth.error);
}
const providerRetrySettings = settingsManager.getProviderRetrySettings();
const timeoutMs =
options?.timeoutMs ??
providerRetrySettings.timeoutMs ??
(model.api === "openai-codex-responses" ? settingsManager.getHttpIdleTimeoutMs() : undefined);
const websocketConnectTimeoutMs =
options?.websocketConnectTimeoutMs ?? settingsManager.getWebSocketConnectTimeoutMs();
const attributionHeaders = getAttributionHeaders(model, settingsManager, options?.sessionId);
return streamSimple(model, context, {
...options,
apiKey: auth.apiKey,
timeoutMs: options?.timeoutMs ?? providerRetrySettings.timeoutMs,
timeoutMs,
websocketConnectTimeoutMs,
maxRetries: options?.maxRetries ?? providerRetrySettings.maxRetries,
maxRetryDelayMs: options?.maxRetryDelayMs ?? providerRetrySettings.maxRetryDelayMs,
headers:

View File

@@ -112,6 +112,7 @@ export interface Settings {
warnings?: WarningSettings;
sessionDir?: string; // Custom session storage directory (same format as --session-dir CLI flag)
httpIdleTimeoutMs?: number; // HTTP header/body idle timeout in milliseconds; 0 disables it
websocketConnectTimeoutMs?: number; // WebSocket connect/open handshake timeout in milliseconds; 0 disables it
}
/** Deep merge settings: project/overrides take precedence, nested objects merge recursively */
@@ -145,6 +146,17 @@ function deepMergeSettings(base: Settings, overrides: Settings): Settings {
return result;
}
function parseTimeoutSetting(value: unknown, settingName: string): number | undefined {
const timeoutMs = parseHttpIdleTimeoutMs(value);
if (timeoutMs !== undefined) {
return timeoutMs;
}
if (value !== undefined) {
throw new Error(`Invalid ${settingName} setting: ${String(value)}`);
}
return undefined;
}
export type SettingsScope = "global" | "project";
export interface SettingsStorage {
@@ -722,15 +734,7 @@ export class SettingsManager {
}
getHttpIdleTimeoutMs(): number {
const value = this.settings.httpIdleTimeoutMs;
const timeoutMs = parseHttpIdleTimeoutMs(value);
if (timeoutMs !== undefined) {
return timeoutMs;
}
if (value !== undefined) {
throw new Error(`Invalid httpIdleTimeoutMs setting: ${String(value)}`);
}
return DEFAULT_HTTP_IDLE_TIMEOUT_MS;
return parseTimeoutSetting(this.settings.httpIdleTimeoutMs, "httpIdleTimeoutMs") ?? DEFAULT_HTTP_IDLE_TIMEOUT_MS;
}
setHttpIdleTimeoutMs(timeoutMs: number): void {
@@ -750,6 +754,10 @@ export class SettingsManager {
};
}
getWebSocketConnectTimeoutMs(): number | undefined {
return parseTimeoutSetting(this.settings.websocketConnectTimeoutMs, "websocketConnectTimeoutMs");
}
getHideThinkingBlock(): boolean {
return this.settings.hideThinkingBlock ?? false;
}

View File

@@ -0,0 +1,153 @@
import { mkdirSync, mkdtempSync, rmSync } from "node:fs";
import { tmpdir } from "node:os";
import { join } from "node:path";
import {
type Api,
type AssistantMessage,
createAssistantMessageEventStream,
type Model,
type SimpleStreamOptions,
} from "@earendil-works/pi-ai";
import { afterEach, beforeEach, describe, expect, it } from "vitest";
import { AuthStorage } from "../src/core/auth-storage.ts";
import { ModelRegistry } from "../src/core/model-registry.ts";
import { createAgentSession } from "../src/core/sdk.ts";
import { SessionManager } from "../src/core/session-manager.ts";
import { SettingsManager } from "../src/core/settings-manager.ts";
describe("createAgentSession stream options", () => {
let tempDir: string;
let cwd: string;
let agentDir: string;
beforeEach(() => {
tempDir = mkdtempSync(join(tmpdir(), "pi-sdk-stream-options-"));
cwd = join(tempDir, "project");
agentDir = join(tempDir, "agent");
mkdirSync(cwd, { recursive: true });
mkdirSync(agentDir, { recursive: true });
});
afterEach(() => {
if (tempDir) {
rmSync(tempDir, { recursive: true, force: true });
}
});
function createModel(api: Api): Model<Api> {
return {
id: "capture-model",
name: "Capture Model",
api,
provider: "capture-provider",
baseUrl: "https://capture.invalid/v1",
reasoning: false,
input: ["text"],
cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 },
contextWindow: 128000,
maxTokens: 4096,
};
}
function createDoneStream(api: Api) {
const stream = createAssistantMessageEventStream();
const message: AssistantMessage = {
role: "assistant",
content: [{ type: "text", text: "ok" }],
api,
provider: "capture-provider",
model: "capture-model",
usage: {
input: 0,
output: 0,
cacheRead: 0,
cacheWrite: 0,
totalTokens: 0,
cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 },
},
stopReason: "stop",
timestamp: Date.now(),
};
stream.end(message);
return stream;
}
async function captureStreamOptions(
api: Api,
settings: { httpIdleTimeoutMs?: number; websocketConnectTimeoutMs?: number },
requestOptions: SimpleStreamOptions = {},
): Promise<SimpleStreamOptions | undefined> {
const model = createModel(api);
const settingsManager = SettingsManager.inMemory(settings);
const authStorage = AuthStorage.create(join(agentDir, "auth.json"));
authStorage.setRuntimeApiKey(model.provider, "test-api-key");
const modelRegistry = ModelRegistry.create(authStorage, join(agentDir, "models.json"));
let capturedOptions: SimpleStreamOptions | undefined;
modelRegistry.registerProvider(model.provider, {
api,
streamSimple: (_model, _context, providerOptions) => {
capturedOptions = providerOptions;
return createDoneStream(api);
},
});
const sessionManager = SessionManager.inMemory(cwd);
const { session } = await createAgentSession({
cwd,
agentDir,
model,
authStorage,
modelRegistry,
settingsManager,
sessionManager,
});
try {
await session.agent.streamFn(model, { messages: [] }, requestOptions);
return capturedOptions;
} finally {
session.dispose();
modelRegistry.unregisterProvider(model.provider);
}
}
it("forwards httpIdleTimeoutMs as timeoutMs for OpenAI Codex", async () => {
const options = await captureStreamOptions("openai-codex-responses", { httpIdleTimeoutMs: 1234 });
expect(options?.timeoutMs).toBe(1234);
});
it("does not default timeoutMs from httpIdleTimeoutMs for other providers", async () => {
const options = await captureStreamOptions("openai-completions", { httpIdleTimeoutMs: 1234 });
expect(options?.timeoutMs).toBeUndefined();
});
it("lets request timeoutMs override httpIdleTimeoutMs for OpenAI Codex", async () => {
const options = await captureStreamOptions(
"openai-codex-responses",
{ httpIdleTimeoutMs: 1234 },
{ timeoutMs: 0 },
);
expect(options?.timeoutMs).toBe(0);
});
it("forwards websocketConnectTimeoutMs from settings", async () => {
const options = await captureStreamOptions("openai-codex-responses", { websocketConnectTimeoutMs: 1234 });
expect(options?.websocketConnectTimeoutMs).toBe(1234);
});
it("lets request websocketConnectTimeoutMs override settings", async () => {
const options = await captureStreamOptions(
"openai-codex-responses",
{ websocketConnectTimeoutMs: 1234 },
{ websocketConnectTimeoutMs: 0 },
);
expect(options?.websocketConnectTimeoutMs).toBe(0);
});
});