PR Pilot

1import { atom, read, update } from 'claude-code'
2import type { EngineInterface, Register } from 'claude-code'
3
4import type { Checks, Pr, Stack } from '../types'
5
6const FIELDS = 'number,url,headRefOid,state,isDraft,mergeable,mergeStateStatus,statusCheckRollup'
7const FAILED = ['FAILURE', 'ERROR', 'TIMED_OUT', 'CANCELLED', 'ACTION_REQUIRED', 'STARTUP_FAILURE']
8
9// Held by the host, so a hot reload keeps the band and does not toast the same result twice.
10const current = atom({ plugin: 'pr-pilot', key: 'pr' } as const, null as Pr | null)
11
12/** Drops control characters, so a check name or gh output cannot carry terminal escapes. */
13export const clean = (text: string) => text.replace(/[\u0000-\u001f\u007f-\u009f]/g, '')
14
15type Rollup = { __typename?: string; name?: string; context?: string; status?: string; conclusion?: string; state?: string }
16
17/** Tallies `statusCheckRollup`, which mixes check runs (`status`, `conclusion`) and commit statuses (`state`). */
18export function tally(rollup: readonly Rollup[]): Checks {
19  const checks: Checks = { total: rollup.length, passed: 0, pending: 0, failed: [] }
20  for (const check of rollup) {
21    const result = check.__typename === 'StatusContext' ? check.state : check.status === 'COMPLETED' ? check.conclusion : 'PENDING'
22    if (FAILED.includes(result ?? '')) checks.failed.push(clean(check.name ?? check.context ?? '?'))
23    else if (result === 'PENDING' || result === 'EXPECTED' || !result) checks.pending += 1
24    else checks.passed += 1
25  }
26
27  return checks
28}
29
30/** True once the pull request can merge as is: open, checks green, nothing blocking it. */
31export function isReady(pr: Pr) {
32  return !pr.isDraft && pr.mergeable === 'MERGEABLE' && pr.checks.failed.length === 0 && pr.checks.pending === 0 && ['CLEAN', 'HAS_HOOKS'].includes(pr.mergeStateStatus)
33}
34
35/** The band's checks entry and its colour. */
36export function checksEntry({ total, passed, pending, failed }: Checks): [string, string] {
37  if (total === 0) return ['no checks', 'gray']
38  if (failed.length) return [`✗ ${failed.length} failing: ${failed.join(', ')}`, 'red']
39  if (pending) return [`◐ ${passed}/${total} checks`, 'yellow']
40
41  return [`✓ ${total}/${total} checks`, 'green']
42}
43
44/** The band's merge entry and its colour. */
45export function mergeEntry(pr: Pr): [string, string] {
46  if (pr.isDraft) return ['draft', 'gray']
47  if (pr.mergeable === 'CONFLICTING') return ['conflicts', 'red']
48  if (pr.mergeStateStatus === 'BEHIND') return ['behind base', 'yellow']
49  if (pr.mergeStateStatus === 'BLOCKED') return ['blocked', 'yellow']
50  if (isReady(pr)) return ['ready · /merge', 'green']
51
52  return ['waiting', 'gray']
53}
54
55type StackView = { trunk: string; currentBranch: string; branches: { name: string; isMerged: boolean; pr?: { number: number } }[] }
56
57/** The current branch's place among the unmerged layers of `gh stack view --json`, or null when it is not one. */
58export function parseStack(view: StackView): Stack | null {
59  const layers = view.branches.filter(b => !b.isMerged).map(b => ({ branch: b.name, pr: b.pr?.number }))
60  const index = layers.findIndex(l => l.branch === view.currentBranch)
61
62  return index < 0 ? null : { position: index + 1, total: layers.length, trunk: view.trunk, layers }
63}
64
65async function fetchStack($: EngineInterface): Promise<Stack | null> {
66  // Exits non-zero off a stack, and where the gh-stack extension is not installed.
67  const ran = await $.process.run(['gh', 'stack', 'view', '--json'], { timeoutMs: 15_000 }).catch(() => null)
68  if (!ran || ran.exitCode !== 0) return null
69
70  try {
71    return parseStack(JSON.parse(ran.stdout))
72  } catch {
73    return null
74  }
75}
76
77async function fetchPr($: EngineInterface, number?: number): Promise<Pr | null> {
78  const argv = ['gh', 'pr', 'view', ...(number ? [String(number)] : []), '--json', FIELDS]
79  const ran = await $.process.run(argv, { timeoutMs: 15_000 }).catch(() => null)
80  if (!ran || ran.exitCode !== 0) return null
81  const view = JSON.parse(ran.stdout)
82  if (view.state !== 'OPEN') return null
83
84  return {
85    number: view.number,
86    url: view.url,
87    head: view.headRefOid,
88    isDraft: view.isDraft,
89    mergeable: view.mergeable,
90    mergeStateStatus: view.mergeStateStatus,
91    checks: tally(view.statusCheckRollup ?? []),
92  }
93}
94
95async function refresh($: EngineInterface) {
96  const pr = await fetchPr($)
97  const stack = pr && (await fetchStack($))
98  if (pr && stack) pr.stack = stack
99  const was = await read($, current)
100  await update($, current, () => pr)
101
102  // Toast when the same pull request's checks finish, not on first sight.
103  if (!pr || was?.number !== pr.number || !was.checks.pending || pr.checks.pending) return
104  $.ui.toast(
105    pr.checks.failed.length
106      ? `PR #${pr.number} checks failed: ${pr.checks.failed.join(', ')}`
107      : `PR #${pr.number} checks passed${isReady(pr) ? ', ready to /merge' : ''}`,
108    { timeoutMs: 8000 },
109  )
110}
111
112/** True when `branch` is at `head` locally, so merging and deleting it loses no unpushed commits. */
113async function isPushed($: EngineInterface, branch: string, head: string) {
114  const local = await $.process.run(['git', 'rev-parse', '--verify', `refs/heads/${branch}^{commit}`]).catch(() => null)
115
116  return local?.exitCode === 0 && local.stdout.trim() === head
117}
118
119/**
120 * Merges the whole stack with `gh stack merge`, after checking every unmerged layer, since it is all
121 * or nothing and has no head pin. It runs without an argument, as a number would be read as a stack
122 * number before a pull request number.
123 */
124async function mergeStack($: EngineInterface, pr: Pr, stack: Stack, method: string) {
125  const numbers: number[] = []
126  for (const layer of stack.layers) {
127    if (!layer.pr) return `${layer.branch} has no pull request yet: run gh stack submit first.`
128    const one = layer.pr === pr.number ? pr : await fetchPr($, layer.pr)
129    if (!one) return `PR #${layer.pr} (${layer.branch}) is not open.`
130    if (!isReady(one)) return `PR #${one.number} (${layer.branch}) is not ready to merge: ${checksEntry(one.checks)[0]}, ${mergeEntry(one)[0]}.`
131    if (!(await isPushed($, layer.branch, one.head))) return `${layer.branch} is not at PR #${one.number}'s head: push or pull first.`
132    numbers.push(one.number)
133  }
134
135  const ran = await $.process.run(['gh', 'stack', 'merge', '--yes', `--${method}`], { timeoutMs: 300_000 })
136  if (ran.exitCode !== 0) return `gh stack merge failed: ${clean((ran.stderr || ran.stdout).trim())}`
137
138  return `Merged stack: PR #${numbers.join(', #')}.${await deleteLayers($, stack)}`
139}
140
141/**
142 * Deletes every layer's branch, local and remote, and switches to the trunk, as `gh pr merge -d` does
143 * for one pull request, which `gh stack merge` has no flag for. It waits until every pull request has
144 * merged, since deleting a head branch while the stack sits in a merge queue would close it.
145 */
146async function deleteLayers($: EngineInterface, stack: Stack) {
147  for (const layer of stack.layers) {
148    const state = await $.process.run(['gh', 'pr', 'view', String(layer.pr), '--json', 'state', '--jq', '.state'], { timeoutMs: 15_000 }).catch(() => null)
149    if (state?.stdout.trim() !== 'MERGED') return ' Branches kept until every pull request has merged.'
150  }
151  const checkout = await $.process.run(['git', 'checkout', stack.trunk]).catch(() => null)
152  if (checkout?.exitCode !== 0) return ` Could not switch to ${stack.trunk}, so the branches were kept.`
153  await $.process.run(['git', 'pull', '--ff-only'], { timeoutMs: 60_000 }).catch(() => null)
154
155  const kept: string[] = []
156  for (const { branch } of stack.layers) {
157    // Fails harmlessly when the repository already deleted the head branch on merge.
158    await $.process.run(['gh', 'api', '-X', 'DELETE', `repos/{owner}/{repo}/git/refs/heads/${branch}`], { timeoutMs: 15_000 }).catch(() => null)
159    const deleted = await $.process.run(['git', 'branch', '-D', branch]).catch(() => null)
160    if (deleted?.exitCode !== 0) kept.push(branch)
161  }
162
163  return kept.length ? ` Switched to ${stack.trunk}, but could not delete ${kept.join(', ')} locally.` : ` Deleted ${stack.layers.map(l => l.branch).join(', ')} and switched to ${stack.trunk}.`
164}
165
166export const register: Register = (on, options) => {
167  const method = typeof options.method === 'string' ? options.method : 'merge'
168
169  on('session.start', async ($, e, next) => {
170    const result = await next(e)
171    await $.command.register({ name: 'merge', description: "Merge this branch's pull request, or its whole stack, once checks pass" })
172    await refresh($)
173    $.clock.every(60_000, () => refresh($))
174
175    return result
176  })
177
178  // A turn may have pushed or opened a pull request.
179  on('turn.complete', async ($, e, next) => {
180    const result = await next(e)
181    if (e.agentId === undefined) void refresh($)
182
183    return result
184  })
185
186  on('command.run', { command: 'merge' }, async $ => {
187    await refresh($)
188    const pr = await read($, current)
189    if (!pr) return { text: 'No open pull request for this branch.' }
190    if (pr.stack) {
191      const text = await mergeStack($, pr, pr.stack, method)
192      await refresh($)
193
194      return { text }
195    }
196    if (!isReady(pr)) return { text: `PR #${pr.number} is not ready to merge: ${checksEntry(pr.checks)[0]}, ${mergeEntry(pr)[0]}.` }
197
198    // --delete-branch drops the local branch, so refuse while it holds commits the PR does not.
199    const local = await $.process.run(['git', 'rev-parse', 'HEAD']).catch(() => null)
200    if (local?.stdout.trim() !== pr.head) return { text: `This branch is not at PR #${pr.number}'s head: push or pull first.` }
201
202    // Pinned to the head that was checked, so a push since then is not merged unchecked.
203    const ran = await $.process.run(['gh', 'pr', 'merge', String(pr.number), `--${method}`, '--delete-branch', '--match-head-commit', pr.head], { timeoutMs: 120_000 })
204    await refresh($)
205
206    return { text: ran.exitCode === 0 ? `Merged PR #${pr.number}.` : `gh pr merge failed: ${clean((ran.stderr || ran.stdout).trim())}` }
207  })
208
209  on('ui.render', { component: 'AbovePrompt' }, async ($, e, next) => {
210    const pr = await read($, current)
211
212    if (e.props.hasSurvey || !pr) {
213      return next(e)
214    }
215
216    const { Box, Text } = $.ui.resolve(e)
217    const [checks, checksColour] = checksEntry(pr.checks)
218    const [merge, mergeColour] = mergeEntry(pr)
219    // Stack above whatever the plugins beneath draw, so other bands stay.
220    const below = await next(e)
221
222    return (
223      <Box flexDirection="column">
224        <Box flexDirection="row" paddingX={1} gap={2}>
225          <Text dimColor>PR #{pr.number}</Text>
226          {pr.stack && <Text dimColor>stack {pr.stack.position}/{pr.stack.total}</Text>}
227          <Text color={checksColour} wrap="truncate-end">{checks}</Text>
228          <Text color={mergeColour}>{merge}</Text>
229        </Box>
230        {below}
231      </Box>
232    )
233  })
234}
235

1/** How the pull request's checks stand: `failed` names the failing ones. */
2export type Checks = { total: number; passed: number; pending: number; failed: string[] }
3
4/** The unmerged layers of a gh-stack stack, bottom first, where the current branch sits among them, and the trunk they merge into. */
5export type Stack = { position: number; total: number; trunk: string; layers: { branch: string; pr?: number }[] }
6
7/** The open pull request for the current branch, as `gh pr view` last reported it. */
8export type Pr = {
9  number: number
10  url: string
11  /** The head commit the checks ran on. */
12  head: string
13  isDraft: boolean
14  mergeable: string
15  mergeStateStatus: string
16  checks: Checks
17  /** Set when the branch is a layer of a gh-stack stack. */
18  stack?: Stack
19}
20
21declare module 'claude-code' {
22  interface PluginState {
23    'pr-pilot': { pr: Pr | null }
24  }
25}
26