auto-worktree

1import type { EngineInterface, Register } from 'claude-code'
2
3import { excludedText, isIsolated, shieldGit, type Situation, skipReasonOf } from './judge'
4
5async function situationOf($: EngineInterface): Promise<Situation> {
6  const [turns, surfaces, entrypoint, cwd, repo] = await Promise.all([
7    $.session.turns(),
8    $.session.surfaces(),
9    $.env.get('CLAUDE_CODE_ENTRYPOINT'),
10    $.session.cwd(),
11    $.session.repo(),
12  ])
13
14  return { turns, surfaces, entrypoint: entrypoint ?? null, cwd, root: repo?.root ?? null }
15}
16
17const toastOf = ($: EngineInterface, text: string) => $.ui.toast(text, { timeoutMs: 10_000 })
18
19/**
20 * メインのチェックアウトの git status にワークツリーが出ないよう、このクローンの除外ファイル
21 * （.git/info/exclude。コミットされない）に足す。うまくいかなくても作業は止めない。
22 */
23async function excludeWorktrees($: EngineInterface, root: string) {
24  const ran = await $.process
25    .run(['git', 'rev-parse', '--git-path', 'info/exclude'], { cwd: root, timeoutMs: 5_000 })
26    .catch(() => undefined)
27
28  if (ran?.exitCode !== 0) {
29    return
30  }
31
32  const found = ran.stdout.trim()
33  const path = found.startsWith('/') ? found : `${root}/${found}`
34  const text = excludedText(await $.fs.read(path).catch(() => ''))
35
36  if (text !== null) {
37    await $.fs.write(path, text).catch(() => undefined)
38  }
39}
40
41export const register: Register = on => {
42  on('session.start', async ($, e, next) => {
43    await $.command.register({
44      name: 'auto-worktree',
45      description: 'このセッションが最初のプロンプトでワークツリーに入るかどうかと、その理由を出す',
46    })
47
48    return next(e)
49  })
50
51  on('prompt.submit', async ($, e, next) => {
52    // 様子が読めないときも、プロンプトは止めずにそのまま送る
53    const situation = await situationOf($).catch(() => null)
54
55    if (situation === null || skipReasonOf(situation) !== null) {
56      return next(e)
57    }
58
59    // モデルが EnterWorktree を呼ぶのと同じ呼び出し。権限の確認もそのまま通る。
60    // 入れなくてもプロンプトは止めず、元の場所で始める
61    let entered
62    try {
63      entered = await $.tool.call({ tool: 'EnterWorktree' })
64    } catch (error) {
65      toastOf($, `ワークツリーに入れませんでした: ${String(error)}`)
66
67      return next(e)
68    }
69
70    if (entered.deny !== undefined || entered.isError) {
71      toastOf($, `ワークツリーに入れませんでした: ${entered.deny ?? entered.text ?? '理由は不明'}`)
72
73      return next(e)
74    }
75
76    const { worktreePath, worktreeBranch } = entered.result
77    toastOf($, `ワークツリーで始めます: ${worktreePath}`)
78
79    if (situation.root !== null) {
80      await excludeWorktrees($, situation.root)
81    }
82
83    // プラグインからのツール呼び出しの結果はモデルに届かないので、移ったことを添えて伝える。
84    // VS Code 拡張機能のセッションには画面がなくトーストが出ないため、返答でも人に知らせてもらう
85    const branch = worktreeBranch === undefined ? '' : `（ブランチ ${worktreeBranch}）`
86    const note = `このセッションは最初のプロンプトの前に、mod auto-worktree によってワークツリー ${worktreePath}${branch} へ移った。ファイルの読み書きとコマンドはこのワークツリーで行う。最初の返答の冒頭で、このワークツリーとブランチに移ったことを利用者に一言伝える。`
87
88    return next({ ...e, context: [...(e.context ?? []), note] })
89  })
90
91  // ワークツリーに隔離されたセッションでは、どこで動くか読めない git を Claude Code が拒否する。
92  // rtk のフック（設定ファイルの PreToolUse）が git を `rtk git` に書き換えると、それに当たって git が使えなくなる。
93  // rtk のフックは改変すると rtk の整合性検査で全部止まり、classic.PreToolUse は組み込みのセキュリティ用プラグインが
94  // user の mod を通さない。そこで rtk のフックより上の tool.call で、git を rtk が書き換えない `\git` にしておく
95  on('tool.call', { tool: 'Bash' }, async ($, e, next) => {
96    const command = e.tool === 'Bash' && /\bgit\b/.test(e.command) ? shieldGit(e.command) : undefined
97
98    if (e.tool !== 'Bash' || command === undefined || command === e.command) {
99      return next(e)
100    }
101
102    const [cwd, repo] = await Promise.all([$.session.cwd(), $.session.repo()]).catch(() => [null, null] as const)
103
104    return cwd !== null && isIsolated(cwd, repo?.root ?? null) ? next({ ...e, command }) : next(e)
105  })
106
107  on('command.run', { command: 'auto-worktree' }, async $ => {
108    const situation = await situationOf($)
109    const why = skipReasonOf(situation)
110    const lines = [
111      why === null ? '次のプロンプトでワークツリーに入ります。' : `ワークツリーには入りません（${why}）。`,
112      `ディレクトリ: ${situation.cwd}`,
113      `リポジトリ: ${situation.root ?? 'なし'}`,
114      `画面: ${situation.surfaces.join(', ') || 'なし'}`,
115      `起動元: ${situation.entrypoint ?? 'なし'}`,
116      `これまでのプロンプト: ${situation.turns} 通`,
117    ]
118
119    return { text: lines.join('\n') }
120  })
121}
122

