auto-model-router 0.1.0

package/test/wire-request.test.ts

@@ -0,0 +1,297 @@
+import { describe, expect, test } from "bun:test";
+import type { UpstreamMutations } from "../src/wire/types.ts";
+import { WireErrorException } from "../src/wire/openai/errors.ts";
+import { parseChatRequest } from "../src/wire/openai/request.ts";
+
+const HEADERS = new Headers();
+
+function mutations(overrides: Partial<UpstreamMutations> = {}): UpstreamMutations {
+	return {
+		slug: "openai/gpt-5.5",
+		fallbacks: [],
+		sessionId: "omp-test",
+		cacheBreakpointMessageIndices: [],
+		reasoning: undefined,
+		maxTokens: undefined,
+		stripAssistantReasoning: false,
+		...overrides,
+	};
+}
+
+function userBody(content: unknown): Record<string, unknown> {
+	return { model: "auto", messages: [{ role: "user", content }] };
+}
+
+describe("parseChatRequest normalization", () => {
+	test("string content and equivalent content-part array normalize to the same text", () => {
+		const fromString = parseChatRequest(userBody("hello world"), HEADERS);
+		const fromParts = parseChatRequest(
+			userBody([{ type: "text", text: "hello world" }]),
+			HEADERS,
+		);
+		expect(fromParts.messages[0]!.text).toBe(fromString.messages[0]!.text);
+		expect(fromParts.messages[0]!.textBytes).toBe(fromString.messages[0]!.textBytes);
+		expect(fromParts.messages[0]!.images).toBe(0);
+	});
+
+	test("image parts are counted and flag hasImages", () => {
+		const req = parseChatRequest(
+			userBody([
+				{ type: "text", text: "look at these" },
+				{ type: "image_url", image_url: { url: "data:image/png;base64,AA==" } },
+				{ type: "image_url", image_url: { url: "https://example.com/x.png" } },
+			]),
+			HEADERS,
+		);
+		expect(req.messages[0]!.images).toBe(2);
+		expect(req.messages[0]!.text).toBe("look at these");
+		expect(req.hasImages).toBe(true);
+		expect(parseChatRequest(userBody("plain"), HEADERS).hasImages).toBe(false);
+	});
+
+	test("provider prefix is stripped from model", () => {
+		expect(parseChatRequest(userBody("hi"), HEADERS).requestedModel).toBe("auto");
+		const prefixed = parseChatRequest(
+			{ model: "auto-model-router/auto", messages: [{ role: "user", content: "hi" }] },
+			HEADERS,
+		);
+		expect(prefixed.requestedModel).toBe("auto");
+	});
+
+	test("tool schemas, names, and descriptions contribute to promptBytes", () => {
+		const parameters = { type: "object", properties: { path: { type: "string" } } };
+		const withTools = parseChatRequest(
+			{
+				model: "auto",
+				messages: [{ role: "user", content: "hi" }],
+				tools: [
+					{
+						type: "function",
+						function: { name: "read", description: "Read a file", parameters },
+					},
+				],
+			},
+			HEADERS,
+		);
+		const withoutTools = parseChatRequest(userBody("hi"), HEADERS);
+		const tool = withTools.tools[0]!;
+		expect(tool.schemaBytes).toBe(new TextEncoder().encode(JSON.stringify(parameters)).length);
+		expect(withTools.promptBytes).toBe(
+			withoutTools.promptBytes + tool.schemaBytes + 4 + "Read a file".length,
+		);
+	});
+
+	test("tool calls and tool results are carried into NormMessage", () => {
+		const req = parseChatRequest(
+			{
+				model: "auto",
+				messages: [
+					{
+						role: "assistant",
+						content: null,
+						tool_calls: [
+							{ id: "call_1", type: "function", function: { name: "read", arguments: "{\"path\":\"x\"}" } },
+						],
+					},
+					{ role: "tool", tool_call_id: "call_1", name: "read", content: "file bytes" },
+				],
+			},
+			HEADERS,
+		);
+		expect(req.messages[0]!.toolCalls).toEqual([{ id: "call_1", name: "read", argsJson: "{\"path\":\"x\"}" }]);
+		expect(req.messages[1]!.toolCallId).toBe("call_1");
+		expect(req.messages[1]!.toolName).toBe("read");
+	});
+
+	test("forcedToolChoice only for objects and non-auto/none strings", () => {
+		const base = userBody("hi");
+		expect(parseChatRequest(base, HEADERS).forcedToolChoice).toBe(false);
+		expect(parseChatRequest({ ...base, tool_choice: "auto" }, HEADERS).forcedToolChoice).toBe(false);
+		expect(parseChatRequest({ ...base, tool_choice: "none" }, HEADERS).forcedToolChoice).toBe(false);
+		expect(parseChatRequest({ ...base, tool_choice: "required" }, HEADERS).forcedToolChoice).toBe(true);
+		expect(
+			parseChatRequest(
+				{ ...base, tool_choice: { type: "function", function: { name: "read" } } },
+				HEADERS,
+			).forcedToolChoice,
+		).toBe(true);
+	});
+
+	test("reasoning accepted from both spellings, omitted when absent", () => {
+		expect(
+			parseChatRequest({ ...userBody("hi"), reasoning_effort: "high" }, HEADERS).reasoning,
+		).toBe("high");
+		expect(
+			parseChatRequest({ ...userBody("hi"), reasoning: { effort: "low" } }, HEADERS).reasoning,
+		).toBe("low");
+		expect(
+			parseChatRequest({ ...userBody("hi"), reasoning: { enabled: false } }, HEADERS).reasoning,
+		).toBe("off");
+		const plain = parseChatRequest(userBody("hi"), HEADERS);
+		expect("reasoning" in plain).toBe(false);
+	});
+
+	test("malformed input throws WireErrorException", () => {
+		expect(() => parseChatRequest(null, HEADERS)).toThrow(WireErrorException);
+		expect(() => parseChatRequest({ model: "auto", messages: [] }, HEADERS)).toThrow(WireErrorException);
+		expect(() => parseChatRequest({ messages: [{ role: "user", content: "hi" }] }, HEADERS)).toThrow(
+			WireErrorException,
+		);
+		try {
+			parseChatRequest({ model: "auto", messages: [{ role: "user", content: 42 }] }, HEADERS);
+			expect.unreachable();
+		} catch (e) {
+			expect(e).toBeInstanceOf(WireErrorException);
+			expect((e as WireErrorException).wireError.status).toBe(400);
+			expect((e as WireErrorException).wireError.code).toBe("invalid_request");
+		}
+	});
+});
+
+describe("conversationKey", () => {
+	const system = { role: "system", content: "You are a coding agent." };
+	const first = { role: "user", content: "Fix the bug in main.ts" };
+
+	test("stable across later turns of the same conversation", () => {
+		const turn1 = parseChatRequest({ model: "auto", messages: [system, first] }, HEADERS);
+		const turn3 = parseChatRequest(
+			{
+				model: "auto",
+				messages: [
+					system,
+					first,
+					{ role: "assistant", content: "Done." },
+					{ role: "user", content: "Now add a test" },
+				],
+			},
+			HEADERS,
+		);
+		expect(turn1.conversationKey).toBe(turn3.conversationKey);
+		expect(turn1.conversationKey).toMatch(/^[0-9a-f]{32}$/);
+	});
+
+	test("differs when the first non-system message differs", () => {
+		const a = parseChatRequest({ model: "auto", messages: [system, first] }, HEADERS);
+		const b = parseChatRequest(
+			{ model: "auto", messages: [system, { role: "user", content: "Write a poem" }] },
+			HEADERS,
+		);
+		expect(a.conversationKey).not.toBe(b.conversationKey);
+	});
+});
+
+describe("renderUpstreamBody", () => {
+	test("two renders are independent and never mutate the original body", () => {
+		const original = {
+			model: "auto",
+			messages: [{ role: "user", content: "hi" }],
+			stream: false,
+			max_tokens: 100,
+			reasoning_effort: "low",
+		};
+		const req = parseChatRequest(original, HEADERS);
+		const before = JSON.stringify(original);
+
+		const a = req.renderUpstreamBody(mutations({ slug: "openai/aaa" }));
+		const b = req.renderUpstreamBody(
+			mutations({
+				slug: "openai/bbb",
+				fallbacks: ["openai/ccc"],
+				sessionId: "omp-2",
+				maxTokens: 50,
+				reasoning: "high",
+				cacheBreakpointMessageIndices: [0],
+			}),
+		);
+
+		expect(a.model).toBe("openai/aaa");
+		expect(b.model).toBe("openai/bbb");
+		expect(a.models).toBeUndefined();
+		expect(b.models).toEqual(["openai/bbb", "openai/ccc"]);
+		expect(b.session_id).toBe("omp-2");
+		expect(a.stream).toBe(true);
+		expect(b.stream).toBe(true);
+		expect("stream_options" in b).toBe(false);
+		expect(a.max_tokens).toBe(100);
+		expect(b.max_tokens).toBe(50);
+		expect("reasoning_effort" in a).toBe(false);
+		expect("reasoning_effort" in b).toBe(false);
+		expect("reasoning" in a).toBe(false);
+		expect(b.reasoning).toEqual({ effort: "high" });
+
+		// Mutating one render must not leak into the other or the original.
+		(b.messages as Record<string, unknown>[])[0]!.content = "mutated";
+		expect((a.messages as Record<string, unknown>[])[0]!.content).toBe("hi");
+		expect(JSON.stringify(original)).toBe(before);
+	});
+
+	test("maxTokens lands on whichever spelling the client used", () => {
+		const req = parseChatRequest(
+			{ model: "auto", messages: [{ role: "user", content: "hi" }], max_completion_tokens: 200 },
+			HEADERS,
+		);
+		const out = req.renderUpstreamBody(mutations({ maxTokens: 64 }));
+		expect(out.max_completion_tokens).toBe(64);
+		expect("max_tokens" in out).toBe(false);
+	});
+
+	test("reasoning off renders as { enabled: false }", () => {
+		const req = parseChatRequest(userBody("hi"), HEADERS);
+		const out = req.renderUpstreamBody(mutations({ reasoning: "off" }));
+		expect(out.reasoning).toEqual({ enabled: false });
+	});
+
+	test("cache breakpoints land on named messages and promote string content to parts", () => {
+		const req = parseChatRequest(
+			{
+				model: "auto",
+				messages: [
+					{ role: "system", content: "system prompt" },
+					{ role: "user", content: "first question" },
+					{ role: "assistant", content: [{ type: "text", text: "answer" }, { type: "text", text: "more" }] },
+					{ role: "user", content: "follow up" },
+				],
+			},
+			HEADERS,
+		);
+		const out = req.renderUpstreamBody(mutations({ cacheBreakpointMessageIndices: [0, 2] }));
+		const msgs = out.messages as Array<Record<string, unknown>>;
+
+		expect(msgs[0]!.content).toEqual([
+			{ type: "text", text: "system prompt", cache_control: { type: "ephemeral" } },
+		]);
+		// Untouched message keeps its original shape.
+		expect(msgs[1]!.content).toBe("first question");
+		// Breakpoint lands on the LAST text part of the message.
+		expect(msgs[2]!.content).toEqual([
+			{ type: "text", text: "answer" },
+			{ type: "text", text: "more", cache_control: { type: "ephemeral" } },
+		]);
+	});
+
+	test("stripAssistantReasoning removes all three spellings from assistant messages only", () => {
+		const req = parseChatRequest(
+			{
+				model: "auto",
+				messages: [
+					{
+						role: "assistant",
+						content: "a1",
+						reasoning: "r",
+						reasoning_content: "rc",
+						reasoning_details: [{ type: "reasoning.text", text: "rc" }],
+					},
+					{ role: "user", content: "u1", reasoning: "keep-me" },
+				],
+			},
+			HEADERS,
+		);
+		const out = req.renderUpstreamBody(mutations({ stripAssistantReasoning: true }));
+		const msgs = out.messages as Array<Record<string, unknown>>;
+		expect("reasoning" in msgs[0]!).toBe(false);
+		expect("reasoning_content" in msgs[0]!).toBe(false);
+		expect("reasoning_details" in msgs[0]!).toBe(false);
+		expect(msgs[0]!.content).toBe("a1");
+		expect(msgs[1]!.reasoning).toBe("keep-me");
+	});
+});

