From cbfdc0a21254d3ac7a79a13dabfe805cebbaf9ba Mon Sep 17 00:00:00 2001 From: can1357 Date: Thu, 7 May 2026 04:16:11 +0200 Subject: [PATCH] feat(coding-agent): added notebook-aware `.ipynb` edit/read path helpers - Added `.ipynb` detection and editable-cell conversion utilities, including merge/serialize helpers in `edit/notebook`. - Rerouted hashline, patch, and replace edit flows, plus read/write paths, through notebook-aware helpers before persistence. - Removed the dedicated `notebook` tool, its schema flags, renderer, and built-in registration/settings checks. - Updated notebook read behavior, docs, and tests so `.ipynb` reads return editable `# %%` cells and edits reserialize to JSON. --- packages/coding-agent/CHANGELOG.md | 13 + .../src/config/settings-schema.ts | 5 - .../coding-agent/src/edit/modes/hashline.ts | 28 +- packages/coding-agent/src/edit/modes/patch.ts | 10 +- .../coding-agent/src/edit/modes/replace.ts | 8 +- packages/coding-agent/src/edit/notebook.ts | 222 ++++++++++++++ packages/coding-agent/src/edit/read-file.ts | 7 + .../coding-agent/src/prompts/tools/read.md | 2 +- packages/coding-agent/src/tools/index.ts | 4 - packages/coding-agent/src/tools/notebook.ts | 286 ------------------ packages/coding-agent/src/tools/read.ts | 17 +- packages/coding-agent/src/tools/renderers.ts | 2 - .../test/tool-discovery/initial-tools.test.ts | 2 - packages/coding-agent/test/tools.test.ts | 61 +++- .../coding-agent/test/tools/index.test.ts | 3 - 15 files changed, 340 insertions(+), 330 deletions(-) create mode 100644 packages/coding-agent/src/edit/notebook.ts delete mode 100644 packages/coding-agent/src/tools/notebook.ts diff --git a/packages/coding-agent/CHANGELOG.md b/packages/coding-agent/CHANGELOG.md index 682291f08..8e92d3198 100644 --- a/packages/coding-agent/CHANGELOG.md +++ b/packages/coding-agent/CHANGELOG.md @@ -1,6 +1,19 @@ # Changelog ## [Unreleased] +### Breaking Changes + +- Removed the dedicated `notebook` tool; `.ipynb` reads and edits now go through `read` and `edit`. + +### Changed + +- Changed `read` tool behavior for `.ipynb:raw` requests to return raw notebook content instead of converting via markit +- Changed `.ipynb` edit and read handling to route through notebook serialization helpers +- Changed `.ipynb` reads to return an editable cell text representation and apply edits back to notebook JSON while preserving cell metadata and outputs where possible. + +### Removed + +- Removed the `notebook.enabled` configuration option from tool settings ## [14.7.2] - 2026-05-06 ### Breaking Changes diff --git a/packages/coding-agent/src/config/settings-schema.ts b/packages/coding-agent/src/config/settings-schema.ts index 63bd8a0a2..173ce47ac 100644 --- a/packages/coding-agent/src/config/settings-schema.ts +++ b/packages/coding-agent/src/config/settings-schema.ts @@ -1756,11 +1756,6 @@ export const SETTINGS_SCHEMA = { }, // Optional tools - "notebook.enabled": { - type: "boolean", - default: true, - ui: { tab: "tools", label: "Notebook", description: "Enable the notebook tool for notebook editing" }, - }, "renderMermaid.enabled": { type: "boolean", diff --git a/packages/coding-agent/src/edit/modes/hashline.ts b/packages/coding-agent/src/edit/modes/hashline.ts index ca2e6670d..7348bf052 100644 --- a/packages/coding-agent/src/edit/modes/hashline.ts +++ b/packages/coding-agent/src/edit/modes/hashline.ts @@ -52,6 +52,7 @@ import { HL_HASH_CAPTURE_RE_RAW, } from "../line-hash"; import { detectLineEnding, normalizeToLF, restoreLineEndings, stripBom } from "../normalize"; +import { readEditFileText, serializeEditFileText } from "../read-file"; import type { EditToolDetails, LspBatchRequest } from "../renderer"; // ─────────────────────────────────────────────────────────────────────────── @@ -1656,11 +1657,14 @@ export function splitHashlineInputs(input: string, options: SplitHashlineOptions // 13. Diff computation (for streaming preview) // ─────────────────────────────────────────────────────────────────────────── -async function readHashlineFileText(file: { text(): Promise }, pathText: string): Promise { +async function readHashlineFileText( + _file: { text(): Promise }, + absolutePath: string, + pathText: string, +): Promise { try { - return await file.text(); + return await readEditFileText(absolutePath, pathText); } catch (error) { - if (isEnoent(error)) throw new Error(`File not found: ${pathText}`); const message = error instanceof Error ? error.message : String(error); throw new Error(message || `Unable to read ${pathText}`); } @@ -1679,7 +1683,7 @@ export async function computeHashlineDiff( const [section] = sections; const absolutePath = resolveToCwd(section.path, cwd); - const rawContent = await readHashlineFileText(Bun.file(absolutePath), section.path); + const rawContent = await readHashlineFileText(Bun.file(absolutePath), absolutePath, section.path); const { text: content } = stripBom(rawContent); const normalized = normalizeToLF(content); const result = applyHashlineEdits(normalized, parseHashline(section.diff), options); @@ -1699,11 +1703,13 @@ interface ReadHashlineFileResult { rawContent: string; } -async function readHashlineFile(absolutePath: string): Promise { +async function readHashlineFile(absolutePath: string, pathText: string): Promise { try { - return { exists: true, rawContent: await Bun.file(absolutePath).text() }; + return { exists: true, rawContent: await readEditFileText(absolutePath, pathText) }; } catch (error) { if (isEnoent(error)) return { exists: false, rawContent: "" }; + if (error instanceof Error && error.message === `File not found: ${pathText}`) + return { exists: false, rawContent: "" }; throw error; } } @@ -1746,7 +1752,7 @@ async function preflightHashlineSection(options: ExecuteHashlineSingleOptions & const { edits } = parseHashlineWithWarnings(diff); enforcePlanModeWrite(session, sectionPath, { op: "update" }); - const source = await readHashlineFile(absolutePath); + const source = await readHashlineFile(absolutePath, sectionPath); if (!source.exists && hasAnchorScopedEdit(edits)) throw new Error(`File not found: ${sectionPath}`); if (source.exists) assertEditableFileContent(source.rawContent, sectionPath); @@ -1773,7 +1779,7 @@ async function executeHashlineSection( const { edits, warnings: parseWarnings } = parseHashlineWithWarnings(diff); enforcePlanModeWrite(session, sourcePath, { op: "update" }); - const source = await readHashlineFile(absolutePath); + const source = await readHashlineFile(absolutePath, sourcePath); if (!source.exists && hasAnchorScopedEdit(edits)) throw new Error(`File not found: ${sourcePath}`); if (source.exists) assertEditableFileContent(source.rawContent, sourcePath); @@ -1789,7 +1795,11 @@ async function executeHashlineSection( }; } - const finalContent = bom + restoreLineEndings(result.lines, originalEnding); + const finalContent = await serializeEditFileText( + absolutePath, + sourcePath, + bom + restoreLineEndings(result.lines, originalEnding), + ); const diagnostics = await writethrough( absolutePath, finalContent, diff --git a/packages/coding-agent/src/edit/modes/patch.ts b/packages/coding-agent/src/edit/modes/patch.ts index 937703a2e..1b6036714 100644 --- a/packages/coding-agent/src/edit/modes/patch.ts +++ b/packages/coding-agent/src/edit/modes/patch.ts @@ -44,6 +44,7 @@ import { restoreLineEndings, stripBom, } from "../normalize"; +import { readEditFileText, serializeEditFileText } from "../read-file"; import type { EditToolDetails, LspBatchRequest } from "../renderer"; import { type ContextLineResult, @@ -104,13 +105,13 @@ export const defaultFileSystem: FileSystem = { return Bun.file(path).exists(); }, async read(path: string): Promise { - return Bun.file(path).text(); + return readEditFileText(path, path); }, async readBinary(path: string): Promise { return fs.promises.readFile(path); }, async write(path: string, content: string): Promise { - await Bun.write(path, content); + await Bun.write(path, await serializeEditFileText(path, path, content)); }, async delete(path: string): Promise { await fs.promises.unlink(path); @@ -1637,7 +1638,7 @@ class LspFileSystem implements FileSystem { } async read(path: string): Promise { - return this.#getFile(path).text(); + return readEditFileText(path, path); } async readBinary(path: string): Promise { @@ -1647,10 +1648,11 @@ class LspFileSystem implements FileSystem { async write(path: string, content: string): Promise { const file = this.#getFile(path); + const finalContent = await serializeEditFileText(path, path, content); const deferredForPath = this.deferredForPath; const result = await this.writethrough( path, - content, + finalContent, this.signal, file, this.batchRequest, diff --git a/packages/coding-agent/src/edit/modes/replace.ts b/packages/coding-agent/src/edit/modes/replace.ts index 08cd75f9d..d0568cc82 100644 --- a/packages/coding-agent/src/edit/modes/replace.ts +++ b/packages/coding-agent/src/edit/modes/replace.ts @@ -21,7 +21,7 @@ import { restoreLineEndings, stripBom, } from "../normalize"; -import { readEditFileText } from "../read-file"; +import { readEditFileText, serializeEditFileText } from "../read-file"; import type { EditToolDetails, LspBatchRequest } from "../renderer"; export interface FuzzyMatch { @@ -1065,7 +1065,11 @@ export async function executeReplaceSingle( throw new Error(`Edits to ${path} resulted in no changes being made.`); } - const finalContent = bom + restoreLineEndings(result.content, originalEnding); + const finalContent = await serializeEditFileText( + absolutePath, + path, + bom + restoreLineEndings(result.content, originalEnding), + ); const diagnostics = await writethrough( absolutePath, finalContent, diff --git a/packages/coding-agent/src/edit/notebook.ts b/packages/coding-agent/src/edit/notebook.ts new file mode 100644 index 000000000..5383eef72 --- /dev/null +++ b/packages/coding-agent/src/edit/notebook.ts @@ -0,0 +1,222 @@ +import * as path from "node:path"; +import { isEnoent } from "@oh-my-pi/pi-utils"; + +export type NotebookCellType = "code" | "markdown" | "raw"; + +export interface NotebookCell { + cell_type: NotebookCellType; + source?: string | string[]; + metadata?: Record; + execution_count?: number | null; + outputs?: unknown[]; + [key: string]: unknown; +} + +export interface NotebookDocument { + cells: NotebookCell[]; + metadata: Record; + nbformat: number; + nbformat_minor: number; + [key: string]: unknown; +} + +const CELL_MARKER_RE = /^# %% \[(code|markdown|raw)\](?: cell:(\d+))?$/; + +export function isNotebookPath(filePath: string): boolean { + return path.extname(filePath).toLowerCase() === ".ipynb"; +} + +function isRecord(value: unknown): value is Record { + return typeof value === "object" && value !== null && !Array.isArray(value); +} + +function isCellType(value: unknown): value is NotebookCellType { + return value === "code" || value === "markdown" || value === "raw"; +} + +function sourceToText(source: string | string[] | undefined): string { + if (source === undefined) return ""; + if (typeof source === "string") return source; + return source.join(""); +} + +export function splitNotebookSource(content: string): string[] { + if (content.length === 0) return []; + return content.match(/[^\n]*\n|[^\n]+$/g) ?? []; +} + +function cloneCell(cell: NotebookCell): NotebookCell { + return structuredClone(cell); +} + +function createNotebookCell(cellType: NotebookCellType, source: string): NotebookCell { + const cell: NotebookCell = { + cell_type: cellType, + metadata: {}, + source: splitNotebookSource(source), + }; + if (cellType === "code") { + cell.execution_count = null; + cell.outputs = []; + } + return cell; +} + +function createEmptyNotebook(): NotebookDocument { + return { + cells: [], + metadata: {}, + nbformat: 4, + nbformat_minor: 5, + }; +} + +function validateNotebook(value: unknown, displayPath: string): NotebookDocument { + if (!isRecord(value)) { + throw new Error(`Invalid notebook structure (expected object): ${displayPath}`); + } + if (!Array.isArray(value.cells)) { + throw new Error(`Invalid notebook structure (missing cells array): ${displayPath}`); + } + for (let index = 0; index < value.cells.length; index++) { + const cell = value.cells[index]; + if (!isRecord(cell) || !isCellType(cell.cell_type)) { + throw new Error(`Invalid notebook cell ${index} in ${displayPath}`); + } + } + return value as unknown as NotebookDocument; +} + +export async function readNotebookDocument(absolutePath: string, displayPath: string): Promise { + try { + return validateNotebook(await Bun.file(absolutePath).json(), displayPath); + } catch (error) { + if (isEnoent(error)) throw new Error(`File not found: ${displayPath}`); + if (error instanceof SyntaxError) throw new Error(`Invalid JSON in notebook: ${displayPath}`); + throw error; + } +} + +export function notebookToEditableText(notebook: NotebookDocument): string { + return notebook.cells + .map((cell, index) => { + const source = sourceToText(cell.source); + return source.length > 0 + ? `# %% [${cell.cell_type}] cell:${index}\n${source}` + : `# %% [${cell.cell_type}] cell:${index}`; + }) + .join("\n"); +} + +interface ParsedVirtualCell { + cellType: NotebookCellType; + cellIndex?: number; + source: string; +} + +function parseVirtualCellMarker(line: string): { cellType: NotebookCellType; cellIndex?: number } | undefined { + const match = CELL_MARKER_RE.exec(line); + if (!match) return undefined; + const cellType = match[1] as NotebookCellType; + const cellIndexText = match[2]; + return { + cellType, + cellIndex: cellIndexText === undefined ? undefined : Number.parseInt(cellIndexText, 10), + }; +} + +function linesToSourceText(lines: string[]): string { + if (lines.length === 0) return ""; + return lines.join("\n"); +} + +function parseNotebookEditableText(text: string, displayPath: string): ParsedVirtualCell[] { + const lines = text.length === 0 ? [] : text.split("\n"); + const cells: ParsedVirtualCell[] = []; + let current: { cellType: NotebookCellType; cellIndex?: number; lines: string[] } | undefined; + + const flush = () => { + if (!current) return; + cells.push({ + cellType: current.cellType, + cellIndex: current.cellIndex, + source: linesToSourceText(current.lines), + }); + }; + + for (const line of lines) { + const marker = parseVirtualCellMarker(line); + if (marker) { + flush(); + current = { ...marker, lines: [] }; + continue; + } + if (!current) { + throw new Error( + `Invalid notebook editable representation for ${displayPath}: expected first line to be "# %% [code] cell:0", "# %% [markdown] cell:0", or "# %% [raw] cell:0".`, + ); + } + current.lines.push(line); + } + flush(); + return cells; +} + +export function applyNotebookEditableText( + notebook: NotebookDocument, + text: string, + displayPath: string, +): NotebookDocument { + const parsedCells = parseNotebookEditableText(text, displayPath); + const usedOriginalCells = new Set(); + const nextNotebook = structuredClone(notebook); + nextNotebook.cells = parsedCells.map(parsedCell => { + const originalIndex = parsedCell.cellIndex; + const originalCell = + originalIndex !== undefined && + originalIndex >= 0 && + originalIndex < notebook.cells.length && + !usedOriginalCells.has(originalIndex) + ? notebook.cells[originalIndex] + : undefined; + if (originalCell) { + usedOriginalCells.add(originalIndex!); + const cell = cloneCell(originalCell); + cell.cell_type = parsedCell.cellType; + cell.source = splitNotebookSource(parsedCell.source); + if (parsedCell.cellType === "code") { + cell.execution_count ??= null; + cell.outputs ??= []; + } else { + delete cell.execution_count; + delete cell.outputs; + } + return cell; + } + return createNotebookCell(parsedCell.cellType, parsedCell.source); + }); + return nextNotebook; +} + +export async function readEditableNotebookText(absolutePath: string, displayPath: string): Promise { + return notebookToEditableText(await readNotebookDocument(absolutePath, displayPath)); +} + +export async function serializeEditedNotebookText( + absolutePath: string, + displayPath: string, + text: string, +): Promise { + let notebook: NotebookDocument; + try { + notebook = await readNotebookDocument(absolutePath, displayPath); + } catch (error) { + if (error instanceof Error && error.message === `File not found: ${displayPath}`) { + notebook = createEmptyNotebook(); + } else { + throw error; + } + } + const nextNotebook = applyNotebookEditableText(notebook, text, displayPath); + return JSON.stringify(nextNotebook, null, 1); +} diff --git a/packages/coding-agent/src/edit/read-file.ts b/packages/coding-agent/src/edit/read-file.ts index 29b56e23f..d31dd36cc 100644 --- a/packages/coding-agent/src/edit/read-file.ts +++ b/packages/coding-agent/src/edit/read-file.ts @@ -5,9 +5,11 @@ * error referencing the display path. */ import { isEnoent } from "@oh-my-pi/pi-utils"; +import { isNotebookPath, readEditableNotebookText, serializeEditedNotebookText } from "./notebook"; export async function readEditFileText(absolutePath: string, path: string): Promise { try { + if (isNotebookPath(absolutePath)) return await readEditableNotebookText(absolutePath, path); return await Bun.file(absolutePath).text(); } catch (error) { if (isEnoent(error)) { @@ -16,3 +18,8 @@ export async function readEditFileText(absolutePath: string, path: string): Prom throw error; } } + +export async function serializeEditFileText(absolutePath: string, path: string, content: string): Promise { + if (isNotebookPath(absolutePath)) return serializeEditedNotebookText(absolutePath, path, content); + return content; +} diff --git a/packages/coding-agent/src/prompts/tools/read.md b/packages/coding-agent/src/prompts/tools/read.md index d40e51355..636ff1561 100644 --- a/packages/coding-agent/src/prompts/tools/read.md +++ b/packages/coding-agent/src/prompts/tools/read.md @@ -33,7 +33,7 @@ The `read` tool is multi-purpose and more capable than it looks — inspects fil # Inspection -Extracts text from PDF, Word, PowerPoint, Excel, RTF, EPUB, and Jupyter notebook files. Can inspect images. +Extracts text from PDF, Word, PowerPoint, Excel, RTF, EPUB, and Jupyter notebook files. Notebooks are shown as editable `# %% [type] cell:N` text; edits to that text are applied back to the underlying `.ipynb` JSON while preserving notebook metadata where possible. Can inspect images. # Directories & Archives diff --git a/packages/coding-agent/src/tools/index.ts b/packages/coding-agent/src/tools/index.ts index b8ab8e5e5..59e2b62f7 100644 --- a/packages/coding-agent/src/tools/index.ts +++ b/packages/coding-agent/src/tools/index.ts @@ -37,7 +37,6 @@ import { HindsightRetainTool } from "./hindsight-retain"; import { InspectImageTool } from "./inspect-image"; import { IrcTool } from "./irc"; import { JobTool } from "./job"; -import { NotebookTool } from "./notebook"; import { wrapToolWithMetaNotice } from "./output-meta"; import { ReadTool } from "./read"; import { RecipeTool } from "./recipe"; @@ -80,7 +79,6 @@ export * from "./image-gen"; export * from "./inspect-image"; export * from "./irc"; export * from "./job"; -export * from "./notebook"; export * from "./read"; export * from "./recipe"; export * from "./render-mermaid"; @@ -271,7 +269,6 @@ export const BUILTIN_TOOLS: Record = { find: s => new FindTool(s), search: s => new SearchTool(s), lsp: LspTool.createIf, - notebook: s => new NotebookTool(s), inspect_image: s => new InspectImageTool(s), browser: s => new BrowserTool(s), checkpoint: CheckpointTool.createIf, @@ -430,7 +427,6 @@ export async function createTools(session: ToolSession, toolNames?: string[]): P if (name === "ast_grep") return session.settings.get("astGrep.enabled"); if (name === "ast_edit") return session.settings.get("astEdit.enabled"); if (name === "render_mermaid") return session.settings.get("renderMermaid.enabled"); - if (name === "notebook") return session.settings.get("notebook.enabled"); if (name === "inspect_image") return session.settings.get("inspect_image.enabled"); if (name === "web_search") return session.settings.get("web_search.enabled"); // search_tool_bm25 is allowed when either legacy mcp.discoveryMode or new tools.discoveryMode is active. diff --git a/packages/coding-agent/src/tools/notebook.ts b/packages/coding-agent/src/tools/notebook.ts deleted file mode 100644 index 77be1ef00..000000000 --- a/packages/coding-agent/src/tools/notebook.ts +++ /dev/null @@ -1,286 +0,0 @@ -import type { AgentTool, AgentToolContext, AgentToolResult, AgentToolUpdateCallback } from "@oh-my-pi/pi-agent-core"; -import { StringEnum } from "@oh-my-pi/pi-ai"; -import type { Component } from "@oh-my-pi/pi-tui"; -import { Text } from "@oh-my-pi/pi-tui"; -import { isEnoent, untilAborted } from "@oh-my-pi/pi-utils"; -import { type Static, Type } from "@sinclair/typebox"; -import type { RenderResultOptions } from "../extensibility/custom-tools/types"; -import type { Theme } from "../modes/theme/theme"; -import type { ToolSession } from "../sdk"; -import { Hasher, type RenderCache, renderCodeCell, renderStatusLine } from "../tui"; -import { resolveToCwd } from "./path-utils"; -import { formatCount, formatErrorMessage, PREVIEW_LIMITS } from "./render-utils"; - -const notebookSchema = Type.Object({ - action: StringEnum(["edit", "insert", "delete"], { - description: "cell action", - examples: ["edit", "insert", "delete"], - }), - notebook_path: Type.String({ description: "notebook path", examples: ["analysis.ipynb"] }), - cell_index: Type.Number({ description: "cell index", examples: [0, 1] }), - content: Type.Optional(Type.String({ description: "new cell content" })), - cell_type: Type.Optional( - StringEnum(["code", "markdown"], { - description: "cell type", - examples: ["code", "markdown"], - }), - ), -}); - -export interface NotebookToolDetails { - /** Action performed */ - action: "edit" | "insert" | "delete"; - /** Cell index operated on */ - cellIndex: number; - /** Cell type */ - cellType?: string; - /** Total cell count after operation */ - totalCells: number; - /** Cell content lines after operation (or removed content for delete) */ - cellSource?: string[]; -} - -interface NotebookCell { - cell_type: "code" | "markdown" | "raw"; - source: string[]; - metadata: Record; - execution_count?: number | null; - outputs?: unknown[]; -} - -interface Notebook { - cells: NotebookCell[]; - metadata: Record; - nbformat: number; - nbformat_minor: number; -} - -function splitIntoLines(content: string): string[] { - return content.split("\n").map((line, i, arr) => (i < arr.length - 1 ? `${line}\n` : line)); -} - -type NotebookParams = Static; - -export class NotebookTool implements AgentTool { - readonly name = "notebook"; - readonly label = "Notebook"; - readonly loadMode = "discoverable"; - readonly summary = "Read and execute Jupyter notebooks"; - readonly description = "Edit, insert, or delete cells in Jupyter notebooks (.ipynb). cell_index is 0-based."; - readonly parameters = notebookSchema; - readonly strict = true; - readonly concurrency = "exclusive"; - - constructor(private readonly session: ToolSession) {} - - async execute( - _toolCallId: string, - params: NotebookParams, - signal?: AbortSignal, - _onUpdate?: AgentToolUpdateCallback, - _context?: AgentToolContext, - ): Promise> { - const { action, notebook_path, cell_index, content, cell_type } = params; - const absolutePath = resolveToCwd(notebook_path, this.session.cwd); - - return untilAborted(signal, async () => { - // Read and parse notebook - let notebook: Notebook; - try { - notebook = await Bun.file(absolutePath).json(); - } catch (err) { - if (isEnoent(err)) throw new Error(`Notebook not found: ${notebook_path}`); - throw new Error(`Invalid JSON in notebook: ${notebook_path}`); - } - - // Validate notebook structure - if (!notebook.cells || !Array.isArray(notebook.cells)) { - throw new Error(`Invalid notebook structure (missing cells array): ${notebook_path}`); - } - - const cellCount = notebook.cells.length; - - // Validate cell_index based on action - if (action === "insert") { - if (cell_index < 0 || cell_index > cellCount) { - throw new Error(`Cell index ${cell_index} out of range for insert (0-${cellCount}) in ${notebook_path}`); - } - } else { - if (cell_index < 0 || cell_index >= cellCount) { - throw new Error(`Cell index ${cell_index} out of range (0-${cellCount - 1}) in ${notebook_path}`); - } - } - - // Validate content for edit/insert - if ((action === "edit" || action === "insert") && content === undefined) { - throw new Error(`Content is required for ${action} action`); - } - - // Perform the action - let resultMessage: string; - let finalCellType: string | undefined; - let cellSource: string[] | undefined; - - switch (action) { - case "edit": { - const sourceLines = splitIntoLines(content!); - notebook.cells[cell_index].source = sourceLines; - finalCellType = notebook.cells[cell_index].cell_type; - cellSource = sourceLines; - resultMessage = `Replaced cell ${cell_index} (${finalCellType})`; - break; - } - case "insert": { - const sourceLines = splitIntoLines(content!); - const newCellType = (cell_type as "code" | "markdown") || "code"; - const newCell: NotebookCell = { - cell_type: newCellType, - source: sourceLines, - metadata: {}, - }; - if (newCellType === "code") { - newCell.execution_count = null; - newCell.outputs = []; - } - notebook.cells.splice(cell_index, 0, newCell); - finalCellType = newCellType; - cellSource = sourceLines; - resultMessage = `Inserted ${newCellType} cell at position ${cell_index}`; - break; - } - case "delete": { - const removedCell = notebook.cells[cell_index]; - finalCellType = removedCell.cell_type; - cellSource = removedCell.source; - notebook.cells.splice(cell_index, 1); - resultMessage = `Deleted cell ${cell_index} (${finalCellType})`; - break; - } - default: { - throw new Error(`Invalid action: ${action}`); - } - } - - // Write back with single-space indentation - await Bun.write(absolutePath, JSON.stringify(notebook, null, 1)); - - const newCellCount = notebook.cells.length; - return { - content: [ - { - type: "text", - text: `${resultMessage}. Notebook now has ${newCellCount} cells.`, - }, - ], - details: { - action: action as "edit" | "insert" | "delete", - cellIndex: cell_index, - cellType: finalCellType, - totalCells: newCellCount, - cellSource, - }, - }; - }); - } -} - -// ============================================================================= -// TUI Renderer -// ============================================================================= - -interface NotebookRenderArgs { - action: string; - notebookPath?: string; - notebook_path?: string; - cellNumber?: number; - cell_index?: number; - cellType?: string; - cell_type?: string; - content?: string; -} - -const COLLAPSED_TEXT_LIMIT = PREVIEW_LIMITS.COLLAPSED_LINES * 2; - -export const notebookToolRenderer = { - renderCall(args: NotebookRenderArgs, _options: RenderResultOptions, uiTheme: Theme): Component { - const meta: string[] = []; - const notebookPath = args.notebookPath ?? args.notebook_path; - const cellNumber = args.cellNumber ?? args.cell_index; - const cellType = args.cellType ?? args.cell_type; - meta.push(`in ${notebookPath || "?"}`); - if (cellNumber !== undefined) meta.push(`cell:${cellNumber}`); - if (cellType) meta.push(`type:${cellType}`); - - const text = renderStatusLine( - { icon: "pending", title: "Notebook", description: args.action || "?", meta }, - uiTheme, - ); - return new Text(text, 0, 0); - }, - - renderResult( - result: { content: Array<{ type: string; text?: string }>; details?: NotebookToolDetails }, - options: RenderResultOptions, - uiTheme: Theme, - args?: NotebookRenderArgs, - ): Component { - const content = result.content?.[0]; - if (content?.type === "text" && content.text?.startsWith("Error:")) { - const notebookPath = args?.notebookPath ?? args?.notebook_path ?? "?"; - const header = renderStatusLine({ icon: "error", title: "Notebook", description: notebookPath }, uiTheme); - return new Text([header, formatErrorMessage(content.text, uiTheme)].join("\n"), 0, 0); - } - - const details = result.details; - const action = details?.action ?? "edit"; - const cellIndex = details?.cellIndex; - const cellType = details?.cellType; - const totalCells = details?.totalCells; - const cellSource = details?.cellSource ?? []; - const lineCount = cellSource.length; - - const actionLabel = action === "insert" ? "Inserted" : action === "delete" ? "Deleted" : "Edited"; - const cellLabel = cellType || "cell"; - const summaryParts = [`${actionLabel} ${cellLabel} ${cellIndex ?? "?"}`]; - if (lineCount > 0) summaryParts.push(formatCount("line", lineCount)); - if (totalCells !== undefined) summaryParts.push(`${totalCells} total`); - - const outputLines = summaryParts.map(part => uiTheme.fg("dim", part)); - const codeText = cellSource.join(""); - const language = cellType === "markdown" ? "markdown" : undefined; - - const notebookPath = args?.notebookPath ?? args?.notebook_path; - const notebookLabel = notebookPath ? `${actionLabel} ${notebookPath}` : "Notebook"; - let cached: RenderCache | undefined; - - return { - render: (width: number): string[] => { - // REACTIVE: read mutable options at render time - const { expanded } = options; - const key = new Hasher().bool(expanded).u32(width).digest(); - if (cached?.key === key) return cached.lines; - - const lines = renderCodeCell( - { - code: codeText, - language, - title: notebookLabel, - status: "complete", - output: outputLines.join("\n"), - codeMaxLines: expanded ? Number.POSITIVE_INFINITY : COLLAPSED_TEXT_LIMIT, - expanded, - width, - }, - uiTheme, - ); - - cached = { key, lines }; - return lines; - }, - invalidate: () => { - cached = undefined; - }, - }; - }, - mergeCallAndResult: true, -}; diff --git a/packages/coding-agent/src/tools/read.ts b/packages/coding-agent/src/tools/read.ts index 0fe9cad79..bf7f57b67 100644 --- a/packages/coding-agent/src/tools/read.ts +++ b/packages/coding-agent/src/tools/read.ts @@ -9,6 +9,7 @@ import { Text } from "@oh-my-pi/pi-tui"; import { getRemoteDir, prompt, readImageMetadata, untilAborted } from "@oh-my-pi/pi-utils"; import { type Static, Type } from "@sinclair/typebox"; import { formatHashLine, formatHashLines, formatLineHash, HL_BODY_SEP } from "../edit/line-hash"; +import { isNotebookPath, readEditableNotebookText } from "../edit/notebook"; import type { RenderResultOptions } from "../extensibility/custom-tools/types"; import { parseInternalUrl } from "../internal-urls/parse"; import type { InternalUrl } from "../internal-urls/types"; @@ -1196,7 +1197,7 @@ export class ReadTool implements AgentTool { const ext = path.extname(absolutePath).toLowerCase(); const _hasEditTool = this.session.hasEditTool ?? true; const _language = getLanguageFromPath(absolutePath); - const shouldConvertWithMarkit = CONVERTIBLE_EXTENSIONS.has(ext) || (ext === ".ipynb" && parsed.kind === "raw"); + const shouldConvertWithMarkit = CONVERTIBLE_EXTENSIONS.has(ext); // Read the file based on type let content: Array | undefined; let details: ReadToolDetails = {}; @@ -1263,8 +1264,20 @@ export class ReadTool implements AgentTool { throw error; } } + } else if (isNotebookPath(absolutePath) && parsed.kind !== "raw") { + const { offset, limit } = selToOffsetLimit(parsed); + return this.#buildInMemoryTextResult( + await readEditableNotebookText(absolutePath, localReadPath), + offset, + limit, + { + details: { resolvedPath: absolutePath }, + sourcePath: absolutePath, + entityLabel: "notebook", + }, + ); } else if (shouldConvertWithMarkit) { - // Convert document or notebook via markit. + // Convert document via markit. const result = await convertFileWithMarkit(absolutePath, signal); if (result.ok) { // Apply truncation to converted content diff --git a/packages/coding-agent/src/tools/renderers.ts b/packages/coding-agent/src/tools/renderers.ts index 22cd8543a..62b0cd883 100644 --- a/packages/coding-agent/src/tools/renderers.ts +++ b/packages/coding-agent/src/tools/renderers.ts @@ -22,7 +22,6 @@ import { findToolRenderer } from "./find"; import { githubToolRenderer } from "./gh-renderer"; import { inspectImageToolRenderer } from "./inspect-image-renderer"; import { jobToolRenderer } from "./job"; -import { notebookToolRenderer } from "./notebook"; import { readToolRenderer } from "./read"; import { recipeToolRenderer } from "./recipe/render"; import { resolveToolRenderer } from "./resolve"; @@ -60,7 +59,6 @@ export const toolRenderers: Record = { find: findToolRenderer as ToolRenderer, search: searchToolRenderer as ToolRenderer, lsp: lspToolRenderer as ToolRenderer, - notebook: notebookToolRenderer as ToolRenderer, inspect_image: inspectImageToolRenderer as ToolRenderer, read: readToolRenderer as ToolRenderer, job: jobToolRenderer as ToolRenderer, diff --git a/packages/coding-agent/test/tool-discovery/initial-tools.test.ts b/packages/coding-agent/test/tool-discovery/initial-tools.test.ts index c1af11c88..50fbca9ea 100644 --- a/packages/coding-agent/test/tool-discovery/initial-tools.test.ts +++ b/packages/coding-agent/test/tool-discovery/initial-tools.test.ts @@ -22,7 +22,6 @@ const allToolsSettings = Settings.isolated({ "search.enabled": true, "github.enabled": true, "lsp.enabled": true, - "notebook.enabled": true, "inspect_image.enabled": true, "web_search.enabled": true, "calc.enabled": true, @@ -99,7 +98,6 @@ describe("built-in tool loadMode annotations", () => { "find", "search", "lsp", - "notebook", "inspect_image", "browser", "checkpoint", diff --git a/packages/coding-agent/test/tools.test.ts b/packages/coding-agent/test/tools.test.ts index cb814304b..b72931be9 100644 --- a/packages/coding-agent/test/tools.test.ts +++ b/packages/coding-agent/test/tools.test.ts @@ -17,7 +17,6 @@ import { wrapToolWithMetaNotice } from "@oh-my-pi/pi-coding-agent/tools/output-m import { ReadTool } from "@oh-my-pi/pi-coding-agent/tools/read"; import { DEFAULT_MATCH_LIMIT, SearchTool } from "@oh-my-pi/pi-coding-agent/tools/search"; import { WriteTool } from "@oh-my-pi/pi-coding-agent/tools/write"; -import * as markitUtils from "@oh-my-pi/pi-coding-agent/utils/markit"; import { $which, Snowflake } from "@oh-my-pi/pi-utils"; import { unzipSync } from "fflate"; @@ -280,15 +279,22 @@ describe("Coding Agent Tools", () => { expect(result.details?.truncation).toBeUndefined(); }); - it("should convert ipynb files through markit for raw reads", async () => { + it("should read ipynb files as editable cell text", async () => { const notebookPath = path.join(testDir, "notebook.ipynb"); const notebook = { cells: [ { cell_type: "markdown", - metadata: {}, + metadata: { keep: true }, source: ["# Notebook Title\n", "\n", "Notebook body\n"], }, + { + cell_type: "code", + metadata: {}, + source: ["print('hello')\n"], + execution_count: 7, + outputs: [{ output_type: "stream", name: "stdout", text: "hello\n" }], + }, ], metadata: {}, nbformat: 4, @@ -296,20 +302,55 @@ describe("Coding Agent Tools", () => { }; fs.writeFileSync(notebookPath, JSON.stringify(notebook)); - const convertSpy = vi.spyOn(markitUtils, "convertFileWithMarkit").mockResolvedValue({ - ok: true, - content: "# Notebook Title\n\nNotebook body\n", - }); - - const result = await readTool.execute("test-call-ipynb", { path: `${notebookPath}:raw` }); + const result = await readTool.execute("test-call-ipynb", { path: notebookPath }); const output = getTextOutput(result); - expect(convertSpy).toHaveBeenCalledTimes(1); + expect(output).toContain("# %% [markdown] cell:0"); expect(output).toContain("# Notebook Title"); expect(output).toContain("Notebook body"); + expect(output).toContain("# %% [code] cell:1"); + expect(output).toContain("print('hello')"); expect(output).not.toContain('"cell_type"'); }); + it("should apply edits to the ipynb editable representation", async () => { + const notebookPath = path.join(testDir, "editable.ipynb"); + const notebook = { + cells: [ + { + cell_type: "markdown", + metadata: { keep: true }, + source: ["Original title\n"], + }, + { + cell_type: "code", + metadata: { trusted: true }, + source: ["print('old')\n"], + execution_count: 3, + outputs: [{ output_type: "stream", name: "stdout", text: "old\n" }], + }, + ], + metadata: { kernelspec: { name: "python3" } }, + nbformat: 4, + nbformat_minor: 5, + }; + fs.writeFileSync(notebookPath, JSON.stringify(notebook)); + const noLspEditTool = wrapToolWithMetaNotice(new EditTool({ ...session, enableLsp: false })); + + await noLspEditTool.execute("test-edit-ipynb", { + path: notebookPath, + edits: [{ old_text: "print('old')", new_text: "print('new')" }], + }); + + const updated = JSON.parse(fs.readFileSync(notebookPath, "utf-8")); + expect(updated.cells).toHaveLength(2); + expect(updated.cells[1].source).toEqual(["print('new')\n"]); + expect(updated.cells[1].metadata).toEqual({ trusted: true }); + expect(updated.cells[1].execution_count).toBe(3); + expect(updated.cells[1].outputs).toEqual([{ output_type: "stream", name: "stdout", text: "old\n" }]); + expect(updated.metadata).toEqual({ kernelspec: { name: "python3" } }); + }); + it("should handle non-existent files", async () => { const testFile = path.join(testDir, "nonexistent.txt"); diff --git a/packages/coding-agent/test/tools/index.test.ts b/packages/coding-agent/test/tools/index.test.ts index d64450051..ce295ed08 100644 --- a/packages/coding-agent/test/tools/index.test.ts +++ b/packages/coding-agent/test/tools/index.test.ts @@ -61,7 +61,6 @@ describe("createTools", () => { expect(names).toContain("search"); expect(names).toContain("find"); expect(names).toContain("lsp"); - expect(names).toContain("notebook"); expect(names).toContain("task"); expect(names).toContain("todo_write"); expect(names).toContain("web_search"); @@ -212,7 +211,6 @@ describe("createTools", () => { "astEdit.enabled": false, "renderMermaid.enabled": false, "web_search.enabled": false, - "notebook.enabled": false, "browser.enabled": false, "inspect_image.enabled": false, "calc.enabled": false, @@ -227,7 +225,6 @@ describe("createTools", () => { expect(names).not.toContain("ast_edit"); expect(names).not.toContain("render_mermaid"); expect(names).not.toContain("web_search"); - expect(names).not.toContain("notebook"); expect(names).not.toContain("browser"); expect(names).not.toContain("inspect_image"); expect(names).not.toContain("calc");