auto-model-router 0.2.32 → 0.3.1

package/test/turn.test.ts

@@ -1,7 +1,7 @@
 import { describe, expect, test } from "bun:test";
 import { createDisabledBridge } from "../src/context/bridge.ts";
 import type { ContextBridge, TurnRecord } from "../src/context/types.ts";
-import type { CatalogSource } from "../src/catalog/types.ts";
+import type { CatalogModel, CatalogSource } from "../src/catalog/types.ts";
 import type { EscalationConfig, RouterConfig } from "../src/config/types.ts";
 import { EMPTY_USAGE, type Ledger, type LedgerEntry, type UsageCounts } from "../src/cost/types.ts";
 import type {
@@ -31,6 +31,7 @@ function mkConfig(escalation: Partial<EscalationConfig> = {}): RouterConfig {
 	return {
 		server: { host: "127.0.0.1", port: 8787, maxConcurrentTurns: 24 },
 		openrouter: { baseUrl: "https://openrouter.ai/api/v1", apiKey: "", title: "test", timeoutMs: 30_000, catalogTtlMs: 3_600_000, catalogRefreshMs: 0 },
+		ollama: { enabled: false, baseUrl: "http://127.0.0.1:11434/v1", apiKey: "", timeoutMs: 30_000, catalogTtlMs: 300_000, includeLocal: false, prices: {}, twins: {}, costBias: 1, biasUntilUsage: 0.9, usagePollMs: 0, quotaCooldownMs: 0, rateLimitCooldownMs: 0, planCreditsUsd: 0 },
 		benchmarks: { enabled: false, artificialAnalysisApiKey: "", benchlm: true, refreshMs: 86_400_000, timeoutMs: 30_000, useLocalScores: false },
 		tiers: {
 			trivial: { minQuality: 0, maxInputPerMtok: 0.3, qualityExponent: 0, pin: [] },
@@ -45,7 +46,7 @@ function mkConfig(escalation: Partial<EscalationConfig> = {}): RouterConfig {
 			data: { axis: "intelligence", minQuality: 0 },
 			chat: { axis: "intelligence", minQuality: 0 },
 		},
-		filters: { allow: [], deny: [], includeFree: false, requireToolSupport: true, minTrust: 0.6, minTrustSamples: 5, trustScopedByHarness: false, trustWindowDays: 0, contextHeadroom: 1.2, latencyWeight: 0, latencyReferenceMs: 5000, latencyReferenceTokensPerSec: 30, latencyMinSamples: 20 },
+		filters: { allow: [], deny: [], includeFree: false, requireToolSupport: true, minTrust: 0.6, minTrustSamples: 5, trustScopedByHarness: false, trustWindowDays: 0, contextHeadroom: 1.2, latencyWeight: 0, latencyReferenceMs: 5000, latencyReferenceTokensPerSec: 30, latencyMinSamples: 20, escalationCostWeight: 0 },
 		classifier: {
 			ambiguityThreshold: 0,
 			model: "test/adjudicator",
@@ -56,6 +57,7 @@ function mkConfig(escalation: Partial<EscalationConfig> = {}): RouterConfig {
 			toolAxis: "coding",
 			chatAxis: "intelligence",
 			agenticLoopDepth: 3,
+			mechanicalRetryFactor: 0.2,
 			reasoningWeights: { medium: 0.14, high: 0.24, xhigh: 0.3, max: 0.34 },
 		},
 		escalation: {
@@ -68,11 +70,11 @@ function mkConfig(escalation: Partial<EscalationConfig> = {}): RouterConfig {
 			escalateOnLengthStop: false,
 			...escalation,
 		},
-		hysteresis: { holdTurns: 2, holdTurnsAfterEscalation: 4, switchMargin: 1.5, cacheWarmTtlMs: 600_000, maxDowngradePerTurn: 1, breakHoldOnMechanical: false },
+		hysteresis: { holdTurns: 2, holdTurnsAfterEscalation: 4, switchMargin: 1.5, cacheWarmTtlMs: 600_000, maxDowngradePerTurn: 1, breakHoldOnMechanical: false, switchHorizonTurns: 1 },
 		exploration: { enabled: false, rates: {}, stickyPolicy: "never", holdTurns: { enabled: false, values: [2, 3, 4] } },
 		cache: { injectBreakpoints: true, maxBreakpoints: 4, minPromptTokens: 1024, milestoneTokens: 20_000 },
 		context: { enabled: false, baseUrl: "", token: "", defaultScope: "", timeoutMs: 3_000, maxStalenessMs: 900_000, maxBlockChars: 24_000, memoryLimit: 8, docsLimit: 2, sessionLimit: 6, briefChars: 0, recordTurns: false, maxQueue: 64 },
-		compaction: { enabled: false, budgetTokens: 40_000, floorRatio: 1, fitToWindow: true, protectRecentTurns: 4, maxToolResultBytes: 4_096, keepHeadBytes: 512, keepTailBytes: 512, elideSupersededReads: true, collapseDuplicateResults: true },
+		compaction: { enabled: false, budgetTokens: 40_000, floorRatio: 1, fitToWindow: true, protectRecentTurns: 4, maxToolResultBytes: 4_096, keepHeadBytes: 512, keepTailBytes: 512, elideSupersededReads: true, collapseDuplicateResults: true, replanGrowthRatio: 1 },
 		budget: { onExceeded: "downgrade" },
 		profiles: [],
 		ledger: { path: ":memory:", blendWindowDays: 7, blendMinSamples: 20, fallbackBlend: { inputPerMtok: 1, outputPerMtok: 4 }, conversationTtlMs: 86_400_000 },
@@ -147,6 +149,8 @@ function mkDecision(tier: Tier, slug: string, probe: Partial<ProbePlan> = {}): D
 		cacheBreakpointMessageIndices: [],
 		compactionPlan: [],
 		promptTokensSaved: 0,
+		compactionSavedBytes: 0,
+		compactionPlanTokens: 0,
 		reasoning: undefined,
 		maxTokens: undefined,
 		stripAssistantReasoning: false,
@@ -403,10 +407,13 @@ describe("runTurn", () => {
 		expect(entries[1]!.attempt).toBe(1);
 		expect(entries[1]!.slug).toBe("better/model");
 
-		// The escalated attempt re-routed one tier up.
-		expect(calls).toHaveLength(2);
+		// A refusal indicts the provider, so a same-tier sibling is probed first;
+		// this fake router only has the simple-tier decision left, which is the
+		// wrong tier, so the turn then escalates one tier up for real.
+		expect(calls).toHaveLength(3);
 		expect(calls[0]).toEqual({ attempt: 0 });
-		expect(calls[1]).toEqual({ attempt: 1, escalateFrom: "trivial" });
+		expect(calls[1]).toEqual({ attempt: 1 });
+		expect(calls[2]).toEqual({ attempt: 1, escalateFrom: "trivial" });
 
 		// The held refusal text never reached the client.
 		expect(textOut(chunks)).toBe("Here is the answer.");
@@ -729,4 +736,163 @@ describe("latency measurement covers the work the router actually does", () => {
 		expect(entries[0]?.ttftMs).not.toBeNull();
 		expect(entries[0]?.ttftMs ?? -1).toBeGreaterThanOrEqual(0);
 	});
+
+	test("a client hang-up AFTER the finish event is a completed turn, not an error", async () => {
+		// Measured live: 1,842 of 2,068 "request aborted" rows carried a finish
+		// reason and full usage — the generation had finished and the client
+		// closed before the trailing [DONE] was read. Recording that as an error
+		// skipped the state save on 12% of turns.
+		const { router } = mkRouter([mkDecision("trivial", "cheap/model", { escalateTo: "simple" })]);
+		const { upstream } = mkUpstream([
+			{
+				kind: "die",
+				chunks: [
+					startChunk("cheap/model"),
+					textChunk("all done here"),
+					finishChunk("stop"),
+					usageChunk({ promptTokens: 120, cachedTokens: 100, completionTokens: 4 }, 0.0004),
+				],
+				error: new UpstreamError("aborted", 0, "request aborted", false),
+			},
+		]);
+		const { ledger, entries } = mkLedger();
+		const { store, map } = mkConversations();
+		const { sink, errors, finishes } = mkSink();
+
+		await runTurn(mkReq(), sink, { config: mkConfig(), router, upstream, ledger, conversations: store, catalog, context: createDisabledBridge() }, new AbortController().signal);
+
+		expect(errors).toHaveLength(0);
+		expect(entries).toHaveLength(1);
+		expect(entries[0]!.error).toBeNull();
+		expect(entries[0]!.wasted).toBe(false);
+		expect(entries[0]!.finishReason).toBe("stop");
+		expect(entries[0]!.reportedUsd).toBe(0.0004);
+		// The turn settled: hysteresis, cache warmth and the turn counter advance.
+		const state = map.get("conv-test")!;
+		expect(state.turn).toBe(1);
+		expect(state.currentSlug).toBe("cheap/model");
+		expect(state.cacheWarmSlug).toBe("cheap/model");
+		expect(state.lastPromptTokens).toBe(120);
+		expect(finishes).toHaveLength(1);
+	});
+
+	test("a client hang-up BEFORE the finish event is still recorded as aborted", async () => {
+		const { router } = mkRouter([mkDecision("trivial", "cheap/model", { escalateTo: "simple", maxTokens: 1 })]);
+		const { upstream } = mkUpstream([
+			{
+				kind: "die",
+				chunks: [startChunk("cheap/model"), textChunk("partial answer that committed")],
+				error: new UpstreamError("aborted", 0, "request aborted", false),
+			},
+		]);
+		const { ledger, entries } = mkLedger();
+		const { store, map } = mkConversations();
+		const { sink, finishes } = mkSink();
+
+		await runTurn(mkReq(), sink, { config: mkConfig(), router, upstream, ledger, conversations: store, catalog, context: createDisabledBridge() }, new AbortController().signal);
+
+		expect(entries).toHaveLength(1);
+		expect(entries[0]!.error).toBe("request aborted");
+		expect(entries[0]!.finishReason).toBeNull();
+		expect(map.get("conv-test")!.turn).toBe(0);
+		expect(finishes).toHaveLength(0);
+	});
+
+
+	test("a provider that reports usage but no cost gets its actual tokens priced at the catalog rate", async () => {
+		// Ollama returns usage without `cost`. Leaving the forecast (assumed 1,024
+		// completion tokens) as the recorded figure overstated a 26-token turn ~3x.
+		const { router } = mkRouter([mkDecision("trivial", "ollama/glm-5.3-flash")]);
+		const { upstream } = mkUpstream([
+			{
+				kind: "chunks",
+				chunks: [startChunk("ollama/glm-5.3-flash"), textChunk("ok"), finishChunk("stop"), usageChunk({ promptTokens: 1_000_000, completionTokens: 1_000_000 }, null)],
+			},
+		]);
+		const { ledger, entries } = mkLedger();
+		const { store } = mkConversations();
+		const { sink, finishes } = mkSink();
+		const priced = {
+			...catalog,
+			find: (slug: string) =>
+				slug === "ollama/glm-5.3-flash"
+					? ({
+							slug,
+							provider: "ollama",
+							canonicalSlug: slug,
+							name: slug,
+							contextLength: 1_000_000,
+							supportsTools: true,
+							supportsReasoning: true,
+							reasoningMandatory: false,
+							supportsToolChoice: false,
+							inputModalities: ["text"],
+							price: { prompt: 0.15 / 1e6, completion: 0.5 / 1e6 },
+							priceTiers: [],
+							quality: {},
+							tokenizer: "Other",
+							isFree: false,
+							createdAtMs: 0,
+							author: "ollama",
+						} satisfies CatalogModel)
+					: undefined,
+		};
+
+		await runTurn(mkReq(), sink, { config: mkConfig(), router, upstream, ledger, conversations: store, catalog: priced, context: createDisabledBridge() }, new AbortController().signal);
+
+		// 1M prompt tokens at $0.15/M + 1M completion tokens at $0.50/M.
+		expect(entries[0]!.reportedUsd).toBeCloseTo(0.65, 6);
+		expect(entries[0]!.priceModel?.slug).toBe("ollama/glm-5.3-flash");
+		expect(finishes[0]!.reportedUsd).toBeCloseTo(0.65, 6);
+	});
+
+	test("an Ollama-served turn after another on the same model prices the previous prompt as cached", async () => {
+		// ollama.com caches prefixes and bills them at the cached rate without
+		// reporting a count; the router estimates it with its warm-cache rule.
+		const ollamaModel: CatalogModel = {
+			slug: "ollama/glm-5.3-flash",
+			provider: "ollama",
+			canonicalSlug: "ollama/glm-5.3-flash",
+			name: "glm",
+			contextLength: 1_000_000,
+			supportsTools: true,
+			supportsReasoning: true,
+			reasoningMandatory: false,
+			supportsToolChoice: false,
+			inputModalities: ["text"],
+			price: { prompt: 0.15 / 1e6, cacheRead: 0.03 / 1e6, completion: 0.5 / 1e6 },
+			priceTiers: [],
+			quality: {},
+			tokenizer: "Other",
+			isFree: false,
+			createdAtMs: 0,
+			author: "ollama",
+		};
+		const priced = { ...catalog, find: (slug: string) => (slug === ollamaModel.slug ? ollamaModel : undefined) };
+		const turn = (prompt: number) => ({
+			kind: "chunks" as const,
+			chunks: [startChunk(ollamaModel.slug), textChunk("ok"), finishChunk("stop"), usageChunk({ promptTokens: prompt, completionTokens: 0 }, null)],
+		});
+		const { router } = mkRouter([mkDecision("simple", ollamaModel.slug)]);
+		const { upstream } = mkUpstream([turn(100_000), turn(120_000)]);
+		const { ledger, entries } = mkLedger();
+		const { store, map } = mkConversations();
+		const deps = { config: mkConfig(), router, upstream, ledger, conversations: store, catalog: priced, context: createDisabledBridge() };
+
+		await runTurn(mkReq(), mkSink().sink, deps, new AbortController().signal);
+		// First turn: nothing to be cached yet, full input rate.
+		expect(entries[0]!.usage.cachedTokens).toBe(0);
+		expect(entries[0]!.usage.cachedEstimated).toBeUndefined();
+		expect(entries[0]!.reportedUsd).toBeCloseTo(0.015, 6);
+
+		await runTurn(mkReq(), mkSink().sink, deps, new AbortController().signal);
+		// Second turn on the same model: the 100k previous prompt is the cached
+		// prefix at $0.03/M, the 20k of growth is fresh at $0.15/M.
+		expect(entries[1]!.usage.cachedTokens).toBe(100_000);
+		expect(entries[1]!.usage.cachedEstimated).toBe(true);
+		expect(entries[1]!.reportedUsd).toBeCloseTo(0.003 + 0.003, 6);
+		// The estimate is evidence enough to keep the Ollama model warm for the stay/switch comparison.
+		expect(map.get("conv-test")!.cacheWarmSlug).toBe(ollamaModel.slug);
+	});
+
 });

package/tools/recompute-ollama-cache.ts

@@ -0,0 +1,129 @@
+#!/usr/bin/env bun
+/**
+ * One-off backfill: re-price historical Ollama ledger rows with the router's
+ * cache estimate (`src/cost/cache-estimate.ts`).
+ *
+ * Rows recorded before the estimate existed booked every prompt token at the
+ * full input rate, overstating Ollama spend ~3.7x against ollama.com's meter.
+ * This walks each conversation's Ollama rows in order, applies the same rule
+ * the live path applies (same model as the previous kept row within
+ * `hysteresis.cacheWarmTtlMs` ⇒ the previous prompt is the cached prefix),
+ * and rewrites `usage`, `cost_breakdown` and `reported_usd` for rows whose
+ * cost the router itself computed (Ollama never reports a cost). Rows that
+ * already carry a cache count are left alone, so it is safe to re-run.
+ *
+ * Usage:
+ *   bun tools/recompute-ollama-cache.ts            # dry run: prints the delta
+ *   bun tools/recompute-ollama-cache.ts --apply    # backs up router.db, then writes
+ */
+
+import { copyFileSync, mkdirSync } from "node:fs";
+import { dirname, join } from "node:path";
+import { Database } from "bun:sqlite";
+import { ollamaRateFor } from "../src/catalog/ollama-prices.ts";
+import { loadConfig } from "../src/config/load.ts";
+import { estimateUnreportedCache } from "../src/cost/cache-estimate.ts";
+import type { UsageCounts } from "../src/cost/types.ts";
+
+const apply = process.argv.includes("--apply");
+const cfg = loadConfig({});
+const ttl = cfg.hysteresis.cacheWarmTtlMs;
+
+interface Row {
+	id: string;
+	ck: string;
+	t: number;
+	latency: number;
+	slug: string;
+	usage: string;
+	reported: number | null;
+	wasted: number;
+	error: string | null;
+}
+
+const db = new Database(cfg.ledger.path);
+db.exec("PRAGMA busy_timeout = 5000");
+const rows = db
+	.query(
+		`SELECT id, conversation_key ck, created_at_ms t, latency_ms latency, COALESCE(served_slug, slug) slug, usage, reported_usd reported, wasted, error
+		 FROM ledger WHERE COALESCE(served_slug, slug) LIKE 'ollama/%' ORDER BY conversation_key, created_at_ms`,
+	)
+	.all() as Row[];
+
+let before = 0;
+let after = 0;
+let changed = 0;
+const updates: { id: string; usage: string; breakdown: string; usd: number }[] = [];
+let prev: Row | undefined;
+let prevPrompt = 0;
+for (const r of rows) {
+	const usage = JSON.parse(r.usage) as UsageCounts;
+	const sameConv = prev !== undefined && prev.ck === r.ck;
+	const rate = ollamaRateFor(r.slug.slice("ollama/".length), cfg.ollama.prices);
+	before += r.reported ?? 0;
+	if (rate === null || r.reported === null || r.error !== null) {
+		after += r.reported ?? 0;
+		if (r.error === null) {
+			prev = r;
+			prevPrompt = usage.promptTokens;
+		}
+		continue;
+	}
+	const est = estimateUnreportedCache(usage, {
+		previousSlug: sameConv ? (prev?.slug ?? null) : null,
+		previousPromptTokens: sameConv ? prevPrompt : 0,
+		previousAtMs: sameConv ? (prev?.t ?? 0) : 0,
+		servedSlug: r.slug,
+		nowMs: r.t,
+		cacheWarmTtlMs: ttl,
+	});
+	const input = rate.rate.input / 1e6;
+	const cached = (rate.rate.cachedInput ?? rate.rate.input) / 1e6;
+	const output = rate.rate.output / 1e6;
+	const fresh = Math.max(0, est.promptTokens - est.cachedTokens);
+	const breakdown = {
+		freshPrompt: fresh * input,
+		cacheRead: est.cachedTokens * cached,
+		cacheWrite: 0,
+		completion: est.completionTokens * output,
+		reasoning: 0,
+		images: 0,
+		request: 0,
+		total: 0,
+		tierAtPromptTokens: 0,
+	};
+	breakdown.total = breakdown.freshPrompt + breakdown.cacheRead + breakdown.completion;
+	after += breakdown.total;
+	if (est.cachedTokens !== usage.cachedTokens) {
+		changed++;
+		updates.push({ id: r.id, usage: JSON.stringify(est), breakdown: JSON.stringify(breakdown), usd: breakdown.total });
+	}
+	// The next row's "previous" is this row as dispatched, whether or not it was wasted:
+	// a wasted probe still warmed the prefix.
+	prev = r;
+	prevPrompt = est.promptTokens;
+}
+
+console.log(`ollama rows: ${rows.length}, re-priced: ${changed}`);
+console.log(`ledger Ollama spend: $${before.toFixed(2)} → $${after.toFixed(2)}`);
+if (!apply) {
+	console.log("dry run; pass --apply to write (router.db is backed up first)");
+	db.close();
+	process.exit(0);
+}
+
+const backupDir = join(dirname(cfg.ledger.path), "backups");
+mkdirSync(backupDir, { recursive: true });
+const stamp = new Date().toISOString().replace(/[-:]/g, "").replace(/\..+/, "");
+const backup = join(backupDir, `router-pre-ollama-cache-${stamp}.db`);
+db.exec("PRAGMA wal_checkpoint(TRUNCATE)");
+copyFileSync(cfg.ledger.path, backup);
+console.log(`backup: ${backup}`);
+
+const stmt = db.prepare("UPDATE ledger SET usage = $usage, cost_breakdown = $breakdown, reported_usd = $usd WHERE id = $id");
+const tx = db.transaction((list: typeof updates) => {
+	for (const u of list) stmt.run({ $usage: u.usage, $breakdown: u.breakdown, $usd: u.usd, $id: u.id });
+});
+tx(updates);
+console.log(`updated ${updates.length} rows`);
+db.close();