token-pilot

1/**
2 * token-pilot as a Claude Code mod: its hooks run inside Claude Code (2.1.275+)
3 * instead of as one process per tool call. The command hooks in
4 * hooks/hooks.json stay for older Claude Code, for sessions where mods are
5 * off, and for any action this mod does not serve yet.
6 *
7 * Engine rule: one unmatched hook per event per module, so every unmatched
8 * hook lives here and delegates.
9 */
10
11import type { EngineInterface, Register } from 'claude-code'
12import { setPluginInstall } from '../../src/core/tool-names.js'
13import { caught } from './host.js'
14import { BASH_ACTIONS, registerBash } from './bash.js'
15import { GREP_ACTIONS, registerGrep } from './grep.js'
16import { MCP_ACTIONS, createEditPrep, registerMcp } from './mcp.js'
17import { EDIT_ACTIONS, registerEdit } from './edit.js'
18import { READ_ACTIONS, registerRead } from './read.js'
19import { AGENT_ACTIONS, registerAgent } from './agent.js'
20import { SESSION_ACTIONS, registerSession } from './session.js'
21import { registerBand } from './band.js'
22import { STATS_COMMAND, registerStats } from './stats.js'
23
24/**
25 * Command-hook actions this mod handles in-process. hooks/run.sh exits early
26 * for each one, so no call is handled twice. With the mod off, nothing is
27 * set and every command hook runs as before.
28 */
29const SERVED: string[] = [...BASH_ACTIONS, ...GREP_ACTIONS, ...MCP_ACTIONS, ...EDIT_ACTIONS, ...READ_ACTIONS, ...SESSION_ACTIONS]
30
31const MAX_SESSIONS = 8
32
33async function handOff($: EngineInterface, id: string): Promise<void> {
34  // Workflow runs keep agent routing on the command hooks, which attach the
35  // workflow budget note (see agent.ts).
36  const inWorkflow = Boolean(
37    (await $.env.get('TOKEN_PILOT_WORKFLOW_ID')) ??
38      (await $.env.get('CLAUDE_CODE_WORKFLOW_ID')) ??
39      (await $.env.get('LOOM_WORKFLOW_ID')),
40  )
41  const served = inWorkflow ? SERVED : [...SERVED, ...AGENT_ACTIONS]
42  // Escape hatch: TOKEN_PILOT_NO_MOD=1 keeps every hook on the command path
43  // (the handlers check it too), exactly as on a Claude Code without mods.
44  const noMod = (await $.env.get('TOKEN_PILOT_NO_MOD')) === '1'
45
46  await $.env.set('TOKEN_PILOT_MOD', noMod ? '' : served.join(','))
47
48  // A nested claude inherits both; run.sh only steps aside for the sessions
49  // listed here. Earlier ones stay: after /clear their background agents run on.
50  const earlier = ((await $.env.get('TOKEN_PILOT_MOD_SESSION')) ?? '').split(',').filter(s => s && s !== id)
51  await $.env.set('TOKEN_PILOT_MOD_SESSION', [id, ...earlier].slice(0, MAX_SESSIONS).join(','))
52}
53
54export const register: Register = on => {
55  setPluginInstall(true)
56
57  // Mod hooks run before the command hooks, so the flag set here is already
58  // visible to the command SessionStart of this same session.
59  // The payload's session_id, not $.session.id(): on /resume that still
60  // answers the session being left, while the command hooks get the new one.
61  on('classic.SessionStart', async ($, e, next) => {
62    await handOff($, e.session_id)
63    return next(e)
64  }).catch(caught)
65
66  // Again on a hot reload, where classic.SessionStart does not fire.
67  on('session.start', async ($, e, next) => {
68    await handOff($, await $.session.id())
69    await $.command.register(STATS_COMMAND)
70    return next(e)
71  }).catch(caught)
72
73  registerBash(on)
74  registerGrep(on)
75
76  const prep = createEditPrep()
77  registerMcp(on, prep)
78  registerEdit(on, prep)
79  registerRead(on)
80  registerAgent(on)
81  registerSession(on)
82  registerBand(on)
83  registerStats(on)
84}
85

1/**
2 * How this install's MCP tools are named to the model.
3 *
4 * A plugin exposes the server under the plugin namespace
5 * (`mcp__plugin_token-pilot_token-pilot__smart_read`); an npm install
6 * registered in `.mcp.json` exposes it bare (`mcp__token-pilot__smart_read`).
7 * Advice that names the wrong one points the model at a tool that does not
8 * exist on its install — and since Claude Code loads MCP definitions on
9 * demand, it cannot even find the real one by searching for the name we
10 * printed.
11 *
12 * v0.52.0 fixed this for the tools listed in agent frontmatter; the hook
13 * messages kept naming the npm form, which is wrong for every plugin-only
14 * user.
15 */
16
17const PLUGIN_PREFIX = "mcp__plugin_token-pilot_token-pilot__";
18const NPM_PREFIX = "mcp__token-pilot__";
19
20/** Full MCP name of one token-pilot tool, as this install exposes it. */
21export function tpTool(name: string): string {
22  return toolPrefix() + name;
23}
24
25let pluginInstall: boolean | undefined;
26
27/**
28 * The Claude Code mod knows it runs as a plugin but has no process.env to
29 * tell; it declares it here once. `undefined` returns to env detection.
30 */
31export function setPluginInstall(value: boolean | undefined): void {
32  pluginInstall = value;
33}
34
35/** The prefix alone — for messages that list several tools. */
36export function toolPrefix(): string {
37  const isPlugin =
38    pluginInstall ?? Boolean(globalThis.process?.env?.CLAUDE_PLUGIN_ROOT);
39  return isPlugin ? PLUGIN_PREFIX : NPM_PREFIX;
40}
41
42/**
43 * Both spellings of one tool. Written into agent `tools:` lists, where the
44 * name that does not resolve is ignored as long as another entry does.
45 */
46export function tpToolBothNames(name: string): string[] {
47  return [NPM_PREFIX + name, PLUGIN_PREFIX + name];
48}
49

