1  import { z } from 'zod/v4'
  2  import { getSessionId, setOriginalCwd } from '../../bootstrap/state.js'
  3  import { clearSystemPromptSections } from '../../constants/systemPromptSections.js'
  4  import { logEvent } from '../../services/analytics/index.js'
  5  import type { Tool } from '../../Tool.js'
  6  import { buildTool, type ToolDef } from '../../Tool.js'
  7  import { clearMemoryFileCaches } from '../../utils/claudemd.js'
  8  import { getCwd } from '../../utils/cwd.js'
  9  import { findCanonicalGitRoot } from '../../utils/git.js'
 10  import { lazySchema } from '../../utils/lazySchema.js'
 11  import { getPlanSlug, getPlansDirectory } from '../../utils/plans.js'
 12  import { setCwd } from '../../utils/Shell.js'
 13  import { saveWorktreeState } from '../../utils/sessionStorage.js'
 14  import {
 15    createWorktreeForSession,
 16    getCurrentWorktreeSession,
 17    validateWorktreeSlug,
 18  } from '../../utils/worktree.js'
 19  import { ENTER_WORKTREE_TOOL_NAME } from './constants.js'
 20  import { getEnterWorktreeToolPrompt } from './prompt.js'
 21  import { renderToolResultMessage, renderToolUseMessage } from './UI.js'
 22  
 23  const inputSchema = lazySchema(() =>
 24    z.strictObject({
 25      name: z
 26        .string()
 27        .superRefine((s, ctx) => {
 28          try {
 29            validateWorktreeSlug(s)
 30          } catch (e) {
 31            ctx.addIssue({ code: 'custom', message: (e as Error).message })
 32          }
 33        })
 34        .optional()
 35        .describe(
 36          'Optional name for the worktree. Each "/"-separated segment may contain only letters, digits, dots, underscores, and dashes; max 64 chars total. A random name is generated if not provided.',
 37        ),
 38    }),
 39  )
 40  type InputSchema = ReturnType<typeof inputSchema>
 41  
 42  const outputSchema = lazySchema(() =>
 43    z.object({
 44      worktreePath: z.string(),
 45      worktreeBranch: z.string().optional(),
 46      message: z.string(),
 47    }),
 48  )
 49  type OutputSchema = ReturnType<typeof outputSchema>
 50  export type Output = z.infer<OutputSchema>
 51  
 52  export const EnterWorktreeTool: Tool<InputSchema, Output> = buildTool({
 53    name: ENTER_WORKTREE_TOOL_NAME,
 54    searchHint: 'create an isolated git worktree and switch into it',
 55    maxResultSizeChars: 100_000,
 56    async description() {
 57      return 'Creates an isolated worktree (via git or configured hooks) and switches the session into it'
 58    },
 59    async prompt() {
 60      return getEnterWorktreeToolPrompt()
 61    },
 62    get inputSchema(): InputSchema {
 63      return inputSchema()
 64    },
 65    get outputSchema(): OutputSchema {
 66      return outputSchema()
 67    },
 68    userFacingName() {
 69      return 'Creating worktree'
 70    },
 71    shouldDefer: true,
 72    toAutoClassifierInput(input) {
 73      return input.name ?? ''
 74    },
 75    renderToolUseMessage,
 76    renderToolResultMessage,
 77    async call(input) {
 78      // Validate not already in a worktree created by this session
 79      if (getCurrentWorktreeSession()) {
 80        throw new Error('Already in a worktree session')
 81      }
 82  
 83      // Resolve to main repo root so worktree creation works from within a worktree
 84      const mainRepoRoot = findCanonicalGitRoot(getCwd())
 85      if (mainRepoRoot && mainRepoRoot !== getCwd()) {
 86        process.chdir(mainRepoRoot)
 87        setCwd(mainRepoRoot)
 88      }
 89  
 90      const slug = input.name ?? getPlanSlug()
 91  
 92      const worktreeSession = await createWorktreeForSession(getSessionId(), slug)
 93  
 94      process.chdir(worktreeSession.worktreePath)
 95      setCwd(worktreeSession.worktreePath)
 96      setOriginalCwd(getCwd())
 97      saveWorktreeState(worktreeSession)
 98      // Clear cached system prompt sections so env_info_simple recomputes with worktree context
 99      clearSystemPromptSections()
100      // Clear memoized caches that depend on CWD
101      clearMemoryFileCaches()
102      getPlansDirectory.cache.clear?.()
103  
104      logEvent('tengu_worktree_created', {
105        mid_session: true,
106      })
107  
108      const branchInfo = worktreeSession.worktreeBranch
109        ? ` on branch ${worktreeSession.worktreeBranch}`
110        : ''
111  
112      return {
113        data: {
114          worktreePath: worktreeSession.worktreePath,
115          worktreeBranch: worktreeSession.worktreeBranch,
116          message: `Created worktree at ${worktreeSession.worktreePath}${branchInfo}. The session is now working in the worktree. Use ExitWorktree to leave mid-session, or exit the session to be prompted.`,
117        },
118      }
119    },
120    mapToolResultToToolResultBlockParam({ message }, toolUseID) {
121      return {
122        type: 'tool_result',
123        content: message,
124        tool_use_id: toolUseID,
125      }
126    },
127  } satisfies ToolDef<InputSchema, Output>)