branch-guard

1import { atom, read, update } from 'claude-code'
2import type { EngineInterface, Register } from 'claude-code'
3
4import { absolute, isFresh, isProtected, parentDir, parseBranches, targetPath } from './rules'
5
6const isOff = atom({ plugin: 'branch-guard', key: 'isOff' } as const, false)
7
8const TOOLS = ['Edit', 'Write', 'NotebookEdit'] as const
9
10async function git($: EngineInterface, dir: string, args: string[]): Promise<{ ok: boolean; out: string }> {
11  const ran = await $.process.run(['git', '-C', dir, ...args]).catch(() => undefined)
12
13  return { ok: ran?.exitCode === 0, out: ran?.stdout.trim() ?? '' }
14}
15
16// The closest existing folder of the file (a Write may create new folders).
17async function existingDir($: EngineInterface, path: string): Promise<string | null> {
18  let dir = parentDir(path)
19  for (let i = 0; i < 20 && dir !== null; i++) {
20    if (await $.fs.exists(dir).catch(() => false)) return dir
21    dir = parentDir(dir)
22  }
23
24  return null
25}
26
27async function repoRoot($: EngineInterface, dir: string): Promise<string | null> {
28  const top = await git($, dir, ['rev-parse', '--show-toplevel'])
29
30  return top.ok && top.out !== '' ? top.out : null
31}
32
33async function branchOf($: EngineInterface, cache: Map<string, { branch: string; at: number }>, root: string, dir: string): Promise<string | null> {
34  const now = await $.clock.now()
35  const hit = cache.get(root)
36  if (hit !== undefined && isFresh(hit.at, now)) return hit.branch
37  const head = await git($, dir, ['rev-parse', '--abbrev-ref', 'HEAD'])
38  if (!head.ok) return null
39  cache.set(root, { branch: head.out, at: now })
40
41  return head.out
42}
43
44// A repo root to ask about, or null when the edit needs no question.
45async function mustAsk(
46  $: EngineInterface,
47  branches: Set<string>,
48  cache: Map<string, { branch: string; at: number }>,
49  allowed: Set<string>,
50  input: unknown,
51): Promise<{ path: string; root: string; branch: string } | null> {
52  const given = targetPath(input)
53  if (given === null) return null
54  const path = given.startsWith('/') ? given : absolute(given, await $.session.cwd())
55  const dir = await existingDir($, path)
56  if (dir === null) return null
57  const root = await repoRoot($, dir)
58  if (root === null || allowed.has(root)) return null
59  const branch = await branchOf($, cache, root, dir)
60  if (branch === null || !isProtected(branch, branches)) return null
61  if ((await git($, dir, ['check-ignore', '-q', path])).ok) return null
62
63  return { path: given, root, branch }
64}
65
66// An asked edit that ran (the user allowed it) opens the whole repo.
67function settle<T extends { deny?: unknown; isError?: boolean }>(pending: Map<string, string>, allowed: Set<string>, path: string, ran: T): T {
68  const root = pending.get(path)
69  if (root !== undefined) {
70    pending.delete(path)
71    if (ran.deny === undefined && ran.isError !== true) allowed.add(root)
72  }
73
74  return ran
75}
76
77export const register: Register = (on, options) => {
78  const branches = parseBranches(options?.protectedBranches)
79  const cache = new Map<string, { branch: string; at: number }>()
80  const allowed = new Set<string>()
81  const pending = new Map<string, string>()
82
83  on('session.start', async ($, e, next) => {
84    await $.command.register({
85      name: 'branch-guard',
86      description: 'branch-guard on|off|status: ask before editing files on main/master',
87    })
88
89    return next(e)
90  })
91
92  on('command.run', { command: 'branch-guard' }, async ($, e) => {
93    const arg = e.args.trim()
94    if (arg === 'off' || arg === 'on') {
95      await update($, isOff, () => arg === 'off')
96      $.ui.status(arg === 'off' ? 'branch-guard: OFF' : undefined)
97    }
98    const off = await read($, isOff)
99
100    return { text: `branch-guard is ${off ? 'OFF for this session' : 'ON'}. Protected: ${[...branches].join(', ')}.` }
101  })
102
103  for (const tool of TOOLS) {
104    on('tool.check', { tool }, async ($, e, next) => {
105      const verdict = await next(e)
106      if (verdict.decision === 'deny' || (await read($, isOff))) return verdict
107      const found = await mustAsk($, branches, cache, allowed, e.input)
108      if (found === null) return verdict
109      pending.set(found.path, found.root)
110
111      return {
112        decision: 'ask',
113        reason: `branch-guard: this edits a file on ${found.branch}. Confirm if the user wants changes there; it will not ask again for this repo in this session. If refused, ask the user to switch branch or stop the checks with /branch-guard off.`,
114      }
115    })
116  }
117
118  on('tool.call', { tool: 'Edit' }, async (_$, e, next) => settle(pending, allowed, e.file_path, await next(e)))
119  on('tool.call', { tool: 'Write' }, async (_$, e, next) => settle(pending, allowed, e.file_path, await next(e)))
120  on('tool.call', { tool: 'NotebookEdit' }, async (_$, e, next) => settle(pending, allowed, e.notebook_path, await next(e)))
121}
122

1export const DEFAULT_BRANCHES = ['main', 'master']
2export const CACHE_MS = 30_000
3
4// "main, develop" -> Set; an empty or missing option means the defaults.
5export function parseBranches(raw: unknown): Set<string> {
6  const names = typeof raw === 'string' ? raw.split(',').map(s => s.trim()).filter(Boolean) : []
7
8  return new Set(names.length > 0 ? names : DEFAULT_BRANCHES)
9}
10
11export const isProtected = (branch: string, branches: Set<string>) => branches.has(branch)
12
13export const isFresh = (at: number, now: number) => now - at < CACHE_MS
14
15// The file a file-editing tool call targets, or null.
16export function targetPath(input: unknown): string | null {
17  const i = (input ?? {}) as { file_path?: unknown; notebook_path?: unknown }
18  const path = typeof i.file_path === 'string' ? i.file_path : i.notebook_path
19
20  return typeof path === 'string' && path !== '' ? path : null
21}
22
23export const absolute = (path: string, cwd: string) => (path.startsWith('/') ? path : `${cwd.replace(/\/$/, '')}/${path}`)
24
25export const parentDir = (path: string): string | null => {
26  const cut = path.lastIndexOf('/')
27
28  return cut > 0 ? path.slice(0, cut) : cut === 0 && path.length > 1 ? '/' : null
29}
30

1export type BranchGuardSwitch = boolean
2
3declare module 'claude-code' {
4  interface PluginState {
5    'branch-guard': { isOff: BranchGuardSwitch }
6  }
7}
8