1/**
2 * $-free helpers shared by the mod's handler files.
3 *
4 * The engine's validator follows `$` only into functions declared at the top
5 * of the same file — never across an import — so anything that calls `$`
6 * lives in the handler file that needs it. What is shared here takes and
7 * returns plain values.
8 */
9
10import type { EngineInterface, HookFailure } from 'claude-code'
11import { resolveConfig } from '../../src/config/resolve.js'
12import { dirname, isAbsolute, join, normalize, relative } from '../../src/core/portable-path.js'
13import type { TokenPilotConfig } from '../../src/types.js'
14
15export const PREFIX = 'mcp__plugin_token-pilot_token-pilot__'
16
17/** Same rules as loadConfig: no readable `.token-pilot.json` → defaults plus env overrides. */
18export function configFrom(
19  raw: string | null,
20  env: Readonly<Record<string, string | undefined>>,
21): TokenPilotConfig {
22  if (raw === null) return resolveConfig(null, env)
23
24  try {
25    return resolveConfig(JSON.parse(raw), env)
26  } catch {
27    return resolveConfig(null, env)
28  }
29}
30
31/** `path` lies strictly inside `root` (both already resolved). */
32export function isInside(root: string, path: string): boolean {
33  const rel = relative(root, path)
34
35  return rel !== '' && !rel.startsWith('..') && !isAbsolute(rel)
36}
37
38/** Model-facing notes after a tool result — never on a deny or an error. */
39export function withContext<R extends { deny?: unknown; isError?: unknown; context?: readonly string[] }>(
40  ran: R,
41  notes: string[],
42): R {
43  if (!notes.length || ran.deny !== undefined || ran.isError) return ran
44
45  return { ...ran, context: [...(ran.context ?? []), ...notes] }
46}
47
48/**
49 * Nearest ancestor of `dir` holding `.git` — a directory for the main
50 * checkout, a file for a linked worktree. `exists` is the caller's
51 * `p => $.fs.exists(p)`: `$` itself cannot cross a file boundary.
52 */
53export async function findCheckout(exists: (path: string) => Promise<boolean>, dir: string): Promise<string | null> {
54  for (let cur = dir; ; cur = dirname(cur)) {
55    if (await exists(join(cur, '.git'))) return cur
56    if (dirname(cur) === cur) return null
57  }
58}
59
60/**
61 * Key for read_for_edit bookkeeping: the path inside its own checkout. A
62 * subagent prepares `src/a.ts` (the command hook maps it into its worktree)
63 * and edits `<worktree>/src/a.ts`; both land on the same key.
64 */
65export function prepKey(checkout: string | null, absPath: string): string {
66  return checkout ? relative(checkout, absPath) : normalize(absPath)
67}
68
69/** Same cap as error-log.ts, which archives hook-errors.jsonl for the CLI. */
70export const ERROR_LOG_MAX_BYTES = 5 * 1024 * 1024
71
72// Archive by hard link then unlink: the link fails, atomically, when the
73// archive name is taken. Where links are not supported, mv -n. find -size
74// reads the size without reading the file (busybox wc -c reads it all).
75const SH_APPEND = [
76  'f=$1 a="${1%.jsonl}.$3.jsonl"',
77  'mkdir -p "$(dirname "$f")" || exit 1',
78  'if [ -n "$(find "$f" -prune -size +$(($2 - 1))c 2>/dev/null)" ]; then',
79  '  { ln "$f" "$a" 2>/dev/null && rm -f "$f"; } || { [ -e "$a" ] || mv -n "$f" "$a" 2>/dev/null; }',
80  'fi',
81  'cat >> "$f"',
82].join('\n')
83
84const NODE_APPEND = [
85  "const fs = require('fs'), path = require('path'), [f, max, now] = process.argv.slice(1)",
86  'fs.mkdirSync(path.dirname(f), { recursive: true })',
87  "const archive = f.replace(/\\.jsonl$/, '.' + now + '.jsonl')",
88  'try {',
89  '  if (fs.statSync(f).size >= +max) {',
90  '    try { fs.linkSync(f, archive); fs.unlinkSync(f) }',
91  "    catch (e) { if (e.code !== 'EEXIST' && !fs.existsSync(archive)) fs.renameSync(f, archive) }",
92  '  }',
93  '} catch {}',
94  'fs.appendFileSync(f, fs.readFileSync(0))',
95].join('\n')
96
97/**
98 * argv that append stdin to `file`, `sh` first and `node` where there is no
99 * `sh`. A file of `maxBytes` or more is first archived as `<name>.<now>.jsonl`,
100 * as the CLI does; an existing archive is never replaced.
101 */
102export function appendArgvs(file: string, maxBytes: number, now: number): string[][] {
103  return [
104    ['sh', '-c', SH_APPEND, 'sh', file, String(maxBytes), String(now)],
105    ['node', '-e', NODE_APPEND, file, String(maxBytes), String(now)],
106  ]
107}
108
109type Run = (argv: string[], init?: { stdin?: string }) => Promise<unknown>
110
111/**
112 * Append one line to a log. Only ever appends: the file is never read and
113 * rewritten, so a failure loses this line at worst. Never throws — telemetry
114 * must not break a tool call. Callers pass `$.process.run` as a closure.
115 */
116export async function appendLog(run: Run, file: string, line: string, maxBytes: number): Promise<void> {
117  for (const argv of appendArgvs(file, maxBytes, Date.now())) {
118    try {
119      const done = (await run(argv, { stdin: line + '\n' })) as { exitCode?: number } | undefined
120      if (done?.exitCode === 0) return
121    } catch {
122      /* this appender is missing — try the next */
123    }
124  }
125}
126
127/**
128 * One hook-errors.jsonl record, in the shape `token-pilot errors` reads. The
129 * engine reports a failed hook as { kind, message } — no Error, no stack.
130 */
131export function errorLine(hook: string, failure: HookFailure | undefined, now: number): string {
132  return JSON.stringify({
133    ts: now,
134    hook,
135    level: 'error',
136    code: 'mod_hook_failed',
137    msg: `${failure?.kind ?? 'throw'}: ${failure?.message ?? 'no message'}`.slice(0, 500),
138  })
139}
140
141/**
142 * A mod hook that throws or overruns is skipped by the engine and the call
143 * goes ahead; this records why, in the log `token-pilot errors` reads. `next`
144 * here is replay-safe, so nothing beneath runs twice.
145 */
146export async function caught<E, R>(
147  $: EngineInterface,
148  e: E,
149  next: ((e: E) => Promise<R>) & { readonly error?: HookFailure },
150): Promise<R> {
151  try {
152    const home = (await $.env.get('HOME')) ?? (await $.env.get('USERPROFILE')) ?? ''
153    if (home && (await $.env.get('TOKEN_PILOT_NO_ERROR_LOG')) !== '1') {
154      const tool = (e as { tool?: unknown } | null)?.tool
155      void appendLog(
156        (argv, init) => $.process.run(argv, init),
157        `${home}/.token-pilot/hook-errors.jsonl`,
158        errorLine(`mod:${typeof tool === 'string' ? tool : 'hook'}`, next.error, Date.now()),
159        ERROR_LOG_MAX_BYTES,
160      )
161    }
162  } catch {
163    /* reporting must never fail the call */
164  }
165
166  return next(e)
167}
168
169let contextModeTool: string | undefined
170
171/**
172 * Note the session's tools as prompt.compose last saw them. The Bash advice
173 * can then name context-mode's execute tool exactly as this install has it.
174 */
175export function noteTools(tools: readonly string[]): void {
176  contextModeTool = tools.find(tool => tool.includes('context-mode') && /__(ctx_)?execute$/.test(tool)) ?? contextModeTool
177}
178
179export function contextModeToolName(): string | undefined {
180  return contextModeTool
181}
182

