model-router

1import { atom, update } from "claude-code";
2import type { EngineInterface, Register } from "claude-code";
3
4import {
5  configFrom,
6  describeConfig,
7  route,
8  stripOverride,
9  subagentModel,
10} from "./route";
11
12const last = atom({ plugin: "model-router", key: "last" } as const, null);
13
14async function isEnabled($: EngineInterface): Promise<boolean> {
15  return (await $.store.get("enabled")) !== false;
16}
17
18export const register: Register = (on, options) => {
19  const config = configFrom(options);
20  // The model picked for each main-thread turn, by turnId.
21  const picks = new Map<string, { model: string; reason: string }>();
22  let pending: { model: string; reason: string } | null = null;
23
24  on("session.start", async ($, e, next) => {
25    await $.command.register({
26      name: "route",
27      description:
28        "Model router: /route on | off | status (prefix a prompt with !haiku, !sonnet or !opus to force)",
29    });
30
31    return next(e);
32  });
33
34  on("command.run", { command: "route" }, async ($, e) => {
35    const arg = e.args.trim().toLowerCase();
36    if (arg === "on" || arg === "off") {
37      await $.store.set("enabled", arg === "on");
38      $.ui.status(undefined);
39      return { text: `Model router ${arg}.` };
40    }
41    return {
42      text: `Model router is ${(await isEnabled($)) ? "on" : "off"}. ${describeConfig(config)}.`,
43    };
44  });
45
46  on("prompt.submit", async ($, e, next) => {
47    pending = null;
48    if (!config.routeMainThread || !(await isEnabled($))) return next(e);
49
50    const picked = route(e.text);
51    const model = picked ? config.models[picked.tier] : null;
52    if (picked && model) pending = { model, reason: picked.reason };
53    $.ui.status(
54      picked && model
55        ? `↳ ${model.replace(/^claude-/, "")} · ${picked.reason}`
56        : undefined,
57    );
58
59    const text = stripOverride(e.text);
60    return next(text === e.text ? e : { ...e, text });
61  }).catch(($, e, next) => next(e));
62
63  on("turn.start", async ($, e, next) => {
64    if (pending) {
65      picks.set(e.turnId, pending);
66      await update($, last, () => ({ ...pending!, text: e.text.slice(0, 80) }));
67      pending = null;
68    }
69    return next(e);
70  });
71
72  on("turn.step", async function* ($, e, next) {
73    const pick = e.agentId === undefined ? picks.get(e.turnId) : undefined;
74    const isWorthIt =
75      pick &&
76      (pick.reason === "forced" || e.messageCount <= config.maxMessages);
77
78    return yield* next(isWorthIt ? { ...e, model: pick.model } : e);
79  });
80
81  on("turn.complete", ($, e, next) => {
82    if (e.agentId === undefined) picks.delete(e.turnId);
83    return next(e);
84  });
85
86  on("agent.spawn", async ($, e, next) => {
87    const model = subagentModel(config, {
88      subagentType: e.subagentType,
89      model: e.model,
90      fork: e.fork,
91      isWorkflow: e.workflow !== undefined,
92    });
93    if (model === undefined || !(await isEnabled($))) return next(e);
94
95    return next({ ...e, model });
96  }).catch(($, e, next) => next(e));
97};
98

