auto-model-router 0.2.11 → 0.2.12

package/.omp-plugin/marketplace.json

@@ -7,14 +7,14 @@
   },
   "metadata": {
     "description": "auto-model-router: a local cost/complexity-aware model router for Oh My Pi, backed by OpenRouter",
-    "version": "0.2.11",
+    "version": "0.2.12",
     "pluginRoot": "."
   },
   "plugins": [
     {
       "name": "auto-model-router",
       "description": "Local cost/complexity-aware model router for Oh My Pi, backed by OpenRouter. Runs in-process, routes per turn by price and task complexity, with budget caps, mid-stream escalation, and cache-aware hysteresis.",
-      "version": "0.2.11",
+      "version": "0.2.12",
       "author": {
         "name": "drewappling",
         "email": "drewappling@gmail.com"

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "auto-model-router",
-  "version": "0.2.11",
+  "version": "0.2.12",
   "private": false,
   "description": "Local cost/complexity-aware model router for Oh My Pi, backed by OpenRouter",
   "type": "module",

package/src/router/state.ts

@@ -57,13 +57,17 @@ export function createConversationStore(db: Database): ConversationStore {
 	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,
-			escalations, spent_usd, last_prompt_tokens, cache_warm_slug, cache_warm_at_ms,
+			last_prompt_tokens, cache_warm_slug, cache_warm_at_ms,
 			context_version, context_fetched_at_ms, updated_at_ms
 		) VALUES ($key, $sessionId, $turn, $currentSlug, $currentTier, $stickyUntilTurn,
-			$escalations, $spentUsd, $lastPromptTokens, $cacheWarmSlug, $cacheWarmAtMs,
+			$lastPromptTokens, $cacheWarmSlug, $cacheWarmAtMs,
 			$contextVersion, $contextFetchedAtMs, $updatedAtMs)
 		ON CONFLICT(key) DO UPDATE SET
 			session_id = excluded.session_id,
@@ -71,8 +75,6 @@ export function createConversationStore(db: Database): ConversationStore {
 			current_slug = excluded.current_slug,
 			current_tier = excluded.current_tier,
 			sticky_until_turn = excluded.sticky_until_turn,
-			escalations = excluded.escalations,
-			spent_usd = excluded.spent_usd,
 			last_prompt_tokens = excluded.last_prompt_tokens,
 			cache_warm_slug = excluded.cache_warm_slug,
 			cache_warm_at_ms = excluded.cache_warm_at_ms,
@@ -80,6 +82,19 @@ export function createConversationStore(db: Database): ConversationStore {
 			context_fetched_at_ms = excluded.context_fetched_at_ms,
 			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 < ?");
 
 	return {
@@ -103,6 +118,7 @@ export function createConversationStore(db: Database): ConversationStore {
 		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,
@@ -110,8 +126,6 @@ export function createConversationStore(db: Database): ConversationStore {
 				$currentSlug: state.currentSlug,
 				$currentTier: state.currentTier,
 				$stickyUntilTurn: state.stickyUntilTurn,
-				$escalations: state.escalations,
-				$spentUsd: state.spentUsd,
 				$lastPromptTokens: state.lastPromptTokens,
 				$cacheWarmSlug: state.cacheWarmSlug,
 				$cacheWarmAtMs: state.cacheWarmAtMs,
@@ -121,6 +135,15 @@ export function createConversationStore(db: Database): ConversationStore {
 			});
 		},
 
+		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() });
+		},
+
 		prune(maxAgeMs) {
 			return deleteStale.run(Date.now() - maxAgeMs).changes;
 		},

package/src/router/types.ts

@@ -179,7 +179,21 @@ export interface ConversationStore {
 	get(key: string): ConversationState | null;
 	/** Loads existing state or creates a fresh record. */
 	load(key: string): 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;
+	/**
+	 * Adds to the persisted counters, atomically in SQL.
+	 *
+	 * Separate from `save` because a dispatch that never reaches the commit path
+	 * — a client abort, an upstream error — was still BILLED, and its cost must
+	 * reach the per-conversation budget guard anyway. Requires `load` to have
+	 * created the row.
+	 */
+	accrue(key: string, delta: { spentUsd?: number; escalations?: number }): void;
 	/** Drops records untouched for longer than `maxAgeMs`. */
 	prune(maxAgeMs: number): number;
 }

package/src/server/turn.ts

@@ -230,6 +230,13 @@ export async function runTurn(
 				error: fields.error,
 				promptTokensSaved: decision.promptTokensSaved,
 			});
+			// Book the money HERE, beside the ledger row, so the two can never
+			// disagree. Every dispatch that reaches this point was billed —
+			// committed, wasted by an escalation, or aborted mid-stream — but only
+			// the committed path used to reach the state update below, so aborted
+			// dispatches (30% of real spend on live data) stayed invisible to the
+			// per-conversation budget guard.
+			conversations.accrue(req.conversationKey, { spentUsd: reportedUsd ?? decision.forecast.expectedUsd });
 		};
 
 		// "retry" re-enters the attempt loop; "done" means the turn is settled
@@ -422,7 +429,9 @@ export async function runTurn(
 		// hysteresis re-arm below can tell whether this turn changed tier.
 		const prevTier = state.currentTier;
 		state.currentTier = decision.tier;
-		state.escalations += escalations;
+		// Escalations accumulate in SQL for the same reason spend does: `save`
+		// below no longer writes this column, so a snapshot cannot clobber it.
+		conversations.accrue(req.conversationKey, { escalations });
 		// Hysteresis window. Only re-arm when the served tier actually changed
 		// (or this turn escalated). Re-arming on EVERY turn — even a trivial one
 		// served by a held hard model — extends the lock forever: the classifier
@@ -434,9 +443,11 @@ export async function runTurn(
 		if (tierChanged || escalations > 0) {
 			state.stickyUntilTurn = turnNumber + resolveHoldTurns(config, req.conversationKey, escalations > 0).turns;
 		}
-		// Reported cost is authoritative; fall back to the forecast so the
-		// budget guard still works when the provider omits cost.
+		// Spend is already booked in `writeEntry`, beside the ledger row, so it is
+		// deliberately NOT accumulated here — doing both would double-count.
+		// Keep the in-memory copy coherent for anything reading `state` later.
 		state.spentUsd += reportedUsd ?? decision.forecast.expectedUsd;
+		state.escalations += escalations;
 		state.lastPromptTokens = usage.promptTokens;
 		if (usage.cachedTokens > 0 || usage.cacheWriteTokens > 0) {
 			// Non-zero cache traffic is direct evidence the upstream cache exists.

package/test/failover.test.ts

@@ -275,6 +275,7 @@ function mkConversations(): { store: ConversationStore; map: Map<string, Convers
 		save: (s) => {
 			map.set(s.key, s);
 		},
+		accrue: () => {},
 		prune: () => 0,
 	};
 	return { store, map };

package/test/state.test.ts

@@ -0,0 +1,90 @@
+import { describe, expect, test } from "bun:test";
+
+import { createConversationStore } from "../src/router/state.ts";
+import { openDb } from "../src/util/sqlite.ts";
+
+function mkStore() {
+	const db = openDb(":memory:");
+	return { db, store: createConversationStore(db) };
+}
+
+describe("conversation spend accounting", () => {
+	test("accrue accumulates instead of overwriting", () => {
+		const { db, store } = mkStore();
+		try {
+			store.load("k1");
+			store.accrue("k1", { spentUsd: 0.25 });
+			store.accrue("k1", { spentUsd: 0.5 });
+			store.accrue("k1", { escalations: 1 });
+			store.accrue("k1", { escalations: 2 });
+
+			const state = store.get("k1");
+			expect(state?.spentUsd).toBeCloseTo(0.75, 10);
+			expect(state?.escalations).toBe(3);
+		} finally {
+			db.close();
+		}
+	});
+
+	test("save cannot clobber spend booked by a dispatch that never committed", () => {
+		// The live bug: a dispatch is billed by the upstream, aborts mid-stream,
+		// and returns before the commit path. The NEXT dispatch had already loaded
+		// a turn-start snapshot, and `save` wrote that snapshot's stale total back
+		// over the aborted dispatch's cost. 30% of real spend vanished this way.
+		const { db, store } = mkStore();
+		try {
+			const snapshot = store.load("k1");
+			expect(snapshot.spentUsd).toBe(0);
+
+			// An aborted dispatch books its cost while `snapshot` is still in hand.
+			store.accrue("k1", { spentUsd: 0.4, escalations: 1 });
+
+			// The in-flight turn now commits using the state it loaded earlier.
+			snapshot.turn = 1;
+			snapshot.currentSlug = "cheap/model";
+			store.save(snapshot);
+
+			const after = store.get("k1");
+			expect(after?.spentUsd).toBeCloseTo(0.4, 10);
+			expect(after?.escalations).toBe(1);
+			// The latest-wins fields still persist normally.
+			expect(after?.turn).toBe(1);
+			expect(after?.currentSlug).toBe("cheap/model");
+		} finally {
+			db.close();
+		}
+	});
+
+	test("interleaved dispatches both keep their money", () => {
+		const { db, store } = mkStore();
+		try {
+			const a = store.load("k1");
+			const b = store.get("k1");
+			expect(b).not.toBeNull();
+
+			store.accrue("k1", { spentUsd: 0.1 });
+			store.save(a);
+			store.accrue("k1", { spentUsd: 0.2 });
+			if (b !== null) store.save(b);
+
+			expect(store.get("k1")?.spentUsd).toBeCloseTo(0.3, 10);
+		} finally {
+			db.close();
+		}
+	});
+
+	test("a zero delta does not touch the row", () => {
+		// Bumping updated_at_ms for a no-op write would keep a dead conversation
+		// alive against `prune`, which reaps on that timestamp.
+		const { db, store } = mkStore();
+		try {
+			store.load("k1");
+			const before = store.get("k1")?.updatedAtMs ?? 0;
+			expect(before).toBeGreaterThan(0);
+			store.accrue("k1", { spentUsd: 0, escalations: 0 });
+			expect(store.get("k1")?.updatedAtMs).toBe(before);
+		} finally {
+			db.close();
+		}
+	});
+});

package/test/turn.test.ts

@@ -241,8 +241,14 @@ function mkLedger(): { ledger: Ledger; entries: LedgerEntry[] } {
 	return { ledger, entries };
 }
 
-function mkConversations(): { store: ConversationStore; map: Map<string, ConversationState> } {
+function mkConversations(): {
+	store: ConversationStore;
+	map: Map<string, ConversationState>;
+	accrued: Map<string, { spentUsd: number; escalations: number }>;
+} {
 	const map = new Map<string, ConversationState>();
+	// Mirrors the real store: money accumulates here, NOT through `save`.
+	const accrued = new Map<string, { spentUsd: number; escalations: number }>();
 	const store: ConversationStore = {
 		get: (k) => map.get(k) ?? null,
 		load: (k) => {
@@ -270,9 +276,15 @@ function mkConversations(): { store: ConversationStore; map: Map<string, Convers
 		save: (s) => {
 			map.set(s.key, s);
 		},
+		accrue: (k, d) => {
+			const cur = accrued.get(k) ?? { spentUsd: 0, escalations: 0 };
+			cur.spentUsd += d.spentUsd ?? 0;
+			cur.escalations += d.escalations ?? 0;
+			accrued.set(k, cur);
+		},
 		prune: () => 0,
 	};
-	return { store, map };
+	return { store, map, accrued };
 }
 
 function mkSink(): { sink: ResponseSink; chunks: UpstreamChunk[]; errors: WireError[]; finishes: TurnSummary[] } {
@@ -631,3 +643,54 @@ describe("agentdox write-back sees the shape of the turn", () => {
 		expect(records).toHaveLength(0);
 	});
 });
+
+describe("spend reaches the conversation total however the dispatch ends", () => {
+	test("a dispatch that dies mid-stream still books what it was billed", async () => {
+		// Live data: 152 aborted dispatches billed $0.9985 — 30% of all spend —
+		// and none of it reached the conversation's running total, because an abort
+		// returns before the commit path. The ledger row and the per-conversation
+		// budget guard must never disagree about money. Probe maxTokens 1 commits
+		// on the first token, so the retryable error below cannot re-enter the
+		// attempt loop and confuse the accounting.
+		const { router } = mkRouter([mkDecision("simple", "cheap/model", { maxTokens: 1, escalateTo: null })]);
+		const { upstream } = mkUpstream([
+			{
+				kind: "die",
+				chunks: [
+					startChunk("cheap/model"),
+					textChunk("plenty of text here, enough to commit on"),
+					usageChunk({ promptTokens: 47_700, completionTokens: 154 }, 0.0071),
+				],
+				error: new UpstreamError("network", 0, "request aborted", true),
+			},
+		]);
+		const { ledger, entries } = mkLedger();
+		const { store, accrued } = mkConversations();
+		const { sink } = mkSink();
+
+		await runTurn(mkReq(), sink, { config: mkConfig(), router, upstream, ledger, conversations: store, catalog, context: createDisabledBridge() }, new AbortController().signal);
+
+		// The turn never committed: the ledger row carries the error.
+		expect(entries).toHaveLength(1);
+		expect(entries[0]?.error).not.toBeNull();
+		// ...but the money was still booked.
+		expect(accrued.get("conv-test")?.spentUsd).toBeCloseTo(0.0071, 10);
+	});
+
+	test("a committed turn books its cost exactly once", async () => {
+		const { router } = mkRouter([mkDecision("simple", "cheap/model", { escalateTo: null })]);
+		const { upstream } = mkUpstream([
+			{ kind: "chunks", chunks: [startChunk("cheap/model"), textChunk("done"), finishChunk("stop"), usageChunk({}, 0.002)] },
+		]);
+		const { ledger, entries } = mkLedger();
+		const { store, accrued } = mkConversations();
+		const { sink, errors } = 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);
+		// Booked in writeEntry only — the commit path must not add it again.
+		expect(accrued.get("conv-test")?.spentUsd).toBeCloseTo(0.002, 10);
+	});
+});