1/**
2 * Bash, in-process: the pre-check (deny `cat` on code, recursive grep,
3 * unbounded git log/diff) and the post-advice after a large output, in one
4 * hook. Replaces the hook-pre-bash and hook-post-bash command hooks.
5 */
6
7import type { On } from 'claude-code'
8import { decidePreBash } from '../../src/hooks/pre-bash.js'
9import { decidePostBashAdvice } from '../../src/hooks/post-bash.js'
10import { parseEnforcementMode } from '../../src/server/enforcement-mode.js'
11import { diagnosticEvent, ROTATION_THRESHOLD_BYTES, tagEvent } from '../../src/core/hook-event.js'
12import { appendLog, caught, contextModeToolName, withContext } from './host.js'
13
14export const BASH_ACTIONS = ['hook-pre-bash', 'hook-post-bash']
15
16export function registerBash(on: On): void {
17  on('tool.call', { tool: 'Bash' }, async ($, e, next) => {
18    if ((await $.env.get('TOKEN_PILOT_NO_MOD')) === '1') return next(e)
19
20    const command = String(e.command ?? '')
21    const mode = parseEnforcementMode(await $.env.get('TOKEN_PILOT_MODE'))
22    const decision = decidePreBash({ tool_name: 'Bash', tool_input: { command } }, mode, {
23      bypass: (await $.env.get('TOKEN_PILOT_BYPASS')) === '1',
24      projectRoot: await $.session.root(),
25    })
26
27    if (decision.kind === 'deny') {
28      // The reason, never the command: a command line can carry secrets.
29      const event = tagEvent(diagnosticEvent({ code: 'bash_denied', detail: { reason: decision.reason.slice(0, 80) } }, Date.now()), {
30        TOKEN_PILOT_WORKFLOW_ID: await $.env.get('TOKEN_PILOT_WORKFLOW_ID'),
31        CLAUDE_CODE_WORKFLOW_ID: await $.env.get('CLAUDE_CODE_WORKFLOW_ID'),
32        LOOM_WORKFLOW_ID: await $.env.get('LOOM_WORKFLOW_ID'),
33        LOOM_TASK_ID: await $.env.get('LOOM_TASK_ID'),
34      })
35      void appendLog(
36        (argv, init) => $.process.run(argv, init),
37        `${await $.session.root()}/.token-pilot/hook-events.jsonl`,
38        JSON.stringify(event),
39        ROTATION_THRESHOLD_BYTES,
40      )
41
42      return { deny: decision.reason }
43    }
44
45    const ran = await next(e)
46    const contextModeTool = contextModeToolName()
47    const advice = decidePostBashAdvice(
48      { tool_name: 'Bash', tool_response: { stdout: ran.text ?? '' } },
49      { contextModeAvailable: contextModeTool !== undefined, contextModeTool },
50    )
51
52    return withContext(ran, [
53      ...(decision.kind === 'advise' ? [decision.reason] : []),
54      ...(advice.additionalContext ? [advice.additionalContext] : []),
55    ])
56  }).catch(caught)
57}
58

1/**
2 * Grep, in-process: a symbol-like pattern goes to find_usages, a TODO scan
3 * gets a pointer to code_audit. Replaces the hook-pre-grep command hook.
4 */
5
6import type { On } from 'claude-code'
7import { decidePreGrep } from '../../src/hooks/pre-grep.js'
8import { parseEnforcementMode } from '../../src/server/enforcement-mode.js'
9import { caught, withContext } from './host.js'
10
11export const GREP_ACTIONS = ['hook-pre-grep']
12
13export function registerGrep(on: On): void {
14  // Claude Code 2.1.289 has no Grep tool (search goes through Bash, which
15  // bash.ts gates). A pattern, not the name, keeps this hook for the builds
16  // that still have one without naming a tool this build's types lack.
17  on('tool.call', { tool: /^Grep$/ }, async ($, e, next) => {
18    if ((await $.env.get('TOKEN_PILOT_NO_MOD')) === '1') return next(e)
19
20    const mode = parseEnforcementMode(await $.env.get('TOKEN_PILOT_MODE'))
21    const decision = decidePreGrep({ tool_name: 'Grep', tool_input: { ...e } } as never, mode)
22    if (decision.kind === 'deny') return { deny: decision.reason }
23
24    return withContext(await next(e), decision.kind === 'advise' ? [decision.reason] : [])
25  }).catch(caught)
26}
27

