ummaya 0.2.3 → 0.2.5

package/tui/src/tools/NotebookEditTool/NotebookEditTool.ts

@@ -1,22 +1,10 @@
 import { 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 { buildTool, type ToolDef } from '../../Tool.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 { callNotebookEditTool } from './call.js'
 import { NOTEBOOK_EDIT_TOOL_NAME } from './constants.js'
 import { DESCRIPTION, PROMPT } from './prompt.js'
 import {
@@ -26,6 +14,7 @@ import {
   renderToolUseMessage,
   renderToolUseRejectedMessage,
 } from './UI.js'
+import { validateNotebookEditInput } from './validateInput.js'
 
 export const inputSchema = lazySchema(() =>
   z.strictObject({
@@ -173,318 +162,10 @@ export const NotebookEditTool = buildTool({
   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 validateInput(input, toolUseContext) {
+    return validateNotebookEditInput(input, toolUseContext)
   },
-  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,
-      }
-    }
+  async call(input, context, _, parentMessage) {
+    return callNotebookEditTool(input, context, parentMessage)
   },
 } satisfies ToolDef<InputSchema, Output>)

package/tui/src/tools/NotebookEditTool/call.ts

@@ -0,0 +1,254 @@
+import { isAbsolute, resolve } from 'path'
+import type { ToolResult, ToolUseContext } from '../../Tool.js'
+import type { AssistantMessage } from '../../types/message.js'
+import { getCwd } from '../../utils/cwd.js'
+import { getFileModificationTime, writeTextContent } from '../../utils/file.js'
+import { readFileSyncWithMetadata } from '../../utils/fileRead.js'
+import { jsonParse, jsonStringify } from '../../utils/slowOperations.js'
+import { documentDerivativeMutationValidationForResolvedTarget } from '../DocumentPrimitive/documentMutationGuard.js'
+import { isNotebookDocument, type NotebookCellRecord } from './notebookModel.js'
+import type { NotebookEditInput } from './validateInput.js'
+
+export type NotebookEditOutput = {
+  new_source: string
+  cell_id?: string
+  cell_type: 'code' | 'markdown'
+  language: string
+  edit_mode: string
+  error?: string
+  notebook_path: string
+  original_file: string
+  updated_file: string
+}
+
+export async function callNotebookEditTool(
+  {
+    notebook_path,
+    new_source,
+    cell_id,
+    cell_type,
+    edit_mode: originalEditMode,
+  }: NotebookEditInput,
+  { readFileState, updateFileHistoryState }: ToolUseContext,
+  parentMessage: AssistantMessage,
+): Promise<ToolResult<NotebookEditOutput>> {
+  const fullPath = isAbsolute(notebook_path)
+    ? notebook_path
+    : resolve(getCwd(), notebook_path)
+  const documentValidation =
+    documentDerivativeMutationValidationForResolvedTarget(fullPath)
+  if (documentValidation !== null) throw new Error(documentValidation.message)
+
+  const { fileHistoryEnabled, fileHistoryTrackEdit } = await import(
+    '../../utils/fileHistory.js'
+  )
+  if (fileHistoryEnabled()) {
+    await fileHistoryTrackEdit(
+      updateFileHistoryState,
+      fullPath,
+      parentMessage.uuid,
+    )
+  }
+
+  try {
+    const { content, encoding, lineEndings } =
+      readFileSyncWithMetadata(fullPath)
+    const parsed: unknown = jsonParse(content)
+    if (!isNotebookDocument(parsed)) {
+      return invalidNotebookResult(new_source, cell_type, cell_id, fullPath)
+    }
+
+    const notebook = parsed
+    let cellIndex = resolveCellIndex(
+      notebook.cells,
+      cell_id,
+      originalEditMode,
+    )
+
+    let editMode = originalEditMode
+    if (editMode === 'replace' && cellIndex === notebook.cells.length) {
+      editMode = 'insert'
+      if (!cell_type) {
+        cell_type = 'code'
+      }
+    }
+
+    const language = notebook.metadata.language_info?.name ?? 'python'
+    const newCellId = resolveNewCellId(
+      notebook.nbformat,
+      notebook.nbformat_minor,
+      editMode,
+      cell_id,
+    )
+
+    if (editMode === 'delete') {
+      notebook.cells.splice(cellIndex, 1)
+    } else if (editMode === 'insert') {
+      notebook.cells.splice(
+        cellIndex,
+        0,
+        createNotebookCell(cell_type, newCellId, new_source),
+      )
+    } else {
+      const targetCell = notebook.cells[cellIndex]
+      if (targetCell === undefined) {
+        throw new Error('Notebook cell index is out of bounds')
+      }
+      targetCell.source = new_source
+      if (targetCell.cell_type === 'code') {
+        targetCell.execution_count = null
+        targetCell.outputs = []
+      }
+      if (cell_type && cell_type !== targetCell.cell_type) {
+        targetCell.cell_type = cell_type
+      }
+    }
+
+    const updatedContent = jsonStringify(notebook, null, 1)
+    writeTextContent(fullPath, updatedContent, encoding, lineEndings)
+    readFileState.set(fullPath, {
+      content: updatedContent,
+      timestamp: getFileModificationTime(fullPath),
+      offset: undefined,
+      limit: undefined,
+    })
+
+    return {
+      data: {
+        new_source,
+        cell_type: cell_type ?? 'code',
+        language,
+        edit_mode: editMode ?? 'replace',
+        cell_id: newCellId,
+        error: '',
+        notebook_path: fullPath,
+        original_file: content,
+        updated_file: updatedContent,
+      },
+    }
+  } catch (error) {
+    if (error instanceof Error) {
+      return {
+        data: buildErrorOutput(
+          new_source,
+          cell_type,
+          cell_id,
+          fullPath,
+          error.message,
+        ),
+      }
+    }
+    return {
+      data: buildErrorOutput(
+        new_source,
+        cell_type,
+        cell_id,
+        fullPath,
+        'Unknown error occurred while editing notebook',
+      ),
+    }
+  }
+}
+
+function resolveCellIndex(
+  cells: NotebookCellRecord[],
+  cellId: string | undefined,
+  originalEditMode: 'replace' | 'insert' | 'delete' | undefined,
+): number {
+  if (!cellId) {
+    return 0
+  }
+
+  let cellIndex = cells.findIndex(cell => cell.id === cellId)
+  if (cellIndex === -1) {
+    const parsedCellIndex = parseNotebookCellId(cellId)
+    if (parsedCellIndex !== undefined) {
+      cellIndex = parsedCellIndex
+    }
+  }
+
+  return originalEditMode === 'insert' ? cellIndex + 1 : cellIndex
+}
+
+function parseNotebookCellId(cellId: string): number | undefined {
+  const match = cellId.match(/^cell-(\d+)$/u)
+  const indexValue = match?.[1]
+  if (indexValue === undefined) return undefined
+  return Number.parseInt(indexValue, 10)
+}
+
+function resolveNewCellId(
+  nbformat: number,
+  nbformatMinor: number,
+  editMode: 'replace' | 'insert' | 'delete' | undefined,
+  cellId: string | undefined,
+): string | undefined {
+  if (nbformat <= 4 && (nbformat !== 4 || nbformatMinor < 5)) {
+    return undefined
+  }
+  if (editMode === 'insert') {
+    return Math.random().toString(36).substring(2, 15)
+  }
+  return cellId
+}
+
+function createNotebookCell(
+  cellType: 'code' | 'markdown' | undefined,
+  newCellId: string | undefined,
+  newSource: string,
+): NotebookCellRecord {
+  if (cellType === 'markdown') {
+    return {
+      cell_type: 'markdown',
+      id: newCellId,
+      source: newSource,
+      metadata: {},
+    }
+  }
+
+  return {
+    cell_type: 'code',
+    id: newCellId,
+    source: newSource,
+    metadata: {},
+    execution_count: null,
+    outputs: [],
+  }
+}
+
+function invalidNotebookResult(
+  newSource: string,
+  cellType: 'code' | 'markdown' | undefined,
+  cellId: string | undefined,
+  notebookPath: string,
+): ToolResult<NotebookEditOutput> {
+  return {
+    data: buildErrorOutput(
+      newSource,
+      cellType,
+      cellId,
+      notebookPath,
+      'Notebook is not valid JSON.',
+    ),
+  }
+}
+
+function buildErrorOutput(
+  newSource: string,
+  cellType: 'code' | 'markdown' | undefined,
+  cellId: string | undefined,
+  notebookPath: string,
+  error: string,
+): NotebookEditOutput {
+  return {
+    new_source: newSource,
+    cell_type: cellType ?? 'code',
+    language: 'python',
+    edit_mode: 'replace',
+    error,
+    cell_id: cellId,
+    notebook_path: notebookPath,
+    original_file: '',
+    updated_file: '',
+  }
+}

package/tui/src/tools/NotebookEditTool/notebookModel.ts

@@ -0,0 +1,51 @@
+export type NotebookCellRecord = {
+  cell_type: 'code' | 'markdown'
+  id?: string
+  source: string | string[]
+  metadata: Record<string, unknown>
+  execution_count?: number | null
+  outputs?: unknown[]
+}
+
+export type NotebookDocument = {
+  cells: NotebookCellRecord[]
+  metadata: {
+    language_info?: {
+      name?: string
+    }
+  }
+  nbformat: number
+  nbformat_minor: number
+}
+
+export function isNotebookDocument(value: unknown): value is NotebookDocument {
+  if (!isRecord(value)) return false
+
+  const cells = value.cells
+  const metadata = value.metadata
+  return (
+    Array.isArray(cells) &&
+    cells.every(isNotebookCellRecord) &&
+    isRecord(metadata) &&
+    typeof value.nbformat === 'number' &&
+    typeof value.nbformat_minor === 'number'
+  )
+}
+
+function isNotebookCellRecord(value: unknown): value is NotebookCellRecord {
+  if (!isRecord(value)) return false
+
+  const cellType = value.cell_type
+  const source = value.source
+  return (
+    (cellType === 'code' || cellType === 'markdown') &&
+    (typeof source === 'string' ||
+      (Array.isArray(source) &&
+        source.every(item => typeof item === 'string'))) &&
+    isRecord(value.metadata)
+  )
+}
+
+function isRecord(value: unknown): value is Record<string, unknown> {
+  return typeof value === 'object' && value !== null && !Array.isArray(value)
+}