import { describe, it, expect, vi, beforeEach, afterEach } from "vitest";
import { classifyAndRoute, isAutoModel, resolveAutoModel } from "./model-router.js";
// ============================================================================
// Mock fetch — callServerProxy uses global fetch internally
// ============================================================================
let originalFetch: typeof globalThis.fetch;
beforeEach(() => {
originalFetch = globalThis.fetch;
});
afterEach(() => {
globalThis.fetch = originalFetch;
});
/**
* Build a mock SSE stream that returns a single text response.
*/
function mockSSEResponse(text: string): ReadableStream<Uint8Array> {
const encoder = new TextEncoder();
const lines = [
`data: ${JSON.stringify({ type: "message_start", message: { usage: { input_tokens: 50 } } })}`,
`data: ${JSON.stringify({ type: "content_block_start", index: 0, content_block: { type: "text", text: "" } })}`,
`data: ${JSON.stringify({ type: "content_block_delta", index: 0, delta: { type: "text_delta", text } })}`,
`data: ${JSON.stringify({ type: "content_block_stop", index: 0 })}`,
`data: ${JSON.stringify({ type: "message_delta", delta: { stop_reason: "end_turn" }, usage: { output_tokens: 20 } })}`,
"data: [DONE]",
];
return new ReadableStream({
start(controller) {
for (const line of lines) {
controller.enqueue(encoder.encode(line + "\n\n"));
}
controller.close();
},
});
}
function mockFetch(classifierResponse: string) {
globalThis.fetch = vi.fn(async () => {
return new Response(mockSSEResponse(classifierResponse), { status: 200 });
}) as any;
}
// ============================================================================
// isAutoModel
// ============================================================================
describe("isAutoModel", () => {
it('returns true for "auto"', () => {
expect(isAutoModel("auto")).toBe(true);
});
it('returns false for "opus"', () => {
expect(isAutoModel("opus")).toBe(false);
});
it('returns false for "sonnet"', () => {
expect(isAutoModel("sonnet")).toBe(false);
});
it('returns false for "haiku"', () => {
expect(isAutoModel("haiku")).toBe(false);
});
it("returns false for empty string", () => {
expect(isAutoModel("")).toBe(false);
});
it('returns false for "Auto" (case sensitive)', () => {
expect(isAutoModel("Auto")).toBe(false);
});
});
// ============================================================================
// classifyAndRoute
// ============================================================================
describe("classifyAndRoute", () => {
const serverUrl = "https://whale-agent.fly.dev";
const authToken = "test-token";
it("routes simple message to haiku", async () => {
mockFetch(JSON.stringify({
model: "haiku",
reasoning: "Simple status check",
confidence: 0.9,
}));
const result = await classifyAndRoute("what time is it?", serverUrl, authToken);
expect(result.model).toBe("haiku");
expect(result.confidence).toBe(0.9);
expect(result.reasoning).toBe("Simple status check");
});
it("routes complex message to opus", async () => {
mockFetch(JSON.stringify({
model: "opus",
reasoning: "Complex multi-file architecture refactor",
confidence: 0.95,
}));
const result = await classifyAndRoute(
"Refactor the entire authentication system to use OAuth2 with PKCE flow, update all 12 affected files, and add comprehensive test coverage",
serverUrl,
authToken,
);
expect(result.model).toBe("opus");
expect(result.confidence).toBe(0.95);
});
it("routes medium message to sonnet", async () => {
mockFetch(JSON.stringify({
model: "sonnet",
reasoning: "Multi-step debugging task",
confidence: 0.85,
}));
const result = await classifyAndRoute(
"Fix the bug in the login form validation and update the tests",
serverUrl,
authToken,
);
expect(result.model).toBe("sonnet");
expect(result.confidence).toBe(0.85);
});
it("defaults to sonnet on malformed classifier response", async () => {
mockFetch("This is not valid JSON at all");
const result = await classifyAndRoute("hello", serverUrl, authToken);
expect(result.model).toBe("sonnet");
expect(result.confidence).toBe(0.5);
expect(result.reasoning).toBe("default fallback");
});
it("defaults to sonnet when confidence is below 0.3", async () => {
mockFetch(JSON.stringify({
model: "haiku",
reasoning: "Uncertain",
confidence: 0.2,
}));
const result = await classifyAndRoute("do something", serverUrl, authToken);
expect(result.model).toBe("sonnet");
expect(result.confidence).toBe(0.2);
});
it("defaults to sonnet on invalid model name in response", async () => {
mockFetch(JSON.stringify({
model: "gpt-5",
reasoning: "Wrong model name",
confidence: 0.9,
}));
const result = await classifyAndRoute("hello", serverUrl, authToken);
expect(result.model).toBe("sonnet");
});
it("defaults to sonnet on network error", async () => {
globalThis.fetch = vi.fn(async () => {
throw new Error("Network failure");
}) as any;
const result = await classifyAndRoute("hello", serverUrl, authToken);
expect(result.model).toBe("sonnet");
expect(result.reasoning).toBe("default fallback");
});
it("defaults to sonnet on HTTP error response", async () => {
globalThis.fetch = vi.fn(async () => {
return new Response("Bad Request", { status: 400 });
}) as any;
const result = await classifyAndRoute("hello", serverUrl, authToken);
expect(result.model).toBe("sonnet");
});
it("handles JSON embedded in markdown code fences", async () => {
mockFetch('```json\n{"model": "opus", "reasoning": "Complex task", "confidence": 0.88}\n```');
const result = await classifyAndRoute("rewrite the entire codebase", serverUrl, authToken);
expect(result.model).toBe("opus");
expect(result.confidence).toBe(0.88);
});
it("handles missing reasoning field gracefully", async () => {
mockFetch(JSON.stringify({
model: "haiku",
confidence: 0.7,
}));
const result = await classifyAndRoute("ls", serverUrl, authToken);
expect(result.model).toBe("haiku");
expect(result.reasoning).toBe("");
});
it("handles missing confidence field gracefully (falls back to sonnet)", async () => {
mockFetch(JSON.stringify({
model: "haiku",
reasoning: "Simple task",
}));
// confidence defaults to 0, which is < 0.3, so falls back to sonnet
const result = await classifyAndRoute("ls", serverUrl, authToken);
expect(result.model).toBe("sonnet");
});
});
// ============================================================================
// resolveAutoModel
// ============================================================================
describe("resolveAutoModel", () => {
const serverUrl = "https://whale-agent.fly.dev";
const authToken = "test-token";
it("returns haiku model ID for simple requests", async () => {
mockFetch(JSON.stringify({
model: "haiku",
reasoning: "Simple",
confidence: 0.9,
}));
const modelId = await resolveAutoModel("what is 2+2?", serverUrl, authToken);
expect(modelId).toBe("claude-haiku-4-5-20251001");
});
it("returns sonnet model ID for medium requests", async () => {
mockFetch(JSON.stringify({
model: "sonnet",
reasoning: "Medium complexity",
confidence: 0.8,
}));
const modelId = await resolveAutoModel("fix the bug in auth", serverUrl, authToken);
expect(modelId).toBe("claude-sonnet-4-20250514");
});
it("returns opus model ID for complex requests", async () => {
mockFetch(JSON.stringify({
model: "opus",
reasoning: "Complex architecture",
confidence: 0.95,
}));
const modelId = await resolveAutoModel("redesign the entire system", serverUrl, authToken);
expect(modelId).toBe("claude-opus-4-6");
});
it("returns sonnet model ID on classifier failure", async () => {
globalThis.fetch = vi.fn(async () => {
throw new Error("Connection refused");
}) as any;
const modelId = await resolveAutoModel("hello", serverUrl, authToken);
expect(modelId).toBe("claude-sonnet-4-20250514");
});
});