1/**
2 * The mod's one unmatched `tool.call` hook (engine rule: one per event). It
3 * handles token-pilot's own MCP tools:
4 *
5 *  - Worktree paths. The MCP server fixes its root at start; a session that
6 *    `cd`s into a git worktree would have relative paths read from the main
7 *    checkout. `$.session.cwd()` follows the main session's `cd`, so its paths
8 *    are rewritten here, through `next()`, which keeps Claude Code's
9 *    permission prompt as it is. A subagent's `cd` is not visible to the mod
10 *    (verified live, 2.1.289), so subagent calls — and the whole-tree warning —
11 *    stay with the hook-mcp-path command hook, which reads each call's own
12 *    cwd. Paths the mod made absolute reach that hook unchanged.
13 *  - read_for_edit calls, recorded so the Edit gate (edit.ts) knows which
14 *    files were prepared — no tmp file, no hashing.
15 */
16
17import type { EngineInterface, On } from 'claude-code'
18import { decideMcpPath } from '../../src/hooks/mcp-path.js'
19import { dirname, isAbsolute, normalize, resolve } from '../../src/core/portable-path.js'
20import { caught, findCheckout, PREFIX, prepKey } from './host.js'
21
22// hook-mcp-path keeps running: it serves subagent calls and the warnings.
23export const MCP_ACTIONS: string[] = []
24
25const TTL_MS = 30 * 60 * 1000 // same as src/core/edit-prep-state.ts
26
27export type EditPrep = {
28  mark(path: string): void
29  isFresh(path: string, now: number): boolean
30}
31
32/** read_for_edit calls seen in this session. A plugin reload clears it; the Edit gate then asks again. */
33export function createEditPrep(): EditPrep {
34  const seen = new Map<string, number>()
35
36  return {
37    mark: path => {
38      seen.set(normalize(path), Date.now())
39    },
40    isFresh: (path, now) => now - (seen.get(normalize(path)) ?? -Infinity) < TTL_MS,
41  }
42}
43
44const checkouts = new Map<string, string | null>()
45
46async function checkoutOf($: EngineInterface, dir: string): Promise<string | null> {
47  if (!checkouts.has(dir)) checkouts.set(dir, await findCheckout(p => $.fs.exists(p), dir))
48
49  return checkouts.get(dir) ?? null
50}
51
52export function registerMcp(on: On, prep: EditPrep): void {
53  on('tool.call', async ($, e, next) => {
54    if (!e.tool.startsWith(PREFIX)) return next(e)
55    if ((await $.env.get('TOKEN_PILOT_NO_MOD')) === '1') return next(e)
56    if (e.agentId !== undefined) return observe($, e, await next(e), prep)
57
58    const args = e as unknown as Record<string, unknown>
59    const root = await $.session.root()
60    const cwd = await $.session.cwd()
61    const serverCheckout = await checkoutOf($, root)
62    const sessionCheckout = await checkoutOf($, cwd)
63    const decision = decideMcpPath(
64      { tool_name: e.tool, tool_input: args, cwd },
65      {
66        projectRoot: root,
67        checkoutOf: dir => (dir === root ? serverCheckout : dir === cwd ? sessionCheckout : null),
68      },
69    )
70
71    const call = decision.kind === 'rewrite' ? { ...e, ...decision.updatedInput } : e
72
73    return observe($, call, await next(call), prep)
74  }).catch(caught)
75}
76
77/** Record a successful read_for_edit for the Edit gate; hand the result back unchanged. */
78async function observe<R extends { deny?: unknown; isError?: unknown }>(
79  $: EngineInterface,
80  call: { tool: string },
81  ran: R,
82  prep: EditPrep,
83): Promise<R> {
84  const path = (call as unknown as Record<string, unknown>).path
85  if (call.tool === `${PREFIX}read_for_edit` && ran.deny === undefined && !ran.isError && typeof path === 'string') {
86    const abs = isAbsolute(path) ? path : resolve(await $.session.root(), path)
87    prep.mark(prepKey(await checkoutOf($, dirname(abs)), abs))
88  }
89
90  return ran
91}
92

1/**
2 * Edit, in-process: an Edit of an existing code file must follow a
3 * read_for_edit of that file. The mod knows which files were prepared from
4 * the read_for_edit calls it saw (mcp.ts), so there is no tmp-file state and
5 * no hashing. Replaces the hook-edit command hook.
6 *
7 * A file the agent wrote itself (Write) counts as prepared: it knows every
8 * byte. A file outside the project is not gated — read_for_edit refuses it.
9 */
10
11import type { EngineInterface, On } from 'claude-code'
12import { decidePreEdit } from '../../src/hooks/pre-edit.js'
13import { isCodeFile } from '../../src/hooks/read-gate.js'
14import { parseEnforcementMode } from '../../src/server/enforcement-mode.js'
15import type { EditPrep } from './mcp.js'
16import { dirname } from '../../src/core/portable-path.js'
17import { caught, findCheckout, isInside, prepKey, withContext } from './host.js'
18
19export const EDIT_ACTIONS = ['hook-edit']
20
21async function realPath($: EngineInterface, path: string): Promise<string> {
22  const stat = await $.fs.stat(path, { resolve: true }).catch(() => null)
23
24  return stat?.realPath ?? path
25}
26
27async function keyOf($: EngineInterface, filePath: string): Promise<string> {
28  return prepKey(await findCheckout(p => $.fs.exists(p), dirname(filePath)), filePath)
29}
30
31export function registerEdit(on: On, prep: EditPrep): void {
32  on('tool.call', { tool: 'Edit' }, async ($, e, next) => {
33    if ((await $.env.get('TOKEN_PILOT_NO_MOD')) === '1') return next(e)
34
35    const filePath = String(e.file_path ?? '')
36    const fileExists = await $.fs.exists(filePath)
37    const decision = decidePreEdit(
38      { tool_name: 'Edit', tool_input: { file_path: filePath } },
39      {
40        mode: parseEnforcementMode(await $.env.get('TOKEN_PILOT_MODE')),
41        isCodeFile: isCodeFile(filePath),
42        fileExists,
43        isPrepared: prep.isFresh(await keyOf($, filePath), Date.now()),
44        bypassed: (await $.env.get('TOKEN_PILOT_BYPASS')) === '1',
45        outsideProject:
46          fileExists && !isInside(await realPath($, await $.session.root()), await realPath($, filePath)),
47      },
48    )
49    if (decision.kind === 'deny') return { deny: decision.reason }
50
51    return withContext(await next(e), decision.kind === 'advise' ? [decision.message] : [])
52  }).catch(caught)
53
54  on('tool.call', { tool: 'Write' }, async ($, e, next) => {
55    const ran = await next(e)
56    const filePath = String(e.file_path ?? '')
57    if (ran.deny === undefined && !ran.isError && isCodeFile(filePath)) prep.mark(await keyOf($, filePath))
58
59    return ran
60  }).catch(caught)
61}
62

