worktrees

1import { atom, read, update } from 'claude-code'
2import type { EngineInterface, Register } from 'claude-code'
3
4import type { PrState, WorktreeRow } from '../types'
5
6const rows = atom({ plugin: 'worktrees', key: 'rows' } as const, [] as WorktreeRow[])
7const isLoading = atom({ plugin: 'worktrees', key: 'isLoading' } as const, false)
8const error = atom({ plugin: 'worktrees', key: 'error' } as const, null as string | null)
9
10const PANE = 'worktrees'
11const GH_CONCURRENCY = 4
12
13export type Listed = { path: string; branch: string | null; isBare: boolean }
14
15// `git worktree list --porcelain`: blank-line separated records, the first the main worktree.
16export const parseWorktrees = (porcelain: string): Listed[] =>
17  porcelain
18    .split(/\n\s*\n/)
19    .map(block => block.trim())
20    .filter(Boolean)
21    .flatMap(block => {
22      const lines = block.split('\n')
23      const path = lines.find(line => line.startsWith('worktree '))?.slice('worktree '.length)
24      if (path === undefined) {
25        return []
26      }
27      const branch = lines.find(line => line.startsWith('branch '))?.slice('branch '.length).replace(/^refs\/heads\//, '') ?? null
28
29      return [{ path, branch, isBare: lines.includes('bare') }]
30    })
31
32// `git status --porcelain=v2 --branch`: changed entries, and `# branch.ab +A -B` when tracking.
33export const parseStatus = (text: string) => {
34  const lines = text.split('\n').filter(Boolean)
35  const ab = lines.find(line => line.startsWith('# branch.ab '))?.match(/\+(\d+) -(\d+)/)
36
37  return {
38    dirty: lines.filter(line => !line.startsWith('#')).length,
39    ahead: ab ? Number(ab[1]) : undefined,
40    behind: ab ? Number(ab[2]) : undefined,
41  }
42}
43
44export const canRemove = (row: WorktreeRow) =>
45  !row.isMain && !row.isCurrent && row.dirty === 0 && (row.pr?.state === 'MERGED' || row.pr?.state === 'CLOSED')
46
47export const mergedRemovable = (list: readonly WorktreeRow[]) =>
48  list.filter(row => canRemove(row) && row.pr?.state === 'MERGED')
49
50const base = (path: string) => path.split('/').filter(Boolean).at(-1) ?? path
51
52const isInside = (dir: string, path: string) => dir === path || dir.startsWith(`${path}/`)
53
54// Runs `work` over `items`, at most `limit` at a time.
55const pool = async <T,>(items: readonly T[], limit: number, work: (item: T) => Promise<void>) => {
56  const queue = [...items]
57  const workers = Array.from({ length: Math.min(limit, queue.length) }, async () => {
58    for (let item = queue.shift(); item !== undefined; item = queue.shift()) {
59      await work(item)
60    }
61  })
62  await Promise.all(workers)
63}
64
65const load = async ($: EngineInterface) => {
66  await update($, isLoading, () => true)
67  await update($, error, () => null)
68
69  try {
70    const root = await $.session.root()
71    const cwd = await $.session.cwd()
72    const listed = await $.process.run(['git', '-C', root, 'worktree', 'list', '--porcelain'])
73    if (listed.exitCode !== 0) {
74      await update($, rows, () => [])
75      await update($, error, () => (listed.stderr.trim() || 'Not a git repository.'))
76
77      return
78    }
79
80    const worktrees = parseWorktrees(listed.stdout).filter(one => !one.isBare)
81    const current = [...worktrees].sort((a, b) => b.path.length - a.path.length).find(one => isInside(cwd, one.path))
82    const built: WorktreeRow[] = []
83    for (const [i, one] of worktrees.entries()) {
84      const status = await $.process.run(['git', '-C', one.path, 'status', '--porcelain=v2', '--branch'])
85      built.push({
86        path: one.path,
87        name: base(one.path),
88        branch: one.branch,
89        isMain: i === 0,
90        isCurrent: one.path === current?.path,
91        ...(status.exitCode === 0 ? parseStatus(status.stdout) : { dirty: 0 }),
92      })
93    }
94    await update($, rows, () => built)
95
96    const prs = new Map<string, { number: number; state: PrState }>()
97    await pool(built.filter(row => row.branch !== null && !row.isMain), GH_CONCURRENCY, async row => {
98      const found = await $.process
99        .run(['gh', 'pr', 'list', '--head', row.branch ?? '', '--state', 'all', '--json', 'number,state', '--limit', '1'], { cwd: root, timeoutMs: 20_000 })
100        .catch(() => undefined)
101      if (found?.exitCode === 0) {
102        const [pr] = JSON.parse(found.stdout || '[]') as { number: number; state: PrState }[]
103        if (pr !== undefined) {
104          prs.set(row.path, pr)
105        }
106      }
107    })
108    await update($, rows, all => all.map(row => ({ ...row, pr: prs.get(row.path) ?? row.pr })))
109  } finally {
110    await update($, isLoading, () => false)
111  }
112}
113
114const remove = async ($: EngineInterface, row: WorktreeRow) => {
115  const answer = await $.ui
116    .ask(`Remove worktree ${row.name} (${row.branch ?? 'detached'})?`, { options: ['Remove', 'Cancel'], header: 'Worktree' })
117    .catch(() => 'Cancel')
118  if (answer !== 'Remove') {
119    return
120  }
121
122  const root = await $.session.root()
123  const ran = await $.process.run(['git', '-C', root, 'worktree', 'remove', row.path])
124  if (ran.exitCode !== 0) {
125    $.ui.toast(`Could not remove ${row.name}: ${(ran.stderr.trim().split('\n')[0] ?? '').slice(0, 120)}`)
126  }
127  await load($)
128}
129
130const removeMerged = async ($: EngineInterface) => {
131  const targets = mergedRemovable(await read($, rows))
132  if (targets.length === 0) {
133    return
134  }
135
136  const names = targets.map(row => row.name).join(', ')
137  const answer = await $.ui
138    .ask(`Remove ${targets.length} merged worktree${targets.length === 1 ? '' : 's'}?  ${names}`, { options: ['Remove all', 'Cancel'], header: 'Worktrees' })
139    .catch(() => 'Cancel')
140  if (answer !== 'Remove all') {
141    return
142  }
143
144  const root = await $.session.root()
145  const failed: string[] = []
146  for (const row of targets) {
147    const ran = await $.process.run(['git', '-C', root, 'worktree', 'remove', row.path]).catch(() => undefined)
148    if (ran?.exitCode !== 0) {
149      failed.push(row.name)
150    }
151  }
152
153  const removed = targets.length - failed.length
154  $.ui.toast(failed.length === 0 ? `Removed ${removed} worktree${removed === 1 ? '' : 's'}` : `Removed ${removed}; could not remove ${failed.join(', ')}`)
155  await load($)
156}
157
158const PR_COLOR: Record<PrState, string> = { MERGED: 'green', OPEN: 'yellow', CLOSED: 'red' }
159
160export const register: Register = on => {
161  on('session.start', async ($, e, next) => {
162    await $.command.register({ name: 'worktrees', description: 'Git worktrees of this repo, with branch, changes and PR state' })
163
164    return next(e)
165  })
166
167  on('command.run', { command: 'worktrees' }, async $ => {
168    await $.ui.open({ id: PANE, title: 'Worktrees' })
169    load($).catch(reason => update($, error, () => String(reason)))
170
171    return { text: 'Worktrees pane opened.' }
172  })
173
174  on('ui.render', { component: 'Pane', requestId: PANE }, async ($, e) => {
175    const { Box, Text, Button } = $.ui.resolve(e)
176    const list = await read($, rows)
177    const loading = await read($, isLoading)
178    const failed = await read($, error)
179    const nameWidth = Math.min(24, Math.max(8, ...list.map(row => row.name.length)))
180    const branchWidth = Math.max(10, Math.min(40, (e.viewport?.columns ?? 80) - nameWidth - 34))
181    const merged = mergedRemovable(list).length
182    const clip = (text: string, width: number) => (text.length > width ? `${text.slice(0, width - 1)}…` : text.padEnd(width))
183
184    return (
185      <Box flexDirection="column">
186        {failed !== null && <Text color="red">{failed}</Text>}
187        {failed === null && list.length === 0 && <Text dimColor>{loading ? 'Loading…' : 'No worktrees.'}</Text>}
188        {list.map(row => (
189          <Box key={row.path}>
190            <Text bold={row.isCurrent}>{row.isCurrent ? '› ' : '  '}{clip(row.name, nameWidth)} </Text>
191            <Text dimColor>{clip(row.branch ?? '(detached)', branchWidth)} </Text>
192            {row.isMain && <Text dimColor>main </Text>}
193            {row.dirty > 0 && <Text color="yellow">●{row.dirty} </Text>}
194            {(row.ahead ?? 0) > 0 && <Text dimColor>↑{row.ahead} </Text>}
195            {(row.behind ?? 0) > 0 && <Text dimColor>↓{row.behind} </Text>}
196            {row.pr !== undefined && <Text color={PR_COLOR[row.pr.state]}>#{row.pr.number} {row.pr.state.toLowerCase()} </Text>}
197            {canRemove(row) && <Button key={`remove-${row.name}`} label="Remove" onPress={() => remove($, row)} />}
198          </Box>
199        ))}
200        <Box>
201          <Button key="refresh" label={loading ? 'Refreshing…' : 'Refresh'} onPress={() => load($)} />
202          {merged > 0 && !loading && (
203            <Button key="remove-merged" label={`Remove ${merged} merged`} onPress={() => removeMerged($)} />
204          )}
205        </Box>
206      </Box>
207    )
208  })
209}
210

1export type PrState = 'OPEN' | 'MERGED' | 'CLOSED'
2
3export type WorktreeRow = {
4  path: string
5  name: string
6  branch: string | null
7  isMain: boolean
8  isCurrent: boolean
9  dirty: number
10  ahead?: number
11  behind?: number
12  pr?: { number: number; state: PrState }
13}
14
15declare module 'claude-code' {
16  interface PluginState {
17    worktrees: { rows: WorktreeRow[]; isLoading: boolean; error: string | null }
18  }
19}
20