feat(auth-gateway): added diagnostic headers to auth-gateway inference responses

- Added a `gatewayResponseHeaders` helper and an optional `json()` headers parameter to attach gateway diagnostics to HTTP responses.
- Updated auth-gateway format and Pi-native handlers to generate request IDs and include them with model metadata in all responses, while adding cost and timing headers when non-streaming responses include a final assistant message.
- Added tests validating header presence for streaming vs non-streaming responses and documented the new auth-gateway diagnostic headers in the unreleased changelog.

packages/ai/CHANGELOG.md

@@ -2,6 +2,10 @@
 
 ## [Unreleased]
 
+### Added
+
+- Added diagnostic response headers to auth-gateway inference endpoints: `x-request-id`/`request-id` (correlates with gateway logs; surfaced by OpenAI/Anthropic SDKs) and LiteLLM-style `x-litellm-model-id`/`x-litellm-model-api-base` on every response, plus `x-litellm-response-cost`, `x-litellm-response-duration-ms`, and `openai-processing-ms` on non-streaming responses
+
 ### Changed
 
 - Switched Google and Google Vertex providers to always use `streamGenerateContent` requests

packages/ai/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",
 };
 

packages/ai/src/auth-gateway/server.ts

@@ -32,7 +32,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,
@@ -334,6 +342,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 +440,7 @@ async function handleFormatEndpoint(
 	);
 
 	logger.info("auth-gateway request", {
+		requestId,
 		format: route.label,
 		model: parsed.modelId,
 		resolvedProvider: model.provider,
@@ -458,7 +469,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 +508,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 +535,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 +623,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 +652,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 +683,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",

packages/ai/test/auth-gateway-response-headers.test.ts

@@ -0,0 +1,103 @@
+import { afterEach, describe, expect, it } from "bun:test";
+import * as fs from "node:fs/promises";
+import * as os from "node:os";
+import * as path from "node:path";
+import { clearCustomApis } from "@oh-my-pi/pi-ai/api-registry";
+import { startAuthGateway } from "@oh-my-pi/pi-ai/auth-gateway";
+import { AuthStorage } from "@oh-my-pi/pi-ai/auth-storage";
+import { createMockModel, type MockModel, registerMockApi } from "@oh-my-pi/pi-ai/providers/mock";
+
+interface GatewayHarness {
+	url: string;
+	mock: MockModel;
+	close(): Promise<void>;
+}
+
+async function bootGateway(): Promise<GatewayHarness> {
+	registerMockApi();
+	const dir = await fs.mkdtemp(path.join(os.tmpdir(), "gw-response-headers-"));
+	const storage = await AuthStorage.create(path.join(dir, "auth.db"));
+	storage.setRuntimeApiKey("openrouter", "test-key");
+	const mock = createMockModel({ provider: "openrouter", id: "mock/header-model" });
+	const handle = startAuthGateway({
+		bind: "127.0.0.1:0",
+		bearerTokens: ["t"],
+		storage,
+		resolveModel: () => mock.model,
+		version: "test",
+	});
+	return {
+		url: handle.url,
+		mock,
+		close: async () => {
+			await handle.close();
+			storage.close();
+			await fs.rm(dir, { recursive: true, force: true });
+		},
+	};
+}
+
+afterEach(() => {
+	clearCustomApis();
+});
+
+describe("auth-gateway diagnostic response headers", () => {
+	it("non-streaming responses carry cost, model id, request id, and duration", async () => {
+		const gw = await bootGateway();
+		try {
+			gw.mock.push({
+				content: ["hello"],
+				usage: {
+					input: 10,
+					output: 5,
+					totalTokens: 15,
+					cost: { input: 0.001, output: 0.0002, total: 0.0012 },
+				},
+			});
+			const res = await fetch(`${gw.url}/v1/chat/completions`, {
+				method: "POST",
+				headers: { "Content-Type": "application/json", Authorization: "Bearer t" },
+				body: JSON.stringify({
+					model: "mock/header-model",
+					messages: [{ role: "user", content: "hi" }],
+					stream: false,
+				}),
+			});
+			expect(res.status).toBe(200);
+			expect(res.headers.get("x-litellm-response-cost")).toBe("0.0012");
+			expect(res.headers.get("x-litellm-model-id")).toBe("mock/header-model");
+			const duration = res.headers.get("x-litellm-response-duration-ms");
+			expect(duration).not.toBeNull();
+			expect(Number(duration)).toBeGreaterThanOrEqual(0);
+			expect(res.headers.get("openai-processing-ms")).toBe(duration);
+			const requestId = res.headers.get("x-request-id");
+			expect(requestId).toMatch(/^[0-9a-f-]{36}$/);
+			expect(res.headers.get("request-id")).toBe(requestId);
+		} finally {
+			await gw.close();
+		}
+	});
+
+	it("streaming responses carry the model and request ids but no cost (unknown at header time)", async () => {
+		const gw = await bootGateway();
+		try {
+			gw.mock.push({ content: ["hello"] });
+			const res = await fetch(`${gw.url}/v1/chat/completions`, {
+				method: "POST",
+				headers: { "Content-Type": "application/json", Authorization: "Bearer t" },
+				body: JSON.stringify({
+					model: "mock/header-model",
+					messages: [{ role: "user", content: "hi" }],
+					stream: true,
+				}),
+			});
+			expect(res.status).toBe(200);
+			expect(res.headers.get("x-litellm-model-id")).toBe("mock/header-model");
+			expect(res.headers.get("x-request-id")).toMatch(/^[0-9a-f-]{36}$/);
+			expect(res.headers.get("x-litellm-response-cost")).toBeNull();
+			await res.text();
+		} finally {
+			await gw.close();
+		}
+	});
+});