1/**
2 * Read, in-process. A whole-file Read of a big code file gets the file's
3 * structural outline (smart_read) back as a normal Read result, with a header
4 * that says so — instead of the command hook's refusal the model then has to
5 * recover from. Core still reads one line, so Claude Code counts the file as
6 * read and a later Edit works. Replaces the hook-read command hook.
7 */
8
9import type { EngineInterface, On } from 'claude-code'
10import {
11  decideReadGate,
12  decideReadGateFromStats,
13  isCodeFile,
14  outlineHeader,
15  spanCannotExceed,
16} from '../../src/hooks/read-gate.js'
17import { computeEffectiveThreshold } from '../../src/hooks/adaptive-threshold.js'
18import { estimateTokens } from '../../src/core/token-estimator.js'
19import { relative } from '../../src/core/portable-path.js'
20import { ROTATION_THRESHOLD_BYTES, tagEvent } from '../../src/core/hook-event.js'
21import type { HookEvent } from '../../src/core/event-log.js'
22import { appendLog, caught, configFrom, isInside, PREFIX } from './host.js'
23
24export const READ_ACTIONS = ['hook-read']
25
26// $.fs.read rejects files this large; they are measured with wc -l instead.
27const FS_READ_MAX = 4 * 1024 * 1024
28
29/**
30 * Line count of a file $.fs.read will not return; without `wc`, estimated
31 * from its size. null when wc ran and failed: the file cannot be read, and
32 * Read is the one to say so.
33 */
34async function countLines($: EngineInterface, filePath: string, size: number): Promise<number | null> {
35  let run
36  try {
37    run = await $.process.run(['wc', '-l', filePath])
38  } catch {
39    return Math.ceil(size / 40) // no wc (native Windows)
40  }
41
42  const lines = Number.parseInt(run.stdout, 10)
43
44  return run.exitCode === 0 && Number.isFinite(lines) ? lines : null
45}
46
47export function registerRead(on: On): void {
48  let savedThisSession = 0
49
50  on('tool.call', { tool: 'Read' }, async ($, e, next) => {
51    if ((await $.env.get('TOKEN_PILOT_NO_MOD')) === '1') return next(e)
52
53    const filePath = String(e.file_path)
54    if (!isCodeFile(filePath)) return next(e)
55
56    const root = await $.session.root()
57    // Literal names: the engine lists the variables a module reads.
58    const env = {
59      TOKEN_PILOT_DENY_THRESHOLD: await $.env.get('TOKEN_PILOT_DENY_THRESHOLD'),
60      TOKEN_PILOT_ADAPTIVE_THRESHOLD: await $.env.get('TOKEN_PILOT_ADAPTIVE_THRESHOLD'),
61      TOKEN_PILOT_ADAPTIVE_BUDGET: await $.env.get('TOKEN_PILOT_ADAPTIVE_BUDGET'),
62      TOKEN_PILOT_MODE: await $.env.get('TOKEN_PILOT_MODE'),
63      TOKEN_PILOT_BYPASS: await $.env.get('TOKEN_PILOT_BYPASS'),
64    }
65    const config = configFrom(await $.fs.read(`${root}/.token-pilot.json`).then(String, () => null), env)
66    if (config.hooks.mode === 'off') return next(e)
67
68    // Real paths, so a symlink pointing out of the project passes through.
69    const realRoot = (await $.fs.stat(root, { resolve: true })).realPath ?? root
70    // A missing file is Read's to report, not a hook failure.
71    const fileStat = await $.fs.stat(filePath, { resolve: true }).catch(() => null)
72    if (!fileStat || !isInside(realRoot, fileStat.realPath ?? filePath)) return next(e)
73
74    const threshold = config.hooks.adaptiveThreshold
75      ? computeEffectiveThreshold({
76          baseThreshold: config.hooks.denyThreshold,
77          sessionSavedTokens: savedThisSession,
78          sessionBudgetTokens: config.hooks.adaptiveBudgetTokens ?? 100_000,
79          enabled: true,
80        })
81      : config.hooks.denyThreshold
82    const offset = typeof e.offset === 'number' ? e.offset : null
83    const limit = typeof e.limit === 'number' ? e.limit : null
84    if (spanCannotExceed(offset, limit, threshold)) return next(e)
85
86    const size = typeof fileStat.size === 'number' ? fileStat.size : 0
87    const content = size >= FS_READ_MAX ? null : await $.fs.read(filePath).then(String, () => null)
88    const lineCount = content === null ? await countLines($, filePath, size) : 0
89    if (lineCount === null) return next(e)
90    const gate =
91      content === null
92        ? decideReadGateFromStats({ filePath, lineCount, bytes: size, offset, limit, threshold })
93        : decideReadGate({ filePath, content, offset, limit, threshold })
94    if (gate.kind === 'pass') return next(e)
95
96    const pointer =
97      `[token-pilot] ${filePath} has ${gate.lineCount} lines. Use ${PREFIX}smart_read for its structure, ` +
98      `${PREFIX}read_symbol for one symbol, or Read with offset/limit.`
99    if (config.hooks.mode === 'advisory') return { deny: pointer }
100
101    // force: the server de-duplicates files it already sent, and would answer
102    // a reader that never saw this one with a "previously loaded" reminder.
103    let text: string | undefined
104    try {
105      const outline = await $.tool.call({ tool: `${PREFIX}smart_read`, path: filePath, force: true } as never)
106      text = outline.deny === undefined && !outline.isError ? outline.text : undefined
107    } catch {
108      text = undefined
109    }
110    if (!text) return { deny: pointer }
111
112    const ran = await next({ ...e, offset: 1, limit: 1 })
113    if (ran.deny !== undefined || ran.isError) return ran
114
115    const saved = Math.max(0, gate.estTokens - estimateTokens(text))
116    savedThisSession += saved
117    try {
118      const event = tagEvent({
119        ts: Date.now(),
120        session_id: await $.session.id(),
121        agent_type: null,
122        agent_id: e.agentId ?? null,
123        event: 'denied',
124        file: filePath,
125        lines: gate.lineCount,
126        estTokens: gate.estTokens,
127        summaryTokens: gate.estTokens - saved,
128        savedTokens: saved,
129      } as HookEvent, {
130        TOKEN_PILOT_WORKFLOW_ID: await $.env.get('TOKEN_PILOT_WORKFLOW_ID'),
131        CLAUDE_CODE_WORKFLOW_ID: await $.env.get('CLAUDE_CODE_WORKFLOW_ID'),
132        LOOM_WORKFLOW_ID: await $.env.get('LOOM_WORKFLOW_ID'),
133        LOOM_TASK_ID: await $.env.get('LOOM_TASK_ID'),
134      })
135      void appendLog(
136        (argv, init) => $.process.run(argv, init),
137        `${root}/.token-pilot/hook-events.jsonl`,
138        JSON.stringify(event),
139        ROTATION_THRESHOLD_BYTES,
140      )
141    } catch {
142      /* telemetry must never cost the model its outline */
143    }
144
145    const body = outlineHeader(relative(root, filePath), gate.lineCount, gate.estTokens, PREFIX) + text
146
147    // Claude Code de-duplicates a Read of an unchanged file with the range it
148    // already served — our own 1-line read, the second time round — and
149    // answers "unchanged". The model has only ever seen an outline of this
150    // file, never its text, so it gets the outline again.
151    if (ran.result?.type !== 'text') {
152      const file = { filePath, content: body, numLines: 1, startLine: 1, totalLines: gate.lineCount }
153      return { ...ran, result: { type: 'text', file } } as typeof ran
154    }
155
156    return { ...ran, result: { ...ran.result, file: { ...ran.result.file, content: body } } }
157  }).catch(caught)
158}
159

