Something went wrong. Try again.
fork of prime agent with changes i want but are not upstreamable probqbly (and some fixes)
Something went wrong. Try again.
70 kB · 2033 lines
TypeScript
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809810811812813814815816817818819820821822823824825826827828829830831832833834835836837838839840841842843844845846847848849850851852853854855856857858859860861862863864865866867868869870871872873874875876877878879880881882883884885886887888889890891892893894895896897898899900901902903904905906907908909910911912913914915916917918919920921922923924925926927928929930931932933934935936937938939940941942943944945946947948949950951952953954955956957958959960961962963964965966967968969970971972973974975976977978979980981982983984985986987988989990991992993994995996997998999100010011002100310041005100610071008100910101011101210131014101510161017101810191020102110221023102410251026102710281029103010311032103310341035103610371038103910401041104210431044104510461047104810491050105110521053105410551056105710581059106010611062106310641065106610671068106910701071107210731074107510761077107810791080108110821083108410851086108710881089109010911092109310941095109610971098109911001101110211031104110511061107110811091110111111121113111411151116111711181119112011211122112311241125112611271128112911301131113211331134113511361137113811391140114111421143114411451146114711481149115011511152115311541155115611571158115911601161116211631164116511661167116811691170117111721173117411751176117711781179118011811182118311841185118611871188118911901191119211931194119511961197119811991200120112021203120412051206120712081209121012111212121312141215121612171218121912201221122212231224122512261227122812291230123112321233123412351236123712381239124012411242124312441245124612471248124912501251125212531254125512561257125812591260126112621263126412651266126712681269127012711272127312741275127612771278127912801281128212831284128512861287128812891290129112921293129412951296129712981299130013011302130313041305130613071308130913101311131213131314131513161317131813191320132113221323132413251326132713281329133013311332133313341335133613371338133913401341134213431344134513461347134813491350135113521353135413551356135713581359136013611362136313641365136613671368136913701371137213731374137513761377137813791380138113821383138413851386138713881389139013911392139313941395139613971398139914001401140214031404140514061407140814091410141114121413141414151416141714181419142014211422142314241425142614271428142914301431143214331434143514361437143814391440144114421443144414451446144714481449145014511452145314541455145614571458145914601461146214631464146514661467146814691470147114721473147414751476147714781479148014811482148314841485148614871488148914901491149214931494149514961497149814991500150115021503150415051506150715081509151015111512151315141515151615171518151915201521152215231524152515261527152815291530153115321533153415351536153715381539154015411542154315441545154615471548154915501551155215531554155515561557155815591560156115621563156415651566156715681569157015711572157315741575157615771578157915801581158215831584158515861587158815891590159115921593159415951596159715981599160016011602160316041605160616071608160916101611161216131614161516161617161816191620162116221623162416251626162716281629163016311632163316341635163616371638163916401641164216431644164516461647164816491650165116521653165416551656165716581659166016611662166316641665166616671668166916701671167216731674167516761677167816791680168116821683168416851686168716881689169016911692169316941695169616971698169917001701170217031704170517061707170817091710171117121713171417151716171717181719172017211722172317241725172617271728172917301731173217331734173517361737173817391740174117421743174417451746174717481749175017511752175317541755175617571758175917601761176217631764176517661767176817691770177117721773177417751776177717781779178017811782178317841785178617871788178917901791179217931794179517961797179817991800180118021803180418051806180718081809181018111812181318141815181618171818181918201821182218231824182518261827182818291830183118321833183418351836183718381839184018411842184318441845184618471848184918501851185218531854185518561857185818591860186118621863186418651866186718681869187018711872187318741875187618771878187918801881188218831884188518861887188818891890189118921893189418951896189718981899190019011902190319041905190619071908190919101911191219131914191519161917191819191920192119221923192419251926192719281929193019311932193319341935193619371938193919401941194219431944194519461947194819491950195119521953195419551956195719581959196019611962196319641965196619671968196919701971197219731974197519761977197819791980198119821983198419851986198719881989199019911992199319941995199619971998199920002001200220032004200520062007200820092010201120122013201420152016201720182019202020212022202320242025202620272028202920302031203220332034import { existsSync, mkdirSync, readFileSync, rmSync, writeFileSync } from "node:fs";import { tmpdir } from "node:os";import { join } from "node:path";import type { AnthropicMessagesCompat, Api, Context, LogEntry, Model, OpenAICompletionsCompat,} from "@earendil-works/pi-ai";import { getApiProvider, getModels, setLogSink } from "@earendil-works/pi-ai";import { getOAuthProvider, registerOAuthProvider } from "@earendil-works/pi-ai/oauth";import { afterEach, beforeEach, describe, expect, test, vi } from "vitest";import { AuthStorage } from "../src/core/auth-storage.js";import { ModelRegistry, type ProviderConfigInput } from "../src/core/model-registry.js";
describe("ModelRegistry", () => { let tempDir: string; let modelsJsonPath: string; let authStorage: AuthStorage;
beforeEach(() => { tempDir = join(tmpdir(), `pi-test-model-registry-${Date.now()}-${Math.random().toString(36).slice(2)}`); mkdirSync(tempDir, { recursive: true }); modelsJsonPath = join(tempDir, "models.json"); authStorage = AuthStorage.create(join(tempDir, "auth.json")); });
afterEach(() => { vi.unstubAllGlobals(); if (tempDir && existsSync(tempDir)) { rmSync(tempDir, { recursive: true }); } });
function providerConfig( baseUrl: string, models: Array<{ id: string; name?: string }>, api: string = "anthropic-messages", ): ProviderConfigInput { return { baseUrl, apiKey: "TEST_KEY", api: api as Api, models: models.map((m) => ({ id: m.id, name: m.name ?? m.id, reasoning: false, input: ["text"], cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 }, contextWindow: 100000, maxTokens: 8000, })), }; }
function writeModelsJson(providers: Record<string, ReturnType<typeof providerConfig>>) { writeFileSync(modelsJsonPath, JSON.stringify({ providers })); }
function getModelsForProvider(registry: ModelRegistry, provider: string) { return registry.getAll().filter((m) => m.provider === provider); }
function toShPath(value: string): string { return value.replace(/\\/g, "/").replace(/"/g, '\\"'); }
function overrideConfig(baseUrl: string, headers?: Record<string, string>) { return { baseUrl, ...(headers && { headers }) }; }
function writeRawModelsJson(providers: Record<string, unknown>) { writeFileSync(modelsJsonPath, JSON.stringify({ providers })); }
const openAiModel: Model<Api> = { id: "test-openai-model", name: "Test OpenAI Model", api: "openai-completions", provider: "openai", baseUrl: "https://api.openai.com/v1", reasoning: false, input: ["text"], cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 }, contextWindow: 128000, maxTokens: 4096, };
const emptyContext: Context = { messages: [], };
describe("baseUrl override (no custom models)", () => { test("overriding baseUrl keeps all built-in models", () => { writeRawModelsJson({ anthropic: overrideConfig("https://my-proxy.example.com/v1"), });
const registry = ModelRegistry.create(authStorage, modelsJsonPath); const anthropicModels = getModelsForProvider(registry, "anthropic");
expect(anthropicModels.length).toBeGreaterThan(1); expect(anthropicModels.some((m) => m.id.includes("claude"))).toBe(true); });
test("overriding baseUrl changes URL on all built-in models", () => { writeRawModelsJson({ anthropic: overrideConfig("https://my-proxy.example.com/v1"), });
const registry = ModelRegistry.create(authStorage, modelsJsonPath); const anthropicModels = getModelsForProvider(registry, "anthropic");
for (const model of anthropicModels) { expect(model.baseUrl).toBe("https://my-proxy.example.com/v1"); } });
test("overriding headers resolves at request time", async () => { writeRawModelsJson({ anthropic: overrideConfig("https://my-proxy.example.com/v1", { "X-Custom-Header": "custom-value", }), });
const registry = ModelRegistry.create(authStorage, modelsJsonPath); const anthropicModels = getModelsForProvider(registry, "anthropic");
for (const model of anthropicModels) { const auth = await registry.getApiKeyAndHeaders(model); expect(auth.ok).toBe(true); if (auth.ok) { expect(auth.headers?.["X-Custom-Header"]).toBe("custom-value"); } } });
test("headers-only override resolves at request time", async () => { writeRawModelsJson({ anthropic: { headers: { "X-Custom-Header": "custom-value", }, }, });
const registry = ModelRegistry.create(authStorage, modelsJsonPath); expect(registry.getError()).toBeUndefined(); const anthropicModels = getModelsForProvider(registry, "anthropic");
for (const model of anthropicModels) { const auth = await registry.getApiKeyAndHeaders(model); expect(auth.ok).toBe(true); if (auth.ok) { expect(auth.headers?.["X-Custom-Header"]).toBe("custom-value"); } } });
test("prime inference requests include selected Prime Agent team header", async () => { const primeAuthStorage = AuthStorage.inMemory({ "prime-inference": { type: "api_key", key: "agent-key", primeTeam: { teamId: "team-1", name: "Research" }, }, }); const registry = ModelRegistry.create(primeAuthStorage, modelsJsonPath); const model = getModelsForProvider(registry, "prime-inference")[0]; expect(model).toBeDefined();
const auth = await registry.getApiKeyAndHeaders(model!);
expect(auth).toEqual({ ok: true, apiKey: "agent-key", headers: { "X-Prime-Team-ID": "team-1" }, }); });
test("baseUrl-only override does not affect other providers", () => { writeRawModelsJson({ anthropic: overrideConfig("https://my-proxy.example.com/v1"), });
const registry = ModelRegistry.create(authStorage, modelsJsonPath); const googleModels = getModelsForProvider(registry, "google");
expect(googleModels.length).toBeGreaterThan(0); expect(googleModels[0].baseUrl).not.toBe("https://my-proxy.example.com/v1"); });
test("can mix baseUrl override and models merge", () => { writeRawModelsJson({ anthropic: overrideConfig("https://anthropic-proxy.example.com/v1"), google: providerConfig( "https://google-proxy.example.com/v1", [{ id: "gemini-custom" }], "google-generative-ai", ), });
const registry = ModelRegistry.create(authStorage, modelsJsonPath);
const anthropicModels = getModelsForProvider(registry, "anthropic"); expect(anthropicModels.length).toBeGreaterThan(1); expect(anthropicModels[0].baseUrl).toBe("https://anthropic-proxy.example.com/v1");
const googleModels = getModelsForProvider(registry, "google"); expect(googleModels.length).toBeGreaterThan(1); expect(googleModels.some((m) => m.id === "gemini-custom")).toBe(true); });
test("refresh() picks up baseUrl override changes", () => { writeRawModelsJson({ anthropic: overrideConfig("https://first-proxy.example.com/v1"), }); const registry = ModelRegistry.create(authStorage, modelsJsonPath);
expect(getModelsForProvider(registry, "anthropic")[0].baseUrl).toBe("https://first-proxy.example.com/v1");
writeRawModelsJson({ anthropic: overrideConfig("https://second-proxy.example.com/v1"), }); registry.refresh();
expect(getModelsForProvider(registry, "anthropic")[0].baseUrl).toBe("https://second-proxy.example.com/v1"); }); });
describe("custom models merge behavior", () => { test("built-in provider custom models inherit api and baseUrl without explicit fields", () => { writeRawModelsJson({ openrouter: { models: [ { id: "fake-provider/fake-model", name: "Fake model", reasoning: true, input: ["text"], }, ], }, });
const registry = ModelRegistry.create(authStorage, modelsJsonPath); expect(registry.getError()).toBeUndefined();
const model = registry.find("openrouter", "fake-provider/fake-model"); expect(model).toBeDefined(); expect(model?.api).toBe("openai-completions"); expect(model?.baseUrl).toBe("https://openrouter.ai/api/v1"); });
test("non-built-in provider custom models still require baseUrl and apiKey", () => { writeRawModelsJson({ "my-custom-provider": { models: [ { id: "my-model", api: "openai-completions", reasoning: false, input: ["text"], }, ], }, });
const registry = ModelRegistry.create(authStorage, modelsJsonPath); expect(registry.getError()).toContain("baseUrl"); });
test("custom provider with same name as built-in merges with built-in models", () => { writeModelsJson({ anthropic: providerConfig("https://my-proxy.example.com/v1", [{ id: "claude-custom" }]), });
const registry = ModelRegistry.create(authStorage, modelsJsonPath); const anthropicModels = getModelsForProvider(registry, "anthropic");
expect(anthropicModels.length).toBeGreaterThan(1); expect(anthropicModels.some((m) => m.id === "claude-custom")).toBe(true); expect(anthropicModels.some((m) => m.id.includes("claude"))).toBe(true); });
test("custom model with same id replaces built-in model by id", () => { writeModelsJson({ openrouter: providerConfig( "https://my-proxy.example.com/v1", [{ id: "anthropic/claude-sonnet-4" }], "openai-completions", ), });
const registry = ModelRegistry.create(authStorage, modelsJsonPath); const models = getModelsForProvider(registry, "openrouter"); const sonnetModels = models.filter((m) => m.id === "anthropic/claude-sonnet-4");
expect(sonnetModels).toHaveLength(1); expect(sonnetModels[0].baseUrl).toBe("https://my-proxy.example.com/v1"); });
test("custom provider with same name as built-in does not affect other built-in providers", () => { writeModelsJson({ anthropic: providerConfig("https://my-proxy.example.com/v1", [{ id: "claude-custom" }]), });
const registry = ModelRegistry.create(authStorage, modelsJsonPath);
expect(getModelsForProvider(registry, "google").length).toBeGreaterThan(0); expect(getModelsForProvider(registry, "openai").length).toBeGreaterThan(0); });
test("provider-level baseUrl applies to both built-in and custom models", () => { writeModelsJson({ anthropic: providerConfig("https://merged-proxy.example.com/v1", [{ id: "claude-custom" }]), });
const registry = ModelRegistry.create(authStorage, modelsJsonPath); const anthropicModels = getModelsForProvider(registry, "anthropic");
for (const model of anthropicModels) { expect(model.baseUrl).toBe("https://merged-proxy.example.com/v1"); } });
test("provider-level compat applies to custom models", () => { writeRawModelsJson({ demo: { baseUrl: "https://example.com/v1", apiKey: "DEMO_KEY", api: "openai-completions", compat: { supportsUsageInStreaming: false, maxTokensField: "max_tokens", }, models: [ { id: "demo-model", reasoning: false, input: ["text"], cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 }, contextWindow: 1000, maxTokens: 100, }, ], }, });
const registry = ModelRegistry.create(authStorage, modelsJsonPath); const compat = registry.find("demo", "demo-model")?.compat as OpenAICompletionsCompat | undefined;
expect(compat?.supportsUsageInStreaming).toBe(false); expect(compat?.maxTokensField).toBe("max_tokens"); });
test("model-level compat overrides provider-level compat for custom models", () => { writeRawModelsJson({ demo: { baseUrl: "https://example.com/v1", apiKey: "DEMO_KEY", api: "openai-completions", compat: { supportsUsageInStreaming: false, maxTokensField: "max_tokens", }, models: [ { id: "demo-model", reasoning: false, input: ["text"], cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 }, contextWindow: 1000, maxTokens: 100, compat: { supportsUsageInStreaming: true, maxTokensField: "max_completion_tokens", }, }, ], }, });
const registry = ModelRegistry.create(authStorage, modelsJsonPath); const compat = registry.find("demo", "demo-model")?.compat as OpenAICompletionsCompat | undefined;
expect(compat?.supportsUsageInStreaming).toBe(true); expect(compat?.maxTokensField).toBe("max_completion_tokens"); });
test("provider-level compat applies to built-in models", () => { writeRawModelsJson({ openrouter: { compat: { supportsUsageInStreaming: false, supportsStrictMode: false, }, }, });
const registry = ModelRegistry.create(authStorage, modelsJsonPath); const models = getModelsForProvider(registry, "openrouter");
expect(models.length).toBeGreaterThan(0); for (const model of models) { const compat = model.compat as OpenAICompletionsCompat | undefined; expect(compat?.supportsUsageInStreaming).toBe(false); expect(compat?.supportsStrictMode).toBe(false); } });
test("model schema accepts thinkingLevelMap and compat schema accepts supportsStrictMode and cacheControlFormat", () => { writeRawModelsJson({ demo: { baseUrl: "https://example.com/v1", apiKey: "DEMO_KEY", api: "openai-completions", models: [ { id: "demo-model", reasoning: true, input: ["text"], cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 }, contextWindow: 1000, maxTokens: 100, thinkingLevelMap: { minimal: null, high: "max", }, compat: { supportsStrictMode: false, cacheControlFormat: "anthropic", }, }, ], }, });
const registry = ModelRegistry.create(authStorage, modelsJsonPath); const model = registry.find("demo", "demo-model"); const compat = model?.compat as OpenAICompletionsCompat | undefined;
expect(registry.getError()).toBeUndefined(); expect(model?.thinkingLevelMap).toEqual({ minimal: null, high: "max" }); expect(compat?.supportsStrictMode).toBe(false); expect(compat?.cacheControlFormat).toBe("anthropic"); });
test("compat schema accepts Anthropic eager tool input streaming flag", () => { writeRawModelsJson({ demo: { baseUrl: "https://example.com", apiKey: "DEMO_KEY", api: "anthropic-messages", compat: { supportsEagerToolInputStreaming: false, }, models: [ { id: "demo-model", reasoning: true, input: ["text"], cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 }, contextWindow: 1000, maxTokens: 100, }, ], }, });
const registry = ModelRegistry.create(authStorage, modelsJsonPath); const compat = registry.find("demo", "demo-model")?.compat as AnthropicMessagesCompat | undefined;
expect(registry.getError()).toBeUndefined(); expect(compat?.supportsEagerToolInputStreaming).toBe(false); });
test("compat schema accepts long cache retention flag", () => { writeRawModelsJson({ demo: { baseUrl: "https://example.com", apiKey: "DEMO_KEY", api: "anthropic-messages", compat: { supportsLongCacheRetention: false, }, models: [ { id: "demo-model", reasoning: true, input: ["text"], cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 }, contextWindow: 1000, maxTokens: 100, }, ], }, });
const registry = ModelRegistry.create(authStorage, modelsJsonPath); const compat = registry.find("demo", "demo-model")?.compat as AnthropicMessagesCompat | undefined;
expect(registry.getError()).toBeUndefined(); expect(compat?.supportsLongCacheRetention).toBe(false); });
test("model-level baseUrl overrides provider-level baseUrl for custom models", () => { writeRawModelsJson({ "opencode-go": { baseUrl: "https://opencode.ai/zen/go/v1", apiKey: "TEST_KEY", models: [ { id: "minimax-m2.5", api: "anthropic-messages", baseUrl: "https://opencode.ai/zen/go", reasoning: true, input: ["text"], cost: { input: 0.3, output: 1.2, cacheRead: 0.03, cacheWrite: 0 }, contextWindow: 204800, maxTokens: 131072, }, { id: "glm-5", api: "openai-completions", reasoning: true, input: ["text"], cost: { input: 1, output: 3.2, cacheRead: 0.2, cacheWrite: 0 }, contextWindow: 204800, maxTokens: 131072, }, ], }, });
const registry = ModelRegistry.create(authStorage, modelsJsonPath); const m25 = registry.find("opencode-go", "minimax-m2.5"); const glm5 = registry.find("opencode-go", "glm-5");
expect(m25?.baseUrl).toBe("https://opencode.ai/zen/go"); expect(glm5?.baseUrl).toBe("https://opencode.ai/zen/go/v1"); });
test("modelOverrides still apply when provider also defines models", () => { writeRawModelsJson({ openrouter: { baseUrl: "https://my-proxy.example.com/v1", apiKey: "OPENROUTER_API_KEY", api: "openai-completions", models: [ { id: "custom/openrouter-model", name: "Custom OpenRouter Model", reasoning: false, input: ["text"], cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 }, contextWindow: 128000, maxTokens: 16384, }, ], modelOverrides: { "anthropic/claude-sonnet-4": { name: "Overridden Built-in Sonnet", }, }, }, });
const registry = ModelRegistry.create(authStorage, modelsJsonPath); const models = getModelsForProvider(registry, "openrouter");
expect(models.some((m) => m.id === "custom/openrouter-model")).toBe(true); expect( models.some((m) => m.id === "anthropic/claude-sonnet-4" && m.name === "Overridden Built-in Sonnet"), ).toBe(true); });
test("refresh() reloads merged custom models from disk", () => { writeModelsJson({ anthropic: providerConfig("https://first-proxy.example.com/v1", [{ id: "claude-custom" }]), }); const registry = ModelRegistry.create(authStorage, modelsJsonPath); expect(getModelsForProvider(registry, "anthropic").some((m) => m.id === "claude-custom")).toBe(true);
writeModelsJson({ anthropic: providerConfig("https://second-proxy.example.com/v1", [{ id: "claude-custom-2" }]), }); registry.refresh();
const anthropicModels = getModelsForProvider(registry, "anthropic"); expect(anthropicModels.some((m) => m.id === "claude-custom")).toBe(false); expect(anthropicModels.some((m) => m.id === "claude-custom-2")).toBe(true); expect(anthropicModels.some((m) => m.id.includes("claude"))).toBe(true); });
test("removing custom models from models.json keeps built-in provider models", () => { writeModelsJson({ anthropic: providerConfig("https://proxy.example.com/v1", [{ id: "claude-custom" }]), }); const registry = ModelRegistry.create(authStorage, modelsJsonPath); expect(getModelsForProvider(registry, "anthropic").some((m) => m.id === "claude-custom")).toBe(true);
writeModelsJson({}); registry.refresh();
const anthropicModels = getModelsForProvider(registry, "anthropic"); expect(anthropicModels.some((m) => m.id === "claude-custom")).toBe(false); expect(anthropicModels.some((m) => m.id.includes("claude"))).toBe(true); }); });
describe("live Prime Inference models", () => { test("loads the cache without replacing external providers and applies local overrides", () => { const bundled = getModels("prime-inference") as Model<"openai-completions">[]; const catalogEntries = bundled.map((model) => ({ id: model.id, display_name: `Live ${model.name}`, pricing: { input_usd_per_mtok: model.cost.input, output_usd_per_mtok: model.cost.output }, specs: { context_window: model.contextWindow, max_output_tokens: model.maxTokens, modalities: { input: model.input, output: ["text"] }, supports_reasoning: model.reasoning, }, })); catalogEntries.push({ id: "test/live-added", display_name: "Live Added", pricing: { input_usd_per_mtok: 1, output_usd_per_mtok: 2 }, specs: { context_window: 200_000, max_output_tokens: 20_000, modalities: { input: ["text"], output: ["text"] }, supports_reasoning: false, }, }); writeFileSync( join(tempDir, "prime-inference-models-cache.json"), JSON.stringify({ object: "list", data: catalogEntries }), ); writeRawModelsJson({ "prime-inference": { baseUrl: "https://local-proxy.example.com/v1", modelOverrides: { "test/live-added": { name: "Local Added", contextWindow: 123_456 } }, }, });
const registry = ModelRegistry.create(authStorage, modelsJsonPath); expect(registry.find("prime-inference", "test/live-added")).toMatchObject({ name: "Local Added", baseUrl: "https://local-proxy.example.com/v1", contextWindow: 123_456, cost: { input: 1, output: 2 }, }); expect(getModelsForProvider(registry, "openrouter")).toHaveLength(getModels("openrouter").length); });
test("restores cached authorized deployment metadata without waiting for the network", async () => { const privateRoute = { id: "vendor/model:deployment", display_name: "Private Deployment", pricing: { input_usd_per_mtok: 1, output_usd_per_mtok: 2 }, specs: { context_window: 200_000, max_output_tokens: 20_000, modalities: { input: ["text"], output: ["text"] }, supports_reasoning: false, }, }; authStorage.set("prime-inference", { type: "api_key", key: "prime-key", primeTeam: { teamId: "research-team", name: "Research" }, }); vi.stubGlobal( "fetch", vi.fn( async (_url: string | URL | Request, init?: RequestInit) => new Response( JSON.stringify({ data: new Headers(init?.headers).has("Authorization") ? [privateRoute] : [] }), ), ), ); const firstRegistry = ModelRegistry.create(authStorage, modelsJsonPath); expect( (await firstRegistry.refreshAvailableModels()).find((model) => model.id === privateRoute.id), ).toMatchObject({ name: "Private Deployment", contextWindow: 200_000 });
vi.stubGlobal( "fetch", vi.fn(async () => { throw new Error("offline"); }), ); const restoredRegistry = ModelRegistry.create(authStorage, modelsJsonPath); expect( (await restoredRegistry.refreshAvailableModels()).find((model) => model.id === privateRoute.id), ).toMatchObject({ name: "Private Deployment", contextWindow: 200_000 }); }); });
describe("modelOverrides (per-model customization)", () => { test("model override applies to a single built-in model", () => { writeRawModelsJson({ openrouter: { modelOverrides: { "anthropic/claude-sonnet-4": { name: "Custom Sonnet Name", }, }, }, });
const registry = ModelRegistry.create(authStorage, modelsJsonPath); const models = getModelsForProvider(registry, "openrouter");
const sonnet = models.find((m) => m.id === "anthropic/claude-sonnet-4"); expect(sonnet?.name).toBe("Custom Sonnet Name");
const opus = models.find((m) => m.id === "anthropic/claude-opus-4"); expect(opus?.name).not.toBe("Custom Sonnet Name"); });
test("model override with compat.openRouterRouting", () => { writeRawModelsJson({ openrouter: { modelOverrides: { "anthropic/claude-sonnet-4": { compat: { openRouterRouting: { only: ["amazon-bedrock"] }, }, }, }, }, });
const registry = ModelRegistry.create(authStorage, modelsJsonPath); const models = getModelsForProvider(registry, "openrouter");
const sonnet = models.find((m) => m.id === "anthropic/claude-sonnet-4"); const compat = sonnet?.compat as OpenAICompletionsCompat | undefined; expect(compat?.openRouterRouting).toEqual({ only: ["amazon-bedrock"] }); });
test("model override deep merges compat settings", () => { writeRawModelsJson({ openrouter: { modelOverrides: { "anthropic/claude-sonnet-4": { compat: { openRouterRouting: { order: ["anthropic", "together"] }, }, }, }, }, });
const registry = ModelRegistry.create(authStorage, modelsJsonPath); const models = getModelsForProvider(registry, "openrouter"); const sonnet = models.find((m) => m.id === "anthropic/claude-sonnet-4");
const compat = sonnet?.compat as OpenAICompletionsCompat | undefined; expect(compat?.openRouterRouting).toEqual({ order: ["anthropic", "together"] }); });
test("multiple model overrides on same provider", () => { writeRawModelsJson({ openrouter: { modelOverrides: { "anthropic/claude-sonnet-4": { compat: { openRouterRouting: { only: ["amazon-bedrock"] } }, }, "anthropic/claude-opus-4": { compat: { openRouterRouting: { only: ["anthropic"] } }, }, }, }, });
const registry = ModelRegistry.create(authStorage, modelsJsonPath); const models = getModelsForProvider(registry, "openrouter");
const sonnet = models.find((m) => m.id === "anthropic/claude-sonnet-4"); const opus = models.find((m) => m.id === "anthropic/claude-opus-4");
const sonnetCompat = sonnet?.compat as OpenAICompletionsCompat | undefined; const opusCompat = opus?.compat as OpenAICompletionsCompat | undefined; expect(sonnetCompat?.openRouterRouting).toEqual({ only: ["amazon-bedrock"] }); expect(opusCompat?.openRouterRouting).toEqual({ only: ["anthropic"] }); });
test("model override combined with baseUrl override", () => { writeRawModelsJson({ openrouter: { baseUrl: "https://my-proxy.example.com/v1", modelOverrides: { "anthropic/claude-sonnet-4": { name: "Proxied Sonnet", }, }, }, });
const registry = ModelRegistry.create(authStorage, modelsJsonPath); const models = getModelsForProvider(registry, "openrouter"); const sonnet = models.find((m) => m.id === "anthropic/claude-sonnet-4");
expect(sonnet?.baseUrl).toBe("https://my-proxy.example.com/v1"); expect(sonnet?.name).toBe("Proxied Sonnet");
const opus = models.find((m) => m.id === "anthropic/claude-opus-4"); expect(opus?.baseUrl).toBe("https://my-proxy.example.com/v1"); expect(opus?.name).not.toBe("Proxied Sonnet"); });
test("model override for non-existent model ID is ignored", () => { writeRawModelsJson({ openrouter: { modelOverrides: { "nonexistent/model-id": { name: "This should not appear", }, }, }, });
const registry = ModelRegistry.create(authStorage, modelsJsonPath); const models = getModelsForProvider(registry, "openrouter");
expect(models.find((m) => m.id === "nonexistent/model-id")).toBeUndefined(); expect(registry.getError()).toBeUndefined(); });
test("model override can change cost fields partially", () => { writeRawModelsJson({ openrouter: { modelOverrides: { "anthropic/claude-sonnet-4": { cost: { input: 99 }, }, }, }, });
const registry = ModelRegistry.create(authStorage, modelsJsonPath); const models = getModelsForProvider(registry, "openrouter"); const sonnet = models.find((m) => m.id === "anthropic/claude-sonnet-4");
expect(sonnet?.cost.input).toBe(99); expect(sonnet?.cost.output).toBeGreaterThan(0); });
test("model override can add headers at request time", async () => { writeRawModelsJson({ openrouter: { modelOverrides: { "anthropic/claude-sonnet-4": { headers: { "X-Custom-Model-Header": "value" }, }, }, }, });
const registry = ModelRegistry.create(authStorage, modelsJsonPath); const models = getModelsForProvider(registry, "openrouter"); const sonnet = models.find((m) => m.id === "anthropic/claude-sonnet-4"); expect(sonnet).toBeDefined();
const auth = await registry.getApiKeyAndHeaders(sonnet!); expect(auth.ok).toBe(true); if (auth.ok) { expect(auth.headers?.["X-Custom-Model-Header"]).toBe("value"); } });
test("refresh() picks up model override changes", () => { writeRawModelsJson({ openrouter: { modelOverrides: { "anthropic/claude-sonnet-4": { name: "First Name", }, }, }, });
const registry = ModelRegistry.create(authStorage, modelsJsonPath); expect( getModelsForProvider(registry, "openrouter").find((m) => m.id === "anthropic/claude-sonnet-4")?.name, ).toBe("First Name");
writeRawModelsJson({ openrouter: { modelOverrides: { "anthropic/claude-sonnet-4": { name: "Second Name", }, }, }, }); registry.refresh();
expect( getModelsForProvider(registry, "openrouter").find((m) => m.id === "anthropic/claude-sonnet-4")?.name, ).toBe("Second Name"); });
test("removing model override restores built-in values", () => { writeRawModelsJson({ openrouter: { modelOverrides: { "anthropic/claude-sonnet-4": { name: "Custom Name", }, }, }, });
const registry = ModelRegistry.create(authStorage, modelsJsonPath); const customName = getModelsForProvider(registry, "openrouter").find( (m) => m.id === "anthropic/claude-sonnet-4", )?.name; expect(customName).toBe("Custom Name");
writeRawModelsJson({}); registry.refresh();
const restoredName = getModelsForProvider(registry, "openrouter").find( (m) => m.id === "anthropic/claude-sonnet-4", )?.name; expect(restoredName).not.toBe("Custom Name"); }); });
describe("dynamic provider lifecycle", () => { test("getProviderDisplayName resolves registered, OAuth, built-in, and fallback names", () => { const registry = ModelRegistry.create(authStorage, modelsJsonPath);
expect(registry.getProviderDisplayName("openai")).toBe("OpenAI"); expect(registry.getProviderDisplayName("github-copilot")).toBe("GitHub Copilot"); expect(registry.getProviderDisplayName("unknown-provider")).toBe("unknown-provider");
registry.registerProvider("named-provider", { name: "Named Provider", baseUrl: "https://provider.test/v1", apiKey: "TEST_KEY", api: "openai-completions", models: [ { id: "demo-model", name: "Demo Model", reasoning: false, input: ["text"], cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 }, contextWindow: 128000, maxTokens: 4096, }, ], }); expect(registry.getProviderDisplayName("named-provider")).toBe("Named Provider");
registry.registerProvider("oauth-provider", { baseUrl: "https://provider.test/v1", api: "openai-completions", oauth: { name: "OAuth Provider", login: async () => ({ access: "access", refresh: "refresh", expires: Date.now() + 60_000 }), refreshToken: async (credentials) => credentials, getApiKey: (credentials) => credentials.access, }, models: [ { id: "demo-model", name: "Demo Model", reasoning: false, input: ["text"], cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 }, contextWindow: 128000, maxTokens: 4096, }, ], }); expect(registry.getProviderDisplayName("oauth-provider")).toBe("OAuth Provider"); });
test("failed registerProvider does not persist invalid streamSimple config", () => { const registry = ModelRegistry.create(authStorage, modelsJsonPath);
expect(() => registry.registerProvider("broken-provider", { streamSimple: (() => { throw new Error("should not run"); }) as any, }), ).toThrow('Provider broken-provider: "api" is required when registering streamSimple.');
expect(() => registry.refresh()).not.toThrow(); });
test("failed registerProvider does not remove existing provider models", () => { const registry = ModelRegistry.create(authStorage, modelsJsonPath);
registry.registerProvider("demo-provider", { baseUrl: "https://provider.test/v1", apiKey: "TEST_KEY", api: "openai-completions", models: [ { id: "demo-model", name: "Demo Model", reasoning: false, input: ["text"], cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 }, contextWindow: 128000, maxTokens: 4096, }, ], });
expect(registry.find("demo-provider", "demo-model")).toBeDefined();
expect(() => registry.registerProvider("demo-provider", { baseUrl: "https://provider.test/v2", apiKey: "TEST_KEY", models: [ { id: "broken-model", name: "Broken Model", reasoning: false, input: ["text"], cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 }, contextWindow: 128000, maxTokens: 4096, }, ], }), ).toThrow('Provider demo-provider, model broken-model: no "api" specified.');
expect(registry.find("demo-provider", "demo-model")).toBeDefined(); expect(() => registry.refresh()).not.toThrow(); expect(registry.find("demo-provider", "demo-model")).toBeDefined(); });
test("unregisterProvider removes custom OAuth provider and restores built-in OAuth provider", () => { const registry = ModelRegistry.create(authStorage, modelsJsonPath);
registry.registerProvider("anthropic", { oauth: { name: "Custom Anthropic OAuth", login: async () => ({ access: "custom-access-token", refresh: "custom-refresh-token", expires: Date.now() + 60_000, }), refreshToken: async (credentials) => credentials, getApiKey: (credentials) => credentials.access, }, });
expect(getOAuthProvider("anthropic")?.name).toBe("Custom Anthropic OAuth");
registry.unregisterProvider("anthropic");
expect(getOAuthProvider("anthropic")?.name).not.toBe("Custom Anthropic OAuth"); });
test("unregisterProvider removes custom streamSimple override and restores built-in API stream handler", () => { const registry = ModelRegistry.create(authStorage, modelsJsonPath);
registry.registerProvider("stream-override-provider", { api: "openai-completions", streamSimple: () => { throw new Error("custom streamSimple override"); }, });
let threwCustomOverride = false; try { getApiProvider("openai-completions")?.streamSimple(openAiModel, emptyContext); } catch (error) { threwCustomOverride = error instanceof Error && error.message === "custom streamSimple override"; } expect(threwCustomOverride).toBe(true);
registry.unregisterProvider("stream-override-provider");
let threwCustomOverrideAfterUnregister = false; try { getApiProvider("openai-completions")?.streamSimple(openAiModel, emptyContext); } catch (error) { threwCustomOverrideAfterUnregister = error instanceof Error && error.message === "custom streamSimple override"; } expect(threwCustomOverrideAfterUnregister).toBe(false); });
describe("dynamic provider override persistence", () => { test("baseUrl-only override keeps built-in provider models after refresh", () => { const registry = ModelRegistry.create(authStorage, modelsJsonPath);
registry.registerProvider("anthropic", { baseUrl: "https://proxy.test/anthropic" }); registry.refresh();
const anthropicModels = getModelsForProvider(registry, "anthropic"); expect(anthropicModels.length).toBeGreaterThan(1); expect(anthropicModels.every((m) => m.baseUrl === "https://proxy.test/anthropic")).toBe(true); });
test("models-only override replaces built-in provider models after refresh", () => { const registry = ModelRegistry.create(authStorage, modelsJsonPath);
registry.registerProvider("anthropic", { ...providerConfig("https://custom.test/anthropic", [{ id: "custom-claude" }], "anthropic-messages"), baseUrl: "https://custom.test/anthropic", }); registry.refresh();
expect(getModelsForProvider(registry, "anthropic").map((m) => m.id)).toEqual(["custom-claude"]); expect(registry.find("anthropic", "custom-claude")?.baseUrl).toBe("https://custom.test/anthropic"); });
test("models plus baseUrl override replaces built-in provider models after refresh", () => { const registry = ModelRegistry.create(authStorage, modelsJsonPath);
registry.registerProvider("anthropic", { ...providerConfig("https://custom.test/anthropic", [{ id: "custom-claude" }], "anthropic-messages"), baseUrl: "https://custom.test/anthropic", }); registry.registerProvider("anthropic", { baseUrl: "https://proxy.test/anthropic" }); registry.refresh();
expect(getModelsForProvider(registry, "anthropic").map((m) => m.id)).toEqual(["custom-claude"]); expect(registry.find("anthropic", "custom-claude")?.baseUrl).toBe("https://proxy.test/anthropic"); });
test("models-only custom provider registration survives refresh", () => { const registry = ModelRegistry.create(authStorage, modelsJsonPath);
registry.registerProvider( "custom-provider", providerConfig("https://custom.test/v1", [{ id: "custom-a" }, { id: "custom-b" }], "openai-completions"), ); registry.refresh();
expect(getModelsForProvider(registry, "custom-provider").map((m) => m.id)).toEqual([ "custom-a", "custom-b", ]); });
test("baseUrl-only override keeps custom provider models after refresh", () => { const registry = ModelRegistry.create(authStorage, modelsJsonPath);
registry.registerProvider( "custom-provider", providerConfig("https://custom.test/v1", [{ id: "custom-a" }, { id: "custom-b" }], "openai-completions"), ); registry.registerProvider("custom-provider", { baseUrl: "https://proxy.test/custom" }); registry.refresh();
expect(getModelsForProvider(registry, "custom-provider").map((m) => m.id)).toEqual([ "custom-a", "custom-b", ]); expect( getModelsForProvider(registry, "custom-provider").every( (m) => m.baseUrl === "https://proxy.test/custom", ), ).toBe(true); });
test("headers-only override keeps custom provider models after refresh", async () => { const registry = ModelRegistry.create(authStorage, modelsJsonPath);
registry.registerProvider( "custom-provider", providerConfig("https://custom.test/v1", [{ id: "custom-a" }, { id: "custom-b" }], "openai-completions"), ); registry.registerProvider("custom-provider", { headers: { "x-proxy": "enabled" } }); registry.refresh();
const models = getModelsForProvider(registry, "custom-provider"); expect(models.map((m) => m.id)).toEqual(["custom-a", "custom-b"]); expect(models.every((m) => m.baseUrl === "https://custom.test/v1")).toBe(true); expect(await registry.getApiKeyAndHeaders(models[0])).toMatchObject({ ok: true, headers: { "x-proxy": "enabled" }, }); }); }); });
describe("auth refresh across processes", () => { test("model catalog includes unauthenticated public models and hides private Prime routes", async () => { const savedPrimeApiKey = process.env.PRIME_API_KEY; const savedOpenAiApiKey = process.env.OPENAI_API_KEY; delete process.env.PRIME_API_KEY; delete process.env.OPENAI_API_KEY; try { const registry = ModelRegistry.create(authStorage, modelsJsonPath);
const unauthenticated = await registry.refreshModelCatalog(); expect(unauthenticated.models.some((model) => model.provider === "openai")).toBe(true); expect(unauthenticated.configuredProviders).not.toContain("openai"); expect( unauthenticated.models.some( (model) => model.provider === "prime-inference" && model.id.startsWith("internal/"), ), ).toBe(false);
authStorage.setRuntimeApiKey("openai", "test-key"); const authenticated = await registry.refreshModelCatalog(); expect(authenticated.configuredProviders).toContain("openai"); } finally { if (savedPrimeApiKey !== undefined) { process.env.PRIME_API_KEY = savedPrimeApiKey; } if (savedOpenAiApiKey !== undefined) { process.env.OPENAI_API_KEY = savedOpenAiApiKey; } } });
test("refresh() picks up credentials written by another process", () => { const savedEnvKey = process.env.PRIME_API_KEY; delete process.env.PRIME_API_KEY; try { const registry = ModelRegistry.create(authStorage, modelsJsonPath); expect(registry.getAvailable().some((m) => m.provider === "prime-inference")).toBe(false);
const otherProcessAuth = AuthStorage.create(join(tempDir, "auth.json")); otherProcessAuth.set("prime-inference", { type: "api_key", key: "test-key" });
registry.refresh(); expect(registry.getAvailable().some((m) => m.provider === "prime-inference")).toBe(true); } finally { if (savedEnvKey !== undefined) { process.env.PRIME_API_KEY = savedEnvKey; } } }); });
describe("API key resolution", () => { function providerWithApiKey(apiKey: string) { return { baseUrl: "https://example.com/v1", apiKey, api: "anthropic-messages", models: [ { id: "test-model", name: "Test Model", reasoning: false, input: ["text"], cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 }, contextWindow: 100000, maxTokens: 8000, }, ], }; }
test("apiKey with ! prefix executes command and uses stdout", async () => { writeRawModelsJson({ "custom-provider": providerWithApiKey("!echo test-api-key-from-command"), });
const registry = ModelRegistry.create(authStorage, modelsJsonPath); const apiKey = await registry.getApiKeyForProvider("custom-provider");
expect(apiKey).toBe("test-api-key-from-command"); });
test("apiKey with ! prefix trims whitespace from command output", async () => { writeRawModelsJson({ "custom-provider": providerWithApiKey("!echo ' spaced-key '"), });
const registry = ModelRegistry.create(authStorage, modelsJsonPath); const apiKey = await registry.getApiKeyForProvider("custom-provider");
expect(apiKey).toBe("spaced-key"); });
test("apiKey with ! prefix handles multiline output (uses trimmed result)", async () => { writeRawModelsJson({ "custom-provider": providerWithApiKey("!printf 'line1\\nline2'"), });
const registry = ModelRegistry.create(authStorage, modelsJsonPath); const apiKey = await registry.getApiKeyForProvider("custom-provider");
expect(apiKey).toBe("line1\nline2"); });
test("apiKey with ! prefix returns undefined on command failure", async () => { writeRawModelsJson({ "custom-provider": providerWithApiKey("!exit 1"), });
const registry = ModelRegistry.create(authStorage, modelsJsonPath); const apiKey = await registry.getApiKeyForProvider("custom-provider");
expect(apiKey).toBeUndefined(); });
test("apiKey with ! prefix returns undefined on nonexistent command", async () => { writeRawModelsJson({ "custom-provider": providerWithApiKey("!nonexistent-command-12345"), });
const registry = ModelRegistry.create(authStorage, modelsJsonPath); const apiKey = await registry.getApiKeyForProvider("custom-provider");
expect(apiKey).toBeUndefined(); });
test("apiKey with ! prefix returns undefined on empty output", async () => { writeRawModelsJson({ "custom-provider": providerWithApiKey("!printf ''"), });
const registry = ModelRegistry.create(authStorage, modelsJsonPath); const apiKey = await registry.getApiKeyForProvider("custom-provider");
expect(apiKey).toBeUndefined(); });
test("apiKey as environment variable name resolves to env value", async () => { const originalEnv = process.env.TEST_API_KEY_12345; process.env.TEST_API_KEY_12345 = "env-api-key-value";
try { writeRawModelsJson({ "custom-provider": providerWithApiKey("TEST_API_KEY_12345"), });
const registry = ModelRegistry.create(authStorage, modelsJsonPath); const apiKey = await registry.getApiKeyForProvider("custom-provider");
expect(apiKey).toBe("env-api-key-value"); } finally { if (originalEnv === undefined) { delete process.env.TEST_API_KEY_12345; } else { process.env.TEST_API_KEY_12345 = originalEnv; } } });
test("apiKey as literal value is used directly when not an env var", async () => { delete process.env.literal_api_key_value;
writeRawModelsJson({ "custom-provider": providerWithApiKey("literal_api_key_value"), });
const registry = ModelRegistry.create(authStorage, modelsJsonPath); const apiKey = await registry.getApiKeyForProvider("custom-provider");
expect(apiKey).toBe("literal_api_key_value"); });
test("apiKey command can use shell features like pipes", async () => { writeRawModelsJson({ "custom-provider": providerWithApiKey("!echo 'hello world' | tr ' ' '-'"), });
const registry = ModelRegistry.create(authStorage, modelsJsonPath); const apiKey = await registry.getApiKeyForProvider("custom-provider");
expect(apiKey).toBe("hello-world"); });
describe("request-time resolution", () => { test("command is executed on every provider lookup", async () => { const counterFile = join(tempDir, "counter"); writeFileSync(counterFile, "0");
const counterPath = toShPath(counterFile); const command = `!sh -c 'count=$(cat "${counterPath}"); echo $((count + 1)) > "${counterPath}"; echo "key-value"'`; writeRawModelsJson({ "custom-provider": providerWithApiKey(command), });
const registry = ModelRegistry.create(authStorage, modelsJsonPath); await registry.getApiKeyForProvider("custom-provider"); await registry.getApiKeyForProvider("custom-provider"); await registry.getApiKeyForProvider("custom-provider");
const count = parseInt(readFileSync(counterFile, "utf-8").trim(), 10); expect(count).toBe(3); });
test("commands are re-executed across registry instances", async () => { const counterFile = join(tempDir, "counter"); writeFileSync(counterFile, "0");
const counterPath = toShPath(counterFile); const command = `!sh -c 'count=$(cat "${counterPath}"); echo $((count + 1)) > "${counterPath}"; echo "key-value"'`; writeRawModelsJson({ "custom-provider": providerWithApiKey(command), });
const registry1 = ModelRegistry.create(authStorage, modelsJsonPath); await registry1.getApiKeyForProvider("custom-provider");
const registry2 = ModelRegistry.create(authStorage, modelsJsonPath); await registry2.getApiKeyForProvider("custom-provider");
const count = parseInt(readFileSync(counterFile, "utf-8").trim(), 10); expect(count).toBe(2); });
test("different commands resolve independently", async () => { writeRawModelsJson({ "provider-a": providerWithApiKey("!echo key-a"), "provider-b": providerWithApiKey("!echo key-b"), });
const registry = ModelRegistry.create(authStorage, modelsJsonPath);
const keyA = await registry.getApiKeyForProvider("provider-a"); const keyB = await registry.getApiKeyForProvider("provider-b");
expect(keyA).toBe("key-a"); expect(keyB).toBe("key-b"); });
test("failed commands are retried", async () => { const counterFile = join(tempDir, "counter"); writeFileSync(counterFile, "0");
const counterPath = toShPath(counterFile); const command = `!sh -c 'count=$(cat "${counterPath}"); echo $((count + 1)) > "${counterPath}"; exit 1'`; writeRawModelsJson({ "custom-provider": providerWithApiKey(command), });
const registry = ModelRegistry.create(authStorage, modelsJsonPath); const key1 = await registry.getApiKeyForProvider("custom-provider"); const key2 = await registry.getApiKeyForProvider("custom-provider");
expect(key1).toBeUndefined(); expect(key2).toBeUndefined();
const count = parseInt(readFileSync(counterFile, "utf-8").trim(), 10); expect(count).toBe(2); });
test("provider auth status reports apiKey environment variables from models.json", () => { const envVarName = "TEST_API_KEY_STATUS_TEST_98765"; const originalEnv = process.env[envVarName];
try { process.env[envVarName] = "status-test-key";
writeRawModelsJson({ "custom-provider": providerWithApiKey(envVarName), });
const registry = ModelRegistry.create(authStorage, modelsJsonPath);
expect(registry.getProviderAuthStatus("custom-provider")).toEqual({ configured: true, source: "environment", label: envVarName, }); } finally { if (originalEnv === undefined) { delete process.env[envVarName]; } else { process.env[envVarName] = originalEnv; } } });
test("provider auth status reports non-env apiKey values from models.json as a config key", () => { writeRawModelsJson({ "custom-provider": providerWithApiKey("literal_api_key_value"), });
const registry = ModelRegistry.create(authStorage, modelsJsonPath);
expect(registry.getProviderAuthStatus("custom-provider")).toEqual({ configured: true, source: "models_json_key", }); });
test("provider auth status reports models.json auth when stored auth is stale", async () => { authStorage.setRuntimeApiKey("custom-provider", "stale-runtime-key"); expect(authStorage.markAuthStale("custom-provider")).toBe(true); writeRawModelsJson({ "custom-provider": providerWithApiKey("literal_api_key_value"), });
const registry = ModelRegistry.create(authStorage, modelsJsonPath); const model = registry.find("custom-provider", "test-model"); expect(model).toBeDefined();
expect(registry.getProviderAuthStatus("custom-provider")).toEqual({ configured: true, source: "models_json_key", }); await expect(registry.getApiKeyForProvider("custom-provider")).resolves.toBe("literal_api_key_value"); await expect(registry.getApiKeyAndHeaders(model!)).resolves.toMatchObject({ ok: true, apiKey: "literal_api_key_value", }); });
test("stale marking uses the auth source resolved for the last request", async () => { const providerId = `test-oauth-fallback-${Date.now()}-${Math.random().toString(36).slice(2)}`; registerOAuthProvider({ id: providerId, name: "Test OAuth Fallback", async login() { throw new Error("Not used in this test"); }, async refreshToken() { throw new Error("refresh failed"); }, getApiKey(credentials) { return `Bearer ${credentials.access}`; }, }); authStorage.set(providerId, { type: "oauth", refresh: "refresh-token", access: "expired-access-token", expires: Date.now() - 10_000, }); writeRawModelsJson({ [providerId]: providerWithApiKey("literal_api_key_value"), });
const registry = ModelRegistry.create(authStorage, modelsJsonPath);
await expect(registry.getApiKeyForProvider(providerId)).resolves.toBe("literal_api_key_value"); const token = registry.getCurrentProviderAuthSourceToken(providerId); expect(token?.source).toBe("models_json_key"); expect(token).toBeDefined(); expect(registry.markProviderAuthSourceStale(token!)).toBe(true);
await expect(registry.getApiKeyForProvider(providerId)).resolves.toBeUndefined(); expect(authStorage.getAuthStatus(providerId)).toEqual({ configured: true, source: "stored", }); });
test("changed literal models.json apiKey no longer matches stale provider marker", async () => { writeRawModelsJson({ "custom-provider": providerWithApiKey("stale-key"), });
const registry = ModelRegistry.create(authStorage, modelsJsonPath);
await expect(registry.getApiKeyForProvider("custom-provider")).resolves.toBe("stale-key"); expect(registry.markProviderAuthStale("custom-provider")).toBe(true); expect(registry.getProviderAuthStatus("custom-provider")).toEqual({ configured: false, source: "stale", label: "expired", }); expect(registry.getAvailable().some((model) => model.provider === "custom-provider")).toBe(false);
writeRawModelsJson({ "custom-provider": providerWithApiKey("fresh-key"), }); registry.refresh();
expect(registry.getProviderAuthStatus("custom-provider")).toEqual({ configured: true, source: "models_json_key", }); expect(registry.getAvailable().some((model) => model.provider === "custom-provider")).toBe(true); await expect(registry.getApiKeyForProvider("custom-provider")).resolves.toBe("fresh-key"); });
test("clearProviderAuthStale restores availability for explicit model selection", async () => { writeRawModelsJson({ "custom-provider": providerWithApiKey("literal_api_key_value"), });
const registry = ModelRegistry.create(authStorage, modelsJsonPath);
await expect(registry.getApiKeyForProvider("custom-provider")).resolves.toBe("literal_api_key_value"); expect(registry.markProviderAuthStale("custom-provider")).toBe(true); expect(registry.getAvailable().some((model) => model.provider === "custom-provider")).toBe(false);
registry.clearProviderAuthStale("custom-provider");
expect(registry.getAvailable().some((model) => model.provider === "custom-provider")).toBe(true); await expect(registry.getApiKeyForProvider("custom-provider")).resolves.toBe("literal_api_key_value"); });
test.each([ "team change", "logout", "missing team", "missing credentials", "active credentials", "rotated credentials", ] as const)("recovers stale Prime CLI private-model access but invalidates it after %s", async (change) => { vi.stubEnv("PRIME_API_KEY", ""); vi.stubEnv("PRIME_TEAM_ID", ""); vi.stubEnv("PI_OFFLINE", "0"); const configPath = join(tempDir, "prime-config.json"); writeFileSync(configPath, JSON.stringify({ api_key: "prime-test-key", team_id: "team-a" })); const cliAuth = AuthStorage.inMemory({}, { primeCliConfigPath: configPath }); const registry = ModelRegistry.create(cliAuth, modelsJsonPath); const modelId = "internal/live-private-model"; const fetchSpy = vi.spyOn(globalThis, "fetch").mockImplementation( async () => new Response( JSON.stringify({ data: [ { id: modelId, pricing: { input_usd_per_mtok: 1, output_usd_per_mtok: 2 }, specs: { context_window: 200_000, max_output_tokens: 20_000, supports_reasoning: true, modalities: { input: ["text"], output: ["text"] }, }, }, ], }), ), ); try { registry.registerProvider("unrelated-extension", { baseUrl: "https://unused.invalid" }); const model = (await registry.refreshAvailableModels()).find((candidate) => candidate.id === modelId)!; expect(model).toBeDefined(); expect( fetchSpy.mock.calls.filter(([, init]) => new Headers(init?.headers).has("Authorization")), ).toHaveLength(1); expect(registry.markProviderAuthStale("prime-inference")).toBe(true);
registry.unregisterProvider("unrelated-extension"); await expect(registry.canUseModel(model, { assumeAuthConfigured: true })).resolves.toBe(true); await Promise.all([registry.refreshAvailableModels(), registry.refreshAvailableModels()]); expect(registry.find("prime-inference", modelId)).toEqual(model);
expect(cliAuth.getProviderHeaders("prime-inference")).toEqual({ "X-Prime-Team-ID": "team-a" }); expect(registry.hasConfiguredAuth(model)).toBe(false); await expect(cliAuth.getApiKey("prime-inference")).resolves.toBeUndefined(); await expect(registry.canUseModel(model, { assumeAuthConfigured: true })).resolves.toBe(true); registry.clearProviderAuthStale("prime-inference"); await expect(registry.canUseModel(model)).resolves.toBe(true); await expect(cliAuth.getApiKey("prime-inference")).resolves.toBe("prime-test-key");
expect(registry.markProviderAuthStale("prime-inference")).toBe(true); switch (change) { case "team change": cliAuth.setPrimeInferenceTeamSelection({ teamId: "team-b", name: "Other team" }); break; case "logout": cliAuth.logout("prime-inference"); break; case "missing team": cliAuth.setPrimeInferenceTeamSelection(null); break; case "missing credentials": writeFileSync(configPath, JSON.stringify({ team_id: "team-a" })); break; case "active credentials": registry.clearProviderAuthStale("prime-inference"); break; case "rotated credentials": writeFileSync(configPath, JSON.stringify({ api_key: "rotated-key", team_id: "team-a" })); break; } registry.refresh(); expect(registry.find("prime-inference", modelId)).toBeUndefined(); await expect(registry.canUseModel(model, { assumeAuthConfigured: true })).resolves.toBe(false); expect( fetchSpy.mock.calls.filter(([, init]) => new Headers(init?.headers).has("Authorization")), ).toHaveLength(1); } finally { fetchSpy.mockRestore(); vi.unstubAllEnvs(); } });
test("preserves stale-auth entitlements only for the same prime-inference team", async () => { authStorage.setRuntimeApiKey("prime-inference", "prime-key"); const headerSpy = vi .spyOn(authStorage, "getProviderHeaders") .mockReturnValue({ "X-Prime-Team-ID": "team-a" }); const registry = ModelRegistry.create(authStorage, modelsJsonPath); const internals = registry as unknown as { authorizedPrivatePrimeInferenceModelIds: Set<string>; authorizedPrivatePrimeInferenceTeamId: string | undefined; }; internals.authorizedPrivatePrimeInferenceModelIds.add("internal/private-model"); internals.authorizedPrivatePrimeInferenceTeamId = "team-a"; expect(registry.markProviderAuthStale("prime-inference")).toBe(true);
await registry.refreshAvailableModels(); expect(internals.authorizedPrivatePrimeInferenceModelIds.has("internal/private-model")).toBe(true);
// A team switch invalidates entitlements fetched for the old team. headerSpy.mockReturnValue({ "X-Prime-Team-ID": "team-b" }); await registry.refreshAvailableModels(); expect(internals.authorizedPrivatePrimeInferenceModelIds.size).toBe(0); });
test("concurrent stale-auth refreshes do not drop preserved entitlements", async () => { authStorage.setRuntimeApiKey("prime-inference", "prime-key"); const registry = ModelRegistry.create(authStorage, modelsJsonPath); vi.spyOn(authStorage, "getProviderHeaders").mockReturnValue({ "X-Prime-Team-ID": "team-a" }); const internals = registry as unknown as { authorizedPrivatePrimeInferenceModelIds: Set<string>; authorizedPrivatePrimeInferenceTeamId: string | undefined; }; internals.authorizedPrivatePrimeInferenceTeamId = "team-a"; internals.authorizedPrivatePrimeInferenceModelIds.add("internal/private-model"); expect(registry.markProviderAuthStale("prime-inference")).toBe(true);
await Promise.all([registry.refreshAvailableModels(), registry.refreshAvailableModels()]);
expect(internals.authorizedPrivatePrimeInferenceModelIds.has("internal/private-model")).toBe(true); });
test("provider auth status reports command apiKey values from models.json without executing them", () => { const counterFile = join(tempDir, "status-counter"); writeFileSync(counterFile, "0"); const counterPath = toShPath(counterFile); const command = `!sh -c 'echo 1 > "${counterPath}"; echo key-value'`; writeRawModelsJson({ "custom-provider": providerWithApiKey(command), });
const registry = ModelRegistry.create(authStorage, modelsJsonPath);
expect(registry.getProviderAuthStatus("custom-provider")).toEqual({ configured: true, source: "models_json_command", }); expect(readFileSync(counterFile, "utf-8")).toBe("0"); });
test("provider auth status reports stale command auth without executing it", () => { const counterFile = join(tempDir, "stale-status-counter"); writeFileSync(counterFile, "0"); const counterPath = toShPath(counterFile); const command = `!sh -c 'count=$(cat "${counterPath}"); echo $((count + 1)) > "${counterPath}"; echo key-value'`; writeRawModelsJson({ "custom-provider": providerWithApiKey(command), });
const registry = ModelRegistry.create(authStorage, modelsJsonPath);
expect(registry.markProviderAuthStale("custom-provider")).toBe(true); expect(readFileSync(counterFile, "utf-8").trim()).toBe("1"); writeFileSync(counterFile, "0");
expect(registry.getProviderAuthStatus("custom-provider")).toEqual({ configured: false, source: "stale", label: "expired", }); expect(readFileSync(counterFile, "utf-8").trim()).toBe("0"); });
test("environment variables are not cached (changes are picked up)", async () => { const envVarName = "TEST_API_KEY_CACHE_TEST_98765"; const originalEnv = process.env[envVarName];
try { process.env[envVarName] = "first-value";
writeRawModelsJson({ "custom-provider": providerWithApiKey(envVarName), });
const registry = ModelRegistry.create(authStorage, modelsJsonPath);
const key1 = await registry.getApiKeyForProvider("custom-provider"); expect(key1).toBe("first-value");
process.env[envVarName] = "second-value";
const key2 = await registry.getApiKeyForProvider("custom-provider"); expect(key2).toBe("second-value"); } finally { if (originalEnv === undefined) { delete process.env[envVarName]; } else { process.env[envVarName] = originalEnv; } } });
test("getAvailable does not execute command-backed apiKey resolution", async () => { const counterFile = join(tempDir, "counter"); writeFileSync(counterFile, "0");
const counterPath = toShPath(counterFile); const command = `!sh -c 'count=$(cat "${counterPath}"); echo $((count + 1)) > "${counterPath}"; echo "key-value"'`; writeRawModelsJson({ "custom-provider": providerWithApiKey(command), });
const registry = ModelRegistry.create(authStorage, modelsJsonPath); const available = registry.getAvailable();
expect(available.some((m) => m.provider === "custom-provider")).toBe(true); const count = parseInt(readFileSync(counterFile, "utf-8").trim(), 10); expect(count).toBe(0); });
test("changed command-backed apiKey no longer matches stale models.json marker", async () => { const tokenFile = join(tempDir, "models-json-token"); writeFileSync(tokenFile, "stale-key"); const tokenPath = toShPath(tokenFile);
writeRawModelsJson({ "custom-provider": providerWithApiKey(`!sh -c 'cat "${tokenPath}"'`), });
const registry = ModelRegistry.create(authStorage, modelsJsonPath);
await expect(registry.getApiKeyForProvider("custom-provider")).resolves.toBe("stale-key"); expect(registry.markProviderAuthStale("custom-provider")).toBe(true); await expect(registry.getApiKeyForProvider("custom-provider")).resolves.toBeUndefined(); expect(registry.getProviderAuthStatus("custom-provider")).toEqual({ configured: false, source: "stale", label: "expired", });
writeFileSync(tokenFile, "fresh-key");
await expect(registry.getApiKeyForProvider("custom-provider")).resolves.toBe("fresh-key"); expect(registry.getProviderAuthStatus("custom-provider")).toEqual({ configured: true, source: "models_json_command", }); });
test("getApiKeyAndHeaders resolves authHeader on every request", async () => { const tokenFile = join(tempDir, "token"); writeFileSync(tokenFile, "token-1"); const tokenPath = toShPath(tokenFile);
writeRawModelsJson({ "custom-provider": { ...providerWithApiKey(`!sh -c 'cat "${tokenPath}"'`), authHeader: true, }, });
const registry = ModelRegistry.create(authStorage, modelsJsonPath); const model = registry.find("custom-provider", "test-model"); expect(model).toBeDefined();
const auth1 = await registry.getApiKeyAndHeaders(model!); expect(auth1).toEqual({ ok: true, apiKey: "token-1", headers: { Authorization: "Bearer token-1" }, });
writeFileSync(tokenFile, "token-2");
const auth2 = await registry.getApiKeyAndHeaders(model!); expect(auth2).toEqual({ ok: true, apiKey: "token-2", headers: { Authorization: "Bearer token-2" }, }); });
test("getApiKeyAndHeaders returns an error for failed authHeader resolution", async () => { writeRawModelsJson({ "custom-provider": { ...providerWithApiKey("!exit 1"), authHeader: true, }, });
const registry = ModelRegistry.create(authStorage, modelsJsonPath); const model = registry.find("custom-provider", "test-model"); expect(model).toBeDefined();
const auth = await registry.getApiKeyAndHeaders(model!); expect(auth.ok).toBe(false); if (!auth.ok) { expect(auth.error).toContain('Failed to resolve API key for provider "custom-provider"'); } }); });
describe("models.dev live reloading and caching", () => { test("loads live models from models-dev-cache.json synchronously on create", () => { const cachePath = join(tempDir, "models-dev-cache.json"); const cachedModel: Model<Api> = { id: "test-live-model", name: "Test Live Model", api: "openai-responses", provider: "openai", baseUrl: "https://api.openai.com/v1", reasoning: false, input: ["text"], cost: { input: 1, output: 2, cacheRead: 0.1, cacheWrite: 0 }, contextWindow: 128000, maxTokens: 4096, }; writeFileSync( cachePath, JSON.stringify({ updatedAt: Date.now(), models: [cachedModel], }), );
const registry = ModelRegistry.create(authStorage, modelsJsonPath, cachePath); try { const found = registry.find("openai", "test-live-model"); expect(found).toBeDefined(); expect(found?.name).toBe("Test Live Model"); } finally { registry.dispose(); } });
test("respects 1-hour TTL when not forced", async () => { const cachePath = join(tempDir, "models-dev-cache.json"); const cachedModel: Model<Api> = { id: "test-ttl-model", name: "Test TTL Model", api: "openai-responses", provider: "openai", baseUrl: "https://api.openai.com/v1", reasoning: false, input: ["text"], cost: { input: 1, output: 2, cacheRead: 0.1, cacheWrite: 0 }, contextWindow: 128000, maxTokens: 4096, }; const recentTime = Date.now() - 1000 * 60 * 10; // 10 minutes ago (< 1h) writeFileSync( cachePath, JSON.stringify({ updatedAt: recentTime, models: [cachedModel], }), );
const registry = ModelRegistry.create(authStorage, modelsJsonPath, cachePath); try { const result = await registry.refreshModelsDev({ force: false }); expect(result.refreshed).toBe(false); const found = registry.find("openai", "test-ttl-model"); expect(found).toBeDefined(); } finally { registry.dispose(); } });
test("handles corrupted or invalid cache file gracefully without throwing", async () => { const cachePath = join(tempDir, "models-dev-cache.json"); writeFileSync(cachePath, "invalid-not-json{{{");
const registry = ModelRegistry.create(authStorage, modelsJsonPath, cachePath); try { expect(registry.getAll().length).toBeGreaterThan(0); } finally { registry.dispose(); } });
test("logs failure to structured logger instead of console.warn when live fetch fails", async () => { const cachePath = join(tempDir, "models-dev-cache.json"); const registry = ModelRegistry.create(authStorage, modelsJsonPath, cachePath); const consoleWarnSpy = vi.spyOn(console, "warn").mockImplementation(() => {}); const fetchSpy = vi.spyOn(globalThis, "fetch").mockRejectedValue(new Error("connection refused")); const logged: LogEntry[] = []; setLogSink((entry) => logged.push(entry));
try { const result = await registry.refreshModelsDev({ force: true }); expect(result.refreshed).toBe(false); expect(result.error).toBe("connection refused"); expect(consoleWarnSpy).not.toHaveBeenCalled(); const modelDevErrorLog = logged.find( (e) => e.component === "coding-agent.model-registry" && e.msg === "failed to refresh models from models.dev", ); expect(modelDevErrorLog).toBeDefined(); expect(modelDevErrorLog?.error).toBe("connection refused"); } finally { setLogSink(undefined); consoleWarnSpy.mockRestore(); fetchSpy.mockRestore(); registry.dispose(); } });
test("does not attempt to fetch models.dev when offline mode is enabled", async () => { const cachePath = join(tempDir, "models-dev-cache.json"); const originalOffline = process.env.PI_OFFLINE; process.env.PI_OFFLINE = "1"; const fetchSpy = vi.spyOn(globalThis, "fetch");
const registry = ModelRegistry.create(authStorage, modelsJsonPath, cachePath); try { const result = await registry.refreshModelsDev({ force: true }); expect(result.refreshed).toBe(false); expect(fetchSpy).not.toHaveBeenCalled(); } finally { if (originalOffline !== undefined) { process.env.PI_OFFLINE = originalOffline; } else { delete process.env.PI_OFFLINE; } fetchSpy.mockRestore(); registry.dispose(); } }); }); });});