rules-injector

1import type { EngineInterface, ProcessRunResult, Register } from 'claude-code'
2
3export const SECTION_ID = 'rules-injector:rules'
4// About 3k tokens: room for house.md (~3.7 KB) and plenty of user rules.
5export const CAP = 12_000
6const MIN_FORGET = 4
7const USAGE = 'usage: /rules | /rules add <text> | /rules add-project <text> | /rules forget <pattern>'
8
9type Part = { kind: 'project' | 'global' | 'house'; path: string; body: string }
10
11// `zstack rule list` prints `## <path>` then the file, for house.md, global.md and the repo's file.
12export function parse(out: string, rulesDir: string): Part[] {
13  const parts: Part[] = []
14  for (const chunk of out.split(/^## (?=\/)/m).slice(1)) {
15    const nl = chunk.indexOf('\n')
16    const path = (nl < 0 ? chunk : chunk.slice(0, nl)).trim()
17    const body = (nl < 0 ? '' : chunk.slice(nl + 1)).replace(/\s*<!--.*?-->/g, '').trim()
18    if (!body) continue
19    const kind = !path.startsWith(`${rulesDir}/`) ? 'house' : path.endsWith('/global.md') ? 'global' : 'project'
20    parts.push({ kind, path, body })
21  }
22  return parts
23}
24
25// Most specific first, so a cut drops generic house rules before the user's own.
26export function compose(parts: Part[], cwd: string): string | undefined {
27  const order = { project: 0, global: 1, house: 2 } as const
28  const title = { project: `Project rules (${cwd})`, global: 'Global rules', house: 'House rules' } as const
29  const sorted = [...parts].sort((a, b) => order[a.kind] - order[b.kind])
30  if (!sorted.length) return undefined
31  const text = [
32    '# Standing rules (zstack)',
33    'The user set these rules and has had to repeat them. Follow every one without being reminded. The user manages them with /rules.',
34    ...sorted.map(p => `## ${title[p.kind]}\n${p.body.replace(/^# House rules\n+/, '').replace(/^#{1,2} /gm, '### ')}`),
35  ].join('\n\n')
36  if (text.length <= CAP) return text
37  const cut = text.slice(0, text.lastIndexOf('\n', CAP))
38  return `${cut}\n\n[${text.length - cut.length} more characters of rules cut; run /rules to see all]`
39}
40
41type Cache = { cwd: string; sig: string; house?: string; text?: string }
42
43// Module state is fine here: nothing draws from it, and a reload just re-reads the rules.
44let cache: Cache | undefined
45let pending: Promise<string | undefined> | undefined
46let isDirty = false
47let isMissing = false
48
49async function rulesDir($: EngineInterface): Promise<string> {
50  const home = await $.env.get('ZSTACK_HOME')
51  return `${home ?? `${(await $.env.get('HOME')) ?? ''}/.local/state/zstack`}/rules`
52}
53
54// Every rules file's mtime and size: `rule add` appends, `rule forget` rewrites,
55// a new project file appears in the directory, house.md is edited by hand.
56async function signature($: EngineInterface, dir: string, house?: string): Promise<string> {
57  const files = await $.fs.list(dir).catch(() => [])
58  const parts = files.map(f => `${f.name}:${f.mtimeMs}:${f.size}`).sort()
59  if (house) {
60    const st = await $.fs.stat(house).catch(() => undefined)
61    parts.push(`house:${st?.mtimeMs ?? 0}:${st?.size ?? 0}`)
62  }
63  return parts.join('|')
64}
65
66function zstack($: EngineInterface, args: string[]): Promise<ProcessRunResult> {
67  return $.process.run(['zstack', ...args], { timeoutMs: 10_000 })
68}
69
70async function load($: EngineInterface): Promise<string | undefined> {
71  if (isMissing) return undefined
72  const cwd = await $.session.cwd()
73  const dir = await rulesDir($)
74  const sig = await signature($, dir, cache?.house)
75  if (cache && !isDirty && cache.cwd === cwd && cache.sig === sig) return cache.text
76  isDirty = false
77
78  let out: ProcessRunResult
79  try {
80    out = await zstack($, ['rule', 'list', '--repo', cwd])
81  } catch {
82    if (cache?.text === undefined) {
83      isMissing = true
84      $.ui.status('rules-injector: zstack not found, standing rules are not injected')
85      return undefined
86    }
87    cache = { ...cache, sig }
88    return cache.text
89  }
90  // A failed read keeps the last good rules and retries on the next change.
91  if (out.exitCode !== 0) {
92    cache = { cwd, sig, house: cache?.house, text: cache?.text }
93    return cache.text
94  }
95  const parts = parse(out.stdout, dir)
96  const house = parts.find(p => p.kind === 'house')?.path
97  cache = { cwd, sig: house === cache?.house ? sig : await signature($, dir, house), house, text: compose(parts, cwd) }
98  return cache.text
99}
100
101function rules($: EngineInterface): Promise<string | undefined> {
102  pending ??= load($).finally(() => {
103    pending = undefined
104  })
105  return pending
106}
107
108function said(r: ProcessRunResult): string {
109  return (r.exitCode === 0 ? r.stdout : r.stderr || `zstack exited ${r.exitCode}`).trim()
110}
111
112async function run($: EngineInterface, args: string[]): Promise<string> {
113  try {
114    const r = await zstack($, args)
115    if (isMissing) {
116      isMissing = false
117      $.ui.status(undefined)
118    }
119    if (r.exitCode === 0 && args[1] !== 'list') isDirty = true
120    return said(r)
121  } catch {
122    return 'zstack is not on PATH; standing rules are not injected.'
123  }
124}
125
126async function command($: EngineInterface, args: string): Promise<string> {
127  const m = /^(\S*)\s*([\s\S]*)$/.exec(args.trim())
128  const sub = m?.[1] ?? ''
129  const rest = (m?.[2] ?? '').trim()
130  const cwd = await $.session.cwd()
131
132  if (sub === '' || sub === 'list') {
133    const listed = await run($, ['rule', 'list', '--repo', cwd])
134    const text = await rules($)
135    const note = text === undefined ? 'nothing is injected' : `${text.length} characters injected into the system prompt`
136    return `${listed || 'no rules yet'}\n\n(${note})`
137  }
138  if ((sub === 'add' || sub === 'add-project' || sub === 'forget') && !rest) return USAGE
139  if (sub === 'add') return run($, ['rule', 'add', '--', rest])
140  if (sub === 'add-project') return run($, ['rule', 'add', '--project', '--repo', cwd, '--', rest])
141  if (sub === 'forget') {
142    // zstack forgets every line containing the pattern, in every project's file.
143    if (rest.length < MIN_FORGET) return `pattern too short: use at least ${MIN_FORGET} characters (it removes every matching rule line in every project)`
144    return run($, ['rule', 'forget', '--', rest])
145  }
146  return USAGE
147}
148
149export const register: Register = on => {
150  on('session.start', async ($, e, next) => {
151    await $.command.register({
152      name: 'rules',
153      description: 'Show zstack standing rules; add, add-project or forget one',
154      argumentHint: '[add <text> | add-project <text> | forget <pattern>]',
155    })
156    return next(e)
157  })
158
159  on('command.run', { command: 'rules' }, async ($, e) => ({ text: await command($, e.args) }))
160
161  // Fail open: if anything here throws, the prompt goes out as the engine composed it.
162  on('prompt.compose', async ($, e, next) => {
163    const composed = await next(e)
164    const text = await rules($)
165    if (text === undefined) return composed
166    const sections = composed.sections.filter(s => s.id !== SECTION_ID)
167    return { sections: [...sections, { id: SECTION_ID, text, scope: 'session' as const }] }
168  }).catch(($, e, next) => next(e))
169}
170