auto-model-router 0.32.0 → 0.34.0

package/test/upstream-keys.test.ts

@@ -0,0 +1,240 @@
+import { describe, expect, test } from "bun:test";
+
+import { buildUpstreamModels } from "../src/catalog/static-catalog.ts";
+import type { CatalogSnapshot } from "../src/catalog/types.ts";
+import { DEFAULT_CONFIG } from "../src/config/defaults.ts";
+import type { RouterConfig, UpstreamEntry } from "../src/config/types.ts";
+import { completeUpstreamEntry } from "../src/config/upstreams.ts";
+import { buildCandidates } from "../src/router/candidates.ts";
+import { extractFeatures } from "../src/router/features.ts";
+import type { Rejection } from "../src/router/types.ts";
+import { createAnthropicClient } from "../src/upstream/anthropic.ts";
+import { compatEndpoint, createCompatClient } from "../src/upstream/compat.ts";
+import { createOllamaClient } from "../src/upstream/ollama.ts";
+import { createOpenRouterClient } from "../src/upstream/openrouter.ts";
+import { parseMessagesRequest } from "../src/wire/anthropic/messages.ts";
+import { parseChatRequest, parseUpstreamKeysHeader } from "../src/wire/openai/request.ts";
+import { parseResponsesRequest } from "../src/wire/openai/responses.ts";
+
+/**
+ * Per-turn upstream credentials (X-Omp-Upstream-Keys): a front door whose
+ * callers bring their own keys sends the turn's credentials in a header, so one
+ * router fleet serves every tenant instead of one process per credential set.
+ *
+ * The property everything else rests on: the credential belongs to the TURN.
+ * The shared `UpstreamEntry` is never written to, so two concurrent turns
+ * cannot see each other's key.
+ */
+
+const NL = String.fromCharCode(10);
+
+function entry(over: Partial<UpstreamEntry> & { id: string; kind: UpstreamEntry["kind"] }): UpstreamEntry {
+	return completeUpstreamEntry({ baseUrl: "https://api.example/v1", apiKey: "sk-configured", models: [{ id: "m1", input: 1, output: 4 }], ...over });
+}
+
+function cfgWith(upstreams: UpstreamEntry[]): RouterConfig {
+	return { ...structuredClone(DEFAULT_CONFIG), upstreams, logLevel: "silent" };
+}
+
+function sse(frames: string[]): Response {
+	return new Response(
+		new ReadableStream({
+			start: (c) => {
+				for (const f of frames) c.enqueue(new TextEncoder().encode(`${f}${NL}${NL}`));
+				c.close();
+			},
+		}),
+		{ status: 200, headers: { "content-type": "text/event-stream" } },
+	);
+}
+
+/** Drains a dispatch so the fake upstream's stream is consumed like a real turn's. */
+async function drain(d: { chunks: AsyncIterable<unknown> }): Promise<void> {
+	for await (const c of d.chunks) void c;
+}
+
+const OK_FRAMES = [`data: ${JSON.stringify({ id: "gen-1", model: "m1", choices: [{ delta: { content: "hi" } }] })}`, "data: [DONE]"];
+const ANTHROPIC_FRAMES = [
+	`event: message_start${NL}data: ${JSON.stringify({ type: "message_start", message: { id: "msg_1", model: "claude-sonnet-4", usage: { input_tokens: 1, output_tokens: 0 } } })}`,
+	`event: message_stop${NL}data: ${JSON.stringify({ type: "message_stop" })}`,
+];
+
+const SIGNAL = new AbortController().signal;
+
+describe("parseUpstreamKeysHeader", () => {
+	test("accepts an id→credential object, drops junk, and never rejects a turn", async () => {
+		// Same defensive contract as X-Omp-Policy: anything unusable is simply no override.
+		expect(parseUpstreamKeysHeader(null)).toBeUndefined();
+		expect(parseUpstreamKeysHeader("   ")).toBeUndefined();
+		expect(parseUpstreamKeysHeader("not json")).toBeUndefined();
+		expect(parseUpstreamKeysHeader('["sk-x"]')).toBeUndefined();
+		expect(parseUpstreamKeysHeader('"sk-x"')).toBeUndefined();
+		expect(parseUpstreamKeysHeader("{}")).toBeUndefined();
+		// Non-string values are dropped; an empty string is KEPT — it means "no credential this turn".
+		expect(parseUpstreamKeysHeader(JSON.stringify({ openrouter: "sk-or-1", " azure-eu ": "az", bad: 7, worse: null, "": "x", off: "" }))).toEqual({
+			openrouter: "sk-or-1",
+			"azure-eu": "az",
+			off: "",
+		});
+		// The credential is copied verbatim: trimming one would break a key whose bytes matter.
+		expect(parseUpstreamKeysHeader(JSON.stringify({ up: " sk-pad " }))).toEqual({ up: " sk-pad " });
+	});
+
+	test("every wire carries the map, and its absence leaves the property off", async () => {
+		const header = new Headers({ "X-Omp-Upstream-Keys": '{"openrouter":"sk-or-tenant"}' });
+		const chat = parseChatRequest({ model: "auto", messages: [{ role: "user", content: "hi" }] }, header);
+		const messages = parseMessagesRequest({ model: "claude-sonnet-4", messages: [{ role: "user", content: "hi" }], max_tokens: 16 }, header);
+		const responses = parseResponsesRequest({ model: "auto", input: "hi" }, header);
+		for (const req of [chat, messages, responses]) expect(req.upstreamKeys).toEqual({ openrouter: "sk-or-tenant" });
+		expect([chat.protocol, messages.protocol, responses.protocol]).toEqual(["openai-chat", "anthropic-messages", "openai-responses"]);
+		// exactOptionalPropertyTypes: absent means absent, not `undefined`.
+		expect("upstreamKeys" in parseChatRequest({ model: "auto", messages: [{ role: "user", content: "hi" }] }, new Headers())).toBe(false);
+		expect("upstreamKeys" in parseMessagesRequest({ model: "claude-sonnet-4", messages: [{ role: "user", content: "hi" }], max_tokens: 16 }, new Headers())).toBe(false);
+		expect("upstreamKeys" in parseResponsesRequest({ model: "auto", input: "hi" }, new Headers())).toBe(false);
+	});
+});
+
+describe("per-turn credentials at dispatch", () => {
+	test("compatEndpoint prefers the per-turn key without touching the entry", async () => {
+		const e = entry({ id: "openai-direct", kind: "openai" });
+		expect(compatEndpoint(e, "m1").headers.authorization).toBe("Bearer sk-configured");
+		expect(compatEndpoint(e, "m1", "sk-turn").headers.authorization).toBe("Bearer sk-turn");
+		// Azure keys with its own header, and an empty credential sends none at all.
+		const az = entry({ id: "azure-eu", kind: "azure" });
+		expect(compatEndpoint(az, "dep", "az-turn").headers["api-key"]).toBe("az-turn");
+		expect(compatEndpoint(az, "dep", "").headers["api-key"]).toBeUndefined();
+		expect(e.apiKey).toBe("sk-configured");
+	});
+
+	test("an OpenAI-compatible upstream uses the turn's key, else the configured one", async () => {
+		const seen: Array<string | null> = [];
+		const cfg = cfgWith([entry({ id: "openai-direct", kind: "openai" })]);
+		const client = createCompatClient(cfg, "openai-direct", async (_url, init) => {
+			seen.push(new Headers(init?.headers).get("authorization"));
+			return sse(OK_FRAMES);
+		});
+		const body = { model: "openai-direct/m1", messages: [] };
+		await drain(await client.dispatch({ body, sessionId: "s", signal: SIGNAL, upstreamKeys: { "openai-direct": "sk-turn" } }));
+		await drain(await client.dispatch({ body, sessionId: "s", signal: SIGNAL }));
+		// A map that names a DIFFERENT upstream leaves this one on its own key.
+		await drain(await client.dispatch({ body, sessionId: "s", signal: SIGNAL, upstreamKeys: { openrouter: "sk-or" } }));
+		// An empty credential is "none", never a fallback to the configured key.
+		await drain(await client.dispatch({ body, sessionId: "s", signal: SIGNAL, upstreamKeys: { "openai-direct": "" } }));
+		expect(seen).toEqual(["Bearer sk-turn", "Bearer sk-configured", "Bearer sk-configured", null]);
+		expect(cfg.upstreams[0]!.apiKey).toBe("sk-configured");
+	});
+
+	test("an Anthropic upstream uses the turn's key on both auth shapes", async () => {
+		const seen: Array<Record<string, string | null>> = [];
+		const cfg = cfgWith([
+			entry({ id: "anthropic-direct", kind: "anthropic", apiKey: "sk-ant-configured" }),
+			entry({ id: "claude-sub", kind: "anthropic", auth: "oauth-bearer", apiKey: "oauth-configured" }),
+		]);
+		const fetchImpl = async (_url: unknown, init?: RequestInit): Promise<Response> => {
+			const h = new Headers(init?.headers);
+			seen.push({ "x-api-key": h.get("x-api-key"), authorization: h.get("authorization") });
+			return sse(ANTHROPIC_FRAMES);
+		};
+		const direct = createAnthropicClient(cfg, "anthropic-direct", fetchImpl);
+		const sub = createAnthropicClient(cfg, "claude-sub", fetchImpl);
+		await drain(await direct.dispatch({ body: { model: "anthropic-direct/m1", messages: [{ role: "user", content: "hi" }], max_tokens: 16 }, sessionId: "s", signal: SIGNAL, upstreamKeys: { "anthropic-direct": "sk-ant-turn" } }));
+		await drain(await direct.dispatch({ body: { model: "anthropic-direct/m1", messages: [{ role: "user", content: "hi" }], max_tokens: 16 }, sessionId: "s", signal: SIGNAL }));
+		await drain(await sub.dispatch({ body: { model: "claude-sub/m1", messages: [{ role: "user", content: "hi" }], max_tokens: 16 }, sessionId: "s", signal: SIGNAL, upstreamKeys: { "claude-sub": "oauth-turn" } }));
+		expect(seen).toEqual([
+			{ "x-api-key": "sk-ant-turn", authorization: null },
+			{ "x-api-key": "sk-ant-configured", authorization: null },
+			{ "x-api-key": null, authorization: "Bearer oauth-turn" },
+		]);
+		expect(cfg.upstreams.map((u) => u.apiKey)).toEqual(["sk-ant-configured", "oauth-configured"]);
+	});
+
+	test("OpenRouter and Ollama take an override under their own reserved ids", async () => {
+		const cfg = cfgWith([]);
+		cfg.openrouter.apiKey = "sk-or-configured";
+		cfg.ollama = { ...cfg.ollama, apiKey: "sk-ollama-configured" };
+		const seen: Array<string | null> = [];
+		const realFetch = globalThis.fetch;
+		globalThis.fetch = (async (_url: unknown, init?: RequestInit) => {
+			seen.push(new Headers(init?.headers).get("authorization"));
+			return sse(OK_FRAMES);
+		}) as unknown as typeof fetch;
+		try {
+			const or = createOpenRouterClient(cfg);
+			await drain(await or.dispatch({ body: { model: "a/b", messages: [] }, sessionId: "s", signal: SIGNAL, upstreamKeys: { openrouter: "sk-or-turn" } }));
+			await drain(await or.dispatch({ body: { model: "a/b", messages: [] }, sessionId: "s", signal: SIGNAL }));
+		} finally {
+			globalThis.fetch = realFetch;
+		}
+		const ollama = createOllamaClient(cfg, async (_url, init) => {
+			seen.push(new Headers(init?.headers).get("authorization"));
+			return sse(OK_FRAMES);
+		});
+		await drain(await ollama.dispatch({ body: { model: "ollama/m", messages: [] }, sessionId: "s", signal: SIGNAL, upstreamKeys: { ollama: "sk-ollama-turn" } }));
+		await drain(await ollama.dispatch({ body: { model: "ollama/m", messages: [] }, sessionId: "s", signal: SIGNAL }));
+		expect(seen).toEqual(["Bearer sk-or-turn", "Bearer sk-or-configured", "Bearer sk-ollama-turn", "Bearer sk-ollama-configured"]);
+		expect(cfg.openrouter.apiKey).toBe("sk-or-configured");
+		expect(cfg.ollama.apiKey).toBe("sk-ollama-configured");
+	});
+
+	test("two concurrent turns over one entry never see each other's credential", async () => {
+		const cfg = cfgWith([entry({ id: "openai-direct", kind: "openai" })]);
+		let release!: () => void;
+		const gate = new Promise<void>((resolve) => {
+			release = resolve;
+		});
+		const seen: Array<string | null> = [];
+		let inFlight = 0;
+		const client = createCompatClient(cfg, "openai-direct", async (_url, init) => {
+			const auth = new Headers(init?.headers).get("authorization");
+			// Hold both requests open at once: a mutation-based implementation
+			// would have overwritten the first turn's key by the time it reads.
+			inFlight++;
+			if (inFlight === 1) await gate;
+			else release();
+			seen.push(auth);
+			return sse(OK_FRAMES);
+		});
+		const turn = async (key: string): Promise<void> => {
+			const d = await client.dispatch({ body: { model: "openai-direct/m1", messages: [] }, sessionId: "s", signal: SIGNAL, upstreamKeys: { "openai-direct": key } });
+			await drain(d);
+		};
+		await Promise.all([turn("sk-tenant-a"), turn("sk-tenant-b")]);
+		expect(seen.sort()).toEqual(["Bearer sk-tenant-a", "Bearer sk-tenant-b"]);
+		// The shared config is exactly as it was configured.
+		expect(cfg.upstreams[0]!.apiKey).toBe("sk-configured");
+	});
+});
+
+describe("candidate selection with per-turn credentials", () => {
+	const twoUpstreams = [
+		...buildUpstreamModels(entry({ id: "up-a", kind: "openai", models: [{ id: "m1", input: 1, output: 4, quality: { coding: 80, intelligence: 80, agentic: 80 } }] }), []),
+		...buildUpstreamModels(entry({ id: "up-b", kind: "openai", models: [{ id: "m1", input: 1, output: 4, quality: { coding: 80, intelligence: 80, agentic: 80 } }] }), []),
+	];
+	const snapshot: CatalogSnapshot = { models: twoUpstreams, fetchedAtMs: Date.now(), keyScoped: true };
+	const cfg: RouterConfig = { ...DEFAULT_CONFIG, adaptiveTierFloors: false, filters: { ...DEFAULT_CONFIG.filters, minTrust: 0 }, tiers: { ...DEFAULT_CONFIG.tiers, moderate: { ...DEFAULT_CONFIG.tiers.moderate, minQuality: 0, maxInputPerMtok: 10 } } };
+
+	function build(keys: Record<string, string> | undefined): { slugs: string[]; rejected: Rejection[] } {
+		const req = parseChatRequest(
+			{ model: "auto", messages: [{ role: "user", content: "hi" }] },
+			new Headers(keys === undefined ? {} : { "x-omp-upstream-keys": JSON.stringify(keys) }),
+		);
+		const { candidates, rejected } = buildCandidates({ req, features: extractFeatures(req, 100), tier: "moderate", task: "chat", snapshot, cfg, expectedCompletionTokens: 128, warmSlug: null });
+		return { slugs: candidates.map((c) => c.model.slug), rejected };
+	}
+
+	test("an upstream with an empty credential is excluded for that turn only", async () => {
+		expect(build(undefined).slugs.sort()).toEqual(["up-a/m1", "up-b/m1"]);
+		// A real credential changes nothing about who may be selected.
+		expect(build({ "up-a": "sk-turn" }).slugs.sort()).toEqual(["up-a/m1", "up-b/m1"]);
+		const off = build({ "up-a": "" });
+		expect(off.slugs).toEqual(["up-b/m1"]);
+		expect(off.rejected).toContainEqual({ slug: "up-a/m1", reason: "no_credential", detail: "upstream up-a has no credential on this turn" });
+		// The next turn, carrying a key, sees it again: nothing was recorded anywhere.
+		expect(build({ "up-a": "sk-turn" }).slugs.sort()).toEqual(["up-a/m1", "up-b/m1"]);
+	});
+
+	test("no rejection reason repeats the credential", async () => {
+		const { rejected } = build({ "up-a": "", "up-b": "sk-secret-value" });
+		expect(JSON.stringify(rejected)).not.toContain("sk-secret-value");
+	});
+});