1/**
2 * Agent, in-process: routing before the subagent starts (hook-pre-task).
3 *
4 * Budget and task telemetry stay on the command hooks: agents run in the
5 * background by default, so the Agent result here is only the launch
6 * acknowledgement — SubagentStop is what sees the real answer.
7 *
8 * Workflow runs (TOKEN_PILOT_WORKFLOW_ID and friends) stay on the command
9 * hooks too, which attach the workflow budget note; register.ts leaves this
10 * action out of the hand-off flag for such sessions.
11 */
12
13import type { EngineInterface, On } from 'claude-code'
14import { buildAgentIndexFromFiles, type AgentIndex } from '../../src/core/agent-matcher.js'
15import { decidePreTask, subagentNeedsToolGuide, subagentToolGuide } from '../../src/hooks/pre-task.js'
16import { parseEnforcementMode } from '../../src/server/enforcement-mode.js'
17import { join } from '../../src/core/portable-path.js'
18import { caught, withContext } from './host.js'
19
20export const AGENT_ACTIONS = ['hook-pre-task']
21
22type AgentFile = { fileName: string; body: string }
23
24async function inWorkflow($: EngineInterface): Promise<boolean> {
25  return Boolean(
26    (await $.env.get('TOKEN_PILOT_WORKFLOW_ID')) ??
27      (await $.env.get('CLAUDE_CODE_WORKFLOW_ID')) ??
28      (await $.env.get('LOOM_WORKFLOW_ID')),
29  )
30}
31
32async function readAgentFiles($: EngineInterface, dir: string): Promise<AgentFile[]> {
33  const files: AgentFile[] = []
34  try {
35    for (const entry of await $.fs.list(dir)) {
36      const fileName = String((entry as { name: string }).name)
37      if (!fileName.startsWith('tp-') || !fileName.endsWith('.md')) continue
38      files.push({ fileName, body: String(await $.fs.read(join(dir, fileName))) })
39    }
40  } catch {
41    /* no agents dir — nothing to route to */
42  }
43
44  return files
45}
46
47export function registerAgent(on: On): void {
48  let files: AgentFile[] | null = null
49  let index: AgentIndex | null = null
50
51  on('tool.call', { tool: 'Agent' }, async ($, e, next) => {
52    if ((await $.env.get('TOKEN_PILOT_NO_MOD')) === '1') return next(e)
53    if (await inWorkflow($)) return next(e)
54
55    files ??= await readAgentFiles($, join($.plugin.root, 'agents'))
56    index ??= buildAgentIndexFromFiles(files)
57
58    const input = {
59      tool_name: 'Agent',
60      tool_input: { subagent_type: e.subagent_type, description: e.description, prompt: e.prompt },
61    }
62    const decision = decidePreTask(input, {
63      mode: parseEnforcementMode(await $.env.get('TOKEN_PILOT_MODE')),
64      agentIndex: index,
65      force: (await $.env.get('TOKEN_PILOT_FORCE_SUBAGENTS')) === '1',
66      agentNamePrefix: 'token-pilot:',
67    })
68    if (decision.kind === 'deny') return { deny: decision.reason }
69
70    // The tool guide is for the subagent: it goes into its prompt, not into
71    // the parent's context after the launch.
72    const call =
73      subagentNeedsToolGuide(input) && typeof e.prompt === 'string'
74        ? { ...e, prompt: `${e.prompt}\n\n${subagentToolGuide()}` }
75        : e
76
77    return withContext(await next(call), decision.kind === 'advise' ? [decision.message] : [])
78  }).catch(caught)
79}
80

