1  import { Hunk } from 'diff'
  2  import { existsSync, mkdirSync, readFileSync, statSync } from 'fs'
  3  import { Box, Text } from 'ink'
  4  import { dirname, isAbsolute, relative, resolve, sep } from 'path'
  5  import * as React from 'react'
  6  import { z } from 'zod'
  7  import { FileEditToolUpdatedMessage } from '../../components/FileEditToolUpdatedMessage.js'
  8  import { StructuredDiff } from '../../components/StructuredDiff.js'
  9  import { logEvent } from '../../services/statsig.js'
 10  import { Tool, ValidationResult } from '../../Tool.js'
 11  import { intersperse } from '../../utils/array.js'
 12  import {
 13    addLineNumbers,
 14    detectFileEncoding,
 15    detectLineEndings,
 16    findSimilarFile,
 17    writeTextContent,
 18  } from '../../utils/file.js'
 19  import { logError } from '../../utils/log.js'
 20  import { getCwd } from '../../utils/state.js'
 21  import { getTheme } from '../../utils/theme.js'
 22  import { NotebookEditTool } from '../NotebookEditTool/NotebookEditTool.js'
 23  import { DESCRIPTION } from './prompt.js'
 24  import { applyEdit } from './utils.js'
 25  import { hasWritePermission } from '../../utils/permissions/filesystem.js'
 26  
 27  const inputSchema = z.strictObject({
 28    file_path: z.string().describe('The absolute path to the file to modify'),
 29    old_string: z.string().describe('The text to replace'),
 30    new_string: z.string().describe('The text to replace it with'),
 31  })
 32  
 33  export type In = typeof inputSchema
 34  
 35  // Number of lines of context to include before/after the change in our result message
 36  const N_LINES_SNIPPET = 4
 37  
 38  export const FileEditTool = {
 39    name: 'Edit',
 40    async description() {
 41      return 'A tool for editing files'
 42    },
 43    async prompt() {
 44      return DESCRIPTION
 45    },
 46    inputSchema,
 47    userFacingName({ old_string, new_string }) {
 48      if (old_string === '') return 'Create'
 49      if (new_string === '') return 'Delete'
 50      return 'Update'
 51    },
 52    async isEnabled() {
 53      return true
 54    },
 55    needsPermissions({ file_path }) {
 56      return !hasWritePermission(file_path)
 57    },
 58    isReadOnly() {
 59      return false
 60    },
 61    renderToolUseMessage(input, { verbose }) {
 62      return `file_path: ${verbose ? input.file_path : relative(getCwd(), input.file_path)}`
 63    },
 64    renderToolResultMessage({ filePath, structuredPatch }, { verbose }) {
 65      return (
 66        <FileEditToolUpdatedMessage
 67          filePath={filePath}
 68          structuredPatch={structuredPatch}
 69          verbose={verbose}
 70        />
 71      )
 72    },
 73    renderToolUseRejectedMessage(
 74      { file_path, old_string, new_string },
 75      { columns, verbose },
 76    ) {
 77      try {
 78        const { patch } = applyEdit(file_path, old_string, new_string)
 79        return (
 80          <Box flexDirection="column">
 81            <Text>
 82              {'  '}⎿{' '}
 83              <Text color={getTheme().error}>
 84                User rejected {old_string === '' ? 'write' : 'update'} to{' '}
 85              </Text>
 86              <Text bold>
 87                {verbose ? file_path : relative(getCwd(), file_path)}
 88              </Text>
 89            </Text>
 90            {intersperse(
 91              patch.map(patch => (
 92                <Box flexDirection="column" paddingLeft={5} key={patch.newStart}>
 93                  <StructuredDiff patch={patch} dim={true} width={columns - 12} />
 94                </Box>
 95              )),
 96              i => (
 97                <Box paddingLeft={5} key={`ellipsis-${i}`}>
 98                  <Text color={getTheme().secondaryText}>...</Text>
 99                </Box>
100              ),
101            )}
102          </Box>
103        )
104      } catch (e) {
105        // Handle the case where while we were showing the diff, the user manually made the change.
106        // TODO: Find a way to show the diff in this case
107        logError(e)
108        return (
109          <Box flexDirection="column">
110            <Text>{'  '}⎿ (No changes)</Text>
111          </Box>
112        )
113      }
114    },
115    async validateInput(
116      { file_path, old_string, new_string },
117      { readFileTimestamps },
118    ) {
119      if (old_string === new_string) {
120        return {
121          result: false,
122          message:
123            'No changes to make: old_string and new_string are exactly the same.',
124          meta: {
125            old_string,
126          },
127        } as ValidationResult
128      }
129  
130      const fullFilePath = isAbsolute(file_path)
131        ? file_path
132        : resolve(getCwd(), file_path)
133  
134      if (existsSync(fullFilePath) && old_string === '') {
135        return {
136          result: false,
137          message: 'Cannot create new file - file already exists.',
138        }
139      }
140  
141      if (!existsSync(fullFilePath) && old_string === '') {
142        return {
143          result: true,
144        }
145      }
146  
147      if (!existsSync(fullFilePath)) {
148        // Try to find a similar file with a different extension
149        const similarFilename = findSimilarFile(fullFilePath)
150        let message = 'File does not exist.'
151  
152        // If we found a similar file, suggest it to the assistant
153        if (similarFilename) {
154          message += ` Did you mean ${similarFilename}?`
155        }
156  
157        return {
158          result: false,
159          message,
160        }
161      }
162  
163      if (fullFilePath.endsWith('.ipynb')) {
164        return {
165          result: false,
166          message: `File is a Jupyter Notebook. Use the ${NotebookEditTool.name} to edit this file.`,
167        }
168      }
169  
170      const readTimestamp = readFileTimestamps[fullFilePath]
171      if (!readTimestamp) {
172        return {
173          result: false,
174          message:
175            'File has not been read yet. Read it first before writing to it.',
176          meta: {
177            isFilePathAbsolute: String(isAbsolute(file_path)),
178          },
179        }
180      }
181  
182      // Check if file exists and get its last modified time
183      const stats = statSync(fullFilePath)
184      const lastWriteTime = stats.mtimeMs
185      if (lastWriteTime > readTimestamp) {
186        return {
187          result: false,
188          message:
189            'File has been modified since read, either by the user or by a linter. Read it again before attempting to write it.',
190        }
191      }
192  
193      const enc = detectFileEncoding(fullFilePath)
194      const file = readFileSync(fullFilePath, enc)
195      if (!file.includes(old_string)) {
196        return {
197          result: false,
198          message: `String to replace not found in file.`,
199          meta: {
200            isFilePathAbsolute: String(isAbsolute(file_path)),
201          },
202        }
203      }
204  
205      const matches = file.split(old_string).length - 1
206      if (matches > 1) {
207        return {
208          result: false,
209          message: `Found ${matches} matches of the string to replace. For safety, this tool only supports replacing exactly one occurrence at a time. Add more lines of context to your edit and try again.`,
210          meta: {
211            isFilePathAbsolute: String(isAbsolute(file_path)),
212          },
213        }
214      }
215  
216      return { result: true }
217    },
218    async *call({ file_path, old_string, new_string }, { readFileTimestamps }) {
219      const { patch, updatedFile } = applyEdit(file_path, old_string, new_string)
220  
221      const fullFilePath = isAbsolute(file_path)
222        ? file_path
223        : resolve(getCwd(), file_path)
224      const dir = dirname(fullFilePath)
225      mkdirSync(dir, { recursive: true })
226      const enc = existsSync(fullFilePath)
227        ? detectFileEncoding(fullFilePath)
228        : 'utf8'
229      const endings = existsSync(fullFilePath)
230        ? detectLineEndings(fullFilePath)
231        : 'LF'
232      const originalFile = existsSync(fullFilePath)
233        ? readFileSync(fullFilePath, enc)
234        : ''
235      writeTextContent(fullFilePath, updatedFile, enc, endings)
236  
237      // Update read timestamp, to invalidate stale writes
238      readFileTimestamps[fullFilePath] = statSync(fullFilePath).mtimeMs
239  
240      // Log when editing CLAUDE.md
241      if (fullFilePath.endsWith(`${sep}CLAUDE.md`)) {
242        logEvent('tengu_write_claudemd', {})
243      }
244  
245      const data = {
246        filePath: file_path,
247        oldString: old_string,
248        newString: new_string,
249        originalFile,
250        structuredPatch: patch,
251      }
252      yield {
253        type: 'result',
254        data,
255        resultForAssistant: this.renderResultForAssistant(data),
256      }
257    },
258    renderResultForAssistant({ filePath, originalFile, oldString, newString }) {
259      const { snippet, startLine } = getSnippet(
260        originalFile || '',
261        oldString,
262        newString,
263      )
264      return `The file ${filePath} has been updated. Here's the result of running \`cat -n\` on a snippet of the edited file:
265  ${addLineNumbers({
266    content: snippet,
267    startLine,
268  })}`
269    },
270  } satisfies Tool<
271    typeof inputSchema,
272    {
273      filePath: string
274      oldString: string
275      newString: string
276      originalFile: string
277      structuredPatch: Hunk[]
278    }
279  >
280  
281  export function getSnippet(
282    initialText: string,
283    oldStr: string,
284    newStr: string,
285  ): { snippet: string; startLine: number } {
286    const before = initialText.split(oldStr)[0] ?? ''
287    const replacementLine = before.split(/\r?\n/).length - 1
288    const newFileLines = initialText.replace(oldStr, newStr).split(/\r?\n/)
289    // Calculate the start and end line numbers for the snippet
290    const startLine = Math.max(0, replacementLine - N_LINES_SNIPPET)
291    const endLine =
292      replacementLine + N_LINES_SNIPPET + newStr.split(/\r?\n/).length
293    // Get snippet
294    const snippetLines = newFileLines.slice(startLine, endLine + 1)
295    const snippet = snippetLines.join('\n')
296    return { snippet, startLine: startLine + 1 }
297  }