oh-my-pi/packages/coding-agent/test/model-discovery.test.ts

import { afterEach, beforeEach, describe, expect, test } from "bun:test";
import * as fs from "node:fs";
import * as os from "node:os";
import * as path from "node:path";
import { Effort, type FetchImpl, type Model } from "@oh-my-pi/pi-ai";
import { buildModel } from "@oh-my-pi/pi-catalog/build";
import { writeModelCache } from "@oh-my-pi/pi-catalog/model-cache";
import { kNoAuth, ModelRegistry } from "@oh-my-pi/pi-coding-agent/config/model-registry";
import { resetSettingsForTest } from "@oh-my-pi/pi-coding-agent/config/settings";
import { AuthStorage } from "@oh-my-pi/pi-coding-agent/session/auth-storage";
import { Snowflake } from "@oh-my-pi/pi-utils";

describe("ModelRegistry runtime discovery", () => {
	let tempDir: string;
	let modelsJsonPath: string;
	let cacheDbPath: string;
	let authStorage: AuthStorage;
	let originalOllamaBaseUrl: string | undefined;
	let originalOllamaHost: string | undefined;
	let originalOllamaContextLength: string | undefined;

	beforeEach(async () => {
		resetSettingsForTest();
		originalOllamaBaseUrl = Bun.env.OLLAMA_BASE_URL;
		originalOllamaHost = Bun.env.OLLAMA_HOST;
		originalOllamaContextLength = Bun.env.OLLAMA_CONTEXT_LENGTH;
		delete Bun.env.OLLAMA_BASE_URL;
		delete Bun.env.OLLAMA_HOST;
		delete Bun.env.OLLAMA_CONTEXT_LENGTH;
		tempDir = path.join(os.tmpdir(), `pi-test-model-registry-${Snowflake.next()}`);
		fs.mkdirSync(tempDir, { recursive: true });
		modelsJsonPath = path.join(tempDir, "models.json");
		cacheDbPath = path.join(tempDir, "models.db");
		// In-memory auth DB: tests need a fresh, isolated credential store per case but
		// never reopen it from disk, so :memory: avoids the WAL/chmod disk-open cost
		// (~3ms/test) while preserving per-test isolation.
		authStorage = await AuthStorage.create(":memory:");
	});

	afterEach(() => {
		resetSettingsForTest();
		if (originalOllamaBaseUrl === undefined) {
			delete Bun.env.OLLAMA_BASE_URL;
		} else {
			Bun.env.OLLAMA_BASE_URL = originalOllamaBaseUrl;
		}
		if (originalOllamaHost === undefined) {
			delete Bun.env.OLLAMA_HOST;
		} else {
			Bun.env.OLLAMA_HOST = originalOllamaHost;
		}
		if (originalOllamaContextLength === undefined) {
			delete Bun.env.OLLAMA_CONTEXT_LENGTH;
		} else {
			Bun.env.OLLAMA_CONTEXT_LENGTH = originalOllamaContextLength;
		}
		authStorage.close();
		if (tempDir && fs.existsSync(tempDir)) {
			fs.rmSync(tempDir, { recursive: true });
		}
	});

	function writeCachedOllamaModels(models: Model<"openai-completions">[]) {
		writeModelCache("ollama", Date.now(), models, true, "", cacheDbPath);
	}

	function getModelsForProvider(registry: ModelRegistry, provider: string) {
		return registry.getAll().filter(m => m.provider === provider);
	}

	function withEnv(name: "OLLAMA_BASE_URL" | "OLLAMA_CONTEXT_LENGTH" | "OLLAMA_HOST", value: string | undefined) {
		const original = Bun.env[name];
		if (value === undefined) {
			delete Bun.env[name];
		} else {
			Bun.env[name] = value;
		}
		return {
			[Symbol.dispose]() {
				if (original === undefined) {
					delete Bun.env[name];
				} else {
					Bun.env[name] = original;
				}
			},
		};
	}

	/** Write raw providers config (for mixed override/replacement scenarios) */
	function writeRawModelsJson(providers: Record<string, unknown>) {
		fs.writeFileSync(modelsJsonPath, JSON.stringify({ providers }));
	}

	function mockOllamaDiscovery(
		modelNames: string[],
		endpoint = "http://127.0.0.1:11434",
		showPayload: Record<string, unknown> = { capabilities: ["completion"] },
	): FetchImpl {
		return async input => {
			const url = String(input);
			if (url === `${endpoint}/api/tags`) {
				return new Response(JSON.stringify({ models: modelNames.map(name => ({ name })) }), {
					status: 200,
					headers: { "Content-Type": "application/json" },
				});
			}
			if (url === `${endpoint}/api/show`) {
				return new Response(JSON.stringify(showPayload), {
					status: 200,
					headers: { "Content-Type": "application/json" },
				});
			}
			throw new Error(`Unexpected URL: ${url}`);
		};
	}

	test("auto-discovers ollama models without provider config", async () => {
		const fetchMock = mockOllamaDiscovery(["phi4-mini"]);
		const registry = new ModelRegistry(authStorage, modelsJsonPath, { fetch: fetchMock });
		await registry.refresh();
		const ollamaModels = getModelsForProvider(registry, "ollama");
		expect(ollamaModels.some(m => m.id === "phi4-mini")).toBe(true);
		expect(registry.getAvailable().some(m => m.provider === "ollama" && m.id === "phi4-mini")).toBe(true);
		expect(await registry.getApiKey(ollamaModels[0])).toBe(kNoAuth);
	});

	test("uses OLLAMA_HOST for implicit ollama discovery", async () => {
		using _baseUrl = withEnv("OLLAMA_BASE_URL", undefined);
		using _host = withEnv("OLLAMA_HOST", "ollama.lan:12345");
		const fetchMock = mockOllamaDiscovery(["phi4-mini"], "http://ollama.lan:12345");
		const registry = new ModelRegistry(authStorage, modelsJsonPath, { fetch: fetchMock });
		await registry.refresh();

		const model = registry.find("ollama", "phi4-mini");
		expect(model?.baseUrl).toBe("http://ollama.lan:12345/v1");
	});

	test("keeps OLLAMA_BASE_URL precedence over OLLAMA_HOST", async () => {
		using _baseUrl = withEnv("OLLAMA_BASE_URL", "http://omp-ollama.example:2222");
		using _host = withEnv("OLLAMA_HOST", "ollama-host.example:3333");
		const fetchMock = mockOllamaDiscovery(["phi4-mini"], "http://omp-ollama.example:2222");
		const registry = new ModelRegistry(authStorage, modelsJsonPath, { fetch: fetchMock });
		await registry.refresh();

		const model = registry.find("ollama", "phi4-mini");
		expect(model?.baseUrl).toBe("http://omp-ollama.example:2222/v1");
	});

	test("uses OLLAMA_CONTEXT_LENGTH for implicit ollama context accounting", async () => {
		using _contextLength = withEnv("OLLAMA_CONTEXT_LENGTH", "16384");
		const fetchMock = mockOllamaDiscovery(["phi4-mini"]);
		const registry = new ModelRegistry(authStorage, modelsJsonPath, { fetch: fetchMock });
		await registry.refresh();

		const model = registry.find("ollama", "phi4-mini");
		expect(model?.contextWindow).toBe(16384);
		expect(model?.maxTokens).toBe(16384);
	});

	test("lets OLLAMA_CONTEXT_LENGTH override ollama show metadata", async () => {
		using _contextLength = withEnv("OLLAMA_CONTEXT_LENGTH", "32768");
		const fetchMock = mockOllamaDiscovery(["phi4-mini"], "http://127.0.0.1:11434", {
			model_info: {
				"phi4.context_length": 4096,
			},
			capabilities: ["completion"],
		});
		const registry = new ModelRegistry(authStorage, modelsJsonPath, { fetch: fetchMock });
		await registry.refresh();

		const model = registry.find("ollama", "phi4-mini");
		expect(model?.contextWindow).toBe(32768);
		expect(model?.maxTokens).toBe(32768);
	});

	test("discovers ollama-cloud through built-in descriptor flow without regressing local implicit ollama", async () => {
		authStorage.setRuntimeApiKey("ollama-cloud", "cloud-test-key");

		const fetchMock: FetchImpl = async (input, init) => {
			const url = String(input);
			if (url === "http://127.0.0.1:11434/api/tags") {
				return new Response(JSON.stringify({ models: [{ name: "phi4-mini" }] }), {
					status: 200,
					headers: { "Content-Type": "application/json" },
				});
			}
			if (url === "http://127.0.0.1:11434/api/show") {
				return new Response(JSON.stringify({ capabilities: ["completion"] }), {
					status: 200,
					headers: { "Content-Type": "application/json" },
				});
			}
			if (url === "https://ollama.com/api/tags") {
				const headers = new Headers(init?.headers);
				expect(headers.get("Authorization")).toBe("Bearer cloud-test-key");
				return new Response(JSON.stringify({ models: [{ name: "gpt-oss:120b" }] }), {
					status: 200,
					headers: { "Content-Type": "application/json" },
				});
			}
			if (url === "https://ollama.com/api/show") {
				const headers = new Headers(init?.headers);
				expect(headers.get("Authorization")).toBe("Bearer cloud-test-key");
				const body = JSON.parse(String(init?.body ?? "{}")) as { model?: string };
				expect(body.model).toBe("gpt-oss:120b");
				return new Response(
					JSON.stringify({
						capabilities: ["completion", "thinking"],
						model_info: { "gpt-oss.context_length": 262144 },
					}),
					{ status: 200, headers: { "Content-Type": "application/json" } },
				);
			}
			throw new Error(`Unexpected URL: ${url}`);
		};

		const registry = new ModelRegistry(authStorage, modelsJsonPath, { fetch: fetchMock });
		await registry.refresh();

		const local = registry.find("ollama", "phi4-mini");
		const cloud = registry.find("ollama-cloud", "gpt-oss:120b");

		expect(local?.provider).toBe("ollama");
		expect(local?.api).toBe("openai-responses");
		expect(cloud?.provider).toBe("ollama-cloud");
		expect(cloud?.api).toBe("ollama-chat");
		expect(cloud?.baseUrl).toBe("https://ollama.com");
		expect(cloud?.reasoning).toBe(true);
		expect(cloud?.contextWindow).toBe(262144);
		expect(await registry.getApiKey(cloud!)).toBe("cloud-test-key");
		expect(registry.getAvailable().some(model => model.provider === "ollama" && model.id === "phi4-mini")).toBe(true);
		expect(
			registry.getAvailable().some(model => model.provider === "ollama-cloud" && model.id === "gpt-oss:120b"),
		).toBe(true);
	});
	test("discovers ollama models at runtime and treats auth:none providers as available", async () => {
		writeRawModelsJson({
			ollama: {
				baseUrl: "http://127.0.0.1:11434/v1",
				api: "openai-completions",
				auth: "none",
				discovery: { type: "ollama" },
			},
		});

		const fetchMock: FetchImpl = async input => {
			const url = String(input);
			if (url === "http://127.0.0.1:11434/api/tags") {
				return new Response(
					JSON.stringify({
						models: [{ name: "qwen2.5-coder:7b" }, { model: "llama3.2:3b", name: "llama3.2:3b" }],
					}),
					{ status: 200, headers: { "Content-Type": "application/json" } },
				);
			}
			if (url === "http://127.0.0.1:11434/api/show") {
				return new Response(JSON.stringify({ capabilities: ["completion"] }), {
					status: 200,
					headers: { "Content-Type": "application/json" },
				});
			}
			throw new Error(`Unexpected URL: ${url}`);
		};

		const registry = new ModelRegistry(authStorage, modelsJsonPath, { fetch: fetchMock });
		await registry.refresh();

		const ollamaModels = getModelsForProvider(registry, "ollama");
		expect(ollamaModels.some(m => m.id === "qwen2.5-coder:7b")).toBe(true);
		expect(ollamaModels.some(m => m.id === "llama3.2:3b")).toBe(true);

		const available = registry.getAvailable().filter(m => m.provider === "ollama");
		expect(available.length).toBe(2);
		expect(await registry.getApiKey(available[0])).toBe(kNoAuth);
	});

	test("normalizes cached ollama completions rows to responses on load", () => {
		writeRawModelsJson({
			ollama: {
				baseUrl: "http://127.0.0.1:11434/v1",
				api: "openai-responses",
				auth: "none",
				discovery: { type: "ollama" },
			},
		});
		writeCachedOllamaModels([
			buildModel({
				id: "phi4-mini",
				name: "phi4-mini",
				api: "openai-completions",
				provider: "ollama",
				baseUrl: "http://127.0.0.1:11434/v1",
				reasoning: false,
				input: ["text"],
				cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 },
				contextWindow: 128000,
				maxTokens: 8192,
			}),
		]);

		const registry = new ModelRegistry(authStorage, modelsJsonPath);
		const ollama = registry.find("ollama", "phi4-mini");

		expect(ollama?.api).toBe("openai-responses");
		expect(ollama?.baseUrl).toBe("http://127.0.0.1:11434/v1");
		expect(registry.getProviderDiscoveryState("ollama")?.status).toBe("cached");
	});

	test("discovers ollama thinking capabilities from show metadata", async () => {
		writeRawModelsJson({
			ollama: {
				baseUrl: "http://127.0.0.1:11434/v1",
				api: "openai-completions",
				auth: "none",
				discovery: { type: "ollama" },
			},
		});

		const fetchMock: FetchImpl = async (input, init) => {
			const url = String(input);
			if (url === "http://127.0.0.1:11434/api/tags") {
				return new Response(
					JSON.stringify({
						models: [{ name: "qwen3.5:397b-cloud" }, { name: "llama3.2:3b" }],
					}),
					{ status: 200, headers: { "Content-Type": "application/json" } },
				);
			}
			if (url === "http://127.0.0.1:11434/api/show") {
				const body = JSON.parse(String(init?.body ?? "{}")) as { model?: string };
				if (body.model === "qwen3.5:397b-cloud") {
					return new Response(JSON.stringify({ capabilities: ["completion", "thinking"] }), {
						status: 200,
						headers: { "Content-Type": "application/json" },
					});
				}
				if (body.model === "llama3.2:3b") {
					return new Response(JSON.stringify({ capabilities: ["completion"] }), {
						status: 200,
						headers: { "Content-Type": "application/json" },
					});
				}
			}
			throw new Error(`Unexpected request: ${url}`);
		};

		const registry = new ModelRegistry(authStorage, modelsJsonPath, { fetch: fetchMock });
		await registry.refresh();

		const qwen = registry.find("ollama", "qwen3.5:397b-cloud");
		expect(qwen?.reasoning).toBe(true);
		expect(qwen?.thinking).toEqual({
			mode: "effort",
			efforts: [Effort.Minimal, Effort.Low, Effort.Medium, Effort.High],
		});

		const llama = registry.find("ollama", "llama3.2:3b");
		expect(llama?.reasoning).toBe(false);
	});

	test("discovers ollama context window from show model_info", async () => {
		const fetchMock: FetchImpl = async (input, init) => {
			const url = String(input);
			if (url === "http://127.0.0.1:11434/api/tags") {
				return new Response(JSON.stringify({ models: [{ name: "gemma3:4b" }] }), {
					status: 200,
					headers: { "Content-Type": "application/json" },
				});
			}
			if (url === "http://127.0.0.1:11434/api/show") {
				const body = JSON.parse(String(init?.body ?? "{}")) as { model?: string };
				if (body.model === "gemma3:4b") {
					return new Response(
						JSON.stringify({
							model_info: {
								"gemma3.context_length": 131072,
							},
						}),
						{
							status: 200,
							headers: { "Content-Type": "application/json" },
						},
					);
				}
			}
			throw new Error(`Unexpected request: ${url}`);
		};

		const registry = new ModelRegistry(authStorage, modelsJsonPath, { fetch: fetchMock });
		await registry.refresh();

		const gemma = registry.find("ollama", "gemma3:4b");
		expect(gemma?.contextWindow).toBe(131072);
		expect(gemma?.maxTokens).toBe(32_768);
		expect(gemma?.input).toEqual(["text"]);
		expect(gemma?.reasoning).toBe(false);
	});

	test("discovery failure does not fail model registry refresh", async () => {
		writeRawModelsJson({
			ollama: {
				baseUrl: "http://127.0.0.1:11434",
				api: "openai-completions",
				auth: "none",
				discovery: { type: "ollama" },
			},
		});

		const fetchMock: FetchImpl = () => {
			throw new Error("connection refused");
		};

		const registry = new ModelRegistry(authStorage, modelsJsonPath, { fetch: fetchMock });
		await registry.refresh();
		expect(getModelsForProvider(registry, "ollama")).toHaveLength(0);
		expect(registry.getError()).toBeUndefined();
	});
	test("loads cached local models before live refresh and preserves them on failure", async () => {
		writeRawModelsJson({
			ollama: {
				baseUrl: "http://127.0.0.1:11434/v1",
				api: "openai-completions",
				auth: "none",
				discovery: { type: "ollama" },
			},
		});

		{
			const fetchMock = mockOllamaDiscovery(["phi4-mini"]);
			const primedRegistry = new ModelRegistry(authStorage, modelsJsonPath, { fetch: fetchMock });
			await primedRegistry.refresh();
		}

		const failingFetch: FetchImpl = () => {
			throw new Error("connection refused");
		};
		const cachedRegistry = new ModelRegistry(authStorage, modelsJsonPath, { fetch: failingFetch });
		expect(getModelsForProvider(cachedRegistry, "ollama").some(model => model.id === "phi4-mini")).toBe(true);
		expect(cachedRegistry.getProviderDiscoveryState("ollama")?.status).toBe("cached");

		await cachedRegistry.refreshProvider("ollama");

		expect(getModelsForProvider(cachedRegistry, "ollama").some(model => model.id === "phi4-mini")).toBe(true);
		const state = cachedRegistry.getProviderDiscoveryState("ollama");
		expect(state?.status).toBe("cached");
		expect(state?.error).toContain("connection refused");
	});

	test("reports unauthenticated discoverable providers without discarding cached models", async () => {
		writeRawModelsJson({
			"custom-local": {
				baseUrl: "http://127.0.0.1:11434/v1",
				api: "openai-completions",
				discovery: { type: "ollama" },
			},
		});
		authStorage.setRuntimeApiKey("custom-local", "test-key");

		{
			const fetchMock: FetchImpl = async input => {
				const url = String(input);
				if (url === "http://127.0.0.1:11434/api/tags") {
					return new Response(JSON.stringify({ models: [{ name: "local-coder" }] }), {
						status: 200,
						headers: { "Content-Type": "application/json" },
					});
				}
				if (url === "http://127.0.0.1:11434/api/show") {
					return new Response(JSON.stringify({ capabilities: ["completion"] }), {
						status: 200,
						headers: { "Content-Type": "application/json" },
					});
				}
				throw new Error(`Unexpected URL: ${url}`);
			};
			const primedRegistry = new ModelRegistry(authStorage, modelsJsonPath, { fetch: fetchMock });
			await primedRegistry.refreshProvider("custom-local");
		}

		authStorage.setRuntimeApiKey("custom-local", "");
		const cachedRegistry = new ModelRegistry(authStorage, modelsJsonPath);
		await cachedRegistry.refreshProvider("custom-local");

		expect(getModelsForProvider(cachedRegistry, "custom-local").some(model => model.id === "local-coder")).toBe(true);
		const state = cachedRegistry.getProviderDiscoveryState("custom-local");
		expect(state?.status).toBe("unauthenticated");
		expect(state?.models).toContain("local-coder");
	});
	test("llama.cpp discovery honors configured API key", async () => {
		authStorage.setRuntimeApiKey("llama.cpp", "test-llama-key");
		const fetchMock: FetchImpl = async (input, init) => {
			const url = String(input);
			if (url === "http://127.0.0.1:8080/models") {
				const headers = init?.headers as Headers | Record<string, string> | undefined;
				let authHeader: string | null = null;
				if (headers instanceof Headers) {
					authHeader = headers.get("Authorization");
				} else if (typeof headers === "object") {
					authHeader = headers.Authorization;
				}
				expect(String(authHeader ?? "")).toBe("Bearer test-llama-key");
				return new Response(JSON.stringify({ data: [{ id: "llama-3.2:3b" }, { id: "mistral:7b" }] }), {
					status: 200,
					headers: { "Content-Type": "application/json" },
				});
			}
			if (url === "http://127.0.0.1:8080/props") {
				const headers = init?.headers as Headers | Record<string, string> | undefined;
				let authHeader: string | null = null;
				if (headers instanceof Headers) {
					authHeader = headers.get("Authorization");
				} else if (typeof headers === "object") {
					authHeader = headers.Authorization;
				}
				expect(String(authHeader ?? "")).toBe("Bearer test-llama-key");
				return new Response(JSON.stringify({ default_generation_settings: { n_ctx: 262144 } }), {
					status: 200,
					headers: { "Content-Type": "application/json" },
				});
			}
			throw new Error(`Unexpected URL: ${url}`);
		};
		const registry = new ModelRegistry(authStorage, modelsJsonPath, { fetch: fetchMock });
		await registry.refresh();
		const llamaModels = getModelsForProvider(registry, "llama.cpp");
		expect(llamaModels.some(m => m.id === "llama-3.2:3b")).toBe(true);
		const apiKey = await registry.getApiKey(llamaModels[0]);
		expect(apiKey).toBe("test-llama-key");
		expect(apiKey).not.toBe(kNoAuth);
	});
	test("llama.cpp discovery without API key is treated as keyless", async () => {
		const fetchMock: FetchImpl = async (input, init) => {
			const url = String(input);
			if (url === "http://127.0.0.1:8080/models") {
				const headers = init?.headers as Headers | Record<string, string> | undefined;
				let authHeader: string | null = null;
				if (headers instanceof Headers) {
					authHeader = headers.get("Authorization");
				} else if (typeof headers === "object") {
					authHeader = headers.Authorization;
				}
				// When no API key, headers should be empty object or undefined
				expect(authHeader).toBeUndefined();
				return new Response(JSON.stringify({ data: [{ id: "llama-3.2:3b" }] }), {
					status: 200,
					headers: { "Content-Type": "application/json" },
				});
			}
			if (url === "http://127.0.0.1:8080/props") {
				const headers = init?.headers as Headers | Record<string, string> | undefined;
				let authHeader: string | null = null;
				if (headers instanceof Headers) {
					authHeader = headers.get("Authorization");
				} else if (typeof headers === "object") {
					authHeader = headers.Authorization;
				}
				expect(authHeader).toBeUndefined();
				return new Response(JSON.stringify({ default_generation_settings: { n_ctx: 262144 } }), {
					status: 200,
					headers: { "Content-Type": "application/json" },
				});
			}
			throw new Error(`Unexpected URL: ${url}`);
		};
		const registry = new ModelRegistry(authStorage, modelsJsonPath, { fetch: fetchMock });
		await registry.refresh();
		const state = registry.getProviderDiscoveryState("llama.cpp");
		if (state?.status !== "ok") {
			throw new Error(`Discovery failed with status ${state?.status}: ${state?.error}`);
		}
		const llamaModels = getModelsForProvider(registry, "llama.cpp");
		const apiKey = await registry.getApiKey(llamaModels[0]);
		expect(apiKey).toBe(kNoAuth);
	});
	test("llama.cpp discovery reads context window from props n_ctx", async () => {
		const fetchMock: FetchImpl = async input => {
			const url = String(input);
			if (url === "http://127.0.0.1:8080/models") {
				return new Response(JSON.stringify({ data: [{ id: "qwen35-35b-a3b" }] }), {
					status: 200,
					headers: { "Content-Type": "application/json" },
				});
			}
			if (url === "http://127.0.0.1:8080/props") {
				return new Response(
					JSON.stringify({
						default_generation_settings: {
							n_ctx: 262144,
						},
						modalities: {
							vision: true,
							audio: false,
						},
					}),
					{
						status: 200,
						headers: { "Content-Type": "application/json" },
					},
				);
			}
			throw new Error(`Unexpected URL: ${url}`);
		};
		const registry = new ModelRegistry(authStorage, modelsJsonPath, { fetch: fetchMock });
		await registry.refresh();
		const llama = registry.find("llama.cpp", "qwen35-35b-a3b");
		expect(llama?.contextWindow).toBe(262144);
		expect(llama?.maxTokens).toBe(32_768);
		expect(llama?.input).toEqual(["text", "image"]);
	});
});