@oh-my-pi/pi-ai 16.4.8 → 16.5.1

package/package.json

@@ -1,7 +1,7 @@
 {
 	"type": "module",
 	"name": "@oh-my-pi/pi-ai",
-	"version": "16.4.8",
+	"version": "16.5.1",
 	"description": "Unified LLM API with automatic model discovery and provider configuration",
 	"homepage": "https://omp.sh",
 	"author": "Can Boluk",
@@ -37,15 +37,15 @@
 		"fmt": "biome format --write ."
 	},
 	"dependencies": {
-		"@bufbuild/protobuf": "^2.12.0",
-		"@oh-my-pi/pi-catalog": "16.4.8",
-		"@oh-my-pi/pi-utils": "16.4.8",
-		"@oh-my-pi/pi-wire": "16.4.8",
-		"arktype": "2.2.2",
+		"@bufbuild/protobuf": "^2.12.1",
+		"@oh-my-pi/pi-catalog": "16.5.1",
+		"@oh-my-pi/pi-utils": "16.5.1",
+		"@oh-my-pi/pi-wire": "16.5.1",
+		"arktype": "2.2.3",
 		"zod": "^4"
 	},
 	"devDependencies": {
-		"@bufbuild/protoc-gen-es": "^2.12.0",
+		"@bufbuild/protoc-gen-es": "^2.12.1",
 		"@types/bun": "^1.3.14"
 	},
 	"engines": {

package/src/auth-broker/discover.ts

@@ -72,6 +72,26 @@ interface ConfigSnapshot {
 	token?: string;
 }
 
+/**
+ * Resolve a dotted config key (e.g. `auth.broker.url`) against a parsed YAML
+ * record, accepting both nested form (`auth: { broker: { url } }`) and the
+ * legacy flat literal-dot key (`"auth.broker.url": ...`). Nested wins when both
+ * are present. Returns the value only when it is a string.
+ */
+function readDottedString(record: Record<string, unknown>, dottedKey: string): string | undefined {
+	let current: unknown = record;
+	for (const segment of dottedKey.split(".")) {
+		if (current === null || typeof current !== "object" || Array.isArray(current)) {
+			current = undefined;
+			break;
+		}
+		current = (current as Record<string, unknown>)[segment];
+	}
+	if (typeof current === "string") return current;
+	const flat = record[dottedKey];
+	return typeof flat === "string" ? flat : undefined;
+}
+
 async function readConfigYaml(agentDir: string): Promise<ConfigSnapshot> {
 	for (const filename of MAIN_CONFIG_FILENAMES) {
 		const configPath = path.join(agentDir, filename);
@@ -80,9 +100,8 @@ async function readConfigYaml(agentDir: string): Promise<ConfigSnapshot> {
 			const parsed = YAML.parse(raw);
 			if (!parsed || typeof parsed !== "object" || Array.isArray(parsed)) return {};
 			const record = parsed as Record<string, unknown>;
-			const url = typeof record["auth.broker.url"] === "string" ? (record["auth.broker.url"] as string) : undefined;
-			const token =
-				typeof record["auth.broker.token"] === "string" ? (record["auth.broker.token"] as string) : undefined;
+			const url = readDottedString(record, "auth.broker.url");
+			const token = readDottedString(record, "auth.broker.token");
 			return { url, token };
 		} catch (err) {
 			if (isEnoent(err)) continue;

package/src/auth-broker/remote-store.ts

@@ -143,14 +143,24 @@ interface UsageCacheEntry {
 
 function usageOverlayKey(
 	provider: Provider,
-	ids: { accountId?: string; email?: string; projectId?: string },
+	ids: { accountId?: string; email?: string; projectId?: string; orgId?: string },
 ): string | undefined {
+	// Org first: one account email can hold several organizations (Anthropic
+	// Team seat + personal Max), each with its own limit pools. Keying the
+	// overlay by account/email would merge the two pools' header ingests.
+	// But the org alone is not enough either: two Team members share the org
+	// id while drawing on per-user pools, so the key stays qualified by the
+	// member's own base identity whenever one is known.
+	let base: string | undefined;
 	const accountId = ids.accountId?.trim().toLowerCase();
-	if (accountId) return `${provider}\0account:${accountId}`;
 	const email = ids.email?.trim().toLowerCase();
-	if (email) return `${provider}\0email:${email}`;
 	const projectId = ids.projectId?.trim().toLowerCase();
-	if (projectId) return `${provider}\0project:${projectId}`;
+	if (accountId) base = `account:${accountId}`;
+	else if (email) base = `email:${email}`;
+	else if (projectId) base = `project:${projectId}`;
+	const orgId = ids.orgId?.trim().toLowerCase();
+	if (orgId) return base ? `${provider}\0org:${orgId}|${base}` : `${provider}\0org:${orgId}`;
+	if (base) return `${provider}\0${base}`;
 	return undefined;
 }
 
@@ -787,6 +797,13 @@ export class RemoteAuthCredentialStore implements AuthCredentialStore {
 		this.#cache.set(key, { value, expiresAtSec });
 	}
 
+	/** Drop all cache rows whose keys start with the supplied prefix. */
+	deleteCachePrefix(prefix: string): void {
+		for (const key of this.#cache.keys()) {
+			if (key.startsWith(prefix)) this.#cache.delete(key);
+		}
+	}
+
 	cleanExpiredCache(): void {
 		const nowSec = Math.floor(Date.now() / 1000);
 		for (const [key, entry] of this.#cache) {
@@ -987,12 +1004,55 @@ export class RemoteAuthCredentialStore implements AuthCredentialStore {
  * usage data" (ranking proceeds without a usage signal for this credential).
  */
 function matchUsageReport(reports: UsageReport[], provider: Provider, credential: OAuthCredential): UsageReport | null {
-	const candidates = reports.filter(report => report.provider === provider);
-	if (candidates.length === 0) return null;
-	if (candidates.length === 1) return candidates[0];
+	const all = reports.filter(report => report.provider === provider);
+	if (all.length === 0) return null;
+	// Org precedence, decisive on EITHER side: an org-scoped credential may
+	// only take its own org's report, and an org-less (legacy) credential may
+	// only take org-less reports — the shared email/account would otherwise
+	// hand one subscription the OTHER subscription's pool (e.g. mark healthy
+	// Max exhausted via Team's report, or rank a legacy row on a sibling's
+	// numbers).
+	const orgId = credential.orgId?.trim().toLowerCase();
 	const accountId = credential.accountId?.trim().toLowerCase();
 	const email = credential.email?.trim().toLowerCase();
 	const projectId = credential.projectId?.trim().toLowerCase();
+	if (orgId) {
+		const sameOrg: UsageReport[] = [];
+		let sawReportOrg = false;
+		for (const report of all) {
+			const metaOrg = readMetadataString((report.metadata ?? {}) as Record<string, unknown>, "orgId");
+			if (metaOrg) {
+				sawReportOrg = true;
+				if (metaOrg.toLowerCase() === orgId) sameOrg.push(report);
+			}
+		}
+		// Org-attributed reports exist: the shared org is a GATE, not a match.
+		// Two Team members share the org id while drawing on per-user pools,
+		// so the credential's own base identity must still line up inside the
+		// same-org subset — a lone sibling report is NOT ours. An org-only
+		// credential (no base identifiers) takes the lone same-org report and
+		// treats several as ambiguous. None in our org → "no usage data"
+		// rather than mis-attributing another org's pool.
+		if (sawReportOrg) {
+			if (accountId || email || projectId) {
+				for (const report of sameOrg) {
+					if (reportMatchesIdentity(report, accountId, email, projectId)) return report;
+				}
+				return null;
+			}
+			return sameOrg.length === 1 ? sameOrg[0]! : null;
+		}
+		// No surviving report carries an org at all: presence mismatch is a
+		// non-match too — the sole org-less report may be a legacy sibling
+		// row's pool, and handing it to a scoped credential would rank/block
+		// on the wrong quota. "No usage data" degrades gracefully instead.
+		return null;
+	}
+	const candidates = all.filter(
+		report => !readMetadataString((report.metadata ?? {}) as Record<string, unknown>, "orgId"),
+	);
+	if (candidates.length === 0) return null;
+	if (all.length === 1 && candidates.length === 1) return candidates[0];
 	for (const report of candidates) {
 		if (reportMatchesIdentity(report, accountId, email, projectId)) return report;
 	}
@@ -1000,15 +1060,48 @@ function matchUsageReport(reports: UsageReport[], provider: Provider, credential
 }
 
 function findMatchingReportIndex(reports: UsageReport[], overlay: UsageReport): number {
-	const candidates = reports
+	const all = reports
 		.map((report, index) => ({ report, index }))
 		.filter(candidate => candidate.report.provider === overlay.provider);
-	if (candidates.length === 0) return -1;
-	if (candidates.length === 1) return candidates[0]!.index;
+	if (all.length === 0) return -1;
 	const metadata = (overlay.metadata ?? {}) as Record<string, unknown>;
+	// Org precedence — mirror matchUsageReport: an org-attributed overlay may
+	// only merge into a report of the SAME org, and an org-less overlay may
+	// only merge into an org-less report. Within the same org the overlay's
+	// base identity must still match — two Team members' reports share the
+	// org id but must not swallow each other's header ingests.
+	const overlayOrg = readMetadataString(metadata, "orgId")?.toLowerCase();
 	const accountId = readMetadataString(metadata, "accountId")?.toLowerCase();
 	const email = readMetadataString(metadata, "email")?.toLowerCase();
 	const projectId = readMetadataString(metadata, "projectId")?.toLowerCase();
+	if (overlayOrg) {
+		const sameOrg: { report: UsageReport; index: number }[] = [];
+		let sawReportOrg = false;
+		for (const candidate of all) {
+			const candidateOrg = readMetadataString((candidate.report.metadata ?? {}) as Record<string, unknown>, "orgId");
+			if (candidateOrg) {
+				sawReportOrg = true;
+				if (candidateOrg.toLowerCase() === overlayOrg) sameOrg.push(candidate);
+			}
+		}
+		if (sawReportOrg) {
+			if (accountId || email || projectId) {
+				for (const candidate of sameOrg) {
+					if (reportMatchesIdentity(candidate.report, accountId, email, projectId)) return candidate.index;
+				}
+				return -1;
+			}
+			return sameOrg.length === 1 ? sameOrg[0]!.index : -1;
+		}
+		// Presence mismatch — mirror matchUsageReport: an org-scoped overlay
+		// never merges into an org-less report; it becomes its own report row.
+		return -1;
+	}
+	const candidates = all.filter(
+		candidate => !readMetadataString((candidate.report.metadata ?? {}) as Record<string, unknown>, "orgId"),
+	);
+	if (candidates.length === 0) return -1;
+	if (all.length === 1 && candidates.length === 1) return candidates[0]!.index;
 	for (const candidate of candidates) {
 		if (reportMatchesIdentity(candidate.report, accountId, email, projectId)) return candidate.index;
 	}

package/src/auth-broker/wire-schemas.ts

@@ -32,6 +32,8 @@ export const oauthCredentialSchema = type({
 	"projectId?": "string",
 	"email?": "string",
 	"accountId?": "string",
+	"orgId?": "string",
+	"orgName?": "string",
 });
 
 /** OAuth credential as it appears in broker snapshots — refresh replaced with sentinel. */
@@ -45,6 +47,8 @@ export const remoteOauthCredentialSchema = type({
 	"projectId?": "string",
 	"email?": "string",
 	"accountId?": "string",
+	"orgId?": "string",
+	"orgName?": "string",
 });
 
 export const apiKeyCredentialSchema = type({

package/src/auth-gateway/http.ts

@@ -5,19 +5,50 @@
  * and peer-resolution logic.
  */
 import { timingSafeEqual as nodeTimingSafeEqual } from "node:crypto";
+import type { Api, AssistantMessage, Model } from "../types";
 
 const JSON_HEADERS = {
 	"Content-Type": "application/json",
 	"X-Content-Type-Options": "nosniff",
 } as const;
 
-export function json(status: number, body: unknown): Response {
+export function json(status: number, body: unknown, headers?: Record<string, string>): Response {
 	return new Response(JSON.stringify(body) ?? "null", {
 		status,
-		headers: JSON_HEADERS,
+		headers: headers ? { ...JSON_HEADERS, ...headers } : JSON_HEADERS,
 	});
 }
 
+/**
+ * Diagnostic response headers for translated inference requests, mirroring the
+ * names existing gateway-aware clients already parse: `x-request-id` /
+ * `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.
+ */
+export function gatewayResponseHeaders(
+	model: Model<Api>,
+	info: { requestId: string; message?: AssistantMessage; startedAt?: number },
+): Record<string, string> {
+	const headers: Record<string, string> = {
+		"x-request-id": info.requestId,
+		"request-id": info.requestId,
+		"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.startedAt !== undefined) {
+		const elapsed = (performance.now() - info.startedAt).toFixed(0);
+		headers["x-litellm-response-duration-ms"] = elapsed;
+		headers["openai-processing-ms"] = elapsed;
+	}
+	return headers;
+}
+
 export function resolvePeer(req: Request): string {
 	const fwd = req.headers.get("x-forwarded-for");
 	if (fwd) return fwd.split(",")[0].trim();
@@ -165,6 +196,8 @@ const CORS_HEADERS: Record<string, string> = {
 	"Access-Control-Allow-Methods": "GET, POST, OPTIONS",
 	"Access-Control-Allow-Headers":
 		"authorization, content-type, anthropic-version, anthropic-beta, openai-organization, openai-project, x-stainless-*, x-api-key",
+	"Access-Control-Expose-Headers":
+		"x-request-id, request-id, x-litellm-model-id, x-litellm-model-api-base, x-litellm-response-cost, x-litellm-response-duration-ms, openai-processing-ms",
 	"Access-Control-Max-Age": "86400",
 };
 

package/src/auth-gateway/server.ts

@@ -22,6 +22,7 @@ import { Effort } from "@oh-my-pi/pi-catalog/effort";
 import { extractHttpStatusFromError, extractRetryHint, 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 } from "../error/gateway";
 import { isUsageLimitOutcome } from "../error/rate-limit";
 import * as anthropicMessages from "../providers/anthropic-messages-server";
@@ -32,7 +33,15 @@ import { completeSimple, streamSimple } from "../stream";
 import type { Api, AssistantMessageEventStream, Context, Model, SimpleStreamOptions } from "../types";
 import { deterministicUuid } from "../utils/deterministic-id";
 import { parseBind } from "../utils/parse-bind";
-import { captureRequestHeaders, corsHeaders, isAuthorized, json, resolvePeer, withCors } from "./http";
+import {
+	captureRequestHeaders,
+	corsHeaders,
+	gatewayResponseHeaders,
+	isAuthorized,
+	json,
+	resolvePeer,
+	withCors,
+} from "./http";
 import type {
 	AuthGatewayServerHandle,
 	AuthGatewayServerOptions,
@@ -227,7 +236,8 @@ async function refreshGatewayApiKeyAfterAuthError(
 	peer: string,
 ): Promise<string | undefined> {
 	const message = error instanceof Error ? error.message : String(error);
-	if (isUsageLimitOutcome(extractHttpStatusFromError(error), message)) {
+	const status = extractHttpStatusFromError(error);
+	if (AIError.isUsageLimit(error) || isUsageLimitOutcome(status, message)) {
 		const retryAfterMs = extractRetryHint(undefined, message);
 		const { switched, retryAtMs } = await storage.markUsageLimitReached(provider, sessionId, {
 			retryAfterMs,
@@ -334,6 +344,8 @@ async function handleFormatEndpoint(
 	req: Request,
 	peer: string,
 ): Promise<Response> {
+	const startedAt = performance.now();
+	const requestId = crypto.randomUUID();
 	const controller = mirrorRequestAbort(req);
 	if (controller.signal.aborted) return clientClosedResponse(route);
 
@@ -430,6 +442,7 @@ async function handleFormatEndpoint(
 	);
 
 	logger.info("auth-gateway request", {
+		requestId,
 		format: route.label,
 		model: parsed.modelId,
 		resolvedProvider: model.provider,
@@ -458,7 +471,11 @@ async function handleFormatEndpoint(
 				const classified = classifyGatewayError(errorMessage);
 				return route.module.formatError(classified.status, classified.type, errorMessage);
 			}
-			return json(200, route.module.encodeResponse(message, parsed.modelId));
+			return json(
+				200,
+				route.module.encodeResponse(message, parsed.modelId),
+				gatewayResponseHeaders(model, { requestId, message, startedAt }),
+			);
 		} catch (error) {
 			if (controller.signal.aborted) return clientClosedResponse(route);
 			const classified = classifyGatewayError(error);
@@ -493,6 +510,7 @@ async function handleFormatEndpoint(
 	return new Response(sseStream, {
 		status: 200,
 		headers: {
+			...gatewayResponseHeaders(model, { requestId }),
 			"Content-Type": "text/event-stream; charset=utf-8",
 			"Cache-Control": "no-cache",
 			Connection: "keep-alive",
@@ -519,6 +537,8 @@ async function handleFormatEndpoint(
  * path.
  */
 async function handlePiNative(bootOpts: AuthGatewayBootOptions, req: Request, peer: string): Promise<Response> {
+	const startedAt = performance.now();
+	const requestId = crypto.randomUUID();
 	const controller = mirrorRequestAbort(req);
 	const aborted = (): Response => piNative.formatError(499, "request_aborted", "client closed request");
 	if (controller.signal.aborted) return aborted();
@@ -605,6 +625,7 @@ async function handlePiNative(bootOpts: AuthGatewayBootOptions, req: Request, pe
 	streamOpts.sessionId ??= sessionId;
 
 	logger.info("auth-gateway request", {
+		requestId,
 		format: "pi-native",
 		model: parsed.modelId,
 		resolvedProvider: model.provider,
@@ -633,7 +654,7 @@ async function handlePiNative(bootOpts: AuthGatewayBootOptions, req: Request, pe
 				const classified = classifyGatewayError(errorMessage);
 				return piNative.formatError(classified.status, classified.type, errorMessage);
 			}
-			return json(200, { message });
+			return json(200, { message }, gatewayResponseHeaders(model, { requestId, message, startedAt }));
 		} catch (error) {
 			if (controller.signal.aborted) return aborted();
 			const classified = classifyGatewayError(error);
@@ -664,6 +685,7 @@ async function handlePiNative(bootOpts: AuthGatewayBootOptions, req: Request, pe
 	return new Response(sseStream, {
 		status: 200,
 		headers: {
+			...gatewayResponseHeaders(model, { requestId }),
 			"Content-Type": "text/event-stream; charset=utf-8",
 			"Cache-Control": "no-cache",
 			Connection: "keep-alive",