@oh-my-pi/pi-ai 18.2.7 → 18.2.8

package/dist/types/rerank/types.d.ts

@@ -0,0 +1,17 @@
+import type { Usage } from "@oh-my-pi/pi-catalog/types";
+export interface RerankRequest {
+    query: string;
+    documents: string[];
+    topN?: number;
+    returnDocuments?: boolean;
+}
+export interface RerankResultItem {
+    index: number;
+    relevanceScore: number;
+    document?: string;
+}
+export interface RerankResult {
+    results: RerankResultItem[];
+    model: string;
+    usage: Usage;
+}

package/dist/types/speech/index.d.ts

@@ -0,0 +1,13 @@
+import type { Api, Model } from "@oh-my-pi/pi-catalog/types";
+import type { SpeechOptions, SpeechRequest, SpeechResult } from "./types.js";
+export * from "./openai-speech.js";
+export * from "./transport.js";
+export * from "./types.js";
+export * from "./xai-tts.js";
+/** Catalog APIs {@link synthesizeSpeech} serves; the `tts` tool and the gateway gate on these. */
+export declare const SPEECH_APIS: readonly ["xai-tts", "openai-speech"];
+export type SpeechApi = (typeof SPEECH_APIS)[number];
+/** Whether a catalog API synthesizes speech through a cloud transport (local inference is tool-only). */
+export declare function isSpeechApi(api: Api): api is SpeechApi;
+/** Synthesize speech through the transport selected by the catalog model's `api`. */
+export declare function synthesizeSpeech(model: Model<Api>, request: SpeechRequest, options: SpeechOptions): Promise<SpeechResult>;

package/dist/types/speech/openai-speech.d.ts

@@ -0,0 +1,3 @@
+import type { Api, Model } from "@oh-my-pi/pi-catalog/types";
+import type { SpeechOptions, SpeechRequest, SpeechResult } from "./types.js";
+export declare function synthesizeOpenAiSpeech(model: Model<Api>, request: SpeechRequest, options: SpeechOptions): Promise<SpeechResult>;

package/dist/types/speech/transport.d.ts

@@ -0,0 +1,7 @@
+import type { Api, Model } from "@oh-my-pi/pi-catalog/types";
+import * as AIError from "../error/index.js";
+import { type SpeechFormat, type SpeechOptions, type SpeechResult } from "./types.js";
+export declare class SpeechApiError extends AIError.ProviderHttpError {
+    readonly name = "SpeechApiError";
+}
+export declare function postSpeechRequest(model: Model<Api>, path: string, payload: Record<string, unknown>, format: SpeechFormat, options: SpeechOptions): Promise<SpeechResult>;

package/dist/types/speech/types.d.ts

@@ -0,0 +1,24 @@
+import type { FetchImpl, Usage } from "@oh-my-pi/pi-catalog/types";
+import type { ApiKey } from "../auth-retry.js";
+export declare const SPEECH_FORMATS: readonly ["mp3", "wav", "pcm", "opus", "aac", "flac"];
+export type SpeechFormat = (typeof SPEECH_FORMATS)[number];
+export declare const SPEECH_FORMAT_MIME_TYPES: Record<SpeechFormat, string>;
+export interface SpeechRequest {
+    text: string;
+    voice?: string;
+    format: SpeechFormat;
+    speed?: number;
+    sampleRate?: number;
+    bitRate?: number;
+    instructions?: string;
+}
+export interface SpeechResult {
+    audio: Uint8Array;
+    mimeType: string;
+    usage: Usage;
+}
+export interface SpeechOptions {
+    apiKey: ApiKey;
+    fetch?: FetchImpl;
+    signal?: AbortSignal;
+}

package/dist/types/speech/xai-tts.d.ts