1import type { PluginOptions } from "claude-code";
2
3export type Tier = "light" | "standard" | "heavy";
4
5export type Route = { tier: Tier; reason: string };
6
7export type RouterConfig = {
8  // Which model each tier runs on; `heavy` keeps the session's own model.
9  models: Record<Tier, string | null>;
10  // Past this many messages a model switch rewrites the whole prompt cache, which costs more than it saves.
11  maxMessages: number;
12  // Subagent types that only search or read: sent to the light model unless the call names one.
13  lightAgents: ReadonlySet<string>;
14  routeMainThread: boolean;
15  routeSubagents: boolean;
16};
17
18export const DEFAULTS = {
19  lightModel: "claude-haiku-5-5",
20  standardModel: "claude-sonnet-5-5",
21  maxMessages: 40,
22  lightAgents: "Explore,claude-code-guide,statusline-setup",
23  routeMainThread: true,
24  routeSubagents: true,
25} as const;
26
27const str = (v: unknown, fallback: string) =>
28  typeof v === "string" && v.trim() !== "" ? v.trim() : fallback;
29const num = (v: unknown, fallback: number) =>
30  typeof v === "number" && Number.isFinite(v) && v >= 0 ? v : fallback;
31const bool = (v: unknown, fallback: boolean) =>
32  typeof v === "boolean" ? v : fallback;
33
34export const configFrom = (
35  options: PluginOptions | undefined,
36): RouterConfig => {
37  const o = options ?? {};
38  return {
39    models: {
40      light: str(o.lightModel, DEFAULTS.lightModel),
41      standard: str(o.standardModel, DEFAULTS.standardModel),
42      heavy: null,
43    },
44    maxMessages: num(o.maxMessages, DEFAULTS.maxMessages),
45    lightAgents: new Set(
46      str(o.lightAgents, DEFAULTS.lightAgents)
47        .split(",")
48        .map((s) => s.trim())
49        .filter(Boolean),
50    ),
51    routeMainThread: bool(o.routeMainThread, DEFAULTS.routeMainThread),
52    routeSubagents: bool(o.routeSubagents, DEFAULTS.routeSubagents),
53  };
54};
55
56export const describeConfig = (c: RouterConfig) =>
57  [
58    `light → ${c.models.light}`,
59    `standard → ${c.models.standard}`,
60    "heavy → session model",
61    `main thread ${c.routeMainThread ? `routed (up to ${c.maxMessages} messages)` : "not routed"}`,
62    `subagents ${c.routeSubagents ? `routed: ${[...c.lightAgents].join(", ") || "none"}` : "not routed"}`,
63  ].join("; ");
64
65const OVERRIDE = /^\s*!(haiku|sonnet|opus)\b/i;
66const HEAVY =
67  /\b(debug|bug|broken|fail|refactor|architect|design|plan|implement|migrat|security|race condition|perf|optimi[sz]|why does|root cause)/i;
68const LIGHT =
69  /^\s*(git (status|log|diff|pull|push)|commit|push|run (the )?(tests?|lint|format|build)|lint|format|typecheck|rename|bump|open (the )?pr|status)\b/i;
70const QUESTION =
71  /^\s*(what|where|which|who|when|how (do|does|can|many|much)|is|are|does|do|can|list|show|explain|summari[sz]e)\b/i;
72
73export const route = (text: string): Route | null => {
74  const forced = OVERRIDE.exec(text)?.[1]?.toLowerCase();
75  if (forced === "haiku") return { tier: "light", reason: "forced" };
76  if (forced === "sonnet") return { tier: "standard", reason: "forced" };
77  if (forced === "opus") return { tier: "heavy", reason: "forced" };
78
79  if (text.trimStart().startsWith("/")) return null;
80  if (HEAVY.test(text)) return { tier: "heavy", reason: "complex task" };
81  if (text.length <= 160 && LIGHT.test(text))
82    return { tier: "light", reason: "routine command" };
83  if (text.length <= 300 && QUESTION.test(text))
84    return { tier: "standard", reason: "question" };
85
86  return null;
87};
88
89export const stripOverride = (text: string) =>
90  text.replace(OVERRIDE, "").trimStart();
91
92// The model a subagent spawn is rewritten to, or undefined to leave it as called.
93export const subagentModel = (
94  c: RouterConfig,
95  spawn: {
96    subagentType: string;
97    model?: string;
98    fork: boolean;
99    isWorkflow: boolean;
100  },
101): string | undefined =>
102  c.routeSubagents &&
103  c.lightAgents.has(spawn.subagentType) &&
104  spawn.model === undefined &&
105  !spawn.fork &&
106  !spawn.isWorkflow
107    ? (c.models.light ?? undefined)
108    : undefined;
109

1export type LastRoute = { model: string; reason: string; text: string };
2
3declare module "claude-code" {
4  interface PluginState {
5    "model-router": { last: LastRoute | null };
6  }
7}
8