1/**
2 * Session guidance as a system-prompt section. The SessionStart command hook
3 * puts it into the first message, and UserPromptSubmit re-sends an anchor on
4 * every turn; a `session` section is part of every request instead, built
5 * once and kept stable so the prompt cache holds. Replaces hook-session-start,
6 * hook-bootstrap and hook-user-prompt.
7 *
8 * Duplicate hook registrations are a note for the person, not the model:
9 * they go out as a toast. Not ported: the subagent-adoption nudge (it reads
10 * the whole event log at every start) and the bootstrap notes (a plugin
11 * always carries its agents; the MCP server reports a missing ast-index).
12 */
13
14import type { EngineInterface, On } from 'claude-code'
15import {
16  buildReminderMessage,
17  countTokenPilotHooks,
18  duplicateWarning,
19  parseAgentEntry,
20  profileBannerNote,
21  snapshotLine,
22  type AgentEntry,
23} from '../../src/hooks/session-context.js'
24import { parseProfileEnv } from '../../src/server/tool-profiles.js'
25import { join } from '../../src/core/portable-path.js'
26import { caught, configFrom, noteTools, PREFIX } from './host.js'
27
28export const SESSION_ACTIONS = ['hook-session-start', 'hook-bootstrap', 'hook-user-prompt']
29
30async function readText($: EngineInterface, path: string): Promise<string | null> {
31  try {
32    const text = await $.fs.read(path)
33    return typeof text === 'string' ? text : null
34  } catch {
35    return null
36  }
37}
38
39async function homeDir($: EngineInterface): Promise<string> {
40  return (await $.env.get('HOME')) ?? (await $.env.get('USERPROFILE')) ?? ''
41}
42
43async function agentEntries($: EngineInterface, dir: string): Promise<AgentEntry[]> {
44  const entries: AgentEntry[] = []
45  try {
46    for (const entry of await $.fs.list(dir)) {
47      const fileName = String((entry as { name: string }).name)
48      if (!fileName.startsWith('tp-') || !fileName.endsWith('.md')) continue
49      const body = await readText($, join(dir, fileName))
50      if (body !== null) entries.push(parseAgentEntry(fileName, body))
51    }
52  } catch {
53    /* no agents dir */
54  }
55
56  return entries
57}
58
59async function sessionText($: EngineInterface): Promise<string | null> {
60  const root = await $.session.root()
61  // Literal names: the engine lists the variables a module reads.
62  const env = {
63    TOKEN_PILOT_DENY_THRESHOLD: await $.env.get('TOKEN_PILOT_DENY_THRESHOLD'),
64    TOKEN_PILOT_ADAPTIVE_THRESHOLD: await $.env.get('TOKEN_PILOT_ADAPTIVE_THRESHOLD'),
65    TOKEN_PILOT_ADAPTIVE_BUDGET: await $.env.get('TOKEN_PILOT_ADAPTIVE_BUDGET'),
66    TOKEN_PILOT_MODE: await $.env.get('TOKEN_PILOT_MODE'),
67    TOKEN_PILOT_BYPASS: await $.env.get('TOKEN_PILOT_BYPASS'),
68  }
69  const config = configFrom(await readText($, join(root, '.token-pilot.json')), env)
70  if (!config.sessionStart.enabled || (await $.env.get('TOKEN_PILOT_BYPASS')) === '1') return null
71
72  // Project agents first, then home agents, then the plugin's own — named as
73  // Claude Code dispatches them (`token-pilot:tp-*`).
74  const seen = new Set<string>()
75  const agents: AgentEntry[] = []
76  for (const [dir, prefix] of [
77    [join(root, '.claude', 'agents'), ''],
78    [join(await homeDir($), '.claude', 'agents'), ''],
79    [join($.plugin.root, 'agents'), 'token-pilot:'],
80  ]) {
81    for (const agent of await agentEntries($, dir)) {
82      if (seen.has(agent.name)) continue
83      seen.add(agent.name)
84      agents.push({ ...agent, name: prefix + agent.name })
85    }
86  }
87
88  let text =
89    profileBannerNote(parseProfileEnv(await $.env.get('TOKEN_PILOT_PROFILE'))) +
90    buildReminderMessage(agents, config.sessionStart.maxReminderTokens)
91
92  const snapshot = join(root, '.token-pilot', 'snapshots', 'latest.md')
93  const body = await readText($, snapshot)
94  if (body !== null) {
95    try {
96      const { mtimeMs } = await $.fs.stat(snapshot)
97      const line = typeof mtimeMs === 'number' ? snapshotLine(body, Math.max(0, Date.now() - mtimeMs)) : null
98      if (line) text += `\n\n${line}`
99    } catch {
100      /* no snapshot age — skip the line */
101    }
102  }
103
104  return text
105}
106
107async function warnDuplicates($: EngineInterface): Promise<void> {
108  const root = await $.session.root()
109  const sources: Array<{ path: string; count: number }> = []
110
111  for (const path of [
112    join(await homeDir($), '.claude', 'settings.json'),
113    join(root, '.claude', 'settings.json'),
114    join(root, '.claude', 'settings.local.json'),
115  ]) {
116    const text = await readText($, path)
117    if (text === null) continue
118    try {
119      const count = countTokenPilotHooks(JSON.parse(text))
120      if (count > 0) sources.push({ path, count })
121    } catch {
122      /* malformed settings — not ours to judge */
123    }
124  }
125
126  const warning = duplicateWarning(sources)
127  if (warning) $.ui.toast(warning)
128}
129
130// Built once per session: stable text keeps the prompt cache; a /clear starts a new session.
131let section: { session: string; text: Promise<string | null> } | null = null
132let warned = false
133
134export function registerSession(on: On): void {
135  on('prompt.compose', async ($, e, next) => {
136    const out = await next(e)
137    noteTools(e.tools)
138    if (!e.tools.some(tool => tool.startsWith(PREFIX))) return out
139    if ((await $.env.get('TOKEN_PILOT_NO_MOD')) === '1') return out
140
141    if (!warned) {
142      warned = true
143      void warnDuplicates($).catch(() => {})
144    }
145
146    const session = await $.session.id()
147    if (section?.session !== session) section = { session, text: sessionText($).catch(() => null) }
148    const text = await section.text
149
150    return text ? { sections: [...out.sections, { id: 'token-pilot', scope: 'session' as const, text }] } : out
151  }).catch(caught)
152}
153