package/test/wire-sink.test.ts

@@ -0,0 +1,179 @@
+import { describe, expect, test } from "bun:test";
+import { EMPTY_USAGE } from "../src/cost/types.ts";
+import type { TurnSummary } from "../src/wire/types.ts";
+import { createBufferedSink, createStreamingSink } from "../src/wire/openai/sink.ts";
+
+const SUMMARY: TurnSummary = {
+	servedSlug: "openai/gpt-5.5",
+	tier: "simple",
+	attempts: 1,
+	predictedUsd: 0.001,
+	reportedUsd: 0.0012,
+	usage: EMPTY_USAGE,
+	reasons: [],
+	escalated: false,
+};
+
+function chunk(raw: Record<string, unknown>): { raw: Record<string, unknown>; events: [] } {
+	return { raw, events: [] };
+}
+
+function parseFrames(text: string): unknown[] {
+	return text
+		.split("\n\n")
+		.filter((f) => f.startsWith("data: ") && f !== "data: [DONE]")
+		.map((f) => JSON.parse(f.slice("data: ".length)));
+}
+
+describe("createStreamingSink", () => {
+	test("rewrites model to the virtual id and preserves the served slug under x_auto_model_router", async () => {
+		const { sink, response } = createStreamingSink("auto");
+		expect(response.headers.get("content-type")).toBe("text/event-stream");
+		expect(response.headers.get("cache-control")).toBe("no-cache");
+
+		sink.chunk(
+			chunk({
+				id: "gen-1",
+				object: "chat.completion.chunk",
+				created: 1700000000,
+				model: "openai/gpt-5.5",
+				choices: [{ index: 0, delta: { content: "hi" }, finish_reason: null }],
+			}),
+		);
+		sink.finish(SUMMARY);
+
+		const text = await response.text();
+		const frames = parseFrames(text);
+		const first = frames[0] as Record<string, unknown>;
+		expect(first.model).toBe("auto");
+		expect(first.x_auto_model_router).toEqual({ model: "openai/gpt-5.5" });
+		// Unknown upstream fields ride along verbatim.
+		expect(first.id).toBe("gen-1");
+
+		// Headers flushed with the first chunk, so the summary arrives as the
+		// final x_auto_model_router frame before [DONE].
+		const last = frames[frames.length - 1] as Record<string, unknown>;
+		expect(last.x_auto_model_router).toEqual({
+			model: "openai/gpt-5.5",
+			tier: "simple",
+			cost_usd: 0.0012,
+			attempts: 1,
+		});
+		expect(text.trimEnd().endsWith("data: [DONE]")).toBe(true);
+	});
+
+	test("error emits one OpenAI-envelope frame before closing", async () => {
+		const { sink, response } = createStreamingSink("auto");
+		sink.error({ status: 502, code: "upstream_error", message: "boom" });
+		const frames = parseFrames(await response.text());
+		expect(frames).toEqual([
+			{ error: { message: "boom", type: "server_error", code: "upstream_error" } },
+		]);
+	});
+});
+
+describe("createBufferedSink", () => {
+	test("reassembles split tool_calls argument fragments into one valid JSON string", async () => {
+		const { sink, response } = createBufferedSink("auto");
+		sink.chunk(
+			chunk({
+				id: "gen-1",
+				created: 1700000000,
+				model: "anthropic/claude-haiku-4.5",
+				choices: [
+					{
+						index: 0,
+						delta: {
+							role: "assistant",
+							tool_calls: [
+								{
+									index: 0,
+									id: "call_abc",
+									type: "function",
+									function: { name: "read", arguments: "{\"path\":" },
+								},
+							],
+						},
+						finish_reason: null,
+					},
+				],
+			}),
+		);
+		// Second fragment carries neither id nor name; they must survive from the first.
+		sink.chunk(
+			chunk({
+				model: "anthropic/claude-haiku-4.5",
+				choices: [
+					{
+						index: 0,
+						delta: { tool_calls: [{ index: 0, function: { arguments: "\"src/main.ts\"}" } }] },
+						finish_reason: null,
+					},
+				],
+			}),
+		);
+		sink.chunk(
+			chunk({
+				model: "anthropic/claude-haiku-4.5",
+				choices: [{ index: 0, delta: {}, finish_reason: "tool_calls" }],
+				usage: { prompt_tokens: 10, completion_tokens: 5, total_tokens: 15 },
+			}),
+		);
+		sink.finish(SUMMARY);
+
+		const res = await response;
+		expect(res.headers.get("x-auto-model-router-model")).toBe("openai/gpt-5.5");
+		expect(res.headers.get("x-auto-model-router-tier")).toBe("simple");
+		expect(res.headers.get("x-auto-model-router-cost-usd")).toBe("0.0012");
+		expect(res.headers.get("x-auto-model-router-attempts")).toBe("1");
+
+		const body = (await res.json()) as {
+			object: string;
+			model: string;
+			choices: Array<{
+				finish_reason: string;
+				message: { tool_calls: Array<{ id: string; function: { name: string; arguments: string } }> };
+			}>;
+			usage: unknown;
+		};
+		expect(body.object).toBe("chat.completion");
+		expect(body.model).toBe("auto");
+		expect(body.choices[0]!.finish_reason).toBe("tool_calls");
+		const call = body.choices[0]!.message.tool_calls[0]!;
+		expect(call.id).toBe("call_abc");
+		expect(call.function.name).toBe("read");
+		expect(JSON.parse(call.function.arguments)).toEqual({ path: "src/main.ts" });
+		expect(body.usage).toEqual({ prompt_tokens: 10, completion_tokens: 5, total_tokens: 15 });
+	});
+
+	test("merges text and reasoning deltas per choice", async () => {
+		const { sink, response } = createBufferedSink("auto-cheap");
+		sink.chunk(
+			chunk({
+				model: "x/y",
+				choices: [{ index: 0, delta: { role: "assistant", reasoning: "think ", content: "Hel" } }],
+			}),
+		);
+		sink.chunk(
+			chunk({
+				model: "x/y",
+				choices: [{ index: 0, delta: { reasoning: "more", content: "lo" }, finish_reason: "stop" }],
+			}),
+		);
+		sink.finish(SUMMARY);
+		const body = (await (await response).json()) as {
+			choices: Array<{ message: { role: string; content: string; reasoning: string } }>;
+		};
+		expect(body.choices[0]!.message).toEqual({ role: "assistant", content: "Hello", reasoning: "think more" });
+	});
+
+	test("error before finish resolves with the envelope at the WireError status", async () => {
+		const { sink, response } = createBufferedSink("auto");
+		sink.error({ status: 429, code: "rate_limit", message: "slow down" });
+		const res = await response;
+		expect(res.status).toBe(429);
+		expect(await res.json()).toEqual({
+			error: { message: "slow down", type: "rate_limit_error", code: "rate_limit" },
+		});
+	});
+});

