model-router

1// Model Router: sends routine read-only follow-up steps of a turn to a cheaper
2// model, keeps planning and final answers on the session model, and shows an
3// honest estimate of what that saved above the prompt.
4
5import { update } from "claude-code";
6
7// List prices in USD per million tokens, from the claude-api skill (cached
8// 2026-09-25). Cache write is the 5-minute TTL rate (1.25x input); the usage a
9// step reports does not say which TTL was written, so 1-hour writes (2x) are
10// under-counted on both sides alike. First match wins, so keep specific ids
11// above general ones.
12const PRICES = [
13  { match: /opus-5-5/, input: 4, output: 20, cacheWrite: 5, cacheRead: 0.2 },
14  { match: /opus/, input: 5, output: 25, cacheWrite: 6.25, cacheRead: 0.5 },
15  { match: /sonnet-5/, input: 2, output: 10, cacheWrite: 2.5, cacheRead: 0.2 },
16  { match: /sonnet/, input: 3, output: 15, cacheWrite: 3.75, cacheRead: 0.3 },
17  { match: /haiku/, input: 1, output: 5, cacheWrite: 1.25, cacheRead: 0.1 },
18  { match: /fable-5-1|mythos-5-1/, input: 10, output: 50, cacheWrite: 12.5, cacheRead: 0.25 },
19  { match: /fable|mythos/, input: 10, output: 50, cacheWrite: 12.5, cacheRead: 1 },
20];
21
22// What `/router sonnet` and `/router haiku` send the request as.
23const CHEAP_MODELS = { sonnet: "claude-sonnet-5-5", haiku: "claude-haiku-4-5" };
24const HAIKU_CONTEXT_LIMIT = 180_000;
25
26// Tools that only read. The engine's own read-only verdict (tool.call's
27// `isReadOnly`) is used first; this list and the Bash pattern are the fallback.
28const READ_ONLY_TOOLS = new Set([
29  "Read", "Grep", "Glob", "LS", "WebFetch", "WebSearch", "NotebookRead", "ToolSearch",
30]);
31const READ_ONLY_BASH =
32  /^\s*(ls|cat|head|tail|wc|pwd|file|stat|tree|which|echo|find|grep|rg|git\s+(status|log|diff|show|branch|blame|rev-parse))\b[^;&|>`$]*$/;
33
34const MODES = ["subagents", "steps", "both"];
35const DEFAULT_CONFIG = { enabled: true, cheap: "sonnet", mode: "subagents" };
36const EMPTY_STATS = {
37  turns: 0, steps: 0, routed: 0, actualUsd: 0, counterfactualUsd: 0, lastPrefix: 0, sessionModel: "",
38  subSteps: 0, subActualUsd: 0, subCounterfactualUsd: 0, subagents: 0,
39};
40const MAX_DECISIONS = 5;
41
42const config = { plugin: "model-router", key: "config" };
43const stats = { plugin: "model-router", key: "stats" };
44const decisions = { plugin: "model-router", key: "decisions" };
45
46export function register(on) {
47  on("session.start", async ($, e, next) => {
48    const result = await next(e);
49    const enabled = await $.store.get("enabled");
50    const cheap = await $.store.get("cheap");
51    const mode = await $.store.get("mode");
52    await $.state.set(config, {
53      enabled: enabled !== false,
54      cheap: cheap === "haiku" ? "haiku" : "sonnet",
55      mode: MODES.includes(mode) ? mode : DEFAULT_CONFIG.mode,
56    });
57    await $.command.register({
58      name: "router",
59      description: "Model router status, on or off, cheap model, or routing mode",
60      argumentHint: "[on|off|sonnet|haiku|mode subagents|steps|both]",
61    });
62    return result;
63  });
64
65  // Remember the engine's read-only verdict for each main-loop tool call, by id.
66  on("tool.call", async ($, e, next) => {
67    const result = await next(e);
68    if (e.agentId == null && result.isReadOnly === true && e.tool_use_id) {
69      await $.state.set({ plugin: "model-router", key: "readOnly", id: e.tool_use_id }, true);
70    }
71    return result;
72  });
73
74  on("turn.step", async function* ($, e, next) {
75    const { value: cfg = DEFAULT_CONFIG } = await $.state.get(config);
76    const mode = cfg.mode ?? DEFAULT_CONFIG.mode;
77
78    // Subagent mode: every step of a subagent goes to the cheap model, so the
79    // subagent builds and reads its own prompt cache there. Its context is its
80    // own, so the main loop's cache is never broken by it.
81    if (e.agentId != null) {
82      if (!cfg.enabled || mode === "steps" || !isDearer(e.model, cfg.cheap)) return yield* next(e);
83      const target = CHEAP_MODELS[cfg.cheap];
84      if (e.index === 0) {
85        $.ui.log(`subagent ${shortId(e.agentId)} → ${cfg.cheap} (all its steps)`);
86        await update($, decisions, (list = []) =>
87          [...list, { kind: "subagent", turn: 0, step: 0, to: cfg.cheap, isRouted: true, reason: `subagent ${shortId(e.agentId)}, all its steps` }].slice(-MAX_DECISIONS),
88        );
89      }
90      const result = yield* next(routedInput(e, cfg.cheap, target));
91      const usage = result?.usage;
92      const actual = usage ? costAt(priceOf(usage.model) ?? priceOf(target), usage) : 0;
93      const counterfactual = usage ? costAt(priceOf(e.model) ?? priceOf(target), usage) : 0;
94      if (usage) {
95        $.ui.log(
96          `subagent ${shortId(e.agentId)} step ${e.index} on ${usage.model} read ${usage.cache_read_input_tokens} wrote ${usage.cache_creation_input_tokens} cost ${usd(actual)}, ~${usd(counterfactual)} on ${e.model}`,
97          { to: "debug" },
98        );
99      }
100      await update($, stats, (s = EMPTY_STATS) => ({
101        ...EMPTY_STATS,
102        ...s,
103        subSteps: (s.subSteps ?? 0) + 1,
104        subagents: (s.subagents ?? 0) + (e.index === 0 ? 1 : 0),
105        subActualUsd: (s.subActualUsd ?? 0) + actual,
106        subCounterfactualUsd: (s.subCounterfactualUsd ?? 0) + counterfactual,
107      }));
108      return result;
109    }
110
111    const { value: before = EMPTY_STATS } = await $.state.get(stats);
112    const turn = e.index === 0 ? before.turns + 1 : Math.max(before.turns, 1);
113    const isStepMode = cfg.enabled && (mode === "steps" || mode === "both");
114
115    let decision;
116    if (!cfg.enabled) {
117      decision = { isRouted: false, reason: "router off" };
118    } else if (!isStepMode) {
119      decision = { isRouted: false, reason: "main loop not routed in subagents mode" };
120    } else if (e.index === 0) {
121      decision = { isRouted: false, reason: "first step answers the prompt" };
122    } else {
123      const messages = await $.session.messages();
124      const uses = lastToolUses(messages);
125      const verdicts = [];
126      for (const use of uses) {
127        const { value: engineSaysReadOnly } = await $.state.get({
128          plugin: "model-router", key: "readOnly", id: use.tool_use_id,
129        });
130        verdicts.push({ ...use, readOnly: engineSaysReadOnly === true || looksReadOnly(use) });
131      }
132      decision = decide({
133        verdicts,
134        cheap: cfg.cheap,
135        sessionModel: e.model,
136        lastPrefix: before.lastPrefix,
137      });
138    }
139
140    const target = decision.isRouted ? CHEAP_MODELS[cfg.cheap] : e.model;
141    const stepInput = decision.isRouted ? routedInput(e, cfg.cheap, target) : e;
142    if (decision.isRouted) {
143      $.ui.log(`step ${e.index} → ${cfg.cheap} (${decision.reason})`);
144    }
145
146    const result = yield* next(stepInput);
147
148    const usage = result?.usage;
149    const cost = decision.isRouted && usage ? routedCost(usage, e.model, before.lastPrefix) : null;
150    if (cost) {
151      $.ui.log(
152        `step ${e.index} on ${usage.model} cost ${usd(cost.actual)}, ~${usd(cost.counterfactual)} on ${e.model}`,
153        { to: "debug" },
154      );
155    }
156    await update($, stats, (s = EMPTY_STATS) => {
157      const after = { ...EMPTY_STATS, ...s, turns: turn, steps: s.steps + 1 };
158      if (!decision.isRouted) after.sessionModel = e.model;
159      if (decision.isRouted) after.routed = s.routed + 1;
160      if (cost) {
161        after.actualUsd = s.actualUsd + cost.actual;
162        after.counterfactualUsd = s.counterfactualUsd + cost.counterfactual;
163      }
164      if (usage) after.lastPrefix = promptTokens(usage) + usage.output_tokens;
165      return after;
166    });
167
168    if (isStepMode) {
169      const entry = { kind: "step", turn, step: e.index, to: decision.isRouted ? cfg.cheap : e.model, isRouted: decision.isRouted, reason: decision.reason };
170      await update($, decisions, (list = []) => [...list, entry].slice(-MAX_DECISIONS));
171    }
172    return result;
173  });
174
175  on("command.run", { command: "router" }, async ($, e) => {
176    const words = e.args.trim().toLowerCase().split(/\s+/).filter(Boolean);
177    const arg = words.join(" ");
178    const { value: cfg = DEFAULT_CONFIG } = await $.state.get(config);
179    const mode = cfg.mode ?? DEFAULT_CONFIG.mode;
180    if (arg === "on" || arg === "off") {
181      const next = { ...cfg, mode, enabled: arg === "on" };
182      await $.store.set("enabled", next.enabled);
183      await $.state.set(config, next);
184      return { text: `Model router ${arg}${next.enabled ? ` · mode ${mode} · cheap model ${next.cheap}` : ""}.` };
185    }
186    if (arg === "sonnet" || arg === "haiku") {
187      const next = { ...cfg, mode, cheap: arg };
188      await $.store.set("cheap", arg);
189      await $.state.set(config, next);
190      return { text: `Model router cheap model set to ${arg} (${CHEAP_MODELS[arg]}).` };
191    }
192    if (words[0] === "mode") {
193      if (!MODES.includes(words[1])) {
194        return { text: `Model router mode is ${mode}. Pick one of /router mode subagents, steps or both.` };
195      }
196      const next = { ...cfg, mode: words[1] };
197      await $.store.set("mode", words[1]);
198      await $.state.set(config, next);
199      return { text: `Model router mode set to ${words[1]} · ${MODE_TEXT[words[1]]}.` };
200    }
201    if (arg !== "" && arg !== "status") {
202      return { text: "Usage /router [on|off|sonnet|haiku|mode subagents|mode steps|mode both]" };
203    }
204    const { value: s = EMPTY_STATS } = await $.state.get(stats);
205    const { value: recent = [] } = await $.state.get(decisions);
206    return { text: statusText({ ...cfg, mode }, { ...EMPTY_STATS, ...s }, recent) };
207  });
208
209  on("ui.render", { component: "AbovePrompt" }, async ($, e, next) => {
210    const original = await next(e);
211    const { value: cfg = DEFAULT_CONFIG } = await $.state.get(config);
212    const { value: raw = EMPTY_STATS } = await $.state.get(stats);
213    const s = { ...EMPTY_STATS, ...raw };
214    if (e.props.hasSurvey || !cfg.enabled || s.routed + s.subSteps === 0) return original;
215    const { Box, Text } = $.ui.resolve(e);
216    const row = band(Box, Text, { ...cfg, mode: cfg.mode ?? DEFAULT_CONFIG.mode }, s, e.props.bodyColumns);
217    if (!original) return row;
218    return Box({ flexDirection: "column", children: [original, row] });
219  });
220}
221
222const MODE_TEXT = {
223  subagents: "every step of a subagent goes to the cheap model, the main loop stays put",
224  steps: "main-loop steps after read-only tools go to the cheap model, subagents stay put",
225  both: "subagents and read-only main-loop steps both go to the cheap model",
226};
227
228function isDearer(model, cheap) {
229  const own = priceOf(model);
230  const target = priceOf(CHEAP_MODELS[cheap]);
231  return Boolean(own && target && target.output < own.output);
232}
233
234function shortId(id) {
235  return String(id).slice(0, 8);
236}
237
238// The routing rule, on plain values. Steps after index 0 only.
239function decide({ verdicts, cheap, sessionModel, lastPrefix }) {
240  if (verdicts.length === 0) return { isRouted: false, reason: "no tool results to read" };
241  const failed = verdicts.filter((v) => v.isError);
242  if (failed.length > 0) return { isRouted: false, reason: `after failed ${names(failed)}` };
243  const writes = verdicts.filter((v) => !v.readOnly);
244  if (writes.length > 0) return { isRouted: false, reason: `after ${names(writes)}` };
245  const session = priceOf(sessionModel);
246  const target = priceOf(CHEAP_MODELS[cheap]);
247  if (!session || !target || target.output >= session.output) {
248    return { isRouted: false, reason: `session model is not dearer than ${cheap}` };
249  }
250  if (cheap === "haiku" && lastPrefix > HAIKU_CONTEXT_LIMIT) {
251    return { isRouted: false, reason: "context too large for haiku" };
252  }
253  return { isRouted: true, reason: `after ${names(verdicts)}` };
254}
255
256// The tool calls of the newest assistant message: the results this step reads.
257function lastToolUses(messages) {
258  if (!Array.isArray(messages)) return [];
259  for (let i = messages.length - 1; i >= 0; i--) {
260    const m = messages[i];
261    if (m.role === "assistant") return m.toolUses ?? [];
262    if (m.role === "user" && !(m.toolResults?.length > 0) && m.text) return [];
263  }
264  return [];
265}
266
267function looksReadOnly(use) {
268  if (READ_ONLY_TOOLS.has(use.tool)) return true;
269  if (use.tool === "Bash") return READ_ONLY_BASH.test(String(use.input?.command ?? ""));
270  return false;
271}
272
273function names(list) {
274  return [...new Set(list.map((v) => v.tool))].join(", ");
275}
276
277function routedInput(e, cheap, model) {
278  // Haiku takes no effort setting; Sonnet takes the same levels as Opus.
279  if (cheap === "haiku") {
280    const { effort, ...rest } = e;
281    return { ...rest, model };
282  }
283  return { ...e, model };
284}
285
286function priceOf(model) {
287  return PRICES.find((p) => p.match.test(String(model ?? "")));
288}
289
290function promptTokens(u) {
291  return u.input_tokens + u.cache_creation_input_tokens + u.cache_read_input_tokens;
292}
293
294function costAt(p, u) {
295  return (
296    (u.input_tokens * p.input +
297      u.output_tokens * p.output +
298      u.cache_creation_input_tokens * p.cacheWrite +
299      u.cache_read_input_tokens * p.cacheRead) /
300    1_000_000
301  );
302}
303
304// Actual cost of a routed step, and what the session model would have charged
305// for the same tokens. The session model would have had the previous step's
306// whole prompt and reply in its cache (it ran every step in that world), so
307// up to `lastPrefix` cached tokens are priced as its cache reads and the rest
308// as cache writes. Output is assumed the same length on either model.
309function routedCost(u, sessionModel, lastPrefix) {
310  const actualPrice = priceOf(u.model) ?? priceOf(sessionModel);
311  const sessionPrice = priceOf(sessionModel) ?? actualPrice;
312  const cached = u.cache_creation_input_tokens + u.cache_read_input_tokens;
313  const read = Math.min(cached, lastPrefix);
314  const counterfactual = costAt(sessionPrice, {
315    input_tokens: u.input_tokens,
316    output_tokens: u.output_tokens,
317    cache_read_input_tokens: read,
318    cache_creation_input_tokens: cached - read,
319  });
320  return { actual: costAt(actualPrice, u), counterfactual };
321}
322
323function band(Box, Text, cfg, s, columns) {
324  const dim = (children) => Text({ dimColor: true, children });
325  const saved = s.counterfactualUsd - s.actualUsd + s.subCounterfactualUsd - s.subActualUsd;
326  const counts = [];
327  if (cfg.mode !== "steps" && s.subSteps > 0) counts.push(`${s.subSteps} subagent step${s.subSteps === 1 ? "" : "s"}`);
328  if (cfg.mode !== "subagents" || s.routed > 0) counts.push(`${s.routed} of ${s.steps} main steps`);
329  const parts = [
330    Text({ color: "cyan", bold: true, children: `⇄ router ${cfg.mode}` }),
331    dim(" · "),
332    Text({ children: `${counts.join(" + ")} → ${cfg.cheap}` }),
333    dim(" · "),
334  ];
335  if (saved >= 0) {
336    parts.push(Text({ color: "green", bold: true, children: `saved ~${usd(saved)}` }));
337  } else {
338    parts.push(Text({ color: "yellow", bold: true, children: `net ~${usd(-saved)} more` }));
339  }
340  if (columns >= 70) parts.push(dim(saved >= 0 ? " this session" : " this session (cache misses)"));
341  return Box({ flexDirection: "row", paddingX: 1, children: parts });
342}
343
344function statusText(cfg, s, recent) {
345  const lines = [
346    `Model router ${cfg.enabled ? "on" : "off"} · mode ${cfg.mode} · cheap model ${cfg.cheap} (${CHEAP_MODELS[cfg.cheap]})`,
347    `Mode ${cfg.mode} · ${MODE_TEXT[cfg.mode]}`,
348  ];
349  if (s.subSteps > 0) {
350    const subSaved = s.subCounterfactualUsd - s.subActualUsd;
351    lines.push(
352      `Subagents ${s.subagents} routed, ${s.subSteps} steps, cost ${usd(s.subActualUsd)} vs ~${usd(s.subCounterfactualUsd)} on their own model, ` +
353        (subSaved >= 0 ? `saved ~${usd(subSaved)}` : `net ~${usd(-subSaved)} more`),
354    );
355  } else {
356    lines.push("Subagents none routed yet this session");
357  }
358  lines.push(`Main loop ${s.routed} of ${s.steps} steps routed`);
359  if (s.routed > 0) {
360    const saved = s.counterfactualUsd - s.actualUsd;
361    const base = s.sessionModel || "the session model";
362    lines.push(
363      `Routed main steps cost ${usd(s.actualUsd)} vs ~${usd(s.counterfactualUsd)} on ${base}, ` +
364        (saved >= 0 ? `saved ~${usd(saved)}` : `net ~${usd(-saved)} more (cache misses on the switch)`),
365    );
366  }
367  if (recent.length > 0) {
368    lines.push("Recent decisions");
369    for (const d of recent) {
370      if (d.kind === "subagent") lines.push(`  ${d.reason} → ${d.to}`);
371      else lines.push(`  turn ${d.turn} step ${d.step} ${d.isRouted ? "→" : "stays on"} ${d.to} (${d.reason})`);
372    }
373  }
374  return lines.join("\n");
375}
376
377function usd(n) {
378  if (n < 0.01) return `$${n.toFixed(4)}`;
379  return n < 10 ? `$${n.toFixed(2)}` : `$${n.toFixed(1)}`;
380}
381

1export type ModelRouterCheap = "sonnet" | "haiku";
2export type ModelRouterMode = "subagents" | "steps" | "both";
3export type ModelRouterConfig = { enabled: boolean; cheap: ModelRouterCheap; mode: ModelRouterMode };
4export type ModelRouterDecision = { kind: "step" | "subagent"; turn: number; step: number; to: string; isRouted: boolean; reason: string };
5export type ModelRouterStats = {
6  turns: number;
7  steps: number;
8  routed: number;
9  actualUsd: number;
10  counterfactualUsd: number;
11  lastPrefix: number;
12  sessionModel: string;
13  subSteps: number;
14  subActualUsd: number;
15  subCounterfactualUsd: number;
16  subagents: number;
17};
18
19declare module "claude-code" {
20  interface PluginState {
21    "model-router": {
22      config: ModelRouterConfig;
23      stats: ModelRouterStats;
24      decisions: ModelRouterDecision[];
25      readOnly: StateFamily<boolean>;
26    };
27  }
28}
29