1  import { z } from 'zod/v4'
  2  import {
  3    getOriginalCwd,
  4    getProjectRoot,
  5    setOriginalCwd,
  6    setProjectRoot,
  7  } from '../../bootstrap/state.js'
  8  import { clearSystemPromptSections } from '../../constants/systemPromptSections.js'
  9  import { logEvent } from '../../services/analytics/index.js'
 10  import type { Tool } from '../../Tool.js'
 11  import { buildTool, type ToolDef } from '../../Tool.js'
 12  import { count } from '../../utils/array.js'
 13  import { clearMemoryFileCaches } from '../../utils/claudemd.js'
 14  import { execFileNoThrow } from '../../utils/execFileNoThrow.js'
 15  import { updateHooksConfigSnapshot } from '../../utils/hooks/hooksConfigSnapshot.js'
 16  import { lazySchema } from '../../utils/lazySchema.js'
 17  import { getPlansDirectory } from '../../utils/plans.js'
 18  import { setCwd } from '../../utils/Shell.js'
 19  import { saveWorktreeState } from '../../utils/sessionStorage.js'
 20  import {
 21    cleanupWorktree,
 22    getCurrentWorktreeSession,
 23    keepWorktree,
 24    killTmuxSession,
 25  } from '../../utils/worktree.js'
 26  import { EXIT_WORKTREE_TOOL_NAME } from './constants.js'
 27  import { getExitWorktreeToolPrompt } from './prompt.js'
 28  import { renderToolResultMessage, renderToolUseMessage } from './UI.js'
 29  
 30  const inputSchema = lazySchema(() =>
 31    z.strictObject({
 32      action: z
 33        .enum(['keep', 'remove'])
 34        .describe(
 35          '"keep" leaves the worktree and branch on disk; "remove" deletes both.',
 36        ),
 37      discard_changes: z
 38        .boolean()
 39        .optional()
 40        .describe(
 41          'Required true when action is "remove" and the worktree has uncommitted files or unmerged commits. The tool will refuse and list them otherwise.',
 42        ),
 43    }),
 44  )
 45  type InputSchema = ReturnType<typeof inputSchema>
 46  
 47  const outputSchema = lazySchema(() =>
 48    z.object({
 49      action: z.enum(['keep', 'remove']),
 50      originalCwd: z.string(),
 51      worktreePath: z.string(),
 52      worktreeBranch: z.string().optional(),
 53      tmuxSessionName: z.string().optional(),
 54      discardedFiles: z.number().optional(),
 55      discardedCommits: z.number().optional(),
 56      message: z.string(),
 57    }),
 58  )
 59  type OutputSchema = ReturnType<typeof outputSchema>
 60  export type Output = z.infer<OutputSchema>
 61  
 62  type ChangeSummary = {
 63    changedFiles: number
 64    commits: number
 65  }
 66  
 67  /**
 68   * Returns null when state cannot be reliably determined — callers that use
 69   * this as a safety gate must treat null as "unknown, assume unsafe"
 70   * (fail-closed). A silent 0/0 would let cleanupWorktree destroy real work.
 71   *
 72   * Null is returned when:
 73   * - git status or rev-list exit non-zero (lock file, corrupt index, bad ref)
 74   * - originalHeadCommit is undefined but git status succeeded — this is the
 75   *   hook-based-worktree-wrapping-git case (worktree.ts:525-532 doesn't set
 76   *   originalHeadCommit). We can see the working tree is git, but cannot count
 77   *   commits without a baseline, so we cannot prove the branch is clean.
 78   */
 79  async function countWorktreeChanges(
 80    worktreePath: string,
 81    originalHeadCommit: string | undefined,
 82  ): Promise<ChangeSummary | null> {
 83    const status = await execFileNoThrow('git', [
 84      '-C',
 85      worktreePath,
 86      'status',
 87      '--porcelain',
 88    ])
 89    if (status.code !== 0) {
 90      return null
 91    }
 92    const changedFiles = count(status.stdout.split('\n'), l => l.trim() !== '')
 93  
 94    if (!originalHeadCommit) {
 95      // git status succeeded → this is a git repo, but without a baseline
 96      // commit we cannot count commits. Fail-closed rather than claim 0.
 97      return null
 98    }
 99  
100    const revList = await execFileNoThrow('git', [
101      '-C',
102      worktreePath,
103      'rev-list',
104      '--count',
105      `${originalHeadCommit}..HEAD`,
106    ])
107    if (revList.code !== 0) {
108      return null
109    }
110    const commits = parseInt(revList.stdout.trim(), 10) || 0
111  
112    return { changedFiles, commits }
113  }
114  
115  /**
116   * Restore session state to reflect the original directory.
117   * This is the inverse of the session-level mutations in EnterWorktreeTool.call().
118   *
119   * keepWorktree()/cleanupWorktree() handle process.chdir and currentWorktreeSession;
120   * this handles everything above the worktree utility layer.
121   */
122  function restoreSessionToOriginalCwd(
123    originalCwd: string,
124    projectRootIsWorktree: boolean,
125  ): void {
126    setCwd(originalCwd)
127    // EnterWorktree sets originalCwd to the *worktree* path (intentional — see
128    // state.ts getProjectRoot comment). Reset to the real original.
129    setOriginalCwd(originalCwd)
130    // --worktree startup sets projectRoot to the worktree; mid-session
131    // EnterWorktreeTool does not. Only restore when it was actually changed —
132    // otherwise we'd move projectRoot to wherever the user had cd'd before
133    // entering the worktree (session.originalCwd), breaking the "stable project
134    // identity" contract.
135    if (projectRootIsWorktree) {
136      setProjectRoot(originalCwd)
137      // setup.ts's --worktree block called updateHooksConfigSnapshot() to re-read
138      // hooks from the worktree. Restore symmetrically. (Mid-session
139      // EnterWorktreeTool never touched the snapshot, so no-op there.)
140      updateHooksConfigSnapshot()
141    }
142    saveWorktreeState(null)
143    clearSystemPromptSections()
144    clearMemoryFileCaches()
145    getPlansDirectory.cache.clear?.()
146  }
147  
148  export const ExitWorktreeTool: Tool<InputSchema, Output> = buildTool({
149    name: EXIT_WORKTREE_TOOL_NAME,
150    searchHint: 'exit a worktree session and return to the original directory',
151    maxResultSizeChars: 100_000,
152    async description() {
153      return 'Exits a worktree session created by EnterWorktree and restores the original working directory'
154    },
155    async prompt() {
156      return getExitWorktreeToolPrompt()
157    },
158    get inputSchema(): InputSchema {
159      return inputSchema()
160    },
161    get outputSchema(): OutputSchema {
162      return outputSchema()
163    },
164    userFacingName() {
165      return 'Exiting worktree'
166    },
167    shouldDefer: true,
168    isDestructive(input) {
169      return input.action === 'remove'
170    },
171    toAutoClassifierInput(input) {
172      return input.action
173    },
174    async validateInput(input) {
175      // Scope guard: getCurrentWorktreeSession() is null unless EnterWorktree
176      // (specifically createWorktreeForSession) ran in THIS session. Worktrees
177      // created by `git worktree add`, or by EnterWorktree in a previous
178      // session, do not populate it. This is the sole entry gate — everything
179      // past this point operates on a path EnterWorktree created.
180      const session = getCurrentWorktreeSession()
181      if (!session) {
182        return {
183          result: false,
184          message:
185            'No-op: there is no active EnterWorktree session to exit. This tool only operates on worktrees created by EnterWorktree in the current session — it will not touch worktrees created manually or in a previous session. No filesystem changes were made.',
186          errorCode: 1,
187        }
188      }
189  
190      if (input.action === 'remove' && !input.discard_changes) {
191        const summary = await countWorktreeChanges(
192          session.worktreePath,
193          session.originalHeadCommit,
194        )
195        if (summary === null) {
196          return {
197            result: false,
198            message: `Could not verify worktree state at ${session.worktreePath}. Refusing to remove without explicit confirmation. Re-invoke with discard_changes: true to proceed — or use action: "keep" to preserve the worktree.`,
199            errorCode: 3,
200          }
201        }
202        const { changedFiles, commits } = summary
203        if (changedFiles > 0 || commits > 0) {
204          const parts: string[] = []
205          if (changedFiles > 0) {
206            parts.push(
207              `${changedFiles} uncommitted ${changedFiles === 1 ? 'file' : 'files'}`,
208            )
209          }
210          if (commits > 0) {
211            parts.push(
212              `${commits} ${commits === 1 ? 'commit' : 'commits'} on ${session.worktreeBranch ?? 'the worktree branch'}`,
213            )
214          }
215          return {
216            result: false,
217            message: `Worktree has ${parts.join(' and ')}. Removing will discard this work permanently. Confirm with the user, then re-invoke with discard_changes: true — or use action: "keep" to preserve the worktree.`,
218            errorCode: 2,
219          }
220        }
221      }
222  
223      return { result: true }
224    },
225    renderToolUseMessage,
226    renderToolResultMessage,
227    async call(input) {
228      const session = getCurrentWorktreeSession()
229      if (!session) {
230        // validateInput guards this, but the session is module-level mutable
231        // state — defend against a race between validation and execution.
232        throw new Error('Not in a worktree session')
233      }
234  
235      // Capture before keepWorktree/cleanupWorktree null out currentWorktreeSession.
236      const {
237        originalCwd,
238        worktreePath,
239        worktreeBranch,
240        tmuxSessionName,
241        originalHeadCommit,
242      } = session
243  
244      // --worktree startup calls setOriginalCwd(getCwd()) and
245      // setProjectRoot(getCwd()) back-to-back right after setCwd(worktreePath)
246      // (setup.ts:235/239), so both hold the same realpath'd value and BashTool
247      // cd never touches either. Mid-session EnterWorktreeTool sets originalCwd
248      // but NOT projectRoot. (Can't use getCwd() — BashTool mutates it on every
249      // cd. Can't use session.worktreePath — it's join()'d, not realpath'd.)
250      const projectRootIsWorktree = getProjectRoot() === getOriginalCwd()
251  
252      // Re-count at execution time for accurate analytics and output — the
253      // worktree state at validateInput time may not match now. Null (git
254      // failure) falls back to 0/0; safety gating already happened in
255      // validateInput, so this only affects analytics + messaging.
256      const { changedFiles, commits } = (await countWorktreeChanges(
257        worktreePath,
258        originalHeadCommit,
259      )) ?? { changedFiles: 0, commits: 0 }
260  
261      if (input.action === 'keep') {
262        await keepWorktree()
263        restoreSessionToOriginalCwd(originalCwd, projectRootIsWorktree)
264  
265        logEvent('tengu_worktree_kept', {
266          mid_session: true,
267          commits,
268          changed_files: changedFiles,
269        })
270  
271        const tmuxNote = tmuxSessionName
272          ? ` Tmux session ${tmuxSessionName} is still running; reattach with: tmux attach -t ${tmuxSessionName}`
273          : ''
274        return {
275          data: {
276            action: 'keep' as const,
277            originalCwd,
278            worktreePath,
279            worktreeBranch,
280            tmuxSessionName,
281            message: `Exited worktree. Your work is preserved at ${worktreePath}${worktreeBranch ? ` on branch ${worktreeBranch}` : ''}. Session is now back in ${originalCwd}.${tmuxNote}`,
282          },
283        }
284      }
285  
286      // action === 'remove'
287      if (tmuxSessionName) {
288        await killTmuxSession(tmuxSessionName)
289      }
290      await cleanupWorktree()
291      restoreSessionToOriginalCwd(originalCwd, projectRootIsWorktree)
292  
293      logEvent('tengu_worktree_removed', {
294        mid_session: true,
295        commits,
296        changed_files: changedFiles,
297      })
298  
299      const discardParts: string[] = []
300      if (commits > 0) {
301        discardParts.push(`${commits} ${commits === 1 ? 'commit' : 'commits'}`)
302      }
303      if (changedFiles > 0) {
304        discardParts.push(
305          `${changedFiles} uncommitted ${changedFiles === 1 ? 'file' : 'files'}`,
306        )
307      }
308      const discardNote =
309        discardParts.length > 0 ? ` Discarded ${discardParts.join(' and ')}.` : ''
310      return {
311        data: {
312          action: 'remove' as const,
313          originalCwd,
314          worktreePath,
315          worktreeBranch,
316          discardedFiles: changedFiles,
317          discardedCommits: commits,
318          message: `Exited and removed worktree at ${worktreePath}.${discardNote} Session is now back in ${originalCwd}.`,
319        },
320      }
321    },
322    mapToolResultToToolResultBlockParam({ message }, toolUseID) {
323      return {
324        type: 'tool_result',
325        content: message,
326        tool_use_id: toolUseID,
327      }
328    },
329  } satisfies ToolDef<InputSchema, Output>)