forked from
oppi.li/claude-code
source dump of claude code
1import { feature } from 'bun:bundle'
2import { extname, isAbsolute, resolve } from 'path'
3import {
4 fileHistoryEnabled,
5 fileHistoryTrackEdit,
6} from 'src/utils/fileHistory.js'
7import { z } from 'zod/v4'
8import { buildTool, type ToolDef, type ToolUseContext } from '../../Tool.js'
9import type { NotebookCell, NotebookContent } from '../../types/notebook.js'
10import { getCwd } from '../../utils/cwd.js'
11import { isENOENT } from '../../utils/errors.js'
12import { getFileModificationTime, writeTextContent } from '../../utils/file.js'
13import { readFileSyncWithMetadata } from '../../utils/fileRead.js'
14import { safeParseJSON } from '../../utils/json.js'
15import { lazySchema } from '../../utils/lazySchema.js'
16import { parseCellId } from '../../utils/notebook.js'
17import { checkWritePermissionForTool } from '../../utils/permissions/filesystem.js'
18import type { PermissionDecision } from '../../utils/permissions/PermissionResult.js'
19import { jsonParse, jsonStringify } from '../../utils/slowOperations.js'
20import { NOTEBOOK_EDIT_TOOL_NAME } from './constants.js'
21import { DESCRIPTION, PROMPT } from './prompt.js'
22import {
23 getToolUseSummary,
24 renderToolResultMessage,
25 renderToolUseErrorMessage,
26 renderToolUseMessage,
27 renderToolUseRejectedMessage,
28} from './UI.js'
29
30export 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)
58type InputSchema = ReturnType<typeof inputSchema>
59
60export 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)
86type OutputSchema = ReturnType<typeof outputSchema>
87
88export type Output = z.infer<OutputSchema>
89
90export 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>)