tokendoctor

1import type { EngineInterface, Register } from 'claude-code'
2
3// Tools for /tokendoctor:optimize. Reading tools run without a prompt. A change runs only after the
4// person picks "Apply" in a dialog that shows it: a permission `ask` can be settled by the mode
5// (bypass, a -p run) without a person, so the dialog is the mod's own, and with no one to ask
6// ($.ui.ask rejects in -p) nothing changes. The tools wait behind ToolSearch, so they add nothing
7// to context until the optimize skill loads them.
8const PLUGIN = 'tokendoctor'
9const tool = <N extends string>(name: N) => `mcp__${PLUGIN}__${name}` as const
10const READS: string[] = [tool('metrics'), tool('settings'), tool('config_list')]
11const WRITES: string[] = [tool('config_set'), tool('set_setting')]
12
13const text = (value: unknown) => ({ result: typeof value === 'string' ? value : JSON.stringify(value, null, 1) })
14
15export const register: Register = on => {
16  on('session.start', async ($, e, next) => {
17    const result = await next(e)
18    await $.tool.register({
19      name: 'metrics',
20      description:
21        'Token metrics: the current session (context fill, rate limits, cost, context breakdown by category), ' +
22        'daily totals for the last `days` (default 14) and the current audit findings. Read-only.',
23      inputSchema: { type: 'object', properties: { days: { type: 'number' } } },
24    })
25    await $.tool.register({
26      name: 'settings',
27      description:
28        'Every setting tokendoctor may change, with its current merged value (unset ones included), ' +
29        'whether a change applies live or at the next session, and what it does. Read-only.',
30    })
31    await $.tool.register({
32      name: 'config_list',
33      description: 'The rows of the /config menu with their current values, kinds and options. Read-only.',
34    })
35    await $.tool.register({
36      name: 'config_set',
37      description: 'Change one /config row (key from config_list), live, as if the person did it in /config. Asks the person.',
38      inputSchema: { type: 'object', properties: { key: { type: 'string' }, value: {} }, required: ['key', 'value'] },
39    })
40    await $.tool.register({
41      name: 'set_setting',
42      description:
43        'Write one setting from the `settings` list to a settings file, with a backup. `value` omitted unsets it. ' +
44        'scope: user (default), project or local. Asks the person.',
45      inputSchema: {
46        type: 'object',
47        properties: { key: { type: 'string' }, value: {}, scope: { type: 'string', enum: ['user', 'project', 'local'] } },
48        required: ['key'],
49      },
50    })
51    return result
52  })
53
54  on('tool.check', async ($, e, next) => {
55    if (READS.includes(e.tool)) return { decision: 'allow' }
56    if (WRITES.includes(e.tool)) return { decision: 'allow' }
57    return next(e)
58  })
59
60  on('tool.call', { tool: tool('metrics') }, async ($, e) => {
61    const days = typeof e.days === 'number' ? e.days : 14
62    const usage = await $.session.usage({ breakdown: 'summary' })
63    const b = usage.context.breakdown
64    const session = {
65      tokens: usage.context.tokens,
66      window: usage.context.window,
67      compactAt: b?.autoCompactThreshold,
68      autoCompact: b?.isAutoCompactEnabled,
69      model: b?.model,
70      biggest: b?.categories
71        .filter(c => c.kind === 'used')
72        .sort((x, y) => y.tokens - x.tokens)
73        .slice(0, 6)
74        .map(c => ({ name: c.name, tokens: c.tokens })),
75      rateLimits: usage.rateLimits,
76      costUsd: usage.cost?.usd,
77    }
78    const trend = await bin($, 'tokendoctor-sessions', ['--by', 'day', '--days', String(days), '--json'])
79    const audit = await bin($, 'tokendoctor-audit', ['--json'])
80    let findings: unknown = audit
81    try {
82      findings = JSON.parse(audit).findings.map((f: { id: string; severity: string; title: string; fixes?: unknown }) => ({
83        id: f.id,
84        severity: f.severity,
85        title: f.title,
86        hasFix: Boolean(f.fixes),
87      }))
88    } catch {}
89    return text({ session, trend: safeJson(trend), findings })
90  })
91
92  on('tool.call', { tool: tool('settings') }, async ($) => {
93    const out = await bin($, 'tokendoctor-set', ['--list', '--json'])
94    const rows = safeJson(out)
95    if (!Array.isArray(rows)) return text(out)
96    const merged = (await $.settings.read()) as Record<string, unknown>
97    const env = (merged.env ?? {}) as Record<string, unknown>
98    const perModel = (merged.modelSettings ?? {}) as Record<string, Record<string, unknown>>
99    // A per-model row stays as the pattern for a model with nothing saved, and gets one row with the
100    // saved value for each model that has one (/effort saves the level there, not in effortLevel).
101    return text(
102      rows.flatMap((r: { key: string }) => {
103        const field = /^modelSettings\.<model id>\.(\w+)$/.exec(r.key)?.[1]
104        if (field) {
105          const saved = Object.entries(perModel).filter(([, m]) => m?.[field] !== undefined)
106          return [{ ...r, current: null }, ...saved.map(([model, m]) => ({ ...r, key: `modelSettings.${model}.${field}`, current: m[field] }))]
107        }
108        return [{ ...r, current: r.key.startsWith('env.') ? env[r.key.slice(4)] ?? null : merged[r.key] ?? null }]
109      }),
110    )
111  })
112
113  on('tool.call', { tool: tool('config_list') }, async ($) => {
114    const rows = await $.config.list()
115    return text(rows.map(r => ({ key: r.key, label: r.label, kind: r.kind, value: r.value, options: r.options })))
116  })
117
118  on('tool.call', { tool: tool('config_set') }, async ($, e) => {
119    const { key, value } = e as unknown as { key: string; value: never }
120    const row = (await $.config.list()).find(r => r.key === key)
121    if (!row) return text(`Not changed: /config has no row "${key}".`)
122    const change = `/config ${row.label}: ${JSON.stringify(row.value)} → ${JSON.stringify(value)} (applies now)`
123    if (!(await confirm($, change))) return text(`Not changed: the person did not approve (${change}).`)
124    const r = await $.config.set({ key, value })
125    return text(r.deny ? `Not changed: ${r.deny}` : `${key} is now ${JSON.stringify(r.value)} (live).`)
126  })
127
128  on('tool.call', { tool: tool('set_setting') }, async ($, e) => {
129    const { key, value, scope } = e as unknown as { key: string; value?: unknown; scope?: string }
130    const args = [key, ...(value === undefined ? ['--unset'] : [JSON.stringify(value)]), '--scope', scope ?? 'user']
131    const plan = await bin($, 'tokendoctor-set', args)
132    if (!plan.startsWith('Would change')) return text(plan)
133    if (!(await confirm($, plan.replace(/\n+Dry run.*$/s, '')))) return text(`Not changed: the person did not approve.\n${plan}`)
134    return text(await bin($, 'tokendoctor-set', [...args, '--write']))
135  })
136}
137
138// The person's yes, in the engine's own dialog. False when they decline, dismiss it, or no one can
139// be asked (a -p run).
140async function confirm($: EngineInterface, change: string) {
141  try {
142    const answer = await $.ui.ask(`tokendoctor wants to change a setting:\n${change}\nApply it?`, {
143      header: 'tokendoctor',
144      options: ['Apply', 'Cancel'],
145    })
146    return answer === 'Apply'
147  } catch {
148    return false
149  }
150}
151
152async function bin($: EngineInterface, name: string, args: string[]) {
153  const { exitCode, stdout, stderr } = await $.process.run(['node', `${$.plugin.root}/bin/${name}`, ...args], { timeoutMs: 120_000 })
154  return exitCode === 0 ? stdout : `${name} failed (${exitCode}): ${stderr || stdout}`
155}
156
157function safeJson(s: string): unknown {
158  try {
159    return JSON.parse(s)
160  } catch {
161    return s
162  }
163}
164