1/**
2 * The savings line as a band above the prompt input: one line, drawn by the
3 * mod, in place of a second copy beside the user's `statusLine`. The numbers
4 * come from hooks/tp-statusline.sh — the single savings calculator — run
5 * after each main-loop turn, when they can have changed, and when /clear or
6 * /resume starts another session.
7 *
8 * That script prints nothing in a statusLine for a session the mod serves
9 * (TOKEN_PILOT_MOD / TOKEN_PILOT_MOD_SESSION), so the line shows once. Not
10 * drawn at session start: Claude Code runs the statusLine once before the
11 * session's hooks hand those variables off, and again only when the session
12 * changes (a reply, the mode, the model). Until the first reply that run's
13 * line is the one shown; a band drawn earlier would sit beside it.
14 */
15
16import { atom, read, update } from 'claude-code'
17import type { EngineInterface, On } from 'claude-code'
18import { caught } from './host.js'
19
20// tp-statusline.sh colours with bash $'\033[…' quoting — run it with bash, then strip.
21const ANSI = /\u001b\[[0-9;]*m/g
22
23const line = atom({ plugin: 'token-pilot', key: 'band' } as const, null)
24
25async function refresh($: EngineInterface, session: string): Promise<void> {
26  if ((await $.env.get('TOKEN_PILOT_NO_MOD')) === '1') return
27
28  // The statusLine payload's shape, which the script parses.
29  const { rateLimits } = await $.session.usage()
30  const payload = JSON.stringify({
31    session_id: session,
32    cwd: await $.session.cwd(),
33    rate_limits: Object.fromEntries(rateLimits.map(limit => [limit.kind, { used_percentage: limit.percentUsed }])),
34  })
35  // An empty TOKEN_PILOT_MOD lets the script print for this call: its silence is for the statusLine.
36  const run = await $.process.run(['bash', `${$.plugin.root}/hooks/tp-statusline.sh`], {
37    stdin: payload,
38    env: { TOKEN_PILOT_MOD: '' },
39  })
40  const text = (run.stdout.split('\n')[0] ?? '').replace(ANSI, '').trim()
41
42  // The write redraws the band.
43  await update($, line, () => text || null)
44}
45
46export function registerBand(on: On): void {
47  on('session.start', { isInteractive: true }, async ($, e, next) => {
48    // A reload or an update of the plugin keeps the pin 1.0.3 and earlier set.
49    $.ui.status(undefined)
50
51    return next(e)
52  }).catch(caught)
53
54  // Past the first startup the statusLine has gone quiet already: redraw the
55  // line for the new session at once, or the band keeps the last one's figures.
56  // The payload's session_id: $.session.id() still answers the session being left.
57  on('classic.SessionStart', { source: ['clear', 'resume', 'compact', 'fork'] }, async ($, e, next) => {
58    const done = await next(e)
59    await refresh($, e.session_id).catch(() => {})
60
61    return done
62  }).catch(caught)
63
64  on('turn.complete', async ($, e, next) => {
65    const done = await next(e)
66    if (e.agentId === undefined) await refresh($, await $.session.id()).catch(() => {})
67
68    return done
69  }).catch(caught)
70
71  on('ui.render', { component: 'AbovePrompt' }, async ($, e, next) => {
72    const text = await read($, line)
73    if (text === null || e.props.hasSurvey) return next(e)
74
75    const { Box, Text } = $.ui.resolve(e)
76
77    return (
78      <Box width={e.props.bodyColumns}>
79        <Text dimColor wrap="truncate-end">
80          {text}
81        </Text>
82      </Box>
83    )
84  }).catch(caught)
85}
86

1/**
2 * /tp-stats — this session's token-pilot savings in a pane, from the MCP
3 * server's session_analytics. The command itself is registered in
4 * register.ts (its one session.start hook).
5 */
6
7import type { On } from 'claude-code'
8import { PREFIX, caught } from './host.js'
9
10export const STATS_COMMAND = { name: 'tp-stats', description: 'Show token-pilot savings for this session' }
11
12const PANE = 'tp-stats'
13let text = 'No data yet.'
14
15export function registerStats(on: On): void {
16  on('command.run', { command: STATS_COMMAND.name }, async $ => {
17    const r = await $.tool.call({ tool: `${PREFIX}session_analytics` } as never)
18    text = (r.deny === undefined && !r.isError && r.text) || 'No token-pilot data for this session yet.'
19    await $.ui.open({ id: PANE, title: 'token-pilot' })
20
21    return { text: 'token-pilot stats opened.' }
22  }).catch(caught)
23
24  on('ui.render', { component: 'Pane', requestId: PANE }, async ($, e) => {
25    const { Box, Text } = $.ui.resolve(e)
26
27    return (
28      <Box flexDirection="column">
29        {text.split('\n').map(line => (
30          <Text>{line}</Text>
31        ))}
32      </Box>
33    )
34  })
35}
36