@f5-sales-demo/pi-ai 22.6.0 → 22.7.1

package/package.json

@@ -1,7 +1,7 @@
 {
 	"type": "module",
 	"name": "@f5-sales-demo/pi-ai",
-	"version": "22.6.0",
+	"version": "22.7.1",
 	"description": "Unified LLM API with automatic model discovery and provider configuration",
 	"homepage": "https://github.com/f5-sales-demo/xcsh",
 	"author": "Can Boluk",
@@ -46,7 +46,7 @@
 		"@anthropic-ai/sdk": "0.124.0",
 		"@aws-sdk/client-bedrock-runtime": "3.1127.0",
 		"@bufbuild/protobuf": "2.14.1",
-		"@f5-sales-demo/pi-utils": "22.6.0",
+		"@f5-sales-demo/pi-utils": "22.7.1",
 		"@google/genai": "2.21.0",
 		"@sinclair/typebox": "0.34.52",
 		"@smithy/node-http-handler": "4.12.1",

package/src/providers/openai-codex-responses.ts

@@ -26,6 +26,7 @@ import { getEnvApiKey } from "../stream";
 import {
 	type Api,
 	type AssistantMessage,
+	type AssistantMessageEvent,
 	type Context,
 	isSpecialServiceTier,
 	type Model,
@@ -47,8 +48,9 @@ import {
 } from "../utils";
 import { AssistantMessageEventStream } from "../utils/event-stream";
 import { finalizeErrorMessage, type RawHttpRequestDump } from "../utils/http-inspector";
-import { getOpenAIStreamIdleTimeoutMs, iterateWithIdleTimeout } from "../utils/idle-iterator";
+import { getOpenAIStreamIdleTimeoutMs } from "../utils/idle-iterator";
 import { parseStreamingJson } from "../utils/json-parse";
+import { isTransportClosureError } from "../utils/retry";
 import { adaptSchemaForStrict, NO_STRICT } from "../utils/schema";
 import {
 	CODEX_BASE_URL,
@@ -184,7 +186,6 @@ interface CodexRequestContext {
 
 interface CodexRequestSetup {
 	requestSignal: AbortSignal;
-	wrapCodexSseStream: (source: AsyncGenerator<Record<string, unknown>>) => AsyncGenerator<Record<string, unknown>>;
 	requestAbortController: AbortController;
 }
 
@@ -194,6 +195,8 @@ interface CodexWebSocketRequestState {
 }
 
 interface CodexStreamRuntime {
+	stagedEvents: AssistantMessageEvent[];
+	firstTransportFailure?: unknown;
 	websocketRequestState?: CodexWebSocketRequestState;
 	continuationRecoveryAttempt: number;
 	sawResponseAccepted: boolean;
@@ -211,6 +214,8 @@ interface CodexStreamRuntime {
 	nativeOutputItems: Array<Record<string, unknown>>;
 	websocketStreamRetries: number;
 	providerRetryAttempt: number;
+	transportRestartAttempt: number;
+	sawReplayUnsafeActivity: boolean;
 	sawTerminalEvent: boolean;
 	canSafelyReplayWebsocketOverSse: boolean;
 }
@@ -450,15 +455,7 @@ function createRequestSetup(options: OpenAICodexResponsesOptions | undefined): C
 	const requestSignal = options?.signal
 		? AbortSignal.any([options.signal, requestAbortController.signal])
 		: requestAbortController.signal;
-	const wrapCodexSseStream = (
-		source: AsyncGenerator<Record<string, unknown>>,
-	): AsyncGenerator<Record<string, unknown>> =>
-		iterateWithIdleTimeout(source, {
-			idleTimeoutMs: getOpenAIStreamIdleTimeoutMs(),
-			errorMessage: "OpenAI Codex SSE stream stalled while waiting for the next event",
-			onIdle: () => requestAbortController.abort(),
-		});
-	return { requestAbortController, requestSignal, wrapCodexSseStream };
+	return { requestAbortController, requestSignal };
 }
 
 async function buildCodexRequestContext(
@@ -763,20 +760,83 @@ async function openCodexWebSocketTransport(
 	return { eventStream, requestBodyForState, transport: "websocket", websocketRequestState };
 }
 
+class CodexSseIdleTimeoutError extends Error {
+	constructor() {
+		super("OpenAI Codex SSE stream stalled while waiting for the next event");
+		this.name = "CodexSseIdleTimeoutError";
+	}
+}
+
+async function* iterateCodexSseAttempt(
+	source: AsyncGenerator<Record<string, unknown>>,
+	attempt: AbortController,
+	cleanup: () => void,
+): AsyncGenerator<Record<string, unknown>> {
+	const timeoutMs = getOpenAIStreamIdleTimeoutMs();
+	try {
+		while (true) {
+			attempt.signal.throwIfAborted();
+			let timer: NodeJS.Timeout | undefined;
+			try {
+				const next = source.next();
+				const result =
+					timeoutMs === undefined
+						? await next
+						: await Promise.race([
+								next,
+								new Promise<never>((_resolve, reject) => {
+									timer = setTimeout(() => {
+										const error = new CodexSseIdleTimeoutError();
+										reject(error);
+										attempt.abort(error);
+									}, timeoutMs);
+								}),
+							]);
+				if (timer !== undefined) {
+					clearTimeout(timer);
+					timer = undefined;
+				}
+				attempt.signal.throwIfAborted();
+				if (result.done) return;
+				yield result.value;
+			} finally {
+				if (timer !== undefined) clearTimeout(timer);
+			}
+		}
+	} finally {
+		attempt.abort();
+		cleanup();
+		try {
+			await source.return(undefined);
+		} catch {
+			/* Preserve the original stream failure. */
+		}
+	}
+}
+
 async function openCodexSseTransport(
 	requestContext: CodexRequestContext,
 	requestSetup: CodexRequestSetup,
 	options: OpenAICodexResponsesOptions | undefined,
 	state: CodexWebSocketSessionState | undefined,
 	body = requestContext.transformedBody,
+	retryFetch = true,
 ): Promise<{
 	eventStream: AsyncGenerator<Record<string, unknown>>;
 	requestBodyForState: RequestBody;
 	transport: CodexTransport;
 	websocketRequestState?: CodexWebSocketRequestState;
 }> {
-	const eventStream = requestSetup.wrapCodexSseStream(
-		await openCodexSseEventStream(
+	const attempt = new AbortController();
+	const signal = requestSetup.requestSignal;
+	const onAbort = () => attempt.abort(signal.reason);
+	const cleanup = () => signal.removeEventListener("abort", onAbort);
+	signal.addEventListener("abort", onAbort, { once: true });
+	if (signal.aborted) onAbort();
+	let source: AsyncGenerator<Record<string, unknown>>;
+	try {
+		attempt.signal.throwIfAborted();
+		source = await openCodexSseEventStream(
 			requestContext.url,
 			requestContext.requestHeaders,
 			requestContext.accountId,
@@ -784,9 +844,16 @@ async function openCodexSseTransport(
 			options?.sessionId,
 			body,
 			state,
-			requestSetup.requestSignal,
-		),
-	);
+			attempt.signal,
+			retryFetch,
+		);
+	} catch (error) {
+		attempt.abort();
+		cleanup();
+		throw error;
+	}
+	const eventStream = iterateCodexSseAttempt(source, attempt, cleanup);
+
 	return { eventStream, requestBodyForState: structuredCloneJSON(body), transport: "sse" };
 }
 
@@ -832,6 +899,7 @@ function createCodexStreamRuntime(initial: {
 	websocketRequestState?: CodexWebSocketRequestState;
 }): CodexStreamRuntime {
 	return {
+		stagedEvents: [],
 		websocketRequestState: initial.websocketRequestState,
 		continuationRecoveryAttempt: 0,
 		sawResponseAccepted: false,
@@ -848,11 +916,37 @@ function createCodexStreamRuntime(initial: {
 		nativeOutputItems: [],
 		websocketStreamRetries: 0,
 		providerRetryAttempt: 0,
+		transportRestartAttempt: 0,
+		sawReplayUnsafeActivity: false,
 		sawTerminalEvent: false,
 		canSafelyReplayWebsocketOverSse: true,
 	};
 }
 
+// Empty normalized blocks belong to the attempt until replay becomes unsafe.
+class CodexAttemptEventStream extends AssistantMessageEventStream {
+	constructor(
+		private readonly target: AssistantMessageEventStream,
+		private readonly runtime: CodexStreamRuntime,
+	) {
+		super();
+	}
+
+	override push(event: AssistantMessageEvent): void {
+		if (this.runtime.sawReplayUnsafeActivity || this.runtime.sawTerminalEvent) {
+			this.flush();
+			this.target.push(event);
+		} else {
+			this.runtime.stagedEvents.push(structuredCloneJSON(event));
+		}
+	}
+
+	flush(): void {
+		for (const event of this.runtime.stagedEvents) this.target.push(event);
+		this.runtime.stagedEvents.length = 0;
+	}
+}
+
 async function processCodexResponseStream(
 	context: CodexStreamProcessingContext,
 	runtime: CodexStreamRuntime,
@@ -860,6 +954,7 @@ async function processCodexResponseStream(
 	const { output, stream } = context;
 	runtime.grammarToolInputProperties = context.requestContext.grammarToolInputProperties;
 	stream.push({ type: "start", partial: output });
+	const attemptStream = new CodexAttemptEventStream(stream, runtime);
 
 	while (true) {
 		try {
@@ -868,10 +963,17 @@ async function processCodexResponseStream(
 				await context.options?.onProviderStreamEvent?.(rawEvent, context.model);
 				firstTokenTime = handleCodexStreamEvent({
 					...context,
+					stream: attemptStream,
 					runtime,
 					rawEvent,
 					firstTokenTime,
 				});
+				if (runtime.sawReplayUnsafeActivity || runtime.sawTerminalEvent) attemptStream.flush();
+			}
+			if (!runtime.sawTerminalEvent) {
+				throw Object.assign(new Error("Codex stream ended before terminal completion event"), {
+					code: "ERR_STREAM_PREMATURE_CLOSE",
+				});
 			}
 			return { firstTokenTime };
 		} catch (error) {
@@ -894,6 +996,27 @@ function handleCodexStreamEvent(args: {
 	const { model, output, stream, runtime, rawEvent } = args;
 	const eventType = typeof rawEvent.type === "string" ? rawEvent.type : "";
 	if (!eventType) return args.firstTokenTime;
+	// Tool activity is unsafe from its first event, even if no arguments or tool result exists yet.
+	const item = asRecord(rawEvent.item);
+	const part = asRecord(rawEvent.part);
+	const hasText = (value: unknown): boolean => typeof value === "string" && value.length > 0;
+	if (
+		/^response\.(?:function_call|custom_tool_call|.*_call\.)/.test(eventType) ||
+		(item && typeof item.type === "string" && item.type !== "message" && item.type !== "reasoning") ||
+		hasText(rawEvent.delta) ||
+		hasText(rawEvent.text) ||
+		hasText(part?.text) ||
+		hasText(part?.refusal) ||
+		(item &&
+			[item.content, item.summary].some(
+				parts =>
+					Array.isArray(parts) &&
+					parts.some(value => hasText(asRecord(value)?.text) || hasText(asRecord(value)?.refusal)),
+			))
+	) {
+		runtime.sawReplayUnsafeActivity = true;
+	}
+
 	if (
 		eventType === "response.created" ||
 		eventType === "response.in_progress" ||
@@ -1376,6 +1499,9 @@ async function recoverCodexStreamError(
 	runtime: CodexStreamRuntime,
 	error: unknown,
 ): Promise<boolean> {
+	if (runtime.transport === "sse" && (error instanceof CodexSseIdleTimeoutError || isTransportClosureError(error))) {
+		return tryRestartCodexSseTransport(context, runtime, error);
+	}
 	if (isCodexRejectedContinuation(error)) {
 		return tryRecoverCodexContinuation(context, runtime);
 	}
@@ -1391,6 +1517,101 @@ async function recoverCodexStreamError(
 	return false;
 }
 
+async function sleepBeforeCodexSseRestart(delayMs: number, signal: AbortSignal): Promise<void> {
+	signal.throwIfAborted();
+	const { promise, resolve, reject } = Promise.withResolvers<void>();
+	const timer = setTimeout(resolve, delayMs);
+	const onAbort = () => reject(signal.reason);
+	signal.addEventListener("abort", onAbort, { once: true });
+	try {
+		await promise;
+		signal.throwIfAborted();
+	} finally {
+		clearTimeout(timer);
+		signal.removeEventListener("abort", onAbort);
+	}
+}
+
+async function tryRestartCodexSseTransport(
+	context: CodexStreamProcessingContext,
+	runtime: CodexStreamRuntime,
+	originalError: unknown,
+): Promise<boolean> {
+	runtime.firstTransportFailure ??= originalError;
+	let error = originalError;
+	while (true) {
+		const replayEligible =
+			!runtime.sawReplayUnsafeActivity &&
+			!runtime.sawTerminalEvent &&
+			!context.requestSetup.requestSignal.aborted &&
+			context.output.content.every(
+				block =>
+					(block.type === "text" && block.text.length === 0) ||
+					(block.type === "thinking" && block.thinking.length === 0),
+			);
+		logCodexDebug("codex SSE transport recovery", {
+			transport: runtime.transport,
+			model: context.model.id,
+			elapsedMs: Date.now() - context.startTime,
+			attempt: runtime.transportRestartAttempt,
+			terminalEventSeen: runtime.sawTerminalEvent,
+			replayEligible,
+		});
+		if (!replayEligible || !(error instanceof CodexSseIdleTimeoutError || isTransportClosureError(error))) {
+			if (error !== originalError) throw error;
+			return false;
+		}
+		if (runtime.transportRestartAttempt >= 2) throw runtime.firstTransportFailure;
+		runtime.transportRestartAttempt += 1;
+		try {
+			await runtime.eventStream.return(undefined);
+		} catch {
+			/* Preserve the recovery failure. */
+		}
+		runtime.currentItem = null;
+		runtime.currentBlock = null;
+		runtime.currentIndex = -1;
+		runtime.outputSlots.clear();
+		runtime.completedItemIds.clear();
+		runtime.nativeOutputItems.length = 0;
+		runtime.sawResponseAccepted = false;
+		runtime.sawOutputEvent = false;
+		runtime.sawTerminalEvent = false;
+		runtime.stagedEvents.length = 0;
+		resetOutputState(context.output);
+		delete context.output.responseId;
+		delete context.output.providerPayload;
+		delete context.output.duration;
+		delete context.output.ttft;
+		delete context.output.errorMessage;
+		context.output.responseAttribution = { requestedModel: context.model.id };
+		context.firstTokenTime = undefined;
+		const state = context.requestContext.websocketState;
+		if (state) {
+			resetCodexWebSocketAppendState(state);
+			resetCodexSessionMetadata(state);
+		}
+		await sleepBeforeCodexSseRestart(500 * runtime.transportRestartAttempt, context.requestSetup.requestSignal);
+		try {
+			const next = await openCodexSseTransport(
+				context.requestContext,
+				context.requestSetup,
+				context.options,
+				state,
+				context.requestContext.transformedBody,
+				false,
+			);
+			runtime.eventStream = next.eventStream;
+			runtime.requestBodyForState = next.requestBodyForState;
+			runtime.websocketRequestState = undefined;
+			return true;
+		} catch (reopenError) {
+			if (context.requestSetup.requestSignal.aborted) throw reopenError;
+			error = reopenError;
+		}
+	}
+}
+
 function isCodexRejectedContinuation(error: unknown): boolean {
 	return (
 		error instanceof CodexProviderStreamError &&
@@ -1479,6 +1700,7 @@ async function tryReconnectCodexWebSocketOnConnectionLimit(
 		runtime.outputSlots.clear();
 		runtime.completedItemIds.clear();
 		runtime.nativeOutputItems.length = 0;
+		runtime.stagedEvents.length = 0;
 		resetOutputState(context.output);
 		context.firstTokenTime = undefined;
 		recordCodexWebSocketFailure(websocketState, true);
@@ -1540,6 +1762,7 @@ async function tryReplayWebsocketFailureOverSse(
 		runtime.outputSlots.clear();
 		runtime.completedItemIds.clear();
 		runtime.nativeOutputItems.length = 0;
+		runtime.stagedEvents.length = 0;
 		resetOutputState(context.output);
 		context.firstTokenTime = undefined;
 	}
@@ -1581,6 +1804,7 @@ async function tryRetryCodexProviderError(
 	runtime.outputSlots.clear();
 	runtime.completedItemIds.clear();
 	runtime.sawTerminalEvent = false;
+	runtime.stagedEvents.length = 0;
 	resetOutputState(context.output);
 	context.firstTokenTime = undefined;
 	await abortableSleep(CODEX_RETRY_DELAY_MS * runtime.providerRetryAttempt, context.requestSetup.requestSignal);
@@ -2314,6 +2538,7 @@ async function openCodexSseEventStream(
 	body: RequestBody,
 	state: CodexWebSocketSessionState | undefined,
 	signal?: AbortSignal,
+	retryFetch = true,
 ): Promise<AsyncGenerator<Record<string, unknown>>> {
 	const headers = createCodexHeaders(requestHeaders, accountId, apiKey, sessionId, "sse", state);
 	logCodexDebug("codex request", {
@@ -2332,6 +2557,7 @@ async function openCodexSseEventStream(
 			body: JSON.stringify(body),
 		},
 		signal,
+		retryFetch,
 	);
 	logCodexDebug("codex response", {
 		url: response.url,
@@ -2353,7 +2579,39 @@ async function openCodexSseEventStream(
 	if (!response.body) {
 		throw new Error("No response body");
 	}
-	return readSseJson<Record<string, unknown>>(response.body, signal);
+	signal?.throwIfAborted();
+	const reader = response.body.getReader();
+	let closed = false;
+	const close = () => {
+		if (closed) return;
+		closed = true;
+		// Cancel settles pending reads synchronously; do not await an upstream cancel hook.
+		void reader.cancel().catch(() => {});
+		reader.releaseLock();
+	};
+	const onAbort = () => close();
+	signal?.addEventListener("abort", onAbort, { once: true });
+	const owned = new ReadableStream<Uint8Array>({
+		async pull(controller) {
+			try {
+				const result = await reader.read();
+				if (result.done) controller.close();
+				else controller.enqueue(result.value);
+			} catch (error) {
+				controller.error(error);
+			}
+		},
+		cancel: close,
+	});
+	async function* events(): AsyncGenerator<Record<string, unknown>> {
+		try {
+			yield* readSseJson<Record<string, unknown>>(owned);
+		} finally {
+			close();
+			signal?.removeEventListener("abort", onAbort);
+		}
+	}
+	return events();
 }
 
 function createCodexHeaders(
@@ -2437,13 +2695,14 @@ function getRetryDelayMs(
 	return { delay: CODEX_RETRY_DELAY_MS * (attempt + 1), serverProvided: false };
 }
 
-async function fetchWithRetry(url: string, init: RequestInit, signal?: AbortSignal): Promise<Response> {
+async function fetchWithRetry(url: string, init: RequestInit, signal?: AbortSignal, retry = true): Promise<Response> {
 	let attempt = 0;
 	let rateLimitTimeSpent = 0;
 	while (true) {
 		try {
+			signal?.throwIfAborted();
 			const response = await fetch(url, { ...init, signal: signal ?? init.signal });
-			if (!CODEX_RETRYABLE_STATUS.has(response.status)) {
+			if (!retry || !CODEX_RETRYABLE_STATUS.has(response.status)) {
 				return response;
 			}
 			if (signal?.aborted) return response;
@@ -2465,7 +2724,7 @@ async function fetchWithRetry(url: string, init: RequestInit, signal?: AbortSign
 			}
 			await abortableSleep(delay, signal);
 		} catch (error) {
-			if (attempt >= CODEX_MAX_RETRIES || signal?.aborted) {
+			if (!retry || attempt >= CODEX_MAX_RETRIES || signal?.aborted) {
 				throw error;
 			}
 			const delay = CODEX_RETRY_DELAY_MS * (attempt + 1);

package/src/utils/retry.ts

@@ -1,6 +1,7 @@
 type ErrorLike = {
 	message?: string;
 	name?: string;
+	code?: string;
 	status?: number;
 	statusCode?: number;
 	response?: { status?: number };
@@ -13,6 +14,31 @@ const TRANSIENT_MESSAGE_PATTERN =
 const VALIDATION_MESSAGE_PATTERN =
 	/invalid|validation|bad request|unsupported|schema|missing required|not found|unauthorized|forbidden/i;
 
+const BUN_SOCKET_CLOSE_MESSAGE =
+	"The socket connection was closed unexpectedly. For more information, pass `verbose: true` in the second argument to fetch()";
+
+/** Recognize transport closures without interpreting arbitrary provider failure text as replay-safe. */
+export function isTransportClosureError(error: unknown, depth = 0): boolean {
+	if (!error || typeof error !== "object" || depth > 2) return false;
+	const info = error as ErrorLike;
+	const message = info.message ?? "";
+	const status = extractHttpStatusFromError(error);
+	if (
+		info.name === "AbortError" ||
+		/aborted|cancelled|canceled/i.test(message) ||
+		VALIDATION_MESSAGE_PATTERN.test(message) ||
+		(status !== undefined && status >= 400 && status < 500)
+	)
+		return false;
+	return (
+		message === BUN_SOCKET_CLOSE_MESSAGE ||
+		info.code === "ECONNRESET" ||
+		info.code === "ERR_STREAM_PREMATURE_CLOSE" ||
+		/\b(?:ECONNRESET|ERR_STREAM_PREMATURE_CLOSE)\b/.test(message) ||
+		isTransportClosureError(info.cause, depth + 1)
+	);
+}
+
 /**
  * Identify errors that should be retried (timeouts, 5xx, 408, 429, transient network failures).
  */
@@ -31,7 +57,7 @@ export function isRetryableError(error: unknown): boolean {
 
 	if (VALIDATION_MESSAGE_PATTERN.test(message)) return false;
 
-	return TRANSIENT_MESSAGE_PATTERN.test(message);
+	return isTransportClosureError(error) || TRANSIENT_MESSAGE_PATTERN.test(message);
 }
 
 export function extractHttpStatusFromError(error: unknown): number | undefined {