auto-model-router 0.31.0 → 0.32.0

package/src/router/state.ts

@@ -5,11 +5,11 @@
  * survive a restart: an omp session outlives this process, and forgetting
  * which model is warm would cold-start a paid prompt cache for no reason.
  *
- * The `conversations` table is created by `util/sqlite.ts`, the single
- * migration path.
+ * The `conversations` table is created by the store's migration
+ * (`util/schema.ts`, or `util/sqlite.ts` for a file that predates it).
  */
 
-import type { Database, Statement } from "bun:sqlite";
+import { jsonParam, jsonValue, num, type SqlDb } from "../util/sql.ts";
 
 import type { CompactionEdit } from "../wire/types.ts";
 import type { ConversationState, ConversationStore, Tier } from "./types.ts";
@@ -19,147 +19,135 @@ import type { ConversationState, ConversationStore, Tier } from "./types.ts";
 interface Row {
 	key: string;
 	session_id: string;
-	turn: number;
+	turn: unknown;
 	current_slug: string | null;
 	current_tier: string | null;
-	sticky_until_turn: number;
-	escalations: number;
-	spent_usd: number;
-	last_prompt_tokens: number;
+	sticky_until_turn: unknown;
+	escalations: unknown;
+	spent_usd: unknown;
+	last_prompt_tokens: unknown;
 	cache_warm_slug: string | null;
-	cache_warm_at_ms: number;
+	cache_warm_at_ms: unknown;
 	context_version: string | null;
-	context_fetched_at_ms: number;
-	compaction_plan: string | null;
-	compaction_plan_tokens: number;
+	context_fetched_at_ms: unknown;
+	/** JSON: text on sqlite, already parsed on postgres. */
+	compaction_plan: unknown;
+	compaction_plan_tokens: unknown;
 	upgrade_deferred_tier: string | null;
-	updated_at_ms: number;
+	updated_at_ms: unknown;
 }
 
 function toState(row: Row): ConversationState {
 	return {
 		key: row.key,
 		sessionId: row.session_id,
-		turn: row.turn,
+		// Counts and sums arrive as strings from Postgres, and every one of these
+		// feeds arithmetic — the sticky window, the budget guard, cache warmth.
+		turn: num(row.turn),
 		currentSlug: row.current_slug,
 		// Stored as free text; the column is only ever written from a Tier.
 		currentTier: row.current_tier as Tier | null,
-		stickyUntilTurn: row.sticky_until_turn,
-		escalations: row.escalations,
-		spentUsd: row.spent_usd,
-		lastPromptTokens: row.last_prompt_tokens,
+		stickyUntilTurn: num(row.sticky_until_turn),
+		escalations: num(row.escalations),
+		spentUsd: num(row.spent_usd),
+		lastPromptTokens: num(row.last_prompt_tokens),
 		cacheWarmSlug: row.cache_warm_slug,
-		cacheWarmAtMs: row.cache_warm_at_ms,
+		cacheWarmAtMs: num(row.cache_warm_at_ms),
 		contextVersion: row.context_version,
-		contextFetchedAtMs: row.context_fetched_at_ms,
-		compactionPlan: row.compaction_plan === null ? null : (JSON.parse(row.compaction_plan) as CompactionEdit[]),
-		compactionPlanTokens: row.compaction_plan_tokens,
+		contextFetchedAtMs: num(row.context_fetched_at_ms),
+		compactionPlan: jsonValue<CompactionEdit[]>(row.compaction_plan),
+		compactionPlanTokens: num(row.compaction_plan_tokens),
 		upgradeDeferredTier: row.upgrade_deferred_tier as Tier | null,
-		updatedAtMs: row.updated_at_ms,
+		updatedAtMs: num(row.updated_at_ms),
 	};
 }
 
