Something went wrong. Try again.
source dump of claude code forked from oppi.li/claude-code
Something went wrong. Try again.
15 kB · 490 lines
TypeScript
at main
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491import { feature } from 'bun:bundle'import { extname, isAbsolute, resolve } from 'path'import { fileHistoryEnabled, fileHistoryTrackEdit,} from 'src/utils/fileHistory.js'import { z } from 'zod/v4'import { buildTool, type ToolDef, type ToolUseContext } from '../../Tool.js'import type { NotebookCell, NotebookContent } from '../../types/notebook.js'import { getCwd } from '../../utils/cwd.js'import { isENOENT } from '../../utils/errors.js'import { getFileModificationTime, writeTextContent } from '../../utils/file.js'import { readFileSyncWithMetadata } from '../../utils/fileRead.js'import { safeParseJSON } from '../../utils/json.js'import { lazySchema } from '../../utils/lazySchema.js'import { parseCellId } from '../../utils/notebook.js'import { checkWritePermissionForTool } from '../../utils/permissions/filesystem.js'import type { PermissionDecision } from '../../utils/permissions/PermissionResult.js'import { jsonParse, jsonStringify } from '../../utils/slowOperations.js'import { NOTEBOOK_EDIT_TOOL_NAME } from './constants.js'import { DESCRIPTION, PROMPT } from './prompt.js'import { getToolUseSummary, renderToolResultMessage, renderToolUseErrorMessage, renderToolUseMessage, renderToolUseRejectedMessage,} from './UI.js'
export const inputSchema = lazySchema(() => z.strictObject({ notebook_path: z .string() .describe( 'The absolute path to the Jupyter notebook file to edit (must be absolute, not relative)', ), cell_id: z .string() .optional() .describe( '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.', ), new_source: z.string().describe('The new source for the cell'), cell_type: z .enum(['code', 'markdown']) .optional() .describe( '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.', ), edit_mode: z .enum(['replace', 'insert', 'delete']) .optional() .describe( 'The type of edit to make (replace, insert, delete). Defaults to replace.', ), }),)type InputSchema = ReturnType<typeof inputSchema>
export const outputSchema = lazySchema(() => z.object({ new_source: z .string() .describe('The new source code that was written to the cell'), cell_id: z .string() .optional() .describe('The ID of the cell that was edited'), cell_type: z.enum(['code', 'markdown']).describe('The type of the cell'), language: z.string().describe('The programming language of the notebook'), edit_mode: z.string().describe('The edit mode that was used'), error: z .string() .optional() .describe('Error message if the operation failed'), // Fields for attribution tracking notebook_path: z.string().describe('The path to the notebook file'), original_file: z .string() .describe('The original notebook content before modification'), updated_file: z .string() .describe('The updated notebook content after modification'), }),)type OutputSchema = ReturnType<typeof outputSchema>
export type Output = z.infer<OutputSchema>
export const NotebookEditTool = buildTool({ name: NOTEBOOK_EDIT_TOOL_NAME, searchHint: 'edit Jupyter notebook cells (.ipynb)', maxResultSizeChars: 100_000, shouldDefer: true, async description() { return DESCRIPTION }, async prompt() { return PROMPT }, userFacingName() { return 'Edit Notebook' }, getToolUseSummary, getActivityDescription(input) { const summary = getToolUseSummary(input) return summary ? `Editing notebook ${summary}` : 'Editing notebook' }, get inputSchema(): InputSchema { return inputSchema() }, get outputSchema(): OutputSchema { return outputSchema() }, toAutoClassifierInput(input) { if (feature('TRANSCRIPT_CLASSIFIER')) { const mode = input.edit_mode ?? 'replace' return `${input.notebook_path} ${mode}: ${input.new_source}` } return '' }, getPath(input): string { return input.notebook_path }, async checkPermissions(input, context): Promise<PermissionDecision> { const appState = context.getAppState() return checkWritePermissionForTool( NotebookEditTool, input, appState.toolPermissionContext, ) }, mapToolResultToToolResultBlockParam( { cell_id, edit_mode, new_source, error }, toolUseID, ) { if (error) { return { tool_use_id: toolUseID, type: 'tool_result', content: error, is_error: true, } } switch (edit_mode) { case 'replace': return { tool_use_id: toolUseID, type: 'tool_result', content: `Updated cell ${cell_id} with ${new_source}`, } case 'insert': return { tool_use_id: toolUseID, type: 'tool_result', content: `Inserted cell ${cell_id} with ${new_source}`, } case 'delete': return { tool_use_id: toolUseID, type: 'tool_result', content: `Deleted cell ${cell_id}`, } default: return { tool_use_id: toolUseID, type: 'tool_result', content: 'Unknown edit mode', } } }, renderToolUseMessage, renderToolUseRejectedMessage, renderToolUseErrorMessage, renderToolResultMessage, async validateInput( { notebook_path, cell_type, cell_id, edit_mode = 'replace' }, toolUseContext: ToolUseContext, ) { const fullPath = isAbsolute(notebook_path) ? notebook_path : resolve(getCwd(), notebook_path)
// SECURITY: Skip filesystem operations for UNC paths to prevent NTLM credential leaks. if (fullPath.startsWith('\\\\') || fullPath.startsWith('//')) { return { result: true } }
if (extname(fullPath) !== '.ipynb') { return { result: false, message: 'File must be a Jupyter notebook (.ipynb file). For editing other file types, use the FileEdit tool.', errorCode: 2, } }
if ( edit_mode !== 'replace' && edit_mode !== 'insert' && edit_mode !== 'delete' ) { return { result: false, message: 'Edit mode must be replace, insert, or delete.', errorCode: 4, } }
if (edit_mode === 'insert' && !cell_type) { return { result: false, message: 'Cell type is required when using edit_mode=insert.', errorCode: 5, } }
// Require Read-before-Edit (matches FileEditTool/FileWriteTool). Without // this, the model could edit a notebook it never saw, or edit against a // stale view after an external change — silent data loss. const readTimestamp = toolUseContext.readFileState.get(fullPath) if (!readTimestamp) { return { result: false, message: 'File has not been read yet. Read it first before writing to it.', errorCode: 9, } } if (getFileModificationTime(fullPath) > readTimestamp.timestamp) { return { result: false, message: 'File has been modified since read, either by the user or by a linter. Read it again before attempting to write it.', errorCode: 10, } }
let content: string try { content = readFileSyncWithMetadata(fullPath).content } catch (e) { if (isENOENT(e)) { return { result: false, message: 'Notebook file does not exist.', errorCode: 1, } } throw e } const notebook = safeParseJSON(content) as NotebookContent | null if (!notebook) { return { result: false, message: 'Notebook is not valid JSON.', errorCode: 6, } } if (!cell_id) { if (edit_mode !== 'insert') { return { result: false, message: 'Cell ID must be specified when not inserting a new cell.', errorCode: 7, } } } else { // First try to find the cell by its actual ID const cellIndex = notebook.cells.findIndex(cell => cell.id === cell_id)
if (cellIndex === -1) { // If not found, try to parse as a numeric index (cell-N format) const parsedCellIndex = parseCellId(cell_id) if (parsedCellIndex !== undefined) { if (!notebook.cells[parsedCellIndex]) { return { result: false, message: `Cell with index ${parsedCellIndex} does not exist in notebook.`, errorCode: 7, } } } else { return { result: false, message: `Cell with ID "${cell_id}" not found in notebook.`, errorCode: 8, } } } }
return { result: true } }, async call( { notebook_path, new_source, cell_id, cell_type, edit_mode: originalEditMode, }, { readFileState, updateFileHistoryState }, _, parentMessage, ) { const fullPath = isAbsolute(notebook_path) ? notebook_path : resolve(getCwd(), notebook_path)
if (fileHistoryEnabled()) { await fileHistoryTrackEdit( updateFileHistoryState, fullPath, parentMessage.uuid, ) }
try { // readFileSyncWithMetadata gives content + encoding + line endings in // one safeResolvePath + readFileSync pass, replacing the previous // detectFileEncoding + readFile + detectLineEndings chain (each of // which redid safeResolvePath and/or a 4KB readSync). const { content, encoding, lineEndings } = readFileSyncWithMetadata(fullPath) // Must use non-memoized jsonParse here: safeParseJSON caches by content // string and returns a shared object reference, but we mutate the // notebook in place below (cells.splice, targetCell.source = ...). // Using the memoized version poisons the cache for validateInput() and // any subsequent call() with the same file content. let notebook: NotebookContent try { notebook = jsonParse(content) as NotebookContent } catch { return { data: { new_source, cell_type: cell_type ?? 'code', language: 'python', edit_mode: 'replace', error: 'Notebook is not valid JSON.', cell_id, notebook_path: fullPath, original_file: '', updated_file: '', }, } }
let cellIndex if (!cell_id) { cellIndex = 0 // Default to inserting at the beginning if no cell_id is provided } else { // First try to find the cell by its actual ID cellIndex = notebook.cells.findIndex(cell => cell.id === cell_id)
// If not found, try to parse as a numeric index (cell-N format) if (cellIndex === -1) { const parsedCellIndex = parseCellId(cell_id) if (parsedCellIndex !== undefined) { cellIndex = parsedCellIndex } }
if (originalEditMode === 'insert') { cellIndex += 1 // Insert after the cell with this ID } }
// Convert replace to insert if trying to replace one past the end let edit_mode = originalEditMode if (edit_mode === 'replace' && cellIndex === notebook.cells.length) { edit_mode = 'insert' if (!cell_type) { cell_type = 'code' // Default to code if no cell_type specified } }
const language = notebook.metadata.language_info?.name ?? 'python' let new_cell_id = undefined if ( notebook.nbformat > 4 || (notebook.nbformat === 4 && notebook.nbformat_minor >= 5) ) { if (edit_mode === 'insert') { new_cell_id = Math.random().toString(36).substring(2, 15) } else if (cell_id !== null) { new_cell_id = cell_id } }
if (edit_mode === 'delete') { // Delete the specified cell notebook.cells.splice(cellIndex, 1) } else if (edit_mode === 'insert') { let new_cell: NotebookCell if (cell_type === 'markdown') { new_cell = { cell_type: 'markdown', id: new_cell_id, source: new_source, metadata: {}, } } else { new_cell = { cell_type: 'code', id: new_cell_id, source: new_source, metadata: {}, execution_count: null, outputs: [], } } // Insert the new cell notebook.cells.splice(cellIndex, 0, new_cell) } else { // Find the specified cell const targetCell = notebook.cells[cellIndex]! // validateInput ensures cell_number is in bounds targetCell.source = new_source if (targetCell.cell_type === 'code') { // Reset execution count and clear outputs since cell was modified targetCell.execution_count = null targetCell.outputs = [] } if (cell_type && cell_type !== targetCell.cell_type) { targetCell.cell_type = cell_type } } // Write back to file const IPYNB_INDENT = 1 const updatedContent = jsonStringify(notebook, null, IPYNB_INDENT) writeTextContent(fullPath, updatedContent, encoding, lineEndings) // Update readFileState with post-write mtime (matches FileEditTool/ // FileWriteTool). offset:undefined breaks FileReadTool's dedup match — // without this, Read→NotebookEdit→Read in the same millisecond would // return the file_unchanged stub against stale in-context content. readFileState.set(fullPath, { content: updatedContent, timestamp: getFileModificationTime(fullPath), offset: undefined, limit: undefined, }) const data = { new_source, cell_type: cell_type ?? 'code', language, edit_mode: edit_mode ?? 'replace', cell_id: new_cell_id || undefined, error: '', notebook_path: fullPath, original_file: content, updated_file: updatedContent, } return { data, } } catch (error) { if (error instanceof Error) { const data = { new_source, cell_type: cell_type ?? 'code', language: 'python', edit_mode: 'replace', error: error.message, cell_id, notebook_path: fullPath, original_file: '', updated_file: '', } return { data, } } const data = { new_source, cell_type: cell_type ?? 'code', language: 'python', edit_mode: 'replace', error: 'Unknown error occurred while editing notebook', cell_id, notebook_path: fullPath, original_file: '', updated_file: '', } return { data, } } },} satisfies ToolDef<InputSchema, Output>)