@workweave/router 0.2.6 → 0.2.8

package/pi-router/src/compaction.ts

@@ -1,26 +1,150 @@
 /**
- * EXPERIMENTAL, off by default (enable with WEAVE_CHEAP_COMPACTION=1).
+ * Compaction safeguards for Pi's routed provider.
  *
- * Intent: run context compaction through the cheapest routing knobs. The
- * catch is that compaction bypasses provider hooks and reuses the session's
- * static provider headers (the main loop's quality knobs), and
- * `session_before_compact` is the only per-turn lever. Producing a fully
- * router-routed cheap CompactionResult here is not yet validated, so this
- * handler currently defers to pi's built-in compaction (returns no override)
- * and only reserves the flag + surfaces that it's active. Promote to a real
- * cheap path once validated end-to-end.
+ * Pi 0.74 checks automatic compaction only after an entire agent loop. A long
+ * tool loop can therefore cross the virtual model's context window between
+ * turns. Pi then clamps the next request to max_tokens=1, which looks exactly
+ * like an SDK quota probe to the router and produces a one-token response.
+ *
+ * Preserve real probes, but restore a usable budget for a clamped request that
+ * contains an actual tool result. Once the agent finishes, compact using the
+ * highest context reading seen during the run; routed providers tokenize and
+ * report caches differently, so the final response alone is not authoritative.
+ */
+
+import type {
+	AgentEndEvent,
+	BeforeProviderRequestEvent,
+	ExtensionAPI,
+	ExtensionContext,
+	SessionCompactEvent,
+	TurnEndEvent,
+} from "@mariozechner/pi-coding-agent";
+import type { AssistantMessage } from "@mariozechner/pi-ai";
+
+const PROBE_MAX_TOKENS = 4;
+const CONTINUATION_MAX_TOKENS = 16_384;
+const COMPACTION_RESERVE_TOKENS = 16_384;
+const STATUS_KEY = "weave-compaction";
+
+interface ProviderPayload {
+	max_tokens?: unknown;
+	messages?: unknown;
+	tools?: unknown;
+}
+
+interface ProviderMessage {
+	content?: unknown;
+}
+
+type Schedule = (callback: () => void) => unknown;
+
+function isRecord(value: unknown): value is Record<string, unknown> {
+	return typeof value === "object" && value !== null;
+}
+
+function containsToolResult(messages: unknown): boolean {
+	if (!Array.isArray(messages)) return false;
+	return messages.some((message: unknown) => {
+		if (!isRecord(message)) return false;
+		const content = (message as ProviderMessage).content;
+		return Array.isArray(content) && content.some((block: unknown) => isRecord(block) && block.type === "tool_result");
+	});
+}
+
+/**
+ * Repair only the Pi context-exhaustion shape. Genuine SDK probes have no tool
+ * result transcript and remain at max_tokens=1..4 for the router to hard-pin.
  */