package/tools/install.ts

@@ -0,0 +1,140 @@
+#!/usr/bin/env bun
+/**
+ * auto-model-router installer.
+ *
+ * Wires the auto-model-router extensions into omp's `~/.omp/agent/config.yml`
+ * (or `$PI_CODING_AGENT_DIR/config.yml` when that env var relocates the agent
+ * dir). Cross-platform: Windows, macOS, and Linux.
+ *
+ * It adds the three extension paths under `extensions:`:
+ *   - router-embed.ts      (required — runs the router in-process)
+ *   - router-toast.ts      (optional — chosen-model toasts)
+ *   - router-configure.ts  (optional — /router configure command)
+ *
+ * Idempotent: already-present paths are left untouched, so re-running is safe.
+ * The previous config.yml is backed up to a timestamped .bak before writing.
+ *
+ * Usage:
+ *   bun tools/install.ts
+ *   bun tools/install.ts --no-toast --no-configure   # only the embed extension
+ *   PI_CODING_AGENT_DIR=/custom/agent bun tools/install.ts
+ */
+
+import { copyFileSync, existsSync, mkdirSync, readFileSync, writeFileSync } from "node:fs";
+import { homedir } from "node:os";
+import { dirname, join, resolve } from "node:path";
+
+/** The repo's omp-extension directory, resolved from this script's location. */
+const EXT_DIR = resolve(import.meta.dir, "..", "omp-extension");
+
+/** The three extensions, in install order. */
+const EXTENSIONS = [
+	{ file: "router-embed.ts", required: true, label: "router-embed (required)" },
+	{ file: "router-toast.ts", required: false, label: "router-toast (toasts)" },
+	{ file: "router-configure.ts", required: false, label: "router-configure (/router configure)" },
+];
+
+/** omp's agent config.yml path, honoring $PI_CODING_AGENT_DIR. */
+function agentConfigPath(): string {
+	const agentDir = process.env.PI_CODING_AGENT_DIR;
+	const base = agentDir !== undefined && agentDir !== "" ? agentDir : join(homedir(), ".omp", "agent");
+	return join(base, "config.yml");
+}
+
+/** The absolute path of an extension file, using forward slashes for portability. */
+function extensionPath(file: string): string {
+	return join(EXT_DIR, file).replace(/\\/g, "/");
+}
+
+/**
+ * Inserts the extension paths under an `extensions:` key, preserving comments
+ * and existing content. Creates the key if absent; appends to it if present.
+ */
+function spliceExtensions(text: string, paths: string[]): string {
+	const eol = text.includes("\r\n") ? "\r\n" : "\n";
+	const lines = text.split(/\r?\n/);
+
+	// Find the `extensions:` key and its indentation.
+	const extIdx = lines.findIndex((l) => /^extensions\s*:/.test(l));
+	if (extIdx === -1) {
+		// No extensions key: append one at the end (or after the last non-empty line).
+		const indent = "  ";
+		const block = ["extensions:", ...paths.map((p) => `${indent}- ${p}`)];
+		const trimmed = lines.map((l) => l.trimEnd());
+		while (trimmed.length > 0 && trimmed[trimmed.length - 1] === "") trimmed.pop();
+		return [...trimmed, "", ...block, ""].join(eol);
+	}
+
+	// Existing key: find its child indentation (or default to 2 spaces).
+	const extLine = lines[extIdx] ?? "";
+	const indent = /^([ \t]*)/.exec(extLine)?.[1] ?? "";
+	const childIndent = `${indent}  `;
+
+	// Collect existing children until the next top-level key.
+	const children: string[] = [];
+	let i = extIdx + 1;
+	while (i < lines.length) {
+		const line = lines[i] ?? "";
+		if (line.trim() === "") {
+			i++;
+			continue;
+		}
+		// A line at the same or less indentation than the key ends the list.
+		if (/^[ \t]*\S/.test(line) && !line.startsWith(childIndent)) break;
+		children.push(line);
+		i++;
+	}
+
+	const existing = children.map((l) => l.trim());
+	const toAdd = paths.filter((p) => !existing.includes(`- ${p}`));
+	if (toAdd.length === 0) return text;
+
+	const newChildren = [...children, ...toAdd.map((p) => `${childIndent}- ${p}`)];
+	const next = [...lines.slice(0, extIdx + 1), ...newChildren, ...lines.slice(i)];
+	return next.join(eol);
+}
+
+function main(): void {
+	const args = process.argv.slice(2);
+	const noToast = args.includes("--no-toast");
+	const noConfigure = args.includes("--no-configure");
+
+	const selected = EXTENSIONS.filter((e) => {
+		if (e.file === "router-toast.ts" && noToast) return false;
+		if (e.file === "router-configure.ts" && noConfigure) return false;
+		return true;
+	});
+
+	// Verify the required extension files exist before touching anything.
+	const missing = selected.filter((e) => !existsSync(extensionPath(e.file)));
+	if (missing.length > 0) {
+		console.error(`error: missing extension file(s): ${missing.map((m) => extensionPath(m.file)).join(", ")}`);
+		process.exit(1);
+	}
+
+	const target = agentConfigPath();
+	const existing = existsSync(target) ? readFileSync(target, "utf8") : "";
+	const paths = selected.map((e) => extensionPath(e.file));
+	const next = spliceExtensions(existing, paths);
+
+	if (next === existing) {
+		console.log(`already installed: ${target}`);
+		return;
+	}
+
+	// Back up the previous file.
+	if (existing !== "") {
+		const backup = `${target}.${new Date().toISOString().replace(/[:.]/g, "-")}.bak`;
+		copyFileSync(target, backup);
+		console.log(`backup: ${backup}`);
+	}
+
+	mkdirSync(dirname(target), { recursive: true });
+	writeFileSync(target, next, "utf8");
+
+	console.log(`installed auto-model-router extensions into ${target}:`);
+	for (const e of selected) console.log(`  - ${extensionPath(e.file)}  (${e.label})`);
+	console.log("restart the omp session to load them.");
+}
+
+main();