diff --git a/packages/ai/CHANGELOG.md b/packages/ai/CHANGELOG.md index 03fa9035e..712cd2e4f 100644 --- a/packages/ai/CHANGELOG.md +++ b/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 diff --git a/packages/ai/src/auth-gateway/http.ts b/packages/ai/src/auth-gateway/http.ts index 21ea80d61..ec905a308 100644 --- a/packages/ai/src/auth-gateway/http.ts +++ b/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): 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, + info: { requestId: string; message?: AssistantMessage; startedAt?: number }, +): Record { + const headers: Record = { + "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 = { "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", }; diff --git a/packages/ai/src/auth-gateway/server.ts b/packages/ai/src/auth-gateway/server.ts index ae4d47415..2624dd028 100644 --- a/packages/ai/src/auth-gateway/server.ts +++ b/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 { + 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 { + 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", diff --git a/packages/ai/test/auth-gateway-response-headers.test.ts b/packages/ai/test/auth-gateway-response-headers.test.ts new file mode 100644 index 000000000..986d5c084 --- /dev/null +++ b/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; +} + +async function bootGateway(): Promise { + 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(); + } + }); +});