@vib-rato/ai 0.16.0

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

@@ -0,0 +1,488 @@
+/**
+ * Pi-native wire format for the auth-gateway.
+ *
+ * Where the OpenAI / Anthropic / Responses route modules translate foreign
+ * wire shapes through pi-ai's canonical {@link Context}, this module accepts
+ * the canonical shape *directly* — for clients that already speak pi-ai
+ * (containerized Vibrato deployments and sidecar auth gateways).
+ * Skipping the wire-format → Context → wire-format round-trip cuts
+ * per-request CPU but, more importantly, avoids the quantization that those
+ * translations impose on first-class pi-ai fields (service tier, cache
+ * markers, thinking budgets, tool-choice variants, …).
+ *
+ * The streaming wire is {@link AssistantMessageEvent} serialized as SSE. Public
+ * projections omit private raw reasoning and serialized Responses reasoning
+ * signatures while preserving provider-displayable summaries and genuine opaque
+ * signatures. Including `partial: AssistantMessage` on every delta is O(N²) in
+ * turn length on the wire — acceptable for the loopback / sidecar topology this
+ * transport is designed for; provider latency dominates the actual cost.
+ *
+ * Endpoint contract:
+ *   POST /v1/pi/stream
+ *   body:    { modelId, context, options?, stream? }   // `stream` defaults to true
+ *   200 SSE: stream of `AssistantMessageEvent` (terminated by `data: [DONE]`)
+ *   200 JSON (stream=false): { message: AssistantMessage }
+ *   4xx/5xx: { error: { type, message } }
+ */
+import type {
+	AssistantMessage,
+	AssistantMessageEvent,
+	AssistantMessageEventStream,
+	Context,
+	SimpleStreamOptions,
+	ThinkingContent,
+} from "../types";
+
+export interface PiNativeParsedRequest {
+	modelId: string;
+	context: Context;
+	options: SimpleStreamOptions;
+	stream: boolean;
+}
+/**
+ * Subset of {@link SimpleStreamOptions} accepted from the wire. Function-valued
+ * fields (`fetch`, `onPayload`, `onResponse`, `onSseEvent`, exec handlers, the
+ * provider-session map) and gateway-owned controls (`apiKey`, `signal`) are
+ * intentionally absent — those are server-side concerns. Anything outside this
+ * allow-list is dropped silently rather than 400ing, so clients can forward
+ * `SimpleStreamOptions` from older / newer vib builds without per-version
+ * conditionals.
+ */
+const ALLOWED_OPTION_KEYS: ReadonlySet<keyof SimpleStreamOptions> = new Set([
+	"temperature",
+	"topP",
+	"topK",
+	"minP",
+	"presencePenalty",
+	"frequencyPenalty",
+	"repetitionPenalty",
+	"stopSequences",
+	"maxTokens",
+	"cacheRetention",
+	"headers",
+	"initiatorOverride",
+	"maxRetryDelayMs",
+	"fallbackManaged",
+	"metadata",
+	"sessionId",
+	"streamFirstEventTimeoutMs",
+	"streamIdleTimeoutMs",
+	"reasoning",
+	"disableReasoning",
+	"hideThinkingSummary",
+	"thinkingBudgets",
+	"toolChoice",
+	"serviceTier",
+	"kimiApiFormat",
+	"syntheticApiFormat",
+	"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
+// ---------------------------------------------------------------------------
+
+/**
+ * Parse a pi-native request body. Validation is intentionally minimal — only
+ * the shape the gateway itself reads is checked (`modelId`, `context.messages`
+ * array, options is an object). Everything downstream is the canonical pi-ai
+ * type surface; mis-shaped values surface as a `502 upstream_error` from
+ * `streamSimple` rather than being re-validated here.
+ *
+ * Accepts both `{ modelId: string }` and `{ model: { id: string } }` so the
+ * existing `streamProxy` client (which sends the full Model object) can target
+ * the gateway with only a URL swap.
+ */
+export function parseRequest(body: unknown, _headers?: Headers): PiNativeParsedRequest {
+	if (typeof body !== "object" || body === null || Array.isArray(body)) {
+		throw new Error("Request body must be a JSON object");
+	}
+	const obj = body as Record<string, unknown>;
+
+	let modelId: string | undefined;
+	if (typeof obj.modelId === "string" && obj.modelId.length > 0) {
+		modelId = obj.modelId;
+	} else if (typeof obj.model === "string" && obj.model.length > 0) {
+		modelId = obj.model;
+	} else if (typeof obj.model === "object" && obj.model !== null) {
+		const m = obj.model as Record<string, unknown>;
+		if (typeof m.id === "string" && m.id.length > 0) modelId = m.id;
+	}
+	if (!modelId) throw new Error("Missing `modelId` (or `model.id`) field");
+
+	const context = obj.context;
+	if (typeof context !== "object" || context === null || Array.isArray(context)) {
+		throw new Error("Missing `context` object");
+	}
+	const ctxObj = context as Record<string, unknown>;
+	if (!Array.isArray(ctxObj.messages)) {
+		throw new Error("`context.messages` must be an array");
+	}
+	if (ctxObj.systemPrompt !== undefined && !Array.isArray(ctxObj.systemPrompt)) {
+		throw new Error("`context.systemPrompt` must be an array of strings when present");
+	}
+	if (ctxObj.tools !== undefined && !Array.isArray(ctxObj.tools)) {
+		throw new Error("`context.tools` must be an array when present");
+	}
+
+	const options: SimpleStreamOptions = {};
+	const rawOpts = obj.options;
+	if (typeof rawOpts === "object" && rawOpts !== null && !Array.isArray(rawOpts)) {
+		const optsBag = options as Record<string, unknown>;
+		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;
+		}
+	}
+
+	// `stream` defaults to true — pi-native clients overwhelmingly stream, and
+	// matching `streamProxy`'s implicit-stream behavior avoids a one-flag papercut.
+	const stream = typeof obj.stream === "boolean" ? obj.stream : true;
+
+	return {
+		modelId,
+		context: context as Context,
+		options,
+		stream,
+	};
+}
+// ---------------------------------------------------------------------------
+// encodeStream (SSE)
+// ---------------------------------------------------------------------------
+
+const SSE_ENCODER = new TextEncoder();
+const SSE_DONE = SSE_ENCODER.encode("data: [DONE]\n\n");
+
+function isSerializedResponsesReasoningItem(signature: string): boolean {
+	try {
+		const parsed: unknown = JSON.parse(signature);
+		return (
+			typeof parsed === "object" &&
+			parsed !== null &&
+			!Array.isArray(parsed) &&
+			(parsed as { type?: unknown }).type === "reasoning"
+		);
+	} catch {
+		return false;
+	}
+}
+
+/**
+ * Clone a thinking block for public transport. Raw reasoning is private: omit
+ * raw-only blocks, retain only the displayable summary for mixed blocks, and
+ * never forward a serialized Responses reasoning item as a signature.
+ */
+function isResponsesFamilyApi(api: AssistantMessage["api"]): boolean {
+	return api === "openai-responses" || api === "openai-codex-responses";
+}
+
+function sanitizeThinking(content: ThinkingContent, api: AssistantMessage["api"]): ThinkingContent | undefined {
+	if (isResponsesFamilyApi(api) && content.provenance === undefined) return undefined;
+	if (content.provenance === "raw") return undefined;
+
+	let thinking: string;
+	if (content.provenance === "mixed") {
+		if (content.summaryText === undefined) return undefined;
+		thinking = content.summaryText;
+	} else {
+		thinking = content.provenance === "summary" ? (content.summaryText ?? content.thinking) : content.thinking;
+	}
+	const signature =
+		content.thinkingSignature && isSerializedResponsesReasoningItem(content.thinkingSignature)
+			? undefined
+			: content.thinkingSignature;
+	const { rawText: _rawText, thinkingSignature: _thinkingSignature, ...rest } = content;
+	return signature === undefined ? { ...rest, thinking } : { ...rest, thinking, thinkingSignature: signature };
+}
+
+function sanitizeMessage(message: AssistantMessage): AssistantMessage {
+	let changed = false;
+	const content: AssistantMessage["content"] = [];
+	for (const part of message.content) {
+		if (part.type !== "thinking") {
+			content.push(part);
+			continue;
+		}
+		const needsSanitizing =
+			(isResponsesFamilyApi(message.api) && part.provenance === undefined) ||
+			part.provenance !== undefined ||
+			part.rawText !== undefined ||
+			(part.thinkingSignature !== undefined && isSerializedResponsesReasoningItem(part.thinkingSignature));
+		if (!needsSanitizing) {
+			content.push(part);
+			continue;
+		}
+		const sanitized = sanitizeThinking(part, message.api);
+		changed = true;
+		if (sanitized !== undefined) content.push(sanitized);
+	}
+	return changed ? { ...message, content } : message;
+}
+
+function hasRawOrMixedThinking(partial: AssistantMessage, contentIndex: number): boolean {
+	const content = partial.content[contentIndex];
+	return content?.type === "thinking" && (content.provenance === "raw" || content.provenance === "mixed");
+}
+
+type ThinkingEvent = Extract<AssistantMessageEvent, { type: "thinking_start" | "thinking_delta" | "thinking_end" }>;
+
+interface BufferedThinkingEvent {
+	event: ThinkingEvent;
+	sequence: number;
+}
+
+function isFinalSafeThinking(partial: AssistantMessage, contentIndex: number): boolean {
+	const content = partial.content[contentIndex];
+	return (
+		content?.type === "thinking" &&
+		(!isResponsesFamilyApi(partial.api) || content.provenance !== undefined) &&
+		!hasRawOrMixedThinking(partial, contentIndex)
+	);
+}
+
+function maskBufferedThinking(message: AssistantMessage, contentIndexes: ReadonlySet<number>): AssistantMessage {
+	let changed = false;
+	const content = message.content.map((part, contentIndex) => {
+		if (!contentIndexes.has(contentIndex) || part.type !== "thinking") return part;
+		changed = true;
+		return { type: "thinking" as const, thinking: "", ...(part.itemId ? { itemId: part.itemId } : {}) };
+	});
+	return changed ? { ...message, content } : message;
+}
+
+function maskBufferedThinkingInEvent(
+	event: AssistantMessageEvent,
+	contentIndexes: ReadonlySet<number>,
+): AssistantMessageEvent {
+	if (contentIndexes.size === 0) return event;
+	switch (event.type) {
+		case "done":
+		case "error":
+		case "toolChoiceIncapability":
+			return event;
+		case "start":
+		case "text_start":
+		case "text_delta":
+		case "text_end":
+		case "thinking_start":
+		case "thinking_delta":
+		case "thinking_end":
+		case "reasoning_summary_start":
+		case "reasoning_summary_delta":
+		case "reasoning_summary_end":
+		case "toolcall_start":
+		case "toolcall_delta":
+		case "toolcall_end":
+			return { ...event, partial: maskBufferedThinking(event.partial, contentIndexes) };
+	}
+}
+
+function withSummaryPartial<
+	T extends Extract<
+		AssistantMessageEvent,
+		{ type: "reasoning_summary_start" | "reasoning_summary_delta" | "reasoning_summary_end" }
+	>,
+>(event: T, contentIndexes: ReadonlySet<number>, summaryText: string): T {
+	const partial = maskBufferedThinking(event.partial, contentIndexes);
+	const content = [...partial.content];
+	const original = event.partial.content[event.contentIndex];
+	content[event.contentIndex] = {
+		type: "thinking",
+		thinking: summaryText,
+		provenance: "summary",
+		summaryText,
+		...(original?.type === "thinking" && original.itemId ? { itemId: original.itemId } : {}),
+	};
+	return { ...event, partial: { ...partial, content } };
+}
+
+function sanitizeEvent(event: AssistantMessageEvent): AssistantMessageEvent | undefined {
+	switch (event.type) {
+		case "done":
+			return { ...event, message: sanitizeMessage(event.message) };
+		case "error":
+			return { ...event, error: sanitizeMessage(event.error) };
+		case "toolChoiceIncapability":
+			return event;
+		case "thinking_start":
+		case "thinking_delta":
+		case "thinking_end":
+			return hasRawOrMixedThinking(event.partial, event.contentIndex)
+				? undefined
+				: { ...event, partial: sanitizeMessage(event.partial) };
+		case "start":
+		case "text_start":
+		case "text_delta":
+		case "text_end":
+		case "reasoning_summary_start":
+		case "reasoning_summary_delta":
+		case "reasoning_summary_end":
+		case "toolcall_start":
+		case "toolcall_delta":
+		case "toolcall_end":
+			return { ...event, partial: sanitizeMessage(event.partial) };
+	}
+}
+
+/**
+ * Ship only public-safe {@link AssistantMessageEvent} projections. Unknown
+ * thinking blocks remain buffered until their terminal partial establishes that
+ * the provider-native block is safe; raw and mixed blocks never reach SSE.
+ */
+export function encodeStream(events: AssistantMessageEventStream): ReadableStream<Uint8Array> {
+	return new ReadableStream<Uint8Array>({
+		async start(controller) {
+			const bufferedThinking = new Map<number, BufferedThinkingEvent[]>();
+			const summaryTextByIndex = new Map<number, string>();
+			let sequence = 0;
+			const write = (event: AssistantMessageEvent): void => {
+				const sanitized = sanitizeEvent(event);
+				if (sanitized !== undefined) {
+					controller.enqueue(SSE_ENCODER.encode(`data: ${JSON.stringify(sanitized)}\n\n`));
+				}
+			};
+			const emit = (event: AssistantMessageEvent): void => {
+				write(maskBufferedThinkingInEvent(event, new Set(bufferedThinking.keys())));
+			};
+			const flush = (buffered: BufferedThinkingEvent[]): void => {
+				for (const { event } of buffered.sort((a, b) => a.sequence - b.sequence)) emit(event);
+			};
+			const resolveBufferedThinking = (final: AssistantMessage): void => {
+				const ready: BufferedThinkingEvent[] = [];
+				for (const [contentIndex, buffered] of bufferedThinking) {
+					if (isFinalSafeThinking(final, contentIndex)) ready.push(...buffered);
+				}
+				bufferedThinking.clear();
+				flush(ready);
+			};
+			try {
+				for await (const event of events) {
+					switch (event.type) {
+						case "thinking_start":
+						case "thinking_delta": {
+							if (hasRawOrMixedThinking(event.partial, event.contentIndex)) {
+								bufferedThinking.delete(event.contentIndex);
+								break;
+							}
+							const buffered = bufferedThinking.get(event.contentIndex) ?? [];
+							buffered.push({ event, sequence: sequence++ });
+							bufferedThinking.set(event.contentIndex, buffered);
+							break;
+						}
+						case "thinking_end": {
+							const buffered = bufferedThinking.get(event.contentIndex) ?? [];
+							bufferedThinking.delete(event.contentIndex);
+							if (!isFinalSafeThinking(event.partial, event.contentIndex)) break;
+							buffered.push({ event, sequence: sequence++ });
+							flush(buffered);
+							break;
+						}
+						case "done":
+							resolveBufferedThinking(event.message);
+							summaryTextByIndex.clear();
+							emit(event);
+							controller.enqueue(SSE_DONE);
+							controller.close();
+							return;
+						case "error":
+							resolveBufferedThinking(event.error);
+							summaryTextByIndex.clear();
+							emit(event);
+							controller.enqueue(SSE_DONE);
+							controller.close();
+							return;
+						case "reasoning_summary_start": {
+							summaryTextByIndex.set(event.contentIndex, "");
+							write(withSummaryPartial(event, new Set(bufferedThinking.keys()), ""));
+							break;
+						}
+						case "reasoning_summary_delta": {
+							const summaryText = `${summaryTextByIndex.get(event.contentIndex) ?? ""}${event.delta}`;
+							summaryTextByIndex.set(event.contentIndex, summaryText);
+							write(withSummaryPartial(event, new Set(bufferedThinking.keys()), summaryText));
+							break;
+						}
+						case "reasoning_summary_end": {
+							const summaryText = event.content || summaryTextByIndex.get(event.contentIndex) || "";
+							summaryTextByIndex.delete(event.contentIndex);
+							write(withSummaryPartial(event, new Set(bufferedThinking.keys()), summaryText));
+							break;
+						}
+						case "start":
+						case "text_start":
+						case "text_delta":
+						case "text_end":
+						case "toolcall_start":
+						case "toolcall_delta":
+						case "toolcall_end":
+						case "toolChoiceIncapability":
+							emit(event);
+							break;
+					}
+				}
+				controller.enqueue(SSE_DONE);
+				controller.close();
+			} catch (err) {
+				// Best-effort error envelope so the client iterator resolves
+				// instead of hanging on the dropped connection. Shape matches the
+				// canonical `error` event minus the unrecoverable `error:
+				// AssistantMessage` payload (we don't have a usable one here).
+				const message = err instanceof Error ? err.message : String(err);
+				controller.enqueue(
+					SSE_ENCODER.encode(
+						`data: ${JSON.stringify({ type: "error", reason: "error", errorMessage: message })}\n\n`,
+					),
+				);
+				controller.enqueue(SSE_DONE);
+				controller.close();
+			}
+		},
+	});
+}
+
+// ---------------------------------------------------------------------------
+// formatError
+// ---------------------------------------------------------------------------
+
+/**
+ * Pi-native error envelope:
+ *   `{ error: { type, message } }`
+ *
+ * Mirrors OpenAI's outer shape (which clients/SDKs already parse) without the
+ * provider-specific status taxonomy — pi-native callers consume `type`
+ * directly.
+ */
+export function formatError(status: number, type: string, message: string): Response {
+	return new Response(JSON.stringify({ error: { type, message } }), {
+		status,
+		headers: {
+			"Content-Type": "application/json; charset=utf-8",
+			"Cache-Control": "no-store",
+		},
+	});
+}

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

@@ -0,0 +1,59 @@
+/**
+ * Lazy provider module loading.
+ *
+ * Each provider module is loaded only when its stream function is first called.
+ * This avoids eagerly importing heavy SDK dependencies (e.g., @anthropic-ai/sdk,
+ * openai) at startup. The loaded module promise is cached so subsequent calls
+ * reuse the same import.
+ *
+ * stream.ts imports its provider stream functions from this module (see the
+ * lazy wrappers below), so this file IS the main streaming path's provider
+ * loader: heavy SDKs stay out of the CLI startup parse graph.
+ */
+import type { Api, AssistantMessageEventStream, Context, Model, OptionsForApi } from "../types";
+import { AssistantMessageEventStream as EventStreamImpl } from "../utils/event-stream";
+import type { BedrockOptions } from "./amazon-bedrock";
+/**
+ * Lazy runtime descriptor for a built-in provider implementation.
+ *
+ * The registry stores descriptors with an erased module type because each
+ * provider's stream options are intentionally different. Callers narrow the
+ * loaded module at the single API dispatch boundary instead of forcing
+ * distributive variance through the collection type.
+ */
+export interface ProviderRuntimeDescriptor<TApi extends Api = Api, TModule = unknown> {
+    readonly api: TApi;
+    readonly load: () => Promise<TModule>;
+}
+interface BedrockProviderModule {
+    streamBedrock: (model: Model<"bedrock-converse-stream">, context: Context, options: BedrockOptions) => AssistantMessageEventStream;
+}
+export declare function setBedrockProviderModule(module: BedrockProviderModule): void;
+/**
+ * Resolves the first-event timeout fallback for the outer lazy-stream watchdog.
+ * A configured wrapper-specific fallback (from `LazyStreamLimits`) always wins;
+ * otherwise providers known to have slow first events use the same centralized
+ * fallback as their inner provider-level watchdog. Returns `undefined` for
+ * providers that should use the shared default.
+ */
+export declare function resolveLazyStreamFirstEventFallbackMs(provider: string, configuredFallbackMs?: number): number | undefined;
+/**
+ * Lazy provider descriptors used by core consumers that need to inspect or
+ * prewarm a provider without importing its implementation at startup.
+ */
+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">, 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 {};