ummaya 0.2.4 → 0.2.5

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)
+}

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

@@ -0,0 +1,142 @@
+import { extname, isAbsolute, resolve } from 'path'
+import type { ToolUseContext, ValidationResult } from '../../Tool.js'
+import { getCwd } from '../../utils/cwd.js'
+import { isENOENT } from '../../utils/errors.js'
+import { getFileModificationTime } from '../../utils/file.js'
+import { readFileSyncWithMetadata } from '../../utils/fileRead.js'
+import { safeParseJSON } from '../../utils/json.js'
+import { documentDerivativeMutationValidationForResolvedTarget } from '../DocumentPrimitive/documentMutationGuard.js'
+import { isNotebookDocument } from './notebookModel.js'
+
+export type NotebookEditInput = {
+  notebook_path: string
+  cell_id?: string
+  new_source: string
+  cell_type?: 'code' | 'markdown'
+  edit_mode?: 'replace' | 'insert' | 'delete'
+}
+
+export async function validateNotebookEditInput(
+  { notebook_path, cell_type, cell_id, edit_mode = 'replace' }: NotebookEditInput,
+  toolUseContext: ToolUseContext,
+): Promise<ValidationResult> {
+  const fullPath = isAbsolute(notebook_path)
+    ? notebook_path
+    : resolve(getCwd(), notebook_path)
+  const documentValidation =
+    documentDerivativeMutationValidationForResolvedTarget(fullPath)
+  if (documentValidation !== null) return documentValidation
+
+  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,
+    }
+  }
+
+  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)
+  if (!isNotebookDocument(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 {
+    const cellIndex = notebook.cells.findIndex(cell => cell.id === cell_id)
+
+    if (cellIndex === -1) {
+      const parsedCellIndex = parseNotebookCellId(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 }
+}
+
+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)
+}