@@ -0,0 +1,7 @@
+import type { Api, Model } from "@oh-my-pi/pi-catalog/types";
+import type { SpeechOptions, SpeechRequest, SpeechResult } from "./types.js";
+export declare const DEFAULT_XAI_VOICE_ID = "eve";
+export declare const DEFAULT_XAI_SAMPLE_RATE = 24000;
+export declare const DEFAULT_XAI_BIT_RATE = 128000;
+export declare const XAI_MAX_TEXT_LENGTH = 15000;
+export declare function synthesizeXaiSpeech(model: Model<Api>, request: SpeechRequest, options: SpeechOptions): Promise<SpeechResult>;

package/dist/types/transcription/index.d.ts

@@ -0,0 +1,7 @@
+import type { Api, Model } from "@oh-my-pi/pi-catalog/types";
+import { type TranscriptionOptions } from "./openai-transcriptions.js";
+import type { TranscriptionRequest, TranscriptionResult } from "./types.js";
+export * from "./openai-transcriptions.js";
+export * from "./types.js";
+/** Dispatch an audio transcription through the transport selected by the catalog model. */
+export declare function transcribeAudio(model: Model<Api>, request: TranscriptionRequest, options: TranscriptionOptions): Promise<TranscriptionResult>;

package/dist/types/transcription/openai-transcriptions.d.ts

@@ -0,0 +1,15 @@
+import type { Api, FetchImpl, Model } from "@oh-my-pi/pi-catalog/types";
+import { type ApiKey } from "../auth-retry.js";
+import * as AIError from "../error/index.js";
+import type { TranscriptionRequest, TranscriptionResult } from "./types.js";
+export interface TranscriptionOptions {
+    apiKey: ApiKey;
+    fetch?: FetchImpl;
+    signal?: AbortSignal;
+}
+/** Non-2xx response from an OpenAI-compatible transcription endpoint. */
+export declare class TranscriptionApiError extends AIError.ProviderHttpError {
+    readonly name = "TranscriptionApiError";
+}
+/** Call an OpenAI/OpenRouter-compatible multipart transcription endpoint. */
+export declare function transcribeOpenAI(model: Model<Api>, request: TranscriptionRequest, options: TranscriptionOptions): Promise<TranscriptionResult>;

package/dist/types/transcription/types.d.ts

@@ -0,0 +1,41 @@
+import type { Usage } from "@oh-my-pi/pi-catalog/types";
+export type TranscriptionResponseFormat = "json" | "verbose_json";
+export type TranscriptionTimestampGranularity = "word" | "segment";
+export interface TranscriptionRequest {
+    audio: Uint8Array;
+    mimeType: string;
+    fileName?: string;
+    language?: string;
+    prompt?: string;
+    temperature?: number;
+    responseFormat: TranscriptionResponseFormat;
+    timestampGranularities?: TranscriptionTimestampGranularity[];
+}
+/** Provider-specific segment fields are preserved alongside the normalized timestamps and text. */
+export interface TranscriptionSegment {
+    id?: number | string;
+    start: number;
+    end: number;
+    text: string;
+    speaker?: number | string;
+    [key: string]: unknown;
+}
+/** Provider-specific word fields are preserved alongside the normalized timestamps and spelling. */
+export interface TranscriptionWord {
+    word: string;
+    start: number;
+    end: number;
+    speaker?: number | string;
+    confidence?: number;
+    [key: string]: unknown;
+}
+export interface TranscriptionResult {
+    text: string;
+    language?: string;
+    duration?: number;
+    segments?: TranscriptionSegment[];
+    words?: TranscriptionWord[];
+    /** Provider-billed audio duration, which may be present without verbose output. */
+    seconds?: number;
+    usage: Usage;
+}

package/dist/types/video/index.d.ts

@@ -0,0 +1,11 @@
+import type { Api, Model } from "@oh-my-pi/pi-catalog/types";
+import { type VideoOptions } from "./openrouter-video.js";
+import type { VideoContent, VideoGenerationRequest, VideoJob } from "./types.js";
+export * from "./openrouter-video.js";
+export * from "./types.js";
+/** Submit an asynchronous video generation job through the model's transport. */
+export declare function submitVideo(model: Model<Api>, request: VideoGenerationRequest, options: VideoOptions): Promise<VideoJob>;
+/** Poll an asynchronous video generation job through the model's transport. */
+export declare function pollVideo(model: Model<Api>, jobId: string, options: VideoOptions): Promise<VideoJob>;
+/** Stream generated video content through the model's transport. */
+export declare function downloadVideo(model: Model<Api>, jobId: string, options: VideoOptions): Promise<VideoContent>;

