@workweave/router 0.2.6 → 0.2.7

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);
 	});
 }

package/pi-router/src/savings.ts

@@ -0,0 +1,191 @@
+import { MODEL_PRICING, PRICING_VERSION } from "./pricing.generated.js";
+
+export const SAVINGS_ENTRY_TYPE = "weave-router-savings-v1";
+
+export interface TokenUsage {
+	input: number;
+	output: number;
+	cacheRead: number;
+	cacheWrite: number;
+}
+
+export interface RouteDecision {
+	requestedModel: string;
+	routedModel: string;
+	provider?: string;
+	decision?: string;
+}
+
+export interface SavingsEntryData {
+	version: 1;
+	pricingVersion: string;
+	requestedModel: string;
+	routedModel: string;
+	provider?: string;
+	decision?: string;
+	usage: TokenUsage;
+	requestedCostUsd?: number;
+	routedCostUsd?: number;
+	savingsUsd?: number;
+	priced: boolean;
+	unpricedModels: string[];
+}
+
+export interface SavingsAggregate {
+	totalSavingsUsd: number;
+	pricedResponses: number;
+	unpricedResponses: number;
+	lastEntry?: SavingsEntryData;
+}
+
+function finiteNonNegative(value: number): number | undefined {
+	return Number.isFinite(value) && value >= 0 ? value : undefined;
+}
+
+export function normalizeModelId(model: string): string {
+	return model.trim().replace(/^weave\//, "").replace(/\[[^\]]*\]$/, "").replace(/-[0-9]{8}$/, "");
+}
+
+function normalizedUsage(usage: TokenUsage): TokenUsage | undefined {
+	const input = finiteNonNegative(usage.input);
+	const output = finiteNonNegative(usage.output);
+	const cacheRead = finiteNonNegative(usage.cacheRead);
+	const cacheWrite = finiteNonNegative(usage.cacheWrite);
+	if (input === undefined || output === undefined || cacheRead === undefined || cacheWrite === undefined) {
+		return undefined;
+	}
+	return { input, output, cacheRead, cacheWrite };
+}
+
+function modelCostUsd(model: string, usage: TokenUsage): number | undefined {
+	const price = MODEL_PRICING[normalizeModelId(model)];
+	if (!price) return undefined;
+	const inputTokens = usage.input + 1.25 * usage.cacheWrite + 0.1 * usage.cacheRead;
+	return (inputTokens * price.inputUsdPerMillion + usage.output * price.outputUsdPerMillion) / 1_000_000;
+}
+
+export function createSavingsEntry(decision: RouteDecision, rawUsage: TokenUsage): SavingsEntryData {
+	const requestedModel = normalizeModelId(decision.requestedModel);
+	const routedModel = normalizeModelId(decision.routedModel);
+	const usage = normalizedUsage(rawUsage);
+	const base: SavingsEntryData = {
+		version: 1,
+		pricingVersion: PRICING_VERSION,
+		requestedModel,
+		routedModel,
+		...(decision.provider ? { provider: decision.provider } : {}),
+		...(decision.decision ? { decision: decision.decision } : {}),
+		usage: usage ?? { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 },
+		priced: false,
+		unpricedModels: [],
+	};
+
+	if (!usage || !requestedModel || !routedModel) return base;
+
+	// An unchanged route has an exact zero delta even if the catalog does not
+	// know the model. Avoid calling a no-op response "unpriced" when no price
+	// comparison is necessary.
+	if (requestedModel === routedModel) {
+		const cost = modelCostUsd(requestedModel, usage);
+		return {
+			...base,
+			...(cost === undefined ? {} : { requestedCostUsd: cost, routedCostUsd: cost }),
+			savingsUsd: 0,
+			priced: true,
+		};
+	}
+
+	const requestedCostUsd = modelCostUsd(requestedModel, usage);
+	const routedCostUsd = modelCostUsd(routedModel, usage);
+	if (requestedCostUsd === undefined || routedCostUsd === undefined) {
+		const unpricedModels = [
+			...(requestedCostUsd === undefined ? [requestedModel] : []),
+			...(routedCostUsd === undefined ? [routedModel] : []),
+		];
+		return { ...base, unpricedModels };
+	}
+
+	return {
+		...base,
+		requestedCostUsd,
+		routedCostUsd,
+		savingsUsd: requestedCostUsd - routedCostUsd,
+		priced: true,
+	};
+}
+
+function isRecord(value: unknown): value is Record<string, unknown> {
+	return typeof value === "object" && value !== null;
+}
+
+function isTokenUsage(value: unknown): value is TokenUsage {
+	if (!isRecord(value)) return false;
+	return [value.input, value.output, value.cacheRead, value.cacheWrite].every(
+		(token) => typeof token === "number" && Number.isFinite(token) && token >= 0,
+	);
+}
+
+export function isSavingsEntryData(value: unknown): value is SavingsEntryData {
+	if (!isRecord(value)) return false;
+	return (
+		value.version === 1 &&
+		typeof value.pricingVersion === "string" &&
+		typeof value.requestedModel === "string" &&
+		typeof value.routedModel === "string" &&
+		isTokenUsage(value.usage) &&
+		typeof value.priced === "boolean" &&
+		(value.provider === undefined || typeof value.provider === "string") &&
+		(value.decision === undefined || typeof value.decision === "string") &&
+		(value.requestedCostUsd === undefined ||
+			(typeof value.requestedCostUsd === "number" && Number.isFinite(value.requestedCostUsd))) &&
+		(value.routedCostUsd === undefined ||
+			(typeof value.routedCostUsd === "number" && Number.isFinite(value.routedCostUsd))) &&
+		Array.isArray(value.unpricedModels) &&
+		value.unpricedModels.every((model) => typeof model === "string") &&
+		(value.savingsUsd === undefined || (typeof value.savingsUsd === "number" && Number.isFinite(value.savingsUsd)))
+	);
+}
+
+export function aggregateSavings(entries: Iterable<SavingsEntryData>): SavingsAggregate {
+	let totalSavingsUsd = 0;
+	let pricedResponses = 0;
+	let unpricedResponses = 0;
+	let lastEntry: SavingsEntryData | undefined;
+	for (const entry of entries) {
+		lastEntry = entry;
+		if (entry.priced && entry.savingsUsd !== undefined) {
+			totalSavingsUsd += entry.savingsUsd;
+			pricedResponses++;
+		} else {
+			unpricedResponses++;
+		}
+	}
+	return {
+		totalSavingsUsd,
+		pricedResponses,
+		unpricedResponses,
+		...(lastEntry ? { lastEntry } : {}),
+	};
+}
+
+export function formatMoney(amount: number): string {
+	const absolute = Math.abs(amount);
+	if (absolute > 0 && absolute < 0.005) return "<$0.01";
+	return `$${absolute.toFixed(2)}`;
+}
+
+export function formatSavings(aggregate: SavingsAggregate): string {
+	let clause: string;
+	if (aggregate.pricedResponses === 0) {
+		clause = "saved —";
+	} else if (aggregate.totalSavingsUsd < 0) {
+		clause = `extra ${formatMoney(aggregate.totalSavingsUsd)}`;
+	} else {
+		clause = `saved ${formatMoney(aggregate.totalSavingsUsd)}`;
+	}
+	if (aggregate.unpricedResponses > 0) {
+		const suffix = aggregate.unpricedResponses === 1 ? "1 unpriced" : `${aggregate.unpricedResponses} unpriced`;
+		return `${clause} · ${suffix}`;
+	}
+	return clause;
+}

package/pi-router/src/ui.ts

@@ -0,0 +1,80 @@
+import type { ExtensionContext, Theme } from "@mariozechner/pi-coding-agent";
+import type { TUI } from "@mariozechner/pi-tui";
+import type { SavingsAggregate } from "./savings.js";
+import { formatSavings } from "./savings.js";
+import { WoolyComponent } from "./wooly.js";
+
+const STATUS_KEY = "weave-router";
+const WOOLY_WIDGET_KEY = "weave-wooly";
+const BRAND_OPEN = "\x1b[38;2;255;108;71m";
+const BRAND_CLOSE = "\x1b[39m";
+
+const WEAVE_WORDMARK = ["╦ ╦╔═╗╔═╗╦  ╦╔═╗", "║║║║╣ ╠═╣╚╗╔╝║╣ ", "╚╩╝╚═╝╩ ╩ ╚╝ ╚═╝"] as const;
+const LOOM_WORDMARK = [
+	"██╗    ███████╗ ███████╗ ███╗   ███╗",
+	"██║    ██╔═══██╗██╔═══██╗████╗ ████║",
+	"██║    ██║   ██║██║   ██║██╔████╔██║",
+	"██║    ██║   ██║██║   ██║██║╚██╔╝██║",
+	"██████╗╚██████╔╝╚██████╔╝██║ ╚═╝ ██║",
+	"╚═════╝ ╚═════╝  ╚═════╝ ╚═╝     ╚═╝",
+] as const;
+
+export interface RouterUiSnapshot {
+	requestedModel?: string;
+	routedModel?: string;
+	forcedModel?: string;
+	savings: SavingsAggregate;
+}
+
+function brand(text: string): string {
+	return `${BRAND_OPEN}${text}${BRAND_CLOSE}`;
+}
+
+function headerLines(theme: Theme, width: number): string[] {
+	const weave = WEAVE_WORDMARK.map((line) => theme.bold(brand(line)));
+	if (width < 44) {
+		return ["", ...weave, theme.bold("LOOM"), theme.fg("dim", "Weave Router · Loom for Pi"), ""];
+	}
+	return [
+		"",
+		...weave,
+		...LOOM_WORDMARK.map((line) => theme.bold(line)),
+		theme.bold("Weave Router · Loom for Pi"),
+		"",
+	];
+}
+
+export function installLoomUi(ctx: ExtensionContext): void {
+	if (ctx.mode !== "tui") return;
+	ctx.ui.setTitle("Loom · Weave Router");
+	ctx.ui.setHeader((_tui: TUI, theme: Theme) => ({
+		invalidate() {},
+		render(width: number): string[] {
+			return headerLines(theme, width);
+		},
+	}));
+	ctx.ui.setWidget(WOOLY_WIDGET_KEY, (tui: TUI) => new WoolyComponent(tui), { placement: "belowEditor" });
+}
+
+export function clearLoomUi(ctx: ExtensionContext): void {
+	if (!ctx.hasUI) return;
+	ctx.ui.setStatus(STATUS_KEY, undefined);
+	if (ctx.mode === "tui") {
+		ctx.ui.setWidget(WOOLY_WIDGET_KEY, undefined);
+		ctx.ui.setHeader(undefined);
+	}
+}
+
+export function updateRouterStatus(ctx: ExtensionContext, snapshot: RouterUiSnapshot): void {
+	if (!ctx.hasUI) return;
+	const { requestedModel, routedModel, forcedModel, savings } = snapshot;
+	let route: string;
+	if (routedModel && requestedModel) route = `${routedModel} ← ${requestedModel}`;
+	else if (routedModel) route = routedModel;
+	else if (requestedModel) route = `${requestedModel} · awaiting route`;
+	else route = "automatic routing";
+
+	const label = ctx.mode === "tui" ? ctx.ui.theme.bold(brand("WEAVE ROUTER")) : "WEAVE ROUTER";
+	const detail = forcedModel ? `${forcedModel} [forced]` : `${route} · ${formatSavings(savings)}`;
+	ctx.ui.setStatus(STATUS_KEY, `${label} — ${ctx.mode === "tui" ? ctx.ui.theme.fg("dim", detail) : detail}`);
+}