1  import { randomUUID } from 'crypto'
  2  import { basename } from 'path'
  3  import { useEffect, useMemo, useRef, useState } from 'react'
  4  import { logEvent } from 'src/services/analytics/index.js'
  5  import { readFileSync } from 'src/utils/fileRead.js'
  6  import { expandPath } from 'src/utils/path.js'
  7  import type { PermissionOption } from '../components/permissions/FilePermissionDialog/permissionOptions.js'
  8  import type {
  9    MCPServerConnection,
 10    McpSSEIDEServerConfig,
 11    McpWebSocketIDEServerConfig,
 12  } from '../services/mcp/types.js'
 13  import type { ToolUseContext } from '../Tool.js'
 14  import type { FileEdit } from '../tools/FileEditTool/types.js'
 15  import {
 16    getEditsForPatch,
 17    getPatchForEdits,
 18  } from '../tools/FileEditTool/utils.js'
 19  import { getGlobalConfig } from '../utils/config.js'
 20  import { getPatchFromContents } from '../utils/diff.js'
 21  import { isENOENT } from '../utils/errors.js'
 22  import {
 23    callIdeRpc,
 24    getConnectedIdeClient,
 25    getConnectedIdeName,
 26    hasAccessToIDEExtensionDiffFeature,
 27  } from '../utils/ide.js'
 28  import { WindowsToWSLConverter } from '../utils/idePathConversion.js'
 29  import { logError } from '../utils/log.js'
 30  import { getPlatform } from '../utils/platform.js'
 31  
 32  type Props = {
 33    onChange(
 34      option: PermissionOption,
 35      input: {
 36        file_path: string
 37        edits: FileEdit[]
 38      },
 39    ): void
 40    toolUseContext: ToolUseContext
 41    filePath: string
 42    edits: FileEdit[]
 43    editMode: 'single' | 'multiple'
 44  }
 45  
 46  export function useDiffInIDE({
 47    onChange,
 48    toolUseContext,
 49    filePath,
 50    edits,
 51    editMode,
 52  }: Props): {
 53    closeTabInIDE: () => void
 54    showingDiffInIDE: boolean
 55    ideName: string
 56    hasError: boolean
 57  } {
 58    const isUnmounted = useRef(false)
 59    const [hasError, setHasError] = useState(false)
 60  
 61    const sha = useMemo(() => randomUUID().slice(0, 6), [])
 62    const tabName = useMemo(
 63      () => `✻ [Claude Code] ${basename(filePath)} (${sha}) ⧉`,
 64      [filePath, sha],
 65    )
 66  
 67    const shouldShowDiffInIDE =
 68      hasAccessToIDEExtensionDiffFeature(toolUseContext.options.mcpClients) &&
 69      getGlobalConfig().diffTool === 'auto' &&
 70      // Diffs should only be for file edits.
 71      // File writes may come through here but are not supported for diffs.
 72      !filePath.endsWith('.ipynb')
 73  
 74    const ideName =
 75      getConnectedIdeName(toolUseContext.options.mcpClients) ?? 'IDE'
 76  
 77    async function showDiff(): Promise<void> {
 78      if (!shouldShowDiffInIDE) {
 79        return
 80      }
 81  
 82      try {
 83        logEvent('tengu_ext_will_show_diff', {})
 84  
 85        const { oldContent, newContent } = await showDiffInIDE(
 86          filePath,
 87          edits,
 88          toolUseContext,
 89          tabName,
 90        )
 91        // Skip if component has been unmounted
 92        if (isUnmounted.current) {
 93          return
 94        }
 95  
 96        logEvent('tengu_ext_diff_accepted', {})
 97  
 98        const newEdits = computeEditsFromContents(
 99          filePath,
100          oldContent,
101          newContent,
102          editMode,
103        )
104  
105        if (newEdits.length === 0) {
106          // No changes -- edit was rejected (eg. reverted)
107          logEvent('tengu_ext_diff_rejected', {})
108          // We close the tab here because 'no' no longer auto-closes
109          const ideClient = getConnectedIdeClient(
110            toolUseContext.options.mcpClients,
111          )
112          if (ideClient) {
113            // Close the tab in the IDE
114            await closeTabInIDE(tabName, ideClient)
115          }
116          onChange(
117            { type: 'reject' },
118            {
119              file_path: filePath,
120              edits: edits,
121            },
122          )
123          return
124        }
125  
126        // File was modified - edit was accepted
127        onChange(
128          { type: 'accept-once' },
129          {
130            file_path: filePath,
131            edits: newEdits,
132          },
133        )
134      } catch (error) {
135        logError(error as Error)
136        setHasError(true)
137      }
138    }
139  
140    useEffect(() => {
141      void showDiff()
142  
143      // Set flag on unmount
144      return () => {
145        isUnmounted.current = true
146      }
147      // eslint-disable-next-line react-hooks/exhaustive-deps
148    }, [])
149  
150    return {
151      closeTabInIDE() {
152        const ideClient = getConnectedIdeClient(toolUseContext.options.mcpClients)
153  
154        if (!ideClient) {
155          return Promise.resolve()
156        }
157  
158        return closeTabInIDE(tabName, ideClient)
159      },
160      showingDiffInIDE: shouldShowDiffInIDE && !hasError,
161      ideName: ideName,
162      hasError,
163    }
164  }
165  
166  /**
167   * Re-computes the edits from the old and new contents. This is necessary
168   * to apply any edits the user may have made to the new contents.
169   */
170  export function computeEditsFromContents(
171    filePath: string,
172    oldContent: string,
173    newContent: string,
174    editMode: 'single' | 'multiple',
175  ): FileEdit[] {
176    // Use unformatted patches, otherwise the edits will be formatted.
177    const singleHunk = editMode === 'single'
178    const patch = getPatchFromContents({
179      filePath,
180      oldContent,
181      newContent,
182      singleHunk,
183    })
184  
185    if (patch.length === 0) {
186      return []
187    }
188  
189    // For single edit mode, verify we only got one hunk
190    if (singleHunk && patch.length > 1) {
191      logError(
192        new Error(
193          `Unexpected number of hunks: ${patch.length}. Expected 1 hunk.`,
194        ),
195      )
196    }
197  
198    // Re-compute the edits to match the patch
199    return getEditsForPatch(patch)
200  }
201  
202  /**
203   * Done if:
204   *
205   * 1. Tab is closed in IDE
206   * 2. Tab is saved in IDE (we then close the tab)
207   * 3. User selected an option in IDE
208   * 4. User selected an option in terminal (or hit esc)
209   *
210   * Resolves with the new file content.
211   *
212   * TODO: Time out after 5 mins of inactivity?
213   * TODO: Update auto-approval UI when IDE exits
214   * TODO: Close the IDE tab when the approval prompt is unmounted
215   */
216  async function showDiffInIDE(
217    file_path: string,
218    edits: FileEdit[],
219    toolUseContext: ToolUseContext,
220    tabName: string,
221  ): Promise<{ oldContent: string; newContent: string }> {
222    let isCleanedUp = false
223  
224    const oldFilePath = expandPath(file_path)
225    let oldContent = ''
226    try {
227      oldContent = readFileSync(oldFilePath)
228    } catch (e: unknown) {
229      if (!isENOENT(e)) {
230        throw e
231      }
232    }
233  
234    async function cleanup() {
235      // Careful to avoid race conditions, since this
236      // function can be called from multiple places.
237      if (isCleanedUp) {
238        return
239      }
240      isCleanedUp = true
241  
242      // Don't fail if this fails
243      try {
244        await closeTabInIDE(tabName, ideClient)
245      } catch (e) {
246        logError(e as Error)
247      }
248  
249      process.off('beforeExit', cleanup)
250      toolUseContext.abortController.signal.removeEventListener('abort', cleanup)
251    }
252  
253    // Cleanup if the user hits esc to cancel the tool call - or on exit
254    toolUseContext.abortController.signal.addEventListener('abort', cleanup)
255    process.on('beforeExit', cleanup)
256  
257    // Open the diff in the IDE
258    const ideClient = getConnectedIdeClient(toolUseContext.options.mcpClients)
259    try {
260      const { updatedFile } = getPatchForEdits({
261        filePath: oldFilePath,
262        fileContents: oldContent,
263        edits,
264      })
265  
266      if (!ideClient || ideClient.type !== 'connected') {
267        throw new Error('IDE client not available')
268      }
269      let ideOldPath = oldFilePath
270  
271      // Only convert paths if we're in WSL and IDE is on Windows
272      const ideRunningInWindows =
273        (ideClient.config as McpSSEIDEServerConfig | McpWebSocketIDEServerConfig)
274          .ideRunningInWindows === true
275      if (
276        getPlatform() === 'wsl' &&
277        ideRunningInWindows &&
278        process.env.WSL_DISTRO_NAME
279      ) {
280        const converter = new WindowsToWSLConverter(process.env.WSL_DISTRO_NAME)
281        ideOldPath = converter.toIDEPath(oldFilePath)
282      }
283  
284      const rpcResult = await callIdeRpc(
285        'openDiff',
286        {
287          old_file_path: ideOldPath,
288          new_file_path: ideOldPath,
289          new_file_contents: updatedFile,
290          tab_name: tabName,
291        },
292        ideClient,
293      )
294  
295      // Convert the raw RPC result to a ToolCallResponse format
296      const data = Array.isArray(rpcResult) ? rpcResult : [rpcResult]
297  
298      // If the user saved the file then take the new contents and resolve with that.
299      if (isSaveMessage(data)) {
300        void cleanup()
301        return {
302          oldContent: oldContent,
303          newContent: data[1].text,
304        }
305      } else if (isClosedMessage(data)) {
306        void cleanup()
307        return {
308          oldContent: oldContent,
309          newContent: updatedFile,
310        }
311      } else if (isRejectedMessage(data)) {
312        void cleanup()
313        return {
314          oldContent: oldContent,
315          newContent: oldContent,
316        }
317      }
318  
319      // Indicates that the tool call completed with none of the expected
320      // results. Did the user close the IDE?
321      throw new Error('Not accepted')
322    } catch (error) {
323      logError(error as Error)
324      void cleanup()
325      throw error
326    }
327  }
328  
329  async function closeTabInIDE(
330    tabName: string,
331    ideClient?: MCPServerConnection | undefined,
332  ): Promise<void> {
333    try {
334      if (!ideClient || ideClient.type !== 'connected') {
335        throw new Error('IDE client not available')
336      }
337  
338      // Use direct RPC to close the tab
339      await callIdeRpc('close_tab', { tab_name: tabName }, ideClient)
340    } catch (error) {
341      logError(error as Error)
342      // Don't throw - this is a cleanup operation
343    }
344  }
345  
346  function isClosedMessage(data: unknown): data is { text: 'TAB_CLOSED' } {
347    return (
348      Array.isArray(data) &&
349      typeof data[0] === 'object' &&
350      data[0] !== null &&
351      'type' in data[0] &&
352      data[0].type === 'text' &&
353      'text' in data[0] &&
354      data[0].text === 'TAB_CLOSED'
355    )
356  }
357  
358  function isRejectedMessage(data: unknown): data is { text: 'DIFF_REJECTED' } {
359    return (
360      Array.isArray(data) &&
361      typeof data[0] === 'object' &&
362      data[0] !== null &&
363      'type' in data[0] &&
364      data[0].type === 'text' &&
365      'text' in data[0] &&
366      data[0].text === 'DIFF_REJECTED'
367    )
368  }
369  
370  function isSaveMessage(
371    data: unknown,
372  ): data is [{ text: 'FILE_SAVED' }, { text: string }] {
373    return (
374      Array.isArray(data) &&
375      data[0]?.type === 'text' &&
376      data[0].text === 'FILE_SAVED' &&
377      typeof data[1].text === 'string'
378    )
379  }