+export function repairClampedToolContinuation(payload: unknown): boolean {
+	if (!isRecord(payload)) return false;
+	const body = payload as ProviderPayload;
+	if (typeof body.max_tokens !== "number" || body.max_tokens < 1 || body.max_tokens > PROBE_MAX_TOKENS) return false;
+	if (!Array.isArray(body.tools) || body.tools.length === 0) return false;
+	if (!containsToolResult(body.messages)) return false;
+	body.max_tokens = CONTINUATION_MAX_TOKENS;
+	return true;
+}
+
+function contextTokens(message: AssistantMessage): number {
+	const usage = message.usage;
+	if (usage.totalTokens > 0) return usage.totalTokens;
+	return usage.input + usage.output + usage.cacheRead + usage.cacheWrite;
+}
+
+function latestCompactionId(ctx: ExtensionContext): string | undefined {
+	const branch = ctx.sessionManager.getBranch();
+	for (let i = branch.length - 1; i >= 0; i--) {
+		const entry = branch[i];
+		if (entry?.type === "compaction") return entry.id;
+	}
+	return undefined;
+}
+
+export function registerCompaction(pi: ExtensionAPI, schedule: Schedule = (callback) => setTimeout(callback, 0)): void {
+	let highWaterTokens = 0;
+	let lastTurnTokens = 0;
+	let repairedContinuation = false;
+	let compactionScheduled = false;
+
+	const resetRun = () => {
+		highWaterTokens = 0;
+		lastTurnTokens = 0;
+		repairedContinuation = false;
+	};
+
+	const finishCompaction = (ctx: ExtensionContext) => {
+		compactionScheduled = false;
+		resetRun();
+		if (ctx.hasUI) ctx.ui.setStatus(STATUS_KEY, undefined);
+	};
+
+	pi.on("session_start", resetRun);
+	pi.on("agent_start", resetRun);
+	pi.on("session_compact", (_event: SessionCompactEvent, ctx: ExtensionContext) => finishCompaction(ctx));
+
+	pi.on("before_provider_request", (event: BeforeProviderRequestEvent) => {
+		if (repairClampedToolContinuation(event.payload)) repairedContinuation = true;
+	});
+
+	pi.on("turn_end", (event: TurnEndEvent) => {
+		if (event.message.role !== "assistant") return;
+		lastTurnTokens = contextTokens(event.message as AssistantMessage);
+		highWaterTokens = Math.max(highWaterTokens, lastTurnTokens);
+	});
+
+	pi.on("agent_end", (_event: AgentEndEvent, ctx: ExtensionContext) => {
+		if (process.env.WEAVE_PI_AUTO_COMPACTION === "0" || compactionScheduled) return;
+		const contextWindow = ctx.model?.contextWindow ?? ctx.getContextUsage()?.contextWindow ?? 0;
+		if (contextWindow <= COMPACTION_RESERVE_TOKENS) return;
+		const threshold = contextWindow - COMPACTION_RESERVE_TOKENS;
+		// Pi's built-in check runs immediately after this event and owns the
+		// ordinary final-turn threshold case. Starting another compaction while
+		// that async summary is in flight would race it.
+		if (lastTurnTokens > threshold) return;
+		if (!repairedContinuation && highWaterTokens <= threshold) return;
+
+		compactionScheduled = true;
+		const previousCompactionId = latestCompactionId(ctx);
 
-import type { ExtensionAPI, ExtensionContext, SessionBeforeCompactEvent } from "@mariozechner/pi-coding-agent";
-
-export function registerCheapCompaction(pi: ExtensionAPI): void {
-	let announced = false;
-	pi.on("session_before_compact", (_event: SessionBeforeCompactEvent, ctx: ExtensionContext) => {
-		if (ctx.hasUI && !announced) {
-			announced = true;
-			ctx.ui.setStatus("weave-compaction", "cheap compaction: experimental (using built-in)");
-		}
-		// Defer to built-in compaction until the router-routed cheap path is validated.
-		return undefined;
+		// Pi is still settling agent_end while extension handlers run. Schedule
+		// manual compaction for the next event-loop turn to avoid abort/wait
+		// recursion, and let Pi's own auto-compaction win if it already ran.
+		schedule(() => {
+			if (!compactionScheduled) return;
+			const currentCompactionId = latestCompactionId(ctx);
+			if (currentCompactionId !== previousCompactionId) {
+				finishCompaction(ctx);
+				return;
+			}
+			if (ctx.hasUI) ctx.ui.setStatus(STATUS_KEY, "compacting routed context...");
+			ctx.compact({
+				onComplete: () => finishCompaction(ctx),
+				onError: (error) => {
+					compactionScheduled = false;
+					if (ctx.hasUI) ctx.ui.setStatus(STATUS_KEY, `compaction failed: ${error.message}`);
+				},
+			});
+		});
 	});
 }

package/pi-router/src/config.ts

@@ -225,9 +225,10 @@ export const WEAVE_MODELS: ProviderModelConfig[] = [
 	model("claude-opus-4-7", "Claude Opus 4.7 (via Weave Router)", 64000),
 	model("claude-sonnet-4-6", "Claude Sonnet 4.6 (via Weave Router)", 64000),
 	model("claude-haiku-4-5", "Claude Haiku 4.5 (via Weave Router)", 32000),
+	model("grok-4.6", "Grok 4.6 (via Weave Router)", 131072, 500000),
 ];
 
-function model(id: string, name: string, maxTokens: number): ProviderModelConfig {
+function model(id: string, name: string, maxTokens: number, contextWindow: number = 200000): ProviderModelConfig {
 	return {
 		id,
 		name,
@@ -235,7 +236,7 @@ function model(id: string, name: string, maxTokens: number): ProviderModelConfig
 		input: ["text", "image"],
 		// Real cost is decided by the router per request and is unknown client-side.
 		cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 },
-		contextWindow: 200000,
+		contextWindow,
 		maxTokens,
 	};
 }
@@ -243,6 +244,8 @@ function model(id: string, name: string, maxTokens: number): ProviderModelConfig
 // ---------- dispatch / misc tunables ----------
 
 export const ROUTED_MODEL_HEADER = (process.env.WEAVE_ROUTED_MODEL_HEADER || "x-router-model").toLowerCase();
+export const ROUTED_PROVIDER_HEADER = "x-router-provider";
+export const ROUTER_DECISION_HEADER = "x-router-decision";
 /** Marker a headless child prints to stderr so the parent dispatch can read its routed model. */
 export const ROUTED_MODEL_STDERR_PREFIX = "weave-routed-model:";
 

package/pi-router/src/force-model.ts

@@ -0,0 +1,126 @@
+/**
+ * Pi consumes slash commands locally, so expose the router's force-model
+ * directives as extension commands and forward one canonical user turn.
+ *
+ * The router remains authoritative for aliases, validation, canonical model
+ * ids, and pin persistence. UI state is reconstructed from command/response
+ * pairs on the reachable Pi branch; this is necessary because Pi records the
+ * selected model handle on assistant messages rather than the response model.
+ */
+
+import type {
+	ExtensionAPI,
+	ExtensionCommandContext,
+	SessionEntry,
+} from "@mariozechner/pi-coding-agent";
+
+export type ForceModelTransition =
+	| { kind: "applied"; model: string }
+	| { kind: "cleared" }
+	| { kind: "noop" };
+
+type ForceModelDirective = "force" | "clear";
+
+function messageText(message: unknown): string | undefined {
+	if (!message || typeof message !== "object" || !("content" in message)) return undefined;
+	const content = message.content;
+	if (typeof content === "string") return content;
+	if (!Array.isArray(content)) return undefined;
+	return content
+		.filter(
+			(block): block is { type: "text"; text: string } =>
+				Boolean(
+					block &&
+						typeof block === "object" &&
+						"type" in block &&
+						block.type === "text" &&
+						"text" in block &&
+						typeof block.text === "string",
+				),
+		)
+		.map((block) => block.text)
+		.join("\n");
+}
+
+export function parseForceModelDirective(text: string): ForceModelDirective | undefined {
+	const command = text.trim();
+	if (/^\/(?:force-model|fm)\s+\S+/i.test(command)) return "force";
+	if (/^\/(?:unforce-model|ufm)$/i.test(command)) return "clear";
+	return undefined;
+}
+
+export function parseForceModelAcknowledgement(text: string): ForceModelTransition | undefined {
+	const applied = /force-model applied:\s+([^\s()]+)/i.exec(text);
+	if (applied) return { kind: "applied", model: applied[1] };
+	if (/force-model cleared/i.test(text)) return { kind: "cleared" };
+	if (/is(?:n't| not) a recognized model/i.test(text)) return { kind: "noop" };
+	return undefined;
+}
+
+function isSyntheticPinClear(text: string): boolean {
+	return /^(?:✦ \*\*Weave Router\*\* →|Weave Router:)\s+(?:Tool-call|Repetition|No-progress) loop detected\b[\s\S]*\bclearing the session pin\b/i.test(
+		text.trim(),
+	);
+}
+
+/** Reconstruct the effective router pin on the currently reachable branch. */
+export function forcedModelFromBranch(entries: readonly SessionEntry[]): string | undefined {
+	let forcedModel: string | undefined;
+	let pendingDirective: ForceModelDirective | undefined;
+
+	for (const entry of entries) {
+		if (entry.type !== "message") continue;
+		if (entry.message.role === "user") {
+			pendingDirective = parseForceModelDirective(messageText(entry.message) ?? "");
+			continue;
+		}
+		if (entry.message.role !== "assistant") continue;
+
+		const text = messageText(entry.message) ?? "";
+		if (isSyntheticPinClear(text)) {
+			forcedModel = undefined;
+			pendingDirective = undefined;
+			continue;
+		}
+		if (!pendingDirective) continue;
+
+		const transition = parseForceModelAcknowledgement(text);
+		if (transition?.kind === "applied" && pendingDirective === "force") forcedModel = transition.model;
+		else if (transition?.kind === "cleared" && pendingDirective === "clear") forcedModel = undefined;
+		// Rejected force-model commands intentionally retain the previous pin.
+		pendingDirective = undefined;
+	}
+
+	return forcedModel;
+}
+
+function sendRouterCommand(pi: ExtensionAPI, command: string, ctx: ExtensionCommandContext): void {
+	pi.sendUserMessage(command, ctx.isIdle() ? undefined : { deliverAs: "followUp" });
+}
+
+export function registerForceModelCommands(pi: ExtensionAPI): void {
+	const forceModel = async (args: string, ctx: ExtensionCommandContext): Promise<void> => {
+		const modelAndPrompt = args.trim();
+		if (!modelAndPrompt) {
+			ctx.ui.notify("Usage: /fm <model-id>", "warning");
+			return;
+		}
+		sendRouterCommand(pi, `/force-model ${modelAndPrompt}`, ctx);
+	};
+	const clearForceModel = async (_args: string, ctx: ExtensionCommandContext): Promise<void> => {
+		sendRouterCommand(pi, "/unforce-model", ctx);
+	};
+
+	for (const name of ["fm", "force-model"]) {
+		pi.registerCommand(name, {
+			description: "Pin this session to a specific model via the Weave Router",
+			handler: forceModel,
+		});
+	}
+	for (const name of ["ufm", "unforce-model"]) {
+		pi.registerCommand(name, {
+			description: "Clear this session's forced Weave Router model",
+			handler: clearForceModel,
+		});
+	}
+}

package/pi-router/src/index.ts

@@ -7,9 +7,9 @@
  *                   the main loop, speed/cheap in subagents).
  *   - metadata:     stamp body.metadata.user_id for sticky sessions + subagent
  *                   detection.
- *   - routed-model: show which model the router actually picked.
+ *   - Loom UI:      branded header, Wooly animation, actual route, and saved $.
  *   - safety:       block catastrophic bash (unless WEAVE_NO_SAFETY=1).
- *   - compaction:   experimental cheap path (only when WEAVE_CHEAP_COMPACTION=1).
+ *   - compaction:   protect long tool loops, then compact routed context.
  *   - dispatch:     parallel, context-isolated subagents — top-level process
  *                   only (no grandchildren).
  *
@@ -20,8 +20,9 @@
 import { fileURLToPath } from "node:url";
 import type { ExtensionAPI } from "@mariozechner/pi-coding-agent";
 import { isSubagent } from "./config.js";
-import { registerCheapCompaction } from "./compaction.js";
+import { registerCompaction } from "./compaction.js";
 import { registerDispatch } from "./dispatch.js";
+import { registerForceModelCommands } from "./force-model.js";
 import { registerMetadata } from "./metadata.js";
 import { registerRoutedModel } from "./routed-model.js";
 import { registerSafety } from "./safety.js";
@@ -37,10 +38,11 @@ export default function (pi: ExtensionAPI): void {
 	pi.on("session_start", () => registerWeave(pi));
 
 	registerMetadata(pi);
+	registerForceModelCommands(pi);
 	registerRoutedModel(pi);
+	registerCompaction(pi);
 
 	if (process.env.WEAVE_NO_SAFETY !== "1") registerSafety(pi);
-	if (process.env.WEAVE_CHEAP_COMPACTION === "1") registerCheapCompaction(pi);
 
 	// Only the top-level process fans out. Children (WEAVE_PI_SUBAGENT=1) load
 	// this same extension but get no dispatch tool, so subagents can't spawn

package/pi-router/src/pricing.generated.ts

@@ -0,0 +1,81 @@
+// Code generated by cmd/genprices; DO NOT EDIT.
+// Source: internal/router/catalog (USD per 1M tokens).
+
+export interface ModelPricing {
+	inputUsdPerMillion: number;
+	outputUsdPerMillion: number;
+}
+
+export const PRICING_VERSION = "catalog-sha256:f2841a252235d113";
+
+export const MODEL_PRICING: Readonly<Record<string, ModelPricing>> = Object.freeze({
+	"claude-fable-5": { inputUsdPerMillion: 10, outputUsdPerMillion: 50 },
+	"claude-haiku-4-5": { inputUsdPerMillion: 1, outputUsdPerMillion: 5 },
+	"claude-opus-4-0": { inputUsdPerMillion: 15, outputUsdPerMillion: 75 },
+	"claude-opus-4-1": { inputUsdPerMillion: 15, outputUsdPerMillion: 75 },
+	"claude-opus-4-5": { inputUsdPerMillion: 5, outputUsdPerMillion: 25 },
+	"claude-opus-4-6": { inputUsdPerMillion: 5, outputUsdPerMillion: 25 },
+	"claude-opus-4-7": { inputUsdPerMillion: 5, outputUsdPerMillion: 25 },
+	"claude-opus-4-8": { inputUsdPerMillion: 5, outputUsdPerMillion: 25 },
+	"claude-opus-5": { inputUsdPerMillion: 5, outputUsdPerMillion: 25 },
+	"claude-sonnet-4-5": { inputUsdPerMillion: 3, outputUsdPerMillion: 15 },
+	"claude-sonnet-4-6": { inputUsdPerMillion: 3, outputUsdPerMillion: 15 },
+	"claude-sonnet-5": { inputUsdPerMillion: 3, outputUsdPerMillion: 15 },
+	"deepseek/deepseek-v4-flash": { inputUsdPerMillion: 0.1134, outputUsdPerMillion: 0.2791 },
+	"deepseek/deepseek-v4-pro": { inputUsdPerMillion: 1.74, outputUsdPerMillion: 3.48 },
+	"gemini-2.0-flash": { inputUsdPerMillion: 0.1, outputUsdPerMillion: 0.4 },
+	"gemini-2.0-flash-lite": { inputUsdPerMillion: 0.075, outputUsdPerMillion: 0.3 },
+	"gemini-2.5-flash": { inputUsdPerMillion: 0.3, outputUsdPerMillion: 1.2 },
+	"gemini-2.5-flash-lite": { inputUsdPerMillion: 0.1, outputUsdPerMillion: 0.4 },
+	"gemini-2.5-pro": { inputUsdPerMillion: 1.25, outputUsdPerMillion: 5 },
+	"gemini-3-flash-preview": { inputUsdPerMillion: 0.5, outputUsdPerMillion: 2 },
+	"gemini-3-pro-preview": { inputUsdPerMillion: 2, outputUsdPerMillion: 8 },
+	"gemini-3.1-flash-lite-preview": { inputUsdPerMillion: 0.1, outputUsdPerMillion: 0.4 },
+	"gemini-3.1-pro-preview": { inputUsdPerMillion: 2, outputUsdPerMillion: 8 },
+	"gemini-3.5-flash": { inputUsdPerMillion: 1.5, outputUsdPerMillion: 9 },
+	"gemini-3.5-flash-lite": { inputUsdPerMillion: 0.3, outputUsdPerMillion: 2.5 },
+	"gemini-3.6-flash": { inputUsdPerMillion: 1.5, outputUsdPerMillion: 7.5 },
+	"google/gemma-4-26b-a4b-it": { inputUsdPerMillion: 0.15, outputUsdPerMillion: 0.6 },
+	"gpt-4.1": { inputUsdPerMillion: 2, outputUsdPerMillion: 8 },
+	"gpt-4.1-mini": { inputUsdPerMillion: 0.4, outputUsdPerMillion: 1.6 },
+	"gpt-4.1-nano": { inputUsdPerMillion: 0.1, outputUsdPerMillion: 0.4 },
+	"gpt-4o": { inputUsdPerMillion: 2.5, outputUsdPerMillion: 10 },
+	"gpt-4o-mini": { inputUsdPerMillion: 0.15, outputUsdPerMillion: 0.6 },
+	"gpt-5": { inputUsdPerMillion: 2.5, outputUsdPerMillion: 10 },
+	"gpt-5-chat": { inputUsdPerMillion: 2.5, outputUsdPerMillion: 10 },
+	"gpt-5-mini": { inputUsdPerMillion: 0.5, outputUsdPerMillion: 2 },
+	"gpt-5-nano": { inputUsdPerMillion: 0.1, outputUsdPerMillion: 0.4 },
+	"gpt-5.4": { inputUsdPerMillion: 2.5, outputUsdPerMillion: 15 },
+	"gpt-5.4-mini": { inputUsdPerMillion: 0.75, outputUsdPerMillion: 4.5 },
+	"gpt-5.4-nano": { inputUsdPerMillion: 0.2, outputUsdPerMillion: 1.25 },
+	"gpt-5.4-pro": { inputUsdPerMillion: 30, outputUsdPerMillion: 180 },
+	"gpt-5.5": { inputUsdPerMillion: 5, outputUsdPerMillion: 30 },
+	"gpt-5.5-mini": { inputUsdPerMillion: 0.5, outputUsdPerMillion: 2.5 },
+	"gpt-5.5-nano": { inputUsdPerMillion: 0.15, outputUsdPerMillion: 0.6 },
+	"gpt-5.5-pro": { inputUsdPerMillion: 30, outputUsdPerMillion: 180 },
+	"gpt-5.6-luna": { inputUsdPerMillion: 1, outputUsdPerMillion: 6 },
+	"gpt-5.6-sol": { inputUsdPerMillion: 5, outputUsdPerMillion: 30 },
+	"gpt-5.6-terra": { inputUsdPerMillion: 2.5, outputUsdPerMillion: 15 },
+	"grok-4.5": { inputUsdPerMillion: 2, outputUsdPerMillion: 6 },
+	"grok-4.6": { inputUsdPerMillion: 2, outputUsdPerMillion: 6 },
+	"minimax/minimax-m2.7": { inputUsdPerMillion: 0.3, outputUsdPerMillion: 1.2 },
+	"minimax/minimax-m3": { inputUsdPerMillion: 0.3, outputUsdPerMillion: 1.2 },
+	"mistralai/mistral-small-2603": { inputUsdPerMillion: 0.2, outputUsdPerMillion: 0.6 },
+	"moonshotai/kimi-k2.5": { inputUsdPerMillion: 0.6, outputUsdPerMillion: 3 },
+	"moonshotai/kimi-k2.6": { inputUsdPerMillion: 0.95, outputUsdPerMillion: 4 },
+	"moonshotai/kimi-k2.7": { inputUsdPerMillion: 0.95, outputUsdPerMillion: 4 },
+	"moonshotai/kimi-k3": { inputUsdPerMillion: 3, outputUsdPerMillion: 15 },
+	"qwen/qwen3-235b-a22b-2507": { inputUsdPerMillion: 0.2266, outputUsdPerMillion: 0.9064 },
+	"qwen/qwen3-30b-a3b-instruct-2507": { inputUsdPerMillion: 0.15, outputUsdPerMillion: 0.6 },
+	"qwen/qwen3-coder": { inputUsdPerMillion: 0.9, outputUsdPerMillion: 2.7 },
+	"qwen/qwen3-coder-next": { inputUsdPerMillion: 0.5, outputUsdPerMillion: 1.2 },
+	"qwen/qwen3-next-80b-a3b-instruct": { inputUsdPerMillion: 0.15, outputUsdPerMillion: 1.2 },
+	"qwen/qwen3.5-flash-02-23": { inputUsdPerMillion: 0.05, outputUsdPerMillion: 0.15 },
+	"qwen/qwen3.6-35b-a3b": { inputUsdPerMillion: 0.15, outputUsdPerMillion: 1 },
+	"qwen/qwen3.7-plus": { inputUsdPerMillion: 0.4, outputUsdPerMillion: 1.6 },
+	"qwen/qwen3.8-max": { inputUsdPerMillion: 2, outputUsdPerMillion: 6 },
+	"xiaomi/mimo-v2.5-pro": { inputUsdPerMillion: 1, outputUsdPerMillion: 3 },
+	"z-ai/glm-5": { inputUsdPerMillion: 1, outputUsdPerMillion: 3.2 },
+	"z-ai/glm-5.1": { inputUsdPerMillion: 1.4, outputUsdPerMillion: 4.4 },
+	"z-ai/glm-5.2": { inputUsdPerMillion: 1.4, outputUsdPerMillion: 4.4 },
+});

package/pi-router/src/routed-model.ts

@@ -1,31 +1,165 @@
 /**
- * Surfaces which model the router actually picked for each request.
+ * Route attribution, session savings, and the interactive Loom presentation.
  *
- * The router sets `x-router-model` on every response (streaming, non-streaming,
- * and cache hits). In the interactive UI we show it in the status bar and
- * notify on change. In a headless child (print/RPC — e.g. a dispatch subagent)
- * there is no UI, so we print a marker to stderr that the parent dispatch tool
- * parses to attribute each subagent's work to a model.
+ * The selected Pi model is only the comparison baseline. The router's response
+ * headers are authoritative for the model that actually served each response.
+ * We pair those headers with Pi's finalized turn usage, persist an audit entry,
+ * and rebuild the reachable total whenever a session resumes or changes branch.
  */
 
+import type { AssistantMessage } from "@mariozechner/pi-ai";
 import type { ExtensionAPI, ExtensionContext } from "@mariozechner/pi-coding-agent";
-import { ROUTED_MODEL_HEADER, ROUTED_MODEL_STDERR_PREFIX } from "./config.js";
+import {
+	isSubagent,
+	ROUTED_MODEL_HEADER,
+	ROUTED_MODEL_STDERR_PREFIX,
+	ROUTED_PROVIDER_HEADER,
+	ROUTER_DECISION_HEADER,
+} from "./config.js";
+import { forcedModelFromBranch } from "./force-model.js";
+import {
+	aggregateSavings,
+	createSavingsEntry,
+	isSavingsEntryData,
+	normalizeModelId,
+	SAVINGS_ENTRY_TYPE,
+	type RouteDecision,
+	type SavingsAggregate,
+	type SavingsEntryData,
+} from "./savings.js";
+import { clearLoomUi, installLoomUi, updateRouterStatus } from "./ui.js";
 
-const STATUS_KEY = "weave";
+interface PendingRoute {
+	requestedModel?: string;
+	routedModel: string;
+	provider?: string;
+	decision?: string;
+}
+
+function savingsFromBranch(ctx: ExtensionContext): { entries: SavingsEntryData[]; aggregate: SavingsAggregate } {
+	const entries: SavingsEntryData[] = [];
+	for (const entry of ctx.sessionManager.getBranch()) {
+		if (entry.type !== "custom" || entry.customType !== SAVINGS_ENTRY_TYPE || !isSavingsEntryData(entry.data)) continue;
+		entries.push(entry.data);
+	}
+	return { entries, aggregate: aggregateSavings(entries) };
+}
+
+function messageUsage(message: AssistantMessage) {
+	return {
+		input: message.usage.input,
+		output: message.usage.output,
+		cacheRead: message.usage.cacheRead,
+		cacheWrite: message.usage.cacheWrite,
+	};
+}
 
 export function registerRoutedModel(pi: ExtensionAPI): void {
-	let last: string | undefined;
+	let pendingRoutes: PendingRoute[] = [];
+	let entries: SavingsEntryData[] = [];
+	let savings = aggregateSavings(entries);
+	let requestedModel: string | undefined;
+	let routedModel: string | undefined;
+	let forcedModel: string | undefined;
+	let lastNotifiedModel: string | undefined;
+
+	const refresh = (ctx: ExtensionContext) => {
+		if (isSubagent()) return;
+		updateRouterStatus(ctx, { requestedModel, routedModel, forcedModel, savings });
+	};
+
+	const restore = (ctx: ExtensionContext) => {
+		const restored = savingsFromBranch(ctx);
+		entries = restored.entries;
+		savings = restored.aggregate;
+		const lastEntry = restored.aggregate.lastEntry;
+		requestedModel = ctx.model?.id ?? lastEntry?.requestedModel;
+		routedModel =
+			lastEntry && requestedModel && normalizeModelId(requestedModel) === lastEntry.requestedModel
+				? lastEntry.routedModel
+				: undefined;
+		forcedModel = forcedModelFromBranch(ctx.sessionManager.getBranch());
+		pendingRoutes = [];
+		lastNotifiedModel = undefined;
+	};
+
+	pi.on("session_start", (_event, ctx: ExtensionContext) => {
+		restore(ctx);
+		if (!isSubagent()) installLoomUi(ctx);
+		refresh(ctx);
+	});
+
+	pi.on("model_select", (event, ctx: ExtensionContext) => {
+		if (isSubagent()) return;
+		requestedModel = event.model.id;
+		routedModel = undefined;
+		refresh(ctx);
+	});
 
 	pi.on("after_provider_response", (event, ctx: ExtensionContext) => {
+		if (event.status < 200 || event.status >= 300) return;
 		const model = event.headers?.[ROUTED_MODEL_HEADER];
-		if (!model || model === last) return;
-		last = model;
-
-		if (ctx.hasUI) {
-			ctx.ui.setStatus(STATUS_KEY, `routed: ${model}`);
-			ctx.ui.notify(`Weave Router routed to ${model}`, "info");
-		} else {
-			process.stderr.write(`${ROUTED_MODEL_STDERR_PREFIX} ${model}\n`);
+		if (!model) return;
+		const route: PendingRoute = {
+			...(ctx.model?.id ? { requestedModel: ctx.model.id } : {}),
+			routedModel: normalizeModelId(model),
+			...(event.headers[ROUTED_PROVIDER_HEADER] ? { provider: event.headers[ROUTED_PROVIDER_HEADER] } : {}),
+			...(event.headers[ROUTER_DECISION_HEADER] ? { decision: event.headers[ROUTER_DECISION_HEADER] } : {}),
+		};
+		if (!isSubagent()) pendingRoutes.push(route);
+
+		if (!ctx.hasUI || isSubagent()) {
+			if (route.routedModel !== lastNotifiedModel) {
+				process.stderr.write(`${ROUTED_MODEL_STDERR_PREFIX} ${route.routedModel}\n`);
+				lastNotifiedModel = route.routedModel;
+			}
+			return;
+		}
+
+		requestedModel = route.requestedModel ?? requestedModel;
+		routedModel = route.routedModel;
+		refresh(ctx);
+		if (route.routedModel !== lastNotifiedModel) {
+			ctx.ui.notify(`Weave Router routed to ${route.routedModel}`, "info");
+			lastNotifiedModel = route.routedModel;
+		}
+	});
+
+	pi.on("turn_end", (event, ctx: ExtensionContext) => {
+		if (isSubagent() || event.message.role !== "assistant") return;
+		const restoredForcedModel = forcedModelFromBranch(ctx.sessionManager.getBranch());
+		if (restoredForcedModel !== forcedModel) {
+			forcedModel = restoredForcedModel;
+			refresh(ctx);
 		}
+		const pending = pendingRoutes.shift();
+		if (!pending) return;
+		const message = event.message as AssistantMessage;
+		const selected = pending.requestedModel || message.model || ctx.model?.id;
+		if (!selected) return;
+		const decision: RouteDecision = {
+			requestedModel: selected,
+			routedModel: pending.routedModel,
+			...(pending.provider ? { provider: pending.provider } : {}),
+			...(pending.decision ? { decision: pending.decision } : {}),
+		};
+		const entry = createSavingsEntry(decision, messageUsage(message));
+		entries.push(entry);
+		savings = aggregateSavings(entries);
+		requestedModel = entry.requestedModel;
+		routedModel = entry.routedModel;
+		pi.appendEntry(SAVINGS_ENTRY_TYPE, entry);
+		refresh(ctx);
+	});
+
+	pi.on("session_tree", (_event, ctx: ExtensionContext) => {
+		if (isSubagent()) return;
+		restore(ctx);
+		refresh(ctx);
+	});
+
+	pi.on("session_shutdown", (_event, ctx: ExtensionContext) => {
+		pendingRoutes = [];
+		if (!isSubagent()) clearLoomUi(ctx);
 	});
 }