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.
17 kB · 535 lines
TypeScript
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536import { getModels, type KnownProvider, type Model } from "@earendil-works/pi-ai";import { describe, expect, test, vi } from "vitest";import { defaultModelPerProvider, findInitialModel, resolveCliModel, resolveModelScopeFromModels,} from "../src/core/model-resolver.js";
const mockModels: Model<"anthropic-messages">[] = [ { id: "claude-sonnet-4-5", name: "Claude Sonnet 4.5", api: "anthropic-messages", provider: "anthropic", baseUrl: "https://api.anthropic.com", reasoning: true, input: ["text", "image"], cost: { input: 3, output: 15, cacheRead: 0.3, cacheWrite: 3.75 }, contextWindow: 200000, maxTokens: 8192, }, { id: "gpt-4o", name: "GPT-4o", api: "anthropic-messages", // Using same type for simplicity provider: "openai", baseUrl: "https://api.openai.com", reasoning: false, input: ["text", "image"], cost: { input: 5, output: 15, cacheRead: 0.5, cacheWrite: 5 }, contextWindow: 128000, maxTokens: 4096, },];
const mockOpenRouterModels: Model<"anthropic-messages">[] = [ { id: "qwen/qwen3-coder:exacto", name: "Qwen3 Coder Exacto", api: "anthropic-messages", provider: "openrouter", baseUrl: "https://openrouter.ai/api/v1", reasoning: true, input: ["text"], cost: { input: 1, output: 2, cacheRead: 0.1, cacheWrite: 1 }, contextWindow: 128000, maxTokens: 8192, }, { id: "openai/gpt-4o:extended", name: "GPT-4o Extended", api: "anthropic-messages", provider: "openrouter", baseUrl: "https://openrouter.ai/api/v1", reasoning: false, input: ["text", "image"], cost: { input: 5, output: 15, cacheRead: 0.5, cacheWrite: 5 }, contextWindow: 128000, maxTokens: 4096, },];
const allModels = [...mockModels, ...mockOpenRouterModels];
describe("resolveModelScopeFromModels", () => { test("resolves scope patterns against the supplied model list", () => { const daemonModel: Model<"anthropic-messages"> = { id: "daemon-only-model", name: "Daemon Only Model", api: "anthropic-messages", provider: "prime-inference", baseUrl: "https://api.pinference.ai/api/v1", reasoning: true, input: ["text"], cost: { input: 1, output: 2, cacheRead: 0.1, cacheWrite: 1 }, contextWindow: 128000, maxTokens: 8192, };
const result = resolveModelScopeFromModels( ["prime-inference/daemon-only-model:high", "openai/gpt-4o"], [...allModels, daemonModel], );
expect(result).toHaveLength(2); expect(result[0]?.model).toBe(daemonModel); expect(result[0]?.thinkingLevel).toBe("high"); expect(result[1]?.model.provider).toBe("openai"); expect(result[1]?.model.id).toBe("gpt-4o"); });
test("resolves a thinking level after a colon-bearing model id", () => { const result = resolveModelScopeFromModels(["openrouter/qwen/qwen3-coder:exacto:high"], allModels);
expect(result).toEqual([{ model: mockOpenRouterModels[0], thinkingLevel: "high" }]); });
test("keeps the model, warns, and drops an invalid thinking suffix", () => { const warn = vi.spyOn(console, "warn").mockImplementation(() => {}); try { const result = resolveModelScopeFromModels(["sonnet:random"], allModels);
expect(result).toEqual([{ model: mockModels[0], thinkingLevel: undefined }]); expect(warn).toHaveBeenCalledWith(expect.stringContaining('Invalid thinking level "random"')); } finally { warn.mockRestore(); } });
test("preserves provider-qualified selections when model names overlap", () => { const primeInferenceModel: Model<"anthropic-messages"> = { ...mockModels[0]!, id: "z-ai/glm-5.2", name: "GLM 5.2", provider: "prime-inference", baseUrl: "https://api.pinference.ai/api/v1", }; const huggingFaceModel: Model<"anthropic-messages"> = { ...primeInferenceModel, id: "zai-org/GLM-5.2", provider: "huggingface", baseUrl: "https://router.huggingface.co/v1", };
const result = resolveModelScopeFromModels( ["huggingface/zai-org/GLM-5.2"], [primeInferenceModel, huggingFaceModel], );
expect(result).toEqual([{ model: huggingFaceModel, thinkingLevel: undefined }]); });});
describe("resolveCliModel", () => { test("resolves --model provider/id without --provider", () => { const registry = { getAll: () => allModels, } as unknown as Parameters<typeof resolveCliModel>[0]["modelRegistry"];
const result = resolveCliModel({ cliModel: "openai/gpt-4o", modelRegistry: registry, });
expect(result.error).toBeUndefined(); expect(result.model?.provider).toBe("openai"); expect(result.model?.id).toBe("gpt-4o"); });
test("resolves fuzzy patterns within an explicit provider", () => { const registry = { getAll: () => allModels, } as unknown as Parameters<typeof resolveCliModel>[0]["modelRegistry"];
const result = resolveCliModel({ cliProvider: "openai", cliModel: "4o", modelRegistry: registry, });
expect(result.error).toBeUndefined(); expect(result.model?.provider).toBe("openai"); expect(result.model?.id).toBe("gpt-4o"); });
test("supports --model <pattern>:<thinking> (without explicit --thinking)", () => { const registry = { getAll: () => allModels, } as unknown as Parameters<typeof resolveCliModel>[0]["modelRegistry"];
const result = resolveCliModel({ cliModel: "sonnet:high", modelRegistry: registry, });
expect(result.error).toBeUndefined(); expect(result.model?.id).toBe("claude-sonnet-4-5"); expect(result.thinkingLevel).toBe("high"); });
test("prefers exact model id match over provider inference (OpenRouter-style ids)", () => { const registry = { getAll: () => allModels, } as unknown as Parameters<typeof resolveCliModel>[0]["modelRegistry"];
const result = resolveCliModel({ cliModel: "openai/gpt-4o:extended", modelRegistry: registry, });
expect(result.error).toBeUndefined(); expect(result.model?.provider).toBe("openrouter"); expect(result.model?.id).toBe("openai/gpt-4o:extended"); });
test("does not strip invalid :suffix as thinking level in --model (treat as raw id)", () => { const registry = { getAll: () => allModels, } as unknown as Parameters<typeof resolveCliModel>[0]["modelRegistry"];
const result = resolveCliModel({ cliProvider: "openai", cliModel: "gpt-4o:extended", modelRegistry: registry, });
expect(result.error).toBeUndefined(); expect(result.model?.provider).toBe("openai"); expect(result.model?.id).toBe("gpt-4o:extended"); });
test("allows custom model ids for explicit providers without double prefixing", () => { const registry = { getAll: () => allModels, } as unknown as Parameters<typeof resolveCliModel>[0]["modelRegistry"];
const result = resolveCliModel({ cliProvider: "openrouter", cliModel: "openrouter/openai/ghost-model", modelRegistry: registry, });
expect(result.error).toBeUndefined(); expect(result.model?.provider).toBe("openrouter"); expect(result.model?.id).toBe("openai/ghost-model"); });
test("returns a clear error when there are no models", () => { const registry = { getAll: () => [], } as unknown as Parameters<typeof resolveCliModel>[0]["modelRegistry"];
const result = resolveCliModel({ cliProvider: "openai", cliModel: "gpt-4o", modelRegistry: registry, });
expect(result.model).toBeUndefined(); expect(result.error).toContain("No models available"); });
test("prefers provider/model split over gateway model with matching id", () => { const zaiModel: Model<"anthropic-messages"> = { id: "glm-5", name: "GLM-5", api: "anthropic-messages", provider: "zai", baseUrl: "https://open.bigmodel.cn/api/paas/v4", reasoning: true, input: ["text"], cost: { input: 1, output: 2, cacheRead: 0.1, cacheWrite: 1 }, contextWindow: 128000, maxTokens: 8192, }; const gatewayModel: Model<"anthropic-messages"> = { id: "zai/glm-5", name: "GLM-5", api: "anthropic-messages", provider: "vercel-ai-gateway", baseUrl: "https://ai-gateway.vercel.sh", reasoning: true, input: ["text"], cost: { input: 1, output: 2, cacheRead: 0.1, cacheWrite: 1 }, contextWindow: 128000, maxTokens: 8192, }; const registry = { getAll: () => [...allModels, zaiModel, gatewayModel], } as unknown as Parameters<typeof resolveCliModel>[0]["modelRegistry"];
const result = resolveCliModel({ cliModel: "zai/glm-5", modelRegistry: registry, });
expect(result.error).toBeUndefined(); expect(result.model?.provider).toBe("zai"); expect(result.model?.id).toBe("glm-5"); });
test("resolves provider-prefixed fuzzy patterns (openrouter/qwen -> openrouter model)", () => { const registry = { getAll: () => allModels, } as unknown as Parameters<typeof resolveCliModel>[0]["modelRegistry"];
const result = resolveCliModel({ cliModel: "openrouter/qwen", modelRegistry: registry, });
expect(result.error).toBeUndefined(); expect(result.model?.provider).toBe("openrouter"); expect(result.model?.id).toBe("qwen/qwen3-coder:exacto"); });});
describe("default model selection", () => { test("openai defaults track current models", () => { expect(defaultModelPerProvider.openai).toBe("gpt-5.4"); expect(defaultModelPerProvider["openai-codex"]).toBe("gpt-5.5"); expect(defaultModelPerProvider["prime-inference"]).toBe("z-ai/glm-5.2"); });
test("every per-provider default exists in the model catalog", () => { for (const [provider, modelId] of Object.entries(defaultModelPerProvider)) { const models = getModels(provider as KnownProvider); if (models.length === 0) continue; expect( models.map((model) => model.id), `default for ${provider}`, ).toContain(modelId); } });
test("zai, minimax, and cerebras defaults track current models", () => { expect(defaultModelPerProvider.zai).toBe("glm-5.3"); expect(defaultModelPerProvider.minimax).toBe("MiniMax-M2.7"); expect(defaultModelPerProvider["minimax-cn"]).toBe("MiniMax-M2.7"); expect(defaultModelPerProvider.cerebras).toBe("gpt-oss-120b"); });
test("ai-gateway default tracks current model", () => { expect(defaultModelPerProvider["vercel-ai-gateway"]).toBe("zai/glm-5.1"); });
test("findInitialModel accepts explicit provider custom model ids", async () => { const registry = { getAll: () => allModels, } as unknown as Parameters<typeof findInitialModel>[0]["modelRegistry"];
const result = await findInitialModel({ cliProvider: "openrouter", cliModel: "openrouter/openai/ghost-model", scopedModels: [], isContinuing: false, modelRegistry: registry, });
expect(result.model?.provider).toBe("openrouter"); expect(result.model?.id).toBe("openai/ghost-model"); });
test("findInitialModel uses medium as the built-in default thinking level", async () => { const reasoningModel = mockModels[0]; const registry = { refreshAvailableModels: async () => [reasoningModel], } as unknown as Parameters<typeof findInitialModel>[0]["modelRegistry"];
const result = await findInitialModel({ scopedModels: [], isContinuing: false, modelRegistry: registry, });
expect(result.model).toBe(reasoningModel); expect(result.thinkingLevel).toBe("medium"); });
test("findInitialModel prefers GLM 5.2 when Prime Inference is configured", async () => { const anthropicModel: Model<"anthropic-messages"> = { ...mockModels[0], id: "claude-opus-4-7", name: "Claude Opus 4.7", }; const primeModel: Model<"anthropic-messages"> = { id: "z-ai/glm-5.2", name: "GLM 5.2", api: "anthropic-messages", provider: "prime-inference", baseUrl: "https://api.pinference.ai/api/v1", reasoning: true, input: ["text"], cost: { input: 1, output: 2, cacheRead: 0, cacheWrite: 0 }, contextWindow: 1048576, maxTokens: 101376, }; const registry = { refreshAvailableModels: async () => [anthropicModel, primeModel], } as unknown as Parameters<typeof findInitialModel>[0]["modelRegistry"];
const result = await findInitialModel({ scopedModels: [], isContinuing: false, modelRegistry: registry, });
expect(result.model).toBe(primeModel); });
test("findInitialModel uses another provider default when Prime Inference is not configured", async () => { const anthropicModel: Model<"anthropic-messages"> = { ...mockModels[0], id: "claude-opus-4-7", name: "Claude Opus 4.7", }; const registry = { refreshAvailableModels: async () => [anthropicModel], } as unknown as Parameters<typeof findInitialModel>[0]["modelRegistry"];
const result = await findInitialModel({ scopedModels: [], isContinuing: false, modelRegistry: registry, });
expect(result.model).toBe(anthropicModel); });
test("findInitialModel selects ai-gateway default when available", async () => { const aiGatewayModel: Model<"anthropic-messages"> = { id: "anthropic/claude-opus-4-6", name: "Claude Opus 4.6", api: "anthropic-messages", provider: "vercel-ai-gateway", baseUrl: "https://ai-gateway.vercel.sh", reasoning: true, input: ["text", "image"], cost: { input: 5, output: 15, cacheRead: 0.5, cacheWrite: 5 }, contextWindow: 200000, maxTokens: 8192, };
const registry = { refreshAvailableModels: async () => [aiGatewayModel], } as unknown as Parameters<typeof findInitialModel>[0]["modelRegistry"];
const result = await findInitialModel({ scopedModels: [], isContinuing: false, modelRegistry: registry, });
expect(result.model?.provider).toBe("vercel-ai-gateway"); expect(result.model?.id).toBe("anthropic/claude-opus-4-6"); });
test("findInitialModel skips saved defaults without configured auth", async () => { const savedDefault = mockModels[0]; const primeModel: Model<"anthropic-messages"> = { id: "openai/gpt-5.5", name: "GPT 5.5 (Prime Inference)", api: "anthropic-messages", provider: "prime-inference", baseUrl: "https://api.pinference.ai/api/v1", reasoning: true, input: ["text"], cost: { input: 1, output: 2, cacheRead: 0.1, cacheWrite: 1 }, contextWindow: 128000, maxTokens: 8192, }; const registry = { find: (provider: string, modelId: string) => [savedDefault, primeModel].find((model) => model.provider === provider && model.id === modelId), hasConfiguredAuth: (model: Model<"anthropic-messages">) => model.provider === "prime-inference", refreshAvailableModels: async () => [primeModel], } as unknown as Parameters<typeof findInitialModel>[0]["modelRegistry"];
const result = await findInitialModel({ scopedModels: [], isContinuing: false, defaultProvider: savedDefault.provider, defaultModelId: savedDefault.id, modelRegistry: registry, });
expect(result.model?.provider).toBe("prime-inference"); expect(result.model?.id).toBe("openai/gpt-5.5"); });
test("findInitialModel rebuilds a saved default missing from the model snapshot when the provider is authed", async () => { const primeSnapshotModel: Model<"anthropic-messages"> = { id: "openai/gpt-5.5", name: "GPT 5.5 (Prime Inference)", api: "anthropic-messages", provider: "prime-inference", baseUrl: "https://api.pinference.ai/api/v1", reasoning: true, input: ["text"], cost: { input: 1, output: 2, cacheRead: 0.1, cacheWrite: 1 }, contextWindow: 128000, maxTokens: 8192, }; const registry = { find: () => undefined, getAll: () => [primeSnapshotModel], hasConfiguredAuth: (model: Model<"anthropic-messages">) => model.provider === "prime-inference", refreshAvailableModels: async () => [primeSnapshotModel], } as unknown as Parameters<typeof findInitialModel>[0]["modelRegistry"];
const result = await findInitialModel({ scopedModels: [], isContinuing: false, defaultProvider: "prime-inference", defaultModelId: "anthropic/claude-opus-4.6", modelRegistry: registry, });
expect(result.model?.provider).toBe("prime-inference"); expect(result.model?.id).toBe("anthropic/claude-opus-4.6"); });
test("findInitialModel does not rebuild a saved default for an unauthed provider", async () => { const primeSnapshotModel: Model<"anthropic-messages"> = { id: "openai/gpt-5.5", name: "GPT 5.5 (Prime Inference)", api: "anthropic-messages", provider: "prime-inference", baseUrl: "https://api.pinference.ai/api/v1", reasoning: true, input: ["text"], cost: { input: 1, output: 2, cacheRead: 0.1, cacheWrite: 1 }, contextWindow: 128000, maxTokens: 8192, }; const registry = { find: () => undefined, getAll: () => [...mockModels, primeSnapshotModel], hasConfiguredAuth: (model: Model<"anthropic-messages">) => model.provider === "prime-inference", refreshAvailableModels: async () => [primeSnapshotModel], } as unknown as Parameters<typeof findInitialModel>[0]["modelRegistry"];
const result = await findInitialModel({ scopedModels: [], isContinuing: false, defaultProvider: "anthropic", defaultModelId: "claude-ghost-9", modelRegistry: registry, });
expect(result.model?.provider).toBe("prime-inference"); expect(result.model?.id).toBe("openai/gpt-5.5"); });});