package/dist/types/video/openrouter-video.d.ts

@@ -0,0 +1,19 @@
+import type { Api, FetchImpl, Model } from "@oh-my-pi/pi-catalog/types";
+import { type ApiKey } from "../auth-retry.js";
+import * as AIError from "../error/index.js";
+import type { VideoContent, VideoGenerationRequest, VideoJob } from "./types.js";
+export interface VideoOptions {
+    apiKey: ApiKey;
+    fetch?: FetchImpl;
+    signal?: AbortSignal;
+}
+/** Non-2xx response from an OpenRouter video endpoint. */
+export declare class VideoApiError extends AIError.ProviderHttpError {
+    readonly name = "VideoApiError";
+}
+/** Submit an asynchronous OpenRouter video generation job. */
+export declare function submitOpenRouterVideo(model: Model<Api>, request: VideoGenerationRequest, options: VideoOptions): Promise<VideoJob>;
+/** Poll an asynchronous OpenRouter video generation job. */
+export declare function pollOpenRouterVideo(model: Model<Api>, jobId: string, options: VideoOptions): Promise<VideoJob>;
+/** Stream generated video bytes from OpenRouter without buffering them in memory. */
+export declare function downloadOpenRouterVideo(model: Model<Api>, jobId: string, options: VideoOptions): Promise<VideoContent>;

package/dist/types/video/types.d.ts

@@ -0,0 +1,62 @@
+import type { Usage } from "@oh-my-pi/pi-catalog/types";
+export type VideoResolution = "360p" | "480p" | "720p" | "768p" | "1080p" | "1K" | "2K" | "4K";
+export type VideoAspectRatio = "16:9" | "9:16" | "1:1" | "4:3" | "3:4" | "3:2" | "2:3" | "21:9" | "9:21";
+export interface VideoImageReference {
+    type: "image_url";
+    imageUrl: {
+        url: string;
+    };
+}
+export interface VideoFrameImage extends VideoImageReference {
+    frameType: "first_frame" | "last_frame";
+}
+export interface VideoAudioReference {
+    type: "audio_url";
+    audioUrl: {
+        url: string;
+    };
+}
+export interface VideoVideoReference {
+    type: "video_url";
+    videoUrl: {
+        url: string;
+    };
+}
+export type VideoInputReference = VideoImageReference | VideoAudioReference | VideoVideoReference;
+/** Canonical camel-case form of OpenRouter's asynchronous video submit request. */
+export interface VideoGenerationRequest {
+    prompt?: string;
+    duration?: number;
+    resolution?: VideoResolution;
+    aspectRatio?: VideoAspectRatio;
+    size?: string;
+    frameImages?: VideoFrameImage[];
+    inputReferences?: VideoInputReference[];
+    generateAudio?: boolean;
+    seed?: number;
+    callbackUrl?: string;
+    provider?: {
+        options?: Record<string, unknown>;
+    };
+    previousJobId?: string;
+    sessionId?: string;
+    trace?: Record<string, unknown>;
+    user?: string;
+    creativity?: number;
+    upscaleFactor?: number;
+}
+export type VideoJobStatus = "queued" | "processing" | "pending" | "in_progress" | "completed" | "failed" | "cancelled" | "expired";
+export interface VideoJob {
+    id: string;
+    status: VideoJobStatus;
+    pollingUrl?: string;
+    generationId?: string;
+    contentUrls?: string[];
+    error?: string;
+    usage?: Usage;
+}
+export interface VideoContent {
+    body: ReadableStream<Uint8Array>;
+    contentType: string;
+    contentLength?: number;
+}

package/package.json