1/** 最初のプロンプトを送った時点のセッションの様子。 */
2export type Situation = {
3  /** 人がこれまでに送ったプロンプトの数。いま送ろうとしている 1 通目は数えない */
4  turns: number
5  /** セッションが描いている画面。`claude -p` と、SDK 経由で起動する VS Code 拡張機能では空 */
6  surfaces: readonly string[]
7  /** 起動元（環境変数 `CLAUDE_CODE_ENTRYPOINT`）。`claude -p` は `sdk-cli`、VS Code 拡張機能は `claude-vscode` */
8  entrypoint: string | null
9  /** セッションのディレクトリ */
10  cwd: string
11  /** リポジトリのメインの作業ツリーのルート。git の外なら null */
12  root: string | null
13}
14
15/** 画面がなくても人が使っている起動元。 */
16const ATTENDED_ENTRYPOINTS: ReadonlySet<string> = new Set(['claude-vscode'])
17
18const isInside = (dir: string, root: string) => dir === root || dir.startsWith(`${root}/`)
19
20/** メインのチェックアウトの git status にワークツリーが出ないよう、除外ファイルに足す行。 */
21export const EXCLUDE_LINE = '/.claude/worktrees/'
22
23/** これらのどれかがあれば、ワークツリーはもう除外されている。 */
24const COVERING_LINES = ['.claude', '.claude/', '/.claude', '/.claude/', '.claude/worktrees', '.claude/worktrees/', '/.claude/worktrees', EXCLUDE_LINE]
25
26/** 除外ファイルの中身に除外の行を足した結果。もう除外されていれば null。 */
27export function excludedText(current: string): string | null {
28  const lines = current.split('\n').map(line => line.trim())
29
30  if (lines.some(line => COVERING_LINES.includes(line))) {
31    return null
32  }
33
34  const base = current === '' || current.endsWith('\n') ? current : `${current}\n`
35
36  return `${base}${EXCLUDE_LINE}\n`
37}
38
39/** ワークツリーに入らない理由を返す。入るなら null。 */
40export function skipReasonOf(s: Situation): string | null {
41  if (s.turns > 0) {
42    return '最初のプロンプトではない'
43  }
44
45  if (s.surfaces.length === 0 && !ATTENDED_ENTRYPOINTS.has(s.entrypoint ?? '')) {
46    return '人のいないセッション（claude -p など）'
47  }
48
49  if (s.root === null) {
50    return 'git リポジトリの外'
51  }
52
53  // メインの作業ツリーの外にいるなら、別の場所に作ったワークツリーの中
54  if (!isInside(s.cwd, s.root) || isInside(s.cwd, `${s.root}/.claude/worktrees`)) {
55    return 'すでにワークツリーの中'
56  }
57
58  return null
59}
60
61/** Claude Code が作ったワークツリー（`<root>/.claude/worktrees/<名前>`）の中か。そこに入ったセッションは隔離される */
62export const isIsolated = (cwd: string, root: string | null) =>
63  root !== null && cwd.startsWith(`${root}/.claude/worktrees/`)
64
65// コマンドの位置を保つ語。`env git …`、`do git …`、`{ git …; }` の git もコマンドとして見る
66const KEEPS_START: ReadonlySet<string> = new Set(['env', 'time', '!', '{', 'if', 'then', 'elif', 'else', 'do', 'while', 'until'])
67// 変数の代入（`GIT_PAGER=cat git log` の `GIT_PAGER=cat`）もコマンドの位置を保つ
68const ASSIGNMENT = /^[A-Za-z_][A-Za-z0-9_]*=/
69// コマンドの区切り。`&&` `||` `|` `;` `&`、改行、`(` と `$(` の後ろはコマンドの位置になる
70const isOperator = (c: string) => c === ';' || c === '&' || c === '|' || c === '(' || c === '\n'
71const isBlank = (c: string) => c === ' ' || c === '\t'
72
73/** `from` にある引用符を閉じる位置。閉じていなければ -1。二重引用符の中では `\` が次の 1 字を守る */
74function closingOf(text: string, from: number): number {
75  const quote = text[from]
76
77  for (let i = from + 1; i < text.length; i++) {
78    if (quote === '"' && text[i] === '\\') {
79      i++
80    } else if (text[i] === quote) {
81      return i
82    }
83  }
84
85  return -1
86}
87
88/**
89 * コマンドの位置にある git を `\git` にする。シェルにとっては同じ git で、Claude Code の権限ルールも git として当てるが、
90 * rtk のフックはこの形を `rtk git` に書き換えない。引数の git と引用符の中は変えない。
91 * ヒアドキュメント（`<<`）から先は中身がデータなので、そこから後ろには手を出さない。
92 */
93export function shieldGit(command: string): string {
94  let out = ''
95  let isCommandStart = true
96  let i = 0
97
98  while (i < command.length) {
99    const c = command[i] ?? ''
100
101    if (isBlank(c) || isOperator(c)) {
102      out += c
103      i++
104      isCommandStart ||= isOperator(c)
105      continue
106    }
107
108    if (command.startsWith('<<', i)) {
109      return out + command.slice(i)
110    }
111
112    // 語の頭の # から行末まではコメント
113    if (c === '#') {
114      const end = command.indexOf('\n', i)
115      out += end === -1 ? command.slice(i) : command.slice(i, end)
116      i = end === -1 ? command.length : end
117      continue
118    }
119
120    // 1 語を、引用符と `\` を飛ばしながら、空白か区切りまで読む
121    const start = i
122
123    while (i < command.length) {
124      const d = command[i] ?? ''
125
126      if (isBlank(d) || isOperator(d) || command.startsWith('<<', i)) {
127        break
128      }
129
130      if (d === '\\') {
131        i += 2
132      } else if (d === "'" || d === '"') {
133        const end = closingOf(command, i)
134
135        if (end === -1) {
136          return out + command.slice(start)
137        }
138
139        i = end + 1
140      } else {
141        i++
142      }
143    }
144
145    const word = command.slice(start, i)
146    out += isCommandStart && word === 'git' ? '\\git' : word
147    isCommandStart = isCommandStart && (KEEPS_START.has(word) || ASSIGNMENT.test(word))
148  }
149
150  return out
151}
152