1  import { feature } from 'bun:bundle'
  2  import { extname, isAbsolute, resolve } from 'path'
  3  import {
  4    fileHistoryEnabled,
  5    fileHistoryTrackEdit,
  6  } from 'src/utils/fileHistory.js'
  7  import { z } from 'zod/v4'
  8  import { buildTool, type ToolDef, type ToolUseContext } from '../../Tool.js'
  9  import type { NotebookCell, NotebookContent } from '../../types/notebook.js'
 10  import { getCwd } from '../../utils/cwd.js'
 11  import { isENOENT } from '../../utils/errors.js'
 12  import { getFileModificationTime, writeTextContent } from '../../utils/file.js'
 13  import { readFileSyncWithMetadata } from '../../utils/fileRead.js'
 14  import { safeParseJSON } from '../../utils/json.js'
 15  import { lazySchema } from '../../utils/lazySchema.js'
 16  import { parseCellId } from '../../utils/notebook.js'
 17  import { checkWritePermissionForTool } from '../../utils/permissions/filesystem.js'
 18  import type { PermissionDecision } from '../../utils/permissions/PermissionResult.js'
 19  import { jsonParse, jsonStringify } from '../../utils/slowOperations.js'
 20  import { NOTEBOOK_EDIT_TOOL_NAME } from './constants.js'
 21  import { DESCRIPTION, PROMPT } from './prompt.js'
 22  import {
 23    getToolUseSummary,
 24    renderToolResultMessage,
 25    renderToolUseErrorMessage,
 26    renderToolUseMessage,
 27    renderToolUseRejectedMessage,
 28  } from './UI.js'
 29  
 30  export const inputSchema = lazySchema(() =>
 31    z.strictObject({
 32      notebook_path: z
 33        .string()
 34        .describe(
 35          'The absolute path to the Jupyter notebook file to edit (must be absolute, not relative)',
 36        ),
 37      cell_id: z
 38        .string()
 39        .optional()
 40        .describe(
 41          'The ID of the cell to edit. When inserting a new cell, the new cell will be inserted after the cell with this ID, or at the beginning if not specified.',
 42        ),
 43      new_source: z.string().describe('The new source for the cell'),
 44      cell_type: z
 45        .enum(['code', 'markdown'])
 46        .optional()
 47        .describe(
 48          'The type of the cell (code or markdown). If not specified, it defaults to the current cell type. If using edit_mode=insert, this is required.',
 49        ),
 50      edit_mode: z
 51        .enum(['replace', 'insert', 'delete'])
 52        .optional()
 53        .describe(
 54          'The type of edit to make (replace, insert, delete). Defaults to replace.',
 55        ),
 56    }),
 57  )
 58  type InputSchema = ReturnType<typeof inputSchema>
 59  
 60  export const outputSchema = lazySchema(() =>
 61    z.object({
 62      new_source: z
 63        .string()
 64        .describe('The new source code that was written to the cell'),
 65      cell_id: z
 66        .string()
 67        .optional()
 68        .describe('The ID of the cell that was edited'),
 69      cell_type: z.enum(['code', 'markdown']).describe('The type of the cell'),
 70      language: z.string().describe('The programming language of the notebook'),
 71      edit_mode: z.string().describe('The edit mode that was used'),
 72      error: z
 73        .string()
 74        .optional()
 75        .describe('Error message if the operation failed'),
 76      // Fields for attribution tracking
 77      notebook_path: z.string().describe('The path to the notebook file'),
 78      original_file: z
 79        .string()
 80        .describe('The original notebook content before modification'),
 81      updated_file: z
 82        .string()
 83        .describe('The updated notebook content after modification'),
 84    }),
 85  )
 86  type OutputSchema = ReturnType<typeof outputSchema>
 87  
 88  export type Output = z.infer<OutputSchema>
 89  
 90  export const NotebookEditTool = buildTool({
 91    name: NOTEBOOK_EDIT_TOOL_NAME,
 92    searchHint: 'edit Jupyter notebook cells (.ipynb)',
 93    maxResultSizeChars: 100_000,
 94    shouldDefer: true,
 95    async description() {
 96      return DESCRIPTION
 97    },
 98    async prompt() {
 99      return PROMPT
100    },
101    userFacingName() {
102      return 'Edit Notebook'
103    },
104    getToolUseSummary,
105    getActivityDescription(input) {
106      const summary = getToolUseSummary(input)
107      return summary ? `Editing notebook ${summary}` : 'Editing notebook'
108    },
109    get inputSchema(): InputSchema {
110      return inputSchema()
111    },
112    get outputSchema(): OutputSchema {
113      return outputSchema()
114    },
115    toAutoClassifierInput(input) {
116      if (feature('TRANSCRIPT_CLASSIFIER')) {
117        const mode = input.edit_mode ?? 'replace'
118        return `${input.notebook_path} ${mode}: ${input.new_source}`
119      }
120      return ''
121    },
122    getPath(input): string {
123      return input.notebook_path
124    },
125    async checkPermissions(input, context): Promise<PermissionDecision> {
126      const appState = context.getAppState()
127      return checkWritePermissionForTool(
128        NotebookEditTool,
129        input,
130        appState.toolPermissionContext,
131      )
132    },
133    mapToolResultToToolResultBlockParam(
134      { cell_id, edit_mode, new_source, error },
135      toolUseID,
136    ) {
137      if (error) {
138        return {
139          tool_use_id: toolUseID,
140          type: 'tool_result',
141          content: error,
142          is_error: true,
143        }
144      }
145      switch (edit_mode) {
146        case 'replace':
147          return {
148            tool_use_id: toolUseID,
149            type: 'tool_result',
150            content: `Updated cell ${cell_id} with ${new_source}`,
151          }
152        case 'insert':
153          return {
154            tool_use_id: toolUseID,
155            type: 'tool_result',
156            content: `Inserted cell ${cell_id} with ${new_source}`,
157          }
158        case 'delete':
159          return {
160            tool_use_id: toolUseID,
161            type: 'tool_result',
162            content: `Deleted cell ${cell_id}`,
163          }
164        default:
165          return {
166            tool_use_id: toolUseID,
167            type: 'tool_result',
168            content: 'Unknown edit mode',
169          }
170      }
171    },
172    renderToolUseMessage,
173    renderToolUseRejectedMessage,
174    renderToolUseErrorMessage,
175    renderToolResultMessage,
176    async validateInput(
177      { notebook_path, cell_type, cell_id, edit_mode = 'replace' },
178      toolUseContext: ToolUseContext,
179    ) {
180      const fullPath = isAbsolute(notebook_path)
181        ? notebook_path
182        : resolve(getCwd(), notebook_path)
183  
184      // SECURITY: Skip filesystem operations for UNC paths to prevent NTLM credential leaks.
185      if (fullPath.startsWith('\\\\') || fullPath.startsWith('//')) {
186        return { result: true }
187      }
188  
189      if (extname(fullPath) !== '.ipynb') {
190        return {
191          result: false,
192          message:
193            'File must be a Jupyter notebook (.ipynb file). For editing other file types, use the FileEdit tool.',
194          errorCode: 2,
195        }
196      }
197  
198      if (
199        edit_mode !== 'replace' &&
200        edit_mode !== 'insert' &&
201        edit_mode !== 'delete'
202      ) {
203        return {
204          result: false,
205          message: 'Edit mode must be replace, insert, or delete.',
206          errorCode: 4,
207        }
208      }
209  
210      if (edit_mode === 'insert' && !cell_type) {
211        return {
212          result: false,
213          message: 'Cell type is required when using edit_mode=insert.',
214          errorCode: 5,
215        }
216      }
217  
218      // Require Read-before-Edit (matches FileEditTool/FileWriteTool). Without
219      // this, the model could edit a notebook it never saw, or edit against a
220      // stale view after an external change — silent data loss.
221      const readTimestamp = toolUseContext.readFileState.get(fullPath)
222      if (!readTimestamp) {
223        return {
224          result: false,
225          message:
226            'File has not been read yet. Read it first before writing to it.',
227          errorCode: 9,
228        }
229      }
230      if (getFileModificationTime(fullPath) > readTimestamp.timestamp) {
231        return {
232          result: false,
233          message:
234            'File has been modified since read, either by the user or by a linter. Read it again before attempting to write it.',
235          errorCode: 10,
236        }
237      }
238  
239      let content: string
240      try {
241        content = readFileSyncWithMetadata(fullPath).content
242      } catch (e) {
243        if (isENOENT(e)) {
244          return {
245            result: false,
246            message: 'Notebook file does not exist.',
247            errorCode: 1,
248          }
249        }
250        throw e
251      }
252      const notebook = safeParseJSON(content) as NotebookContent | null
253      if (!notebook) {
254        return {
255          result: false,
256          message: 'Notebook is not valid JSON.',
257          errorCode: 6,
258        }
259      }
260      if (!cell_id) {
261        if (edit_mode !== 'insert') {
262          return {
263            result: false,
264            message: 'Cell ID must be specified when not inserting a new cell.',
265            errorCode: 7,
266          }
267        }
268      } else {
269        // First try to find the cell by its actual ID
270        const cellIndex = notebook.cells.findIndex(cell => cell.id === cell_id)
271  
272        if (cellIndex === -1) {
273          // If not found, try to parse as a numeric index (cell-N format)
274          const parsedCellIndex = parseCellId(cell_id)
275          if (parsedCellIndex !== undefined) {
276            if (!notebook.cells[parsedCellIndex]) {
277              return {
278                result: false,
279                message: `Cell with index ${parsedCellIndex} does not exist in notebook.`,
280                errorCode: 7,
281              }
282            }
283          } else {
284            return {
285              result: false,
286              message: `Cell with ID "${cell_id}" not found in notebook.`,
287              errorCode: 8,
288            }
289          }
290        }
291      }
292  
293      return { result: true }
294    },
295    async call(
296      {
297        notebook_path,
298        new_source,
299        cell_id,
300        cell_type,
301        edit_mode: originalEditMode,
302      },
303      { readFileState, updateFileHistoryState },
304      _,
305      parentMessage,
306    ) {
307      const fullPath = isAbsolute(notebook_path)
308        ? notebook_path
309        : resolve(getCwd(), notebook_path)
310  
311      if (fileHistoryEnabled()) {
312        await fileHistoryTrackEdit(
313          updateFileHistoryState,
314          fullPath,
315          parentMessage.uuid,
316        )
317      }
318  
319      try {
320        // readFileSyncWithMetadata gives content + encoding + line endings in
321        // one safeResolvePath + readFileSync pass, replacing the previous
322        // detectFileEncoding + readFile + detectLineEndings chain (each of
323        // which redid safeResolvePath and/or a 4KB readSync).
324        const { content, encoding, lineEndings } =
325          readFileSyncWithMetadata(fullPath)
326        // Must use non-memoized jsonParse here: safeParseJSON caches by content
327        // string and returns a shared object reference, but we mutate the
328        // notebook in place below (cells.splice, targetCell.source = ...).
329        // Using the memoized version poisons the cache for validateInput() and
330        // any subsequent call() with the same file content.
331        let notebook: NotebookContent
332        try {
333          notebook = jsonParse(content) as NotebookContent
334        } catch {
335          return {
336            data: {
337              new_source,
338              cell_type: cell_type ?? 'code',
339              language: 'python',
340              edit_mode: 'replace',
341              error: 'Notebook is not valid JSON.',
342              cell_id,
343              notebook_path: fullPath,
344              original_file: '',
345              updated_file: '',
346            },
347          }
348        }
349  
350        let cellIndex
351        if (!cell_id) {
352          cellIndex = 0 // Default to inserting at the beginning if no cell_id is provided
353        } else {
354          // First try to find the cell by its actual ID
355          cellIndex = notebook.cells.findIndex(cell => cell.id === cell_id)
356  
357          // If not found, try to parse as a numeric index (cell-N format)
358          if (cellIndex === -1) {
359            const parsedCellIndex = parseCellId(cell_id)
360            if (parsedCellIndex !== undefined) {
361              cellIndex = parsedCellIndex
362            }
363          }
364  
365          if (originalEditMode === 'insert') {
366            cellIndex += 1 // Insert after the cell with this ID
367          }
368        }
369  
370        // Convert replace to insert if trying to replace one past the end
371        let edit_mode = originalEditMode
372        if (edit_mode === 'replace' && cellIndex === notebook.cells.length) {
373          edit_mode = 'insert'
374          if (!cell_type) {
375            cell_type = 'code' // Default to code if no cell_type specified
376          }
377        }
378  
379        const language = notebook.metadata.language_info?.name ?? 'python'
380        let new_cell_id = undefined
381        if (
382          notebook.nbformat > 4 ||
383          (notebook.nbformat === 4 && notebook.nbformat_minor >= 5)
384        ) {
385          if (edit_mode === 'insert') {
386            new_cell_id = Math.random().toString(36).substring(2, 15)
387          } else if (cell_id !== null) {
388            new_cell_id = cell_id
389          }
390        }
391  
392        if (edit_mode === 'delete') {
393          // Delete the specified cell
394          notebook.cells.splice(cellIndex, 1)
395        } else if (edit_mode === 'insert') {
396          let new_cell: NotebookCell
397          if (cell_type === 'markdown') {
398            new_cell = {
399              cell_type: 'markdown',
400              id: new_cell_id,
401              source: new_source,
402              metadata: {},
403            }
404          } else {
405            new_cell = {
406              cell_type: 'code',
407              id: new_cell_id,
408              source: new_source,
409              metadata: {},
410              execution_count: null,
411              outputs: [],
412            }
413          }
414          // Insert the new cell
415          notebook.cells.splice(cellIndex, 0, new_cell)
416        } else {
417          // Find the specified cell
418          const targetCell = notebook.cells[cellIndex]! // validateInput ensures cell_number is in bounds
419          targetCell.source = new_source
420          if (targetCell.cell_type === 'code') {
421            // Reset execution count and clear outputs since cell was modified
422            targetCell.execution_count = null
423            targetCell.outputs = []
424          }
425          if (cell_type && cell_type !== targetCell.cell_type) {
426            targetCell.cell_type = cell_type
427          }
428        }
429        // Write back to file
430        const IPYNB_INDENT = 1
431        const updatedContent = jsonStringify(notebook, null, IPYNB_INDENT)
432        writeTextContent(fullPath, updatedContent, encoding, lineEndings)
433        // Update readFileState with post-write mtime (matches FileEditTool/
434        // FileWriteTool). offset:undefined breaks FileReadTool's dedup match —
435        // without this, Read→NotebookEdit→Read in the same millisecond would
436        // return the file_unchanged stub against stale in-context content.
437        readFileState.set(fullPath, {
438          content: updatedContent,
439          timestamp: getFileModificationTime(fullPath),
440          offset: undefined,
441          limit: undefined,
442        })
443        const data = {
444          new_source,
445          cell_type: cell_type ?? 'code',
446          language,
447          edit_mode: edit_mode ?? 'replace',
448          cell_id: new_cell_id || undefined,
449          error: '',
450          notebook_path: fullPath,
451          original_file: content,
452          updated_file: updatedContent,
453        }
454        return {
455          data,
456        }
457      } catch (error) {
458        if (error instanceof Error) {
459          const data = {
460            new_source,
461            cell_type: cell_type ?? 'code',
462            language: 'python',
463            edit_mode: 'replace',
464            error: error.message,
465            cell_id,
466            notebook_path: fullPath,
467            original_file: '',
468            updated_file: '',
469          }
470          return {
471            data,
472          }
473        }
474        const data = {
475          new_source,
476          cell_type: cell_type ?? 'code',
477          language: 'python',
478          edit_mode: 'replace',
479          error: 'Unknown error occurred while editing notebook',
480          cell_id,
481          notebook_path: fullPath,
482          original_file: '',
483          updated_file: '',
484        }
485        return {
486          data,
487        }
488      }
489    },
490  } satisfies ToolDef<InputSchema, Output>)