-export function createConversationStore(db: Database): ConversationStore {
-	// Hoisted: this runs on every turn, twice when an escalation retries.
-	const selectOne: Statement<Row, [string]> = db.query("SELECT * FROM conversations WHERE key = ?");
-	const insertOne: Statement<unknown, [string, string, number]> = db.query(
-		"INSERT INTO conversations (key, session_id, updated_at_ms) VALUES (?, ?, ?)",
-	);
-	// `spent_usd` and `escalations` are ABSENT from this statement on purpose.
-	// They accumulate through `accrueOne` below, so writing a turn-start snapshot
-	// back here would erase whatever a billed-but-uncommitted dispatch added.
-	// The schema defaults both to 0, so the INSERT arm still works.
-	const upsert = db.query(`
-		INSERT INTO conversations (
-			key, session_id, turn, current_slug, current_tier, sticky_until_turn,
-			last_prompt_tokens, cache_warm_slug, cache_warm_at_ms,
-			context_version, context_fetched_at_ms, compaction_plan, compaction_plan_tokens, upgrade_deferred_tier, updated_at_ms
-		) VALUES ($key, $sessionId, $turn, $currentSlug, $currentTier, $stickyUntilTurn,
-			$lastPromptTokens, $cacheWarmSlug, $cacheWarmAtMs,
-			$contextVersion, $contextFetchedAtMs, $compactionPlan, $compactionPlanTokens, $upgradeDeferredTier, $updatedAtMs)
-		ON CONFLICT(key) DO UPDATE SET
-			session_id = excluded.session_id,
-			turn = excluded.turn,
-			current_slug = excluded.current_slug,
-			current_tier = excluded.current_tier,
-			sticky_until_turn = excluded.sticky_until_turn,
-			last_prompt_tokens = excluded.last_prompt_tokens,
-			cache_warm_slug = excluded.cache_warm_slug,
-			cache_warm_at_ms = excluded.cache_warm_at_ms,
-			context_version = excluded.context_version,
-			context_fetched_at_ms = excluded.context_fetched_at_ms,
-			compaction_plan = excluded.compaction_plan,
-			compaction_plan_tokens = excluded.compaction_plan_tokens,
-			upgrade_deferred_tier = excluded.upgrade_deferred_tier,
-			updated_at_ms = excluded.updated_at_ms
-	`);
-	// Read-modify-write in JS lost money: an aborted or failed dispatch is still
-	// billed by the upstream, but it returns before the commit path, so the next
-	// dispatch loaded a stale total and overwrote it. Measured on live data:
-	// 152 aborted dispatches billing $0.9985 — 30% of all spend — never reached
-	// `spent_usd`, leaving the per-conversation budget guard blind to it.
-	// Accumulating in SQL is correct regardless of who raced whom.
-	const accrueOne = db.query(`
-		UPDATE conversations
-		SET spent_usd = spent_usd + $spentUsd,
-			escalations = escalations + $escalations,
-			updated_at_ms = $updatedAtMs
-		WHERE key = $key
-	`);
-	const deleteStale: Statement<unknown, [number]> = db.query("DELETE FROM conversations WHERE updated_at_ms < ?");
+export function createConversationStore(db: SqlDb): ConversationStore {
+	const { sql } = db;
+	const selectOne = async (key: string): Promise<Row | null> =>
+		await db.one<Row>("SELECT * FROM conversations WHERE key = $key", { key });
 
 	return {
-		get(key) {
-			const row = selectOne.get(key);
+		async get(key) {
+			const row = await selectOne(key);
 			return row === null ? null : toState(row);
 		},
 
-		load(key) {
-			const existing = selectOne.get(key);
+		async load(key) {
+			const existing = await selectOne(key);
 			if (existing !== null) return toState(existing);
 			// Session id is derived, not random, so it stays stable if this row is
 			// ever pruned and the same conversation continues afterwards.
 			const sessionId = `omp-${key}`;
-			insertOne.run(key, sessionId, Date.now());
-			const inserted = selectOne.get(key);
+			// Two replicas can reach this at once for the same conversation; the
+			// loser must read the winner's row rather than fail the turn.
+			await sql`INSERT INTO conversations (key, session_id, updated_at_ms) VALUES (${key}, ${sessionId}, ${Date.now()})
+				ON CONFLICT (key) DO NOTHING`;
+			const inserted = await selectOne(key);
 			if (inserted === null) throw new Error(`conversation row vanished immediately after insert: ${key}`);
 			return toState(inserted);
 		},
 
-		save(state) {
-			// bun:sqlite matches named parameters by their literal `$name` key;
-			// bare keys bind nothing at all and every column silently lands NULL.
-			// No $spentUsd / $escalations here — see the statement above.
-			upsert.run({
-				$key: state.key,
-				$sessionId: state.sessionId,
-				$turn: state.turn,
-				$currentSlug: state.currentSlug,
-				$currentTier: state.currentTier,
-				$stickyUntilTurn: state.stickyUntilTurn,
-				$lastPromptTokens: state.lastPromptTokens,
-				$compactionPlan: state.compactionPlan === null ? null : JSON.stringify(state.compactionPlan),
-				$compactionPlanTokens: state.compactionPlanTokens ?? 0,
-				$upgradeDeferredTier: state.upgradeDeferredTier ?? null,
-				$cacheWarmSlug: state.cacheWarmSlug,
-				$cacheWarmAtMs: state.cacheWarmAtMs,
-				$contextVersion: state.contextVersion,
-				$contextFetchedAtMs: state.contextFetchedAtMs,
-				$updatedAtMs: Date.now(),
-			});
+		async save(state) {
+			// `spent_usd` and `escalations` are ABSENT on purpose: they accumulate
+			// through `accrue`, so writing a turn-start snapshot back here would
+			// erase whatever a billed-but-uncommitted dispatch added. The schema
+			// defaults both to 0, so the INSERT arm still works.
+			await sql`
+				INSERT INTO conversations (
+					key, session_id, turn, current_slug, current_tier, sticky_until_turn,
+					last_prompt_tokens, cache_warm_slug, cache_warm_at_ms,
+					context_version, context_fetched_at_ms, compaction_plan, compaction_plan_tokens,
+					upgrade_deferred_tier, updated_at_ms
+				) VALUES (
+					${state.key}, ${state.sessionId}, ${state.turn}, ${state.currentSlug}, ${state.currentTier},
+					${state.stickyUntilTurn}, ${state.lastPromptTokens}, ${state.cacheWarmSlug}, ${state.cacheWarmAtMs},
+					${state.contextVersion}, ${state.contextFetchedAtMs}, ${jsonParam(db, state.compactionPlan)},
+					${state.compactionPlanTokens ?? 0}, ${state.upgradeDeferredTier ?? null}, ${Date.now()}
+				)
+				ON CONFLICT (key) DO UPDATE SET
+					session_id = excluded.session_id,
+					turn = excluded.turn,
+					current_slug = excluded.current_slug,
+					current_tier = excluded.current_tier,
+					sticky_until_turn = excluded.sticky_until_turn,
+					last_prompt_tokens = excluded.last_prompt_tokens,
+					cache_warm_slug = excluded.cache_warm_slug,
+					cache_warm_at_ms = excluded.cache_warm_at_ms,
+					context_version = excluded.context_version,
+					context_fetched_at_ms = excluded.context_fetched_at_ms,
+					compaction_plan = excluded.compaction_plan,
+					compaction_plan_tokens = excluded.compaction_plan_tokens,
+					upgrade_deferred_tier = excluded.upgrade_deferred_tier,
+					updated_at_ms = excluded.updated_at_ms`;
 		},
 
-		accrue(key, delta) {
+		async accrue(key, delta) {
 			const spentUsd = delta.spentUsd ?? 0;
 			const escalations = delta.escalations ?? 0;
 			// Nothing to add: skip the write rather than bump updated_at_ms and
 			// keep a dead conversation alive against `prune`.
 			if (spentUsd === 0 && escalations === 0) return;
-			accrueOne.run({ $key: key, $spentUsd: spentUsd, $escalations: escalations, $updatedAtMs: Date.now() });
+			// Read-modify-write in JS lost money: an aborted or failed dispatch is
+			// still billed by the upstream, but it returns before the commit path,
+			// so the next dispatch loaded a stale total and overwrote it. Measured
+			// on live data: 152 aborted dispatches billing $0.9985 — 30% of all
+			// spend — never reached `spent_usd`. Accumulating in SQL is correct
+			// regardless of who raced whom, and with a shared store the racers can
+			// now be different processes.
+			await sql`
+				UPDATE conversations
+				SET spent_usd = spent_usd + ${spentUsd},
+					escalations = escalations + ${escalations},
+					updated_at_ms = ${Date.now()}
+				WHERE key = ${key}`;
 		},
 
-		prune(maxAgeMs) {
-			return deleteStale.run(Date.now() - maxAgeMs).changes;
+		async prune(maxAgeMs) {
+			const deleted = (await sql`DELETE FROM conversations WHERE updated_at_ms < ${Date.now() - maxAgeMs} RETURNING key`) as {
+				key: string;
+			}[];
+			return deleted.length;
 		},
 	};
 }

package/src/router/types.ts

@@ -226,16 +226,22 @@ export interface ConversationState {
 	updatedAtMs: number;
 }
 
+/**
+ * Every method is asynchronous because the store may be a shared database
+ * rather than a local file: replicas that cannot see each other's conversation
+ * state re-cold-start prompt caches for turns they did not serve, which at an
+ * observed 93% cache hit rate is a cost regression, not a latency one.
+ */
 export interface ConversationStore {
-	get(key: string): ConversationState | null;
+	get(key: string): Promise<ConversationState | null>;
 	/** Loads existing state or creates a fresh record. */
-	load(key: string): ConversationState;
+	load(key: string): Promise<ConversationState>;
 	/**
 	 * Persists the latest-wins fields. Deliberately does NOT write `spentUsd` or
 	 * `escalations` — those accumulate via `accrue`, and writing back a snapshot
 	 * here would clobber what a concurrent or already-billed dispatch added.
 	 */
-	save(state: ConversationState): void;
+	save(state: ConversationState): Promise<void>;
 	/**
 	 * Adds to the persisted counters, atomically in SQL.
 	 *
@@ -244,9 +250,9 @@ export interface ConversationStore {
 	 * reach the per-conversation budget guard anyway. Requires `load` to have
 	 * created the row.
 	 */
-	accrue(key: string, delta: { spentUsd?: number; escalations?: number }): void;
+	accrue(key: string, delta: { spentUsd?: number; escalations?: number }): Promise<void>;
 	/** Drops records untouched for longer than `maxAgeMs`. */
-	prune(maxAgeMs: number): number;
+	prune(maxAgeMs: number): Promise<number>;
 }
 
 /** Guarded-probe configuration for one dispatch. */

package/src/server/advise.ts

@@ -16,7 +16,7 @@
  */
 
 import type { RouterConfig } from "../config/types.ts";
-import type { Ledger } from "../cost/types.ts";
+import type { AsyncLedger } from "../cost/types.ts";
 import { scoreHeuristic } from "../router/classify.ts";
 import { extractFeatures } from "../router/features.ts";
 import { estimateTokens } from "../tokens/estimate.ts";
@@ -69,11 +69,13 @@ function requestOf(req: AdviseRequest): NormRequest {
 }
 
 /** Classifies a prompt the way the first turn of a conversation would be, without dispatching anything. */
-export function advise(cfg: RouterConfig, ledger: Ledger | null, req: AdviseRequest): Advice {
+export async function advise(cfg: RouterConfig, ledger: AsyncLedger | null, req: AdviseRequest): Promise<Advice> {
 	const norm = requestOf(req);
-	const features = extractFeatures(norm, estimateTokens(norm.promptBytes, "unknown", ledger));
+	// Advice is a hint for a client that has not dispatched yet: the default
+	// family ratio is enough, and it keeps this off the store entirely.
+	const features = extractFeatures(norm, estimateTokens(norm.promptBytes, "unknown", null));
 	const cls = scoreHeuristic(features, cfg);
-	const last = req.ompSessionId === "" ? null : (ledger?.latestForSession?.(req.ompSessionId)?.tier ?? null);
+	const last = req.ompSessionId === "" ? null : ((await ledger?.latestForSession(req.ompSessionId))?.tier ?? null);
 	return {
 		tier: cls.tier,
 		task: cls.task,

package/src/server/compaction-digest.ts

@@ -24,7 +24,7 @@ import type { DigestRequest, DigestResult } from "./digest.ts";
 export interface CompactionDigester {
 	digest(req: DigestRequest): Promise<DigestResult>;
 	/** See Digester.noteToolCalls; optional so a fake need not implement it. */
-	noteToolCalls?(ompSessionId: string, calls: readonly { name: string; argsJson: string }[], nowMs?: number): number;
+	noteToolCalls?(ompSessionId: string, calls: readonly { name: string; argsJson: string }[], nowMs?: number): Promise<number>;
 }
 
 export interface DigestCompactionArgs {

package/src/server/digest.ts

@@ -21,7 +21,7 @@
 import type { CatalogModel, CatalogSource } from "../catalog/types.ts";
 import type { DigestConfig, RouterConfig } from "../config/types.ts";
 import { computeCost, forecast } from "../cost/forecast.ts";
-import type { Ledger, LedgerEntry } from "../cost/types.ts";
+import type { AsyncLedger, LedgerEntry } from "../cost/types.ts";
 import { buildCandidates } from "../router/candidates.ts";
 import { primaryArg } from "../router/compaction.ts";
 import { extractFeatures } from "../router/features.ts";
@@ -57,7 +57,7 @@ export type DigestResult =
 export interface DigesterDeps {
 	cfg: RouterConfig;
 	catalog: CatalogSource;
-	ledger: Ledger;
+	ledger: AsyncLedger;
 	upstream: UpstreamClient;
 	log: Logger;
 }
@@ -144,7 +144,7 @@ export interface Digester {
 	 * back for the full output; that digest's ledger row is marked wasted and
 	 * the report shows the re-run rate. Returns how many were marked.
 	 */
-	noteToolCalls(ompSessionId: string, calls: readonly { name: string; argsJson: string }[], nowMs?: number): number;
+	noteToolCalls(ompSessionId: string, calls: readonly { name: string; argsJson: string }[], nowMs?: number): Promise<number>;
 }
 
 export function createDigester(deps: DigesterDeps): Digester {
@@ -163,7 +163,6 @@ export function createDigester(deps: DigesterDeps): Digester {
 				tier: cfg.digest.tier,
 				task: "documentation",
 				snapshot,
-				ledger,
 				cfg,
 				expectedCompletionTokens: cfg.digest.maxOutputTokens,
 				warmSlug: null,
@@ -179,14 +178,14 @@ export function createDigester(deps: DigesterDeps): Digester {
 		async digest(req) {
 			const inputBytes = Buffer.byteLength(req.content);
 			const source = req.source ?? "tool_result";
-			const currentTier = req.tier ?? ledger.latestForSession?.(req.ompSessionId)?.tier ?? null;
+			const currentTier = req.tier ?? (await ledger.latestForSession(req.ompSessionId))?.tier ?? null;
 			const gate = source === "compaction" ? { ...cfg.digest, enabled: cfg.compaction.digestToolResults } : cfg.digest;
 			const applies = digestApplies(gate, req.toolName, inputBytes, false, currentTier);
 			if (!applies.ok) return { digested: false, reason: applies.reason };
 
 			const promptText = `Task: ${req.query === "" ? "(unknown)" : req.query}\nTool: ${req.toolName} ${JSON.stringify(req.input)}\n--- output ---\n${req.content}`;
 			const synthetic = syntheticRequest(req, promptText);
-			const promptTokens = estimateTokens(synthetic.promptBytes, "unknown", ledger);
+			const promptTokens = estimateTokens(synthetic.promptBytes, "unknown", null);
 			const model = await pickModel(synthetic, promptTokens);
 			if (model === null) return { digested: false, reason: "no digest model available" };
 			const est = forecast(model, { promptTokens, completionTokens: cfg.digest.maxOutputTokens, cacheHitRate: 0, images: 0 });
@@ -222,7 +221,7 @@ export function createDigester(deps: DigesterDeps): Digester {
 				clearTimeout(timer);
 			}
 			const ms = Date.now() - startedAt;
-			const completionTokens = estimateTokens(Buffer.byteLength(text), model.tokenizer, ledger);
+			const completionTokens = estimateTokens(Buffer.byteLength(text), model.tokenizer, null);
 			const usage = { promptTokens, cachedTokens: 0, cacheWriteTokens: 0, completionTokens, reasoningTokens: 0, images: 0 };
 			const usd = costUsd ?? computeCost(model, usage).total;
 
@@ -264,7 +263,7 @@ export function createDigester(deps: DigesterDeps): Digester {
 				priceModel: model,
 			};
 			try {
-				ledger.record(entry);
+				await ledger.record(entry);
 			} catch (err) {
 				log.debug("digest ledger record failed", { error: err instanceof Error ? err.message : String(err) });
 			}
@@ -285,7 +284,7 @@ export function createDigester(deps: DigesterDeps): Digester {
 				ms,
 			};
 		},
-		noteToolCalls(ompSessionId, calls, nowMs = Date.now()) {
+		async noteToolCalls(ompSessionId, calls, nowMs = Date.now()) {
 			const list = recent.get(ompSessionId);
 			if (list === undefined || list.length === 0) return 0;
 			let marked = 0;
@@ -303,7 +302,7 @@ export function createDigester(deps: DigesterDeps): Digester {
 					d.rerun = true;
 					marked++;
 					try {
-						ledger.markWasted?.(d.ledgerId);
+						await ledger.markWasted(d.ledgerId);
 					} catch (err) {
 						log.debug("digest re-run mark failed", { error: err instanceof Error ? err.message : String(err) });
 					}