@gajae-code/ai 0.15.5 → 0.15.6

package/src/providers/google-shared.ts

@@ -921,6 +921,7 @@ export function streamGoogleGenAI<T extends "google-generative-ai" | "google-ver
 				options.toolChoice !== "auto" &&
 				options.toolChoice !== "none";
 			const fetchImpl = plan.fetch ?? options?.fetch ?? (globalThis.fetch.bind(globalThis) as FetchImpl);
+			options?.onStreamCreated?.();
 			let response = await fetchImpl(plan.url, {
 				method: "POST",
 				headers: { ...plan.headers, "Content-Type": "application/json", Accept: "text/event-stream" },
@@ -935,6 +936,7 @@ export function streamGoogleGenAI<T extends "google-generative-ai" | "google-ver
 				);
 				if (
 					!options?.fallbackManaged &&
+					!options?.disableProviderRetries &&
 					firstTokenTime === undefined &&
 					isForcedToolChoiceUnsupportedError(error, true)
 				) {
@@ -957,6 +959,7 @@ export function streamGoogleGenAI<T extends "google-generative-ai" | "google-ver
 					params = retryParams;
 					rawRequestDump = { ...rawRequestDump, body: params };
 					wireBody = paramsToWireBody(params);
+					options?.onStreamCreated?.();
 					response = await fetchImpl(plan.url, {
 						method: "POST",
 						headers: { ...plan.headers, "Content-Type": "application/json", Accept: "text/event-stream" },

package/src/providers/kiro-codewhisperer.ts

@@ -29,7 +29,7 @@ import { transportFailureFacts } from "../utils/fallback-transport";
 import { withHttpStatus } from "../utils/http-inspector";
 import { captureUnicodeEscapeEvidence } from "../utils/json-parse";
 import { decodeEventStream } from "./aws-eventstream";
-import { isKiroApiKey, streamKiroApiKey } from "./kiro-api-key";
+import { isKiroApiKey, streamKiroApiKey, toKiroModelId } from "./kiro-api-key";
 
 // ─────────────────────────────────────────────────────────────────────────────
 // Provider options
@@ -67,6 +67,7 @@ interface WireToolResult {
 interface WireUserMessage {
 	userInputMessage: {
 		content: string;
+		modelId?: string;
 		userInputMessageContext?: {
 			tools?: { tools: WireToolSpec[] };
 			toolResults?: { toolResults: WireToolResult[][] };
@@ -353,7 +354,7 @@ export const streamKiroCodeWhisperer: StreamFunction<"kiro-codewhisperer-stream"
 
 function buildConversationState(
 	context: Context,
-	_model: Model<"kiro-codewhisperer-stream">,
+	model: Model<"kiro-codewhisperer-stream">,
 	options: KiroCodeWhispererOptions,
 ): ConversationState {
 	const messages = context.messages;
@@ -361,18 +362,24 @@ function buildConversationState(
 		throw new Error("Kiro CodeWhisperer requires at least one message");
 	}
 
+	// Normalize the local dashed selector/wire id (e.g. "claude-haiku-4-5") to
+	// the canonical dotted upstream Kiro model id (e.g. "claude-haiku-4.5"),
+	// matching the sibling ksk_ API-key transport (kiro-api-key.ts) so both
+	// auth methods send the same wire form for the same catalog entry.
+	const modelId = toKiroModelId(model.wireModelId || model.id);
+
 	// Build history from all messages except the last
 	const history: WireHistoryMessage[] = [];
 	const systemPrompt = context.systemPrompt?.join("\n") ?? "";
 
 	for (let i = 0; i < messages.length - 1; i++) {
 		const msg = messages[i];
-		history.push(convertToWireMessage(msg, i === 0 ? systemPrompt : undefined));
+		history.push(convertToWireMessage(msg, modelId, i === 0 ? systemPrompt : undefined));
 	}
 
 	// Convert the last message as currentMessage
 	const lastMsg = messages[messages.length - 1];
-	const currentMessage = convertToWireUserMessage(lastMsg, systemPrompt);
+	const currentMessage = convertToWireUserMessage(lastMsg, modelId, systemPrompt);
 
 	// Add tools to the current message context
 	if (context.tools && context.tools.length > 0) {
@@ -392,12 +399,16 @@ function buildConversationState(
 	};
 }
 
-function convertToWireMessage(msg: Context["messages"][number], systemPrompt?: string): WireHistoryMessage {
+function convertToWireMessage(
+	msg: Context["messages"][number],
+	modelId: string,
+	systemPrompt?: string,
+): WireHistoryMessage {
 	if (msg.role === "user") {
-		return convertToWireUserMessage(msg, systemPrompt);
+		return convertToWireUserMessage(msg, modelId, systemPrompt);
 	}
 	if (msg.role === "toolResult") {
-		return convertToWireUserMessage(msg, systemPrompt);
+		return convertToWireUserMessage(msg, modelId, systemPrompt);
 	}
 	// assistant → assistant response
 	const textParts: string[] = [];
@@ -417,7 +428,11 @@ function convertToWireMessage(msg: Context["messages"][number], systemPrompt?: s
 	};
 }
 
-function convertToWireUserMessage(msg: Context["messages"][number], systemPrompt?: string): WireUserMessage {
+function convertToWireUserMessage(
+	msg: Context["messages"][number],
+	modelId: string,
+	systemPrompt?: string,
+): WireUserMessage {
 	let content = extractTextContent(msg);
 	if (systemPrompt) {
 		content = `${systemPrompt}\n\n${content}`;
@@ -426,6 +441,7 @@ function convertToWireUserMessage(msg: Context["messages"][number], systemPrompt
 	const userMsg: WireUserMessage = {
 		userInputMessage: {
 			content,
+			modelId,
 		},
 	};
 

package/src/providers/ollama.ts

@@ -419,6 +419,7 @@ export const streamOllama: StreamFunction<"ollama-chat"> = (
 				url: `${baseUrl}/api/chat`,
 				body,
 			};
+			options?.onStreamCreated?.();
 			let response = await fetchWithRetry(`${baseUrl}/api/chat`, {
 				method: "POST",
 				headers: {
@@ -441,6 +442,7 @@ export const streamOllama: StreamFunction<"ollama-chat"> = (
 				if (
 					firstTokenTime === undefined &&
 					!options.fallbackManaged &&
+					!options.disableProviderRetries &&
 					isForcedToolChoiceUnsupportedError(error, true)
 				) {
 					markToolChoiceIncapability(model, "auto", error.message);
@@ -457,6 +459,7 @@ export const streamOllama: StreamFunction<"ollama-chat"> = (
 					body = { ...body };
 					delete (body as { tool_choice?: unknown }).tool_choice;
 					rawRequestDump = { ...rawRequestDump, body };
+					options?.onStreamCreated?.();
 					response = await fetchWithRetry(`${baseUrl}/api/chat`, {
 						method: "POST",
 						headers: {

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

@@ -844,6 +844,7 @@ async function openInitialCodexEventStream(
 	if (websocketState && shouldUseCodexWebSocket(model, websocketState, options?.preferWebsockets)) {
 		const websocketRetryBudget = getCodexWebSocketRetryBudget(options);
 		let websocketRetries = 0;
+		let lastWebsocketError: Error | undefined;
 		while (true) {
 			try {
 				return await openCodexWebSocketTransport(
@@ -855,6 +856,7 @@ async function openInitialCodexEventStream(
 				);
 			} catch (error) {
 				const websocketError = error instanceof Error ? error : new Error(String(error));
+				lastWebsocketError = websocketError;
 				const isFatal = isCodexWebSocketFatalError(websocketError);
 				const activateFallback = isFatal || websocketRetries >= websocketRetryBudget;
 				recordCodexWebSocketFailure(websocketState, activateFallback);
@@ -875,6 +877,9 @@ async function openInitialCodexEventStream(
 				break;
 			}
 		}
+		if (options?.fallbackManaged || options?.disableProviderRetries) {
+			throw lastWebsocketError ?? new Error("Codex websocket transport failed");
+		}
 	}
 	return openCodexSseTransport(model, requestContext, requestSetup, options, websocketState, transformedBody);
 }
@@ -1739,6 +1744,7 @@ async function tryRetryWithoutForcedToolChoice(
 ): Promise<boolean> {
 	if (
 		context.options?.fallbackManaged ||
+		context.options?.disableProviderRetries ||
 		runtime.toolChoiceFallbackAttempted ||
 		context.output.content.length > 0 ||
 		context.firstTokenTime !== undefined ||
@@ -1842,7 +1848,8 @@ async function tryReconnectCodexWebSocketOnConnectionLimit(
 		!websocketState ||
 		runtime.transport !== "websocket" ||
 		context.options?.signal?.aborted ||
-		context.options?.fallbackManaged
+		context.options?.fallbackManaged ||
+		context.options?.disableProviderRetries
 	) {
 		return false;
 	}
@@ -1905,6 +1912,7 @@ async function tryRecoverCodexPreviousResponseNotFound(
 		runtime.previousResponseRecoveryAttempted ||
 		!websocketState ||
 		context.options?.fallbackManaged ||
+		context.options?.disableProviderRetries ||
 		runtime.transport !== "websocket" ||
 		context.output.content.length > 0 ||
 		context.options?.signal?.aborted ||
@@ -1945,7 +1953,8 @@ async function tryReplayWebsocketFailureOverSse(
 		runtime.canSafelyReplayWebsocketOverSse &&
 		!runtime.sawTerminalEvent &&
 		!context.options?.signal?.aborted &&
-		!context.options?.fallbackManaged;
+		!context.options?.fallbackManaged &&
+		!context.options?.disableProviderRetries;
 	if (!canReplay) return false;
 
 	const state = websocketState;
@@ -1999,7 +2008,8 @@ async function tryRetryCodexProviderError(
 		context.output.content.length > 0 ||
 		runtime.providerRetryAttempt >= resolveRetryBudget(context.options?.streamMaxRetries, CODEX_MAX_RETRIES) ||
 		context.options?.signal?.aborted ||
-		context.options?.fallbackManaged
+		context.options?.fallbackManaged ||
+		context.options?.disableProviderRetries
 	) {
 		return false;
 	}
@@ -2119,7 +2129,7 @@ export const streamOpenAICodexResponses: StreamFunction<"openai-codex-responses"
 			try {
 				initialTransport = await openInitialCodexEventStream(model, streamOptions, requestSetup, requestContext);
 			} catch (error) {
-				if (streamOptions.fallbackManaged) throw error;
+				if (streamOptions.fallbackManaged || streamOptions.disableProviderRetries) throw error;
 				initialTransport = await retryCodexInitialTransportWithoutToolChoice(
 					model,
 					streamOptions,
@@ -2633,6 +2643,7 @@ class CodexWebSocketConnection {
 		signal?: AbortSignal,
 		firstEventTimeoutMs?: number,
 		idleTimeoutMs = this.#idleTimeoutMs,
+		onStreamCreated?: () => void,
 	): AsyncGenerator<Record<string, unknown>> {
 		if (!this.#socket || this.#socket.readyState !== WebSocket.OPEN) {
 			throw createCodexWebSocketTransportError("websocket connection is unavailable");
@@ -2654,6 +2665,7 @@ class CodexWebSocketConnection {
 		}
 
 		try {
+			onStreamCreated?.();
 			this.#socket.send(JSON.stringify(request));
 			let sawFirstProgress = false;
 			const startedAt = Date.now();
@@ -2795,7 +2807,7 @@ async function openCodexSseEventStream(
 	signal?: AbortSignal,
 	onSseEvent?: OpenAICodexResponsesOptions["onSseEvent"],
 	fetchOverride?: FetchImpl,
-	options?: Pick<OpenAICodexResponsesOptions, "requestMaxRetries">,
+	options?: Pick<OpenAICodexResponsesOptions, "requestMaxRetries" | "onStreamCreated">,
 ): Promise<AsyncGenerator<Record<string, unknown>>> {
 	const headers = createCodexHeaders(requestHeaders, accountId, apiKey, sessionId, "sse", state);
 	logCodexDebug("codex request", {
@@ -2805,6 +2817,7 @@ async function openCodexSseEventStream(
 		sentTurnStateHeader: headers.has(X_CODEX_TURN_STATE_HEADER),
 		sentModelsEtagHeader: headers.has(X_MODELS_ETAG_HEADER),
 	});
+	options?.onStreamCreated?.();
 	const response = await fetchWithRetry(url, {
 		method: "POST",
 		headers,
@@ -2845,7 +2858,7 @@ async function openCodexWebSocketEventStream(
 	request: Record<string, unknown>,
 	state: CodexWebSocketSessionState,
 	signal?: AbortSignal,
-	options?: Pick<OpenAICodexResponsesOptions, "streamFirstEventTimeoutMs" | "streamIdleTimeoutMs">,
+	options?: Pick<OpenAICodexResponsesOptions, "streamFirstEventTimeoutMs" | "streamIdleTimeoutMs" | "onStreamCreated">,
 	firstEventTimeoutMs?: number,
 ): Promise<AsyncGenerator<Record<string, unknown>>> {
 	const connection = await getOrCreateCodexWebSocketConnection(state, url, headers, signal, options);
@@ -2854,6 +2867,7 @@ async function openCodexWebSocketEventStream(
 		signal,
 		firstEventTimeoutMs,
 		getCodexWebSocketIdleTimeoutMs(options?.streamIdleTimeoutMs),
+		options?.onStreamCreated,
 	);
 }
 

package/src/providers/openai-completions.ts

@@ -636,6 +636,7 @@ export const streamOpenAICompletions: StreamFunction<"openai-completions"> = (
 					headers: requestHeaders,
 					body: params,
 				};
+				options?.onStreamCreated?.();
 				const { data, response, request_id } = await client.chat.completions
 					.create(params as OpenAI.Chat.Completions.ChatCompletionCreateParamsStreaming, { signal: requestSignal })
 					.withResponse();
@@ -648,6 +649,8 @@ export const streamOpenAICompletions: StreamFunction<"openai-completions"> = (
 					provider: model.provider,
 					signal: requestSignal,
 					fallbackManaged: options?.fallbackManaged,
+					requestMaxRetries: options?.requestMaxRetries,
+					disableProviderRetries: options?.disableProviderRetries,
 				});
 			} catch (error) {
 				const capturedErrorResponse = getCapturedErrorResponse();
@@ -656,6 +659,7 @@ export const streamOpenAICompletions: StreamFunction<"openai-completions"> = (
 				);
 				if (
 					!options?.fallbackManaged &&
+					!options?.disableProviderRetries &&
 					firstTokenTime === undefined &&
 					isForcedToolChoiceUnsupportedError(error, sentForcedToolChoice)
 				) {
@@ -675,6 +679,7 @@ export const streamOpenAICompletions: StreamFunction<"openai-completions"> = (
 					openaiStream = await createCompletionsStream();
 				} else if (
 					!options?.fallbackManaged &&
+					!options?.disableProviderRetries &&
 					isOpenRouterAnthropicModel(model) &&
 					!disableStrictTools &&
 					isCompiledGrammarTooLargeStrictError(error, capturedErrorResponse)
@@ -689,6 +694,7 @@ export const streamOpenAICompletions: StreamFunction<"openai-completions"> = (
 				} else {
 					if (
 						options?.fallbackManaged ||
+						options?.disableProviderRetries ||
 						!shouldRetryWithoutStrictTools(error, capturedErrorResponse, appliedToolStrictMode, context.tools)
 					) {
 						throw error;
@@ -733,7 +739,11 @@ export const streamOpenAICompletions: StreamFunction<"openai-completions"> = (
 							continue;
 						}
 						replaySafeChunks.push(chunk);
-						if (hasNetworkErrorFinishReason(chunk) && !options?.fallbackManaged) {
+						if (
+							hasNetworkErrorFinishReason(chunk) &&
+							!options?.fallbackManaged &&
+							!options?.disableProviderRetries
+						) {
 							retryNetworkError = true;
 							break;
 						}

package/src/providers/openai-responses.ts

@@ -402,16 +402,24 @@ export const streamOpenAIResponses: StreamFunction<"openai-responses"> = (
 			};
 			const openaiStream = await callWithCopilotModelRetry(
 				async () => {
+					options?.onStreamCreated?.();
 					const { data, response, request_id } = await client.responses
 						.create(params, { signal: requestSignal })
 						.withResponse();
 					await notifyProviderResponse(options, response, model, request_id);
 					return data;
 				},
-				{ provider: model.provider, signal: requestSignal, fallbackManaged: options?.fallbackManaged },
+				{
+					provider: model.provider,
+					signal: requestSignal,
+					fallbackManaged: options?.fallbackManaged,
+					requestMaxRetries: options?.requestMaxRetries,
+					disableProviderRetries: options?.disableProviderRetries,
+				},
 			).catch(async error => {
 				if (
 					options?.fallbackManaged ||
+					options?.disableProviderRetries ||
 					!isForcedToolChoiceUnsupportedError(error, isForcedOpenAIResponsesToolChoice(params.tool_choice))
 				) {
 					throw error;
@@ -431,6 +439,7 @@ export const streamOpenAIResponses: StreamFunction<"openai-responses"> = (
 				});
 				delete params.tool_choice;
 				if (rawRequestDump) rawRequestDump.body = params;
+				options?.onStreamCreated?.();
 				const { data, response, request_id } = await client.responses
 					.create(params, { signal: requestSignal })
 					.withResponse();

package/src/providers/pi-native-client.ts

@@ -37,6 +37,7 @@ import { AssistantMessageEventStream } from "../utils/event-stream";
 const NON_WIRE_KEYS = new Set<keyof SimpleStreamOptions>([
 	"signal",
 	"apiKey",
+	"onStreamCreated",
 	"fetch",
 	"onPayload",
 	"onResponse",

package/src/providers/pi-native-server.ts

@@ -78,6 +78,21 @@ const ALLOWED_OPTION_KEYS: ReadonlySet<keyof SimpleStreamOptions> = new Set([
 	"preferWebsockets",
 ] as const satisfies readonly (keyof SimpleStreamOptions)[]);
 
+const CREDENTIAL_HEADER_NAMES = new Set([
+	"authorization",
+	"proxy-authorization",
+	"cookie",
+	"set-cookie",
+	"api-key",
+	"x-api-key",
+	"x-goog-api-key",
+	"anthropic-api-key",
+	"x-auth-token",
+	"x-access-token",
+	"x-api-token",
+	"x-client-secret",
+]);
+
 // ---------------------------------------------------------------------------
 // parseRequest
 // ---------------------------------------------------------------------------
@@ -132,6 +147,15 @@ export function parseRequest(body: unknown, _headers?: Headers): PiNativeParsedR
 		for (const [k, v] of Object.entries(rawOpts)) {
 			if (v === undefined || v === null) continue;
 			if (!ALLOWED_OPTION_KEYS.has(k as keyof SimpleStreamOptions)) continue;
+			if (k === "headers" && typeof v === "object" && v !== null && !Array.isArray(v)) {
+				const safeHeaders: Record<string, string> = {};
+				for (const [name, value] of Object.entries(v)) {
+					if (CREDENTIAL_HEADER_NAMES.has(name.toLowerCase())) continue;
+					if (typeof value === "string") safeHeaders[name] = value;
+				}
+				optsBag[k] = safeHeaders;
+				continue;
+			}
 			optsBag[k] = v;
 		}
 	}

package/src/providers/register-builtins.d.ts

@@ -44,16 +44,16 @@ export declare function resolveLazyStreamFirstEventFallbackMs(provider: string,
 export declare const PROVIDER_RUNTIME_DESCRIPTORS: readonly ProviderRuntimeDescriptor<Api, unknown>[];
 /** Return the lazy descriptor for a built-in API, if one is registered. */
 export declare function getProviderRuntimeDescriptor<TApi extends Api>(api: TApi): ProviderRuntimeDescriptor<TApi, unknown> | undefined;
-export declare const streamAnthropic: (model: Model<"anthropic-messages">, context: Context, options: OptionsForApi<"anthropic-messages">) => EventStreamImpl;
-export declare const streamAzureOpenAIResponses: (model: Model<"azure-openai-responses">, context: Context, options: OptionsForApi<"azure-openai-responses">) => EventStreamImpl;
-export declare const streamGoogle: (model: Model<"google-generative-ai">, context: Context, options: OptionsForApi<"google-generative-ai">) => EventStreamImpl;
-export declare const streamGoogleGeminiCli: (model: Model<"google-gemini-cli">, context: Context, options: OptionsForApi<"google-gemini-cli">) => EventStreamImpl;
-export declare const streamGoogleVertex: (model: Model<"google-vertex">, context: Context, options: OptionsForApi<"google-vertex">) => EventStreamImpl;
-export declare const streamOpenAICodexResponses: (model: Model<"openai-codex-responses">, context: Context, options: OptionsForApi<"openai-codex-responses">) => EventStreamImpl;
-export declare const streamOpenAICompletions: (model: Model<"openai-completions">, context: Context, options: OptionsForApi<"openai-completions">) => EventStreamImpl;
-export declare const streamOpenAIResponses: (model: Model<"openai-responses">, context: Context, options: OptionsForApi<"openai-responses">) => EventStreamImpl;
-export declare const streamCursor: (model: Model<"cursor-agent">, context: Context, options: OptionsForApi<"cursor-agent">) => EventStreamImpl;
-export declare const streamOllama: (model: Model<"ollama-chat">, context: Context, options: OptionsForApi<"ollama-chat">) => EventStreamImpl;
-export declare const streamBedrock: (model: Model<"bedrock-converse-stream">, context: Context, options: OptionsForApi<"bedrock-converse-stream">) => EventStreamImpl;
-export declare const streamKiroCodeWhisperer: (model: Model<"kiro-codewhisperer-stream">, context: Context, options: OptionsForApi<"kiro-codewhisperer-stream">) => EventStreamImpl;
+export declare const streamAnthropic: (model: Model<"anthropic-messages">, context: Context, options: OptionsForApi<"anthropic-messages">, onStreamCreated?: () => void) => EventStreamImpl;
+export declare const streamAzureOpenAIResponses: (model: Model<"azure-openai-responses">, context: Context, options: OptionsForApi<"azure-openai-responses">, onStreamCreated?: () => void) => EventStreamImpl;
+export declare const streamGoogle: (model: Model<"google-generative-ai">, context: Context, options: OptionsForApi<"google-generative-ai">, onStreamCreated?: () => void) => EventStreamImpl;
+export declare const streamGoogleGeminiCli: (model: Model<"google-gemini-cli">, context: Context, options: OptionsForApi<"google-gemini-cli">, onStreamCreated?: () => void) => EventStreamImpl;
+export declare const streamGoogleVertex: (model: Model<"google-vertex">, context: Context, options: OptionsForApi<"google-vertex">, onStreamCreated?: () => void) => EventStreamImpl;
+export declare const streamOpenAICodexResponses: (model: Model<"openai-codex-responses">, context: Context, options: OptionsForApi<"openai-codex-responses">, onStreamCreated?: () => void) => EventStreamImpl;
+export declare const streamOpenAICompletions: (model: Model<"openai-completions">, context: Context, options: OptionsForApi<"openai-completions">, onStreamCreated?: () => void) => EventStreamImpl;
+export declare const streamOpenAIResponses: (model: Model<"openai-responses">, context: Context, options: OptionsForApi<"openai-responses">, onStreamCreated?: () => void) => EventStreamImpl;
+export declare const streamCursor: (model: Model<"cursor-agent">, context: Context, options: OptionsForApi<"cursor-agent">, onStreamCreated?: () => void) => EventStreamImpl;
+export declare const streamOllama: (model: Model<"ollama-chat">, context: Context, options: OptionsForApi<"ollama-chat">, onStreamCreated?: () => void) => EventStreamImpl;
+export declare const streamBedrock: (model: Model<"bedrock-converse-stream">, context: Context, options: OptionsForApi<"bedrock-converse-stream">, onStreamCreated?: () => void) => EventStreamImpl;
+export declare const streamKiroCodeWhisperer: (model: Model<"kiro-codewhisperer-stream">, context: Context, options: OptionsForApi<"kiro-codewhisperer-stream">, onStreamCreated?: () => void) => EventStreamImpl;
 export {};

package/src/providers/register-builtins.ts

@@ -252,10 +252,17 @@ function forwardStream<TApi extends Api>(
 	model: Model<TApi>,
 	options: OptionsForApi<TApi>,
 	abortTracker: AbortSourceTracker,
+	onStreamCreated?: () => void,
 	limits?: LazyStreamLimits,
 ): void {
 	(async () => {
 		try {
+			let admitted = false;
+			const markAdmission = (): void => {
+				if (admitted) return;
+				admitted = true;
+				onStreamCreated?.();
+			};
 			let watchedSource = source;
 			if (!limits?.providerOwnsWatchdog) {
 				const idleTimeoutMs = options.streamIdleTimeoutMs ?? getStreamIdleTimeoutMs(limits?.defaultIdleTimeoutMs);
@@ -286,6 +293,7 @@ function forwardStream<TApi extends Api>(
 			}
 
 			for await (const event of watchedSource) {
+				if (event.type !== "start") markAdmission();
 				target.push(event);
 			}
 			if (hasFinalResult(source)) {
@@ -337,8 +345,13 @@ function createLazyLoadErrorMessage<TApi extends Api>(
 function createLazyStream<TApi extends Api>(
 	loadModule: () => Promise<LazyProviderModule<TApi>>,
 	limits?: LazyStreamLimits,
-): (model: Model<TApi>, context: Context, options: OptionsForApi<TApi>) => EventStreamImpl {
-	return (model, context, options) => {
+): (
+	model: Model<TApi>,
+	context: Context,
+	options: OptionsForApi<TApi>,
+	onStreamCreated?: () => void,
+) => EventStreamImpl {
+	return (model, context, options, onStreamCreated) => {
 		let abortTracker: AbortSourceTracker | undefined;
 		const outer = new EventStreamImpl(() =>
 			abortTracker?.abortLocally(new Error("Provider stream consumer stopped before completion")),
@@ -350,7 +363,7 @@ function createLazyStream<TApi extends Api>(
 				abortTracker = createAbortSourceTracker(streamOptions.signal);
 				const providerOptions = { ...streamOptions, signal: abortTracker.requestSignal } as OptionsForApi<TApi>;
 				const inner = module.stream(model, context, providerOptions);
-				forwardStream(outer, inner, model, streamOptions, abortTracker, limits);
+				forwardStream(outer, inner, model, streamOptions, abortTracker, onStreamCreated, limits);
 			})
 			.catch(error => {
 				const message = createLazyLoadErrorMessage(model, error);

package/src/stream.d.ts

@@ -29,12 +29,13 @@ export declare function listProvidersWithEnvKey(): string[];
  * handling, so callers can append it unconditionally.
  */
 export declare function formatProviderCredentialHint(provider: string): string;
+export declare function streamFromLazyImport(createInner: () => Promise<AssistantMessageEventStream>, signal?: AbortSignal, onStreamCreated?: () => void): AssistantMessageEventStream;
 /**
  * Build an actionable "missing API key" error for a provider, used by the
  * low-level `stream`/`complete` entry points (#755).
  */
 export declare function formatMissingApiKeyError(provider: string): string;
-export declare function stream<TApi extends Api>(model: Model<TApi>, context: Context, options?: OptionsForApi<TApi>): AssistantMessageEventStream;
+export declare function stream<TApi extends Api>(model: Model<TApi>, context: Context, options?: OptionsForApi<TApi>, onStreamCreated?: () => void): AssistantMessageEventStream;
 export declare function complete<TApi extends Api>(model: Model<TApi>, context: Context, options?: OptionsForApi<TApi>): Promise<AssistantMessage>;
 export declare function streamSimple<TApi extends Api>(model: Model<TApi>, context: Context, options?: SimpleStreamOptions): AssistantMessageEventStream;
 export declare function completeSimple<TApi extends Api>(model: Model<TApi>, context: Context, options?: SimpleStreamOptions): Promise<AssistantMessage>;