@@ -1,6 +1,6 @@
 {
 	"name": "@oh-my-pi/pi-ai",
-	"version": "18.2.7",
+	"version": "18.2.8",
 	"description": "Unified LLM API with automatic model discovery and provider configuration",
 	"keywords": [
 		"ai",
@@ -45,10 +45,34 @@
 			"types": "./dist/types/error/index.d.ts",
 			"import": "./src/error/index.ts"
 		},
+		"./embeddings": {
+			"types": "./dist/types/embeddings/index.d.ts",
+			"import": "./src/embeddings/index.ts"
+		},
 		"./judgment": {
 			"types": "./dist/types/judgment/index.d.ts",
 			"import": "./src/judgment/index.ts"
 		},
+		"./rerank": {
+			"types": "./dist/types/rerank/index.d.ts",
+			"import": "./src/rerank/index.ts"
+		},
+		"./speech": {
+			"types": "./dist/types/speech/index.d.ts",
+			"import": "./src/speech/index.ts"
+		},
+		"./images": {
+			"types": "./dist/types/images/index.d.ts",
+			"import": "./src/images/index.ts"
+		},
+		"./transcription": {
+			"types": "./dist/types/transcription/index.d.ts",
+			"import": "./src/transcription/index.ts"
+		},
+		"./video": {
+			"types": "./dist/types/video/index.d.ts",
+			"import": "./src/video/index.ts"
+		},
 		"./*": {
 			"types": "./dist/types/*.d.ts",
 			"import": "./src/*.ts"
@@ -128,11 +152,11 @@
 		"fmt": "oxfmt --no-error-on-unmatched-pattern 'src/**/*.{ts,tsx}' '{test,bench,examples,scripts}/**/*.ts' '*.ts'"
 	},
 	"dependencies": {
-		"@oh-my-pi/omptype": "18.2.7",
-		"@oh-my-pi/pi-catalog": "18.2.7",
-		"@oh-my-pi/pi-natives": "18.2.7",
-		"@oh-my-pi/pi-utils": "18.2.7",
-		"@oh-my-pi/pi-wire": "18.2.7"
+		"@oh-my-pi/omptype": "18.2.8",
+		"@oh-my-pi/pi-catalog": "18.2.8",
+		"@oh-my-pi/pi-natives": "18.2.8",
+		"@oh-my-pi/pi-utils": "18.2.8",
+		"@oh-my-pi/pi-wire": "18.2.8"
 	},
 	"devDependencies": {
 		"@types/bun": "^1.3.14"

package/src/auth-gateway/dispatch.ts

@@ -0,0 +1,273 @@
+/**
+ * Per-request plumbing shared by every auth-gateway route.
+ *
+ * A route module (`server.ts` chat/pi-native handlers, `routes/*.ts` for
+ * judgments, images, speech, transcription, …) owns only its wire format:
+ * parse the body, pick a model, call the pi-ai client, encode the reply.
+ * Everything credential-shaped lives here so each route drives the same
+ * broker-backed rotation policy and the same usage ledger.
+ */
+import { extractHttpStatusFromError, logger } from "@oh-my-pi/pi-utils";
+import type { ApiKeyResolver } from "../auth-retry";
+import type { AuthStorage } from "../auth-storage";
+import * as AIError from "../error";
+import { classifyGatewayError, type GatewayErrorClassification } from "../error/gateway";
+import { isUsageLimitOutcome } from "../error/rate-limit";
+import type { Api, FetchImpl, Model, Usage } from "../types";
+import type { ClientUsageIdentity } from "../usage";
+import { extractProviderRetryHint } from "../utils/retry-after";
+import type { AuthGatewayServerOptions } from "./types";
+
+export type ModelResolver = (modelId: string) => Model<Api> | undefined;
+
+export interface AuthGatewayBootOptions extends AuthGatewayServerOptions {
+	/** Source of credentials. Caller wires this to a broker-backed AuthStorage. */
+	storage: AuthStorage;
+	/**
+	 * Resolve a client-requested model id to a pi-ai Model. Caller supplies
+	 * this from a ModelRegistry (lives in `coding-agent` to avoid an inverse
+	 * dependency in `pi-ai`).
+	 */
+	resolveModel: ModelResolver;
+	/** Optional supplier for `/v1/models` listing. Returns the full model array. */
+	listModels?: () => Iterable<Model<Api>>;
+	/** Upstream transport for every provider call; defaults to global `fetch`. Test seam. */
+	fetch?: FetchImpl;
+}
+
+/**
+ * The client's own session key, or `undefined` when it sent none. A blank key
+ * counts as none: honouring it would collapse every caller that sends an empty
+ * key into one shared credential-sticky, prefix-cache and provider-session
+ * bucket.
+ */
+export function normalizeClientSessionKey(clientKey: string | undefined): string | undefined {
+	return clientKey !== undefined && clientKey.trim().length > 0 ? clientKey : undefined;
+}
+
+/**
+ * Stable identity of the account a request's credential belongs to.
+ *
+ * `markUsageLimitReached` and the auth-retry resolver switch a session to a
+ * sibling credential, so the provider state retained for that session can
+ * outlive the account that taught it. OAuth rows expose an account id / email
+ * that survives token refresh — fingerprinting the bearer instead would look
+ * like a rotation every time a token refreshes and discard the retained
+ * lessons for nothing. Key-based rows fall back to a hash of the key, never
+ * the key itself: this value is held for the lifetime of the entry.
+ */
+export function resolveGatewayAccount(
+	storage: AuthStorage,
+	provider: string,
+	sessionId: string,
+	apiKey: string,
+): string {
+	const identity = storage.getOAuthAccountIdentity(provider, sessionId);
+	if (identity) {
+		return `oauth:${JSON.stringify([
+			identity.accountId ?? "",
+			identity.email ?? "",
+			identity.projectId ?? "",
+			identity.orgId ?? "",
+		])}`;
+	}
+	return `key:${Bun.hash(apiKey).toString(36)}`;
+}
+
+/**
+ * Resolve the credential for one request from broker-backed storage.
+ *
+ * pi-ai clients never consult `AuthStorage`; the gateway resolves the bearer
+ * (an OAuth access token refreshed through the broker when needed) and hands
+ * it to the client. Returns the key, or the error classification the route
+ * should encode in its own envelope: storage failures map through
+ * {@link classifyGatewayError}, a provider without any credential is a 401.
+ */
+export async function resolveGatewayApiKey(
+	storage: AuthStorage,
+	model: Model<Api>,
+	sessionId: string,
+	signal: AbortSignal,
+	peer: string,
+): Promise<string | GatewayErrorClassification> {
+	let apiKey: string | undefined;
+	try {
+		apiKey = await storage.getApiKey(model.provider, sessionId, { modelId: model.id, signal });
+	} catch (error) {
+		const classified = classifyGatewayError(error);
+		logger.warn("auth-gateway getApiKey threw", { provider: model.provider, peer, error: classified.message });
+		return classified;
+	}
+	if (apiKey) return apiKey;
+	return {
+		status: 401,
+		type: "authentication_error",
+		message: `No credential available for provider ${model.provider}`,
+	};
+}
+
+/**
+ * Hook fired by a pi-ai client when the upstream request fails in a way
+ * that's rotatable — today that's HTTP 401 (credential is bad) and
+ * usage-limit phrasing matched by {@link isUsageLimitError} (Codex's
+ * `usage_limit_reached`, Anthropic's `usage_limit_reached`, Google's
+ * `resource_exhausted`, …). The two cases need different storage actions:
+ *
+ * - **usage-limit** → {@link AuthStorage.markUsageLimitReached}. Marks just
+ *   the current session's credential as temporarily blocked (honouring
+ *   `retry-after` / `resets_at` hints when present) and returns `true` only
+ *   when a sibling credential is still available. Burning the credential
+ *   with `invalidateCredentialMatching` here would orphan accounts whose
+ *   reset window is several hours away — exactly the bug this helper exists
+ *   to avoid.
+ * - **auth-failure** → {@link AuthStorage.invalidateCredentialMatching}.
+ *   Suspect/delete the row so it doesn't get re-picked next request.
+ *
+ * In both branches we return the next `getApiKey` result (sticky on the
+ * same `sessionId`) so the client can transparently retry the pre-emit
+ * failure with a fresh credential. Returning `undefined` aborts the retry
+ * and surfaces the original error to the caller.
+ */
+async function refreshGatewayApiKeyAfterAuthError(
+	storage: AuthStorage,
+	model: Model<Api>,
+	sessionId: string,
+	provider: string,
+	oldKey: string,
+	error: unknown,
+	signal: AbortSignal,
+	format: string,
+	peer: string,
+): Promise<string | undefined> {
+	const message = error instanceof Error ? error.message : String(error);
+	const status = extractHttpStatusFromError(error);
+	if (AIError.isUsageLimit(error) || isUsageLimitOutcome(status, message)) {
+		const retryAfterMs = extractProviderRetryHint(provider, message);
+		const { switched, retryAtMs } = await storage.markUsageLimitReached(provider, sessionId, {
+			retryAfterMs,
+			providerTimed: retryAfterMs !== undefined,
+			baseUrl: model.baseUrl,
+			modelId: model.id,
+			apiKey: oldKey,
+			signal,
+		});
+		logger.debug("auth-gateway retrying provider request after usage-limit block", {
+			format,
+			provider,
+			peer,
+			switched,
+			retryAfterMs,
+			retryAtMs,
+			error: message,
+		});
+		if (!switched) return undefined;
+		return storage.getApiKey(provider, sessionId, { modelId: model.id, signal });
+	}
+	await storage.invalidateCredentialMatching(provider, oldKey, { sessionId, signal });
+	logger.debug("auth-gateway retrying provider request after credential invalidation", {
+		format,
+		provider,
+		peer,
+		error: message,
+	});
+	return storage.getApiKey(provider, sessionId, { modelId: model.id, signal });
+}
+
+/**
+ * Build the {@link ApiKeyResolver} handed to a pi-ai client for a gateway
+ * request. Drives the central a/b/c auth-retry policy server-side:
+ *
+ * - initial resolve → the credential already resolved for this request.
+ * - step (b) `!lastChance` → force-refresh the SAME session-sticky credential
+ *   (a peer/broker may have rotated its token out from under our cached copy).
+ * - step (c) `lastChance` → {@link refreshGatewayApiKeyAfterAuthError} switches
+ *   to a sibling (usage-limit block vs credential invalidation by error class).
+ *
+ * `lastKey` tracks the most recent bearer so the switch step invalidates the
+ * credential that actually failed. `onResolvedKey` observes every rotation;
+ * routes that retain provider session state use it to re-key the account
+ * lease, one-shot routes pass `undefined`.
+ */
+export function buildGatewayApiKeyResolver(
+	storage: AuthStorage,
+	model: Model<Api>,
+	sessionId: string,
+	initialKey: string,
+	requestSignal: AbortSignal,
+	format: string,
+	peer: string,
+	onResolvedKey?: (apiKey: string) => void,
+): ApiKeyResolver {
+	let lastKey = initialKey;
+	return async ({ lastChance, error, signal }) => {
+		const sig = signal ?? requestSignal;
+		if (error === undefined) {
+			lastKey = initialKey;
+			return initialKey;
+		}
+		if (!lastChance) {
+			const refreshed = await storage.getApiKey(model.provider, sessionId, {
+				modelId: model.id,
+				signal: sig,
+				forceRefresh: true,
+			});
+			lastKey = refreshed ?? lastKey;
+			if (refreshed) onResolvedKey?.(refreshed);
+			return refreshed;
+		}
+		const next = await refreshGatewayApiKeyAfterAuthError(
+			storage,
+			model,
+			sessionId,
+			model.provider,
+			lastKey,
+			error,
+			sig,
+			format,
+			peer,
+		);
+		lastKey = next ?? lastKey;
+		if (next) onResolvedKey?.(next);
+		return next;
+	};
+}
+
+/**
+ * Attribute one settled upstream request to the originating client via the
+ * broker's observed-usage channel (`AuthStorage.recordObservedUsage`, batched
+ * by the remote store). Error/aborted turns still record — the provider
+ * billed whatever tokens the partial turn consumed; zero-usage results
+ * (pre-flight failures) are skipped. `at` defaults to now.
+ */
+export function recordGatewayUsage(
+	storage: AuthStorage,
+	model: Model<Api>,
+	client: ClientUsageIdentity,
+	usage: Usage,
+	at?: number,
+): void {
+	if (usage.input + usage.output + usage.cacheRead + usage.cacheWrite === 0) return;
+	storage.recordObservedUsage({
+		provider: model.provider,
+		model: model.id,
+		at,
+		usage: { input: usage.input, output: usage.output, cacheRead: usage.cacheRead, cacheWrite: usage.cacheWrite },
+		costUsd: usage.cost.total,
+		client,
+	});
+}
+
+/**
+ * An `AbortController` that follows the inbound request's abort signal. Routes
+ * abort it themselves when the response body is cancelled mid-stream, which
+ * `req.signal` alone does not observe.
+ */
+export function mirrorRequestAbort(req: Request): AbortController {
+	const controller = new AbortController();
+	if (req.signal.aborted) {
+		controller.abort(req.signal.reason);
+	} else {
+		req.signal.addEventListener("abort", () => controller.abort(req.signal.reason), { once: true });
+	}
+	return controller;
+}

package/src/auth-gateway/http.ts

@@ -7,7 +7,7 @@
 import { timingSafeEqual as nodeTimingSafeEqual } from "node:crypto";
 import * as os from "node:os";
 import { getInstallId } from "@oh-my-pi/pi-utils";
-import type { Api, AssistantMessage, Model } from "../types";
+import type { Api, Model } from "../types";
 import type { ClientUsageIdentity } from "../usage";
 
 const JSON_HEADERS = {
@@ -28,14 +28,13 @@ export function json(status: number, body: unknown, headers?: Record<string, str
  * `request-id` (surfaced as `_request_id` by the OpenAI and Anthropic SDKs,
  * matches the gateway log line), LiteLLM's model-resolution and cost headers,
  * and OpenAI's `openai-processing-ms`. Model/request-id headers are always
- * present; `message` — the final assistant message, available only on
- * non-streaming responses — adds the computed cost, and `startedAt` the wall
- * time. Streaming responses send headers before usage exists, so they carry
- * only the identity headers.
+ * present; `costUsd` — known only once a non-streaming response has settled —
+ * adds the computed cost, and `startedAt` the wall time. Streaming responses
+ * send headers before usage exists, so they carry only the identity headers.
  */
 export function gatewayResponseHeaders(
 	model: Model<Api>,
-	info: { requestId: string; message?: AssistantMessage; startedAt?: number },
+	info: { requestId: string; costUsd?: number; startedAt?: number },
 ): Record<string, string> {
 	const headers: Record<string, string> = {
 		"x-request-id": info.requestId,
@@ -43,7 +42,7 @@ export function gatewayResponseHeaders(
 		"x-litellm-model-id": model.id,
 	};
 	if (model.baseUrl) headers["x-litellm-model-api-base"] = model.baseUrl;
-	if (info.message) headers["x-litellm-response-cost"] = info.message.usage.cost.total.toString();
+	if (info.costUsd !== undefined) headers["x-litellm-response-cost"] = info.costUsd.toString();
 	if (info.startedAt !== undefined) {
 		const elapsed = (performance.now() - info.startedAt).toFixed(0);
 		headers["x-litellm-response-duration-ms"] = elapsed;

package/src/auth-gateway/index.ts

@@ -1,3 +1,4 @@
+export * from "./dispatch";
 export * from "./http";
 export * from "./session-state";
 export * from "./server";