feat: notebook_edit tool — cell-aware Jupyter (.ipynb) editing
Ported from origin/v0.6.0, adapted to resolveWithinCwd (no setLastEdit). - src/tools/notebookEdit.ts: replace/insert/delete a cell by cell_id or cell_index. Converts the model's single-string new_source to/from nbformat's source line-array (trailing-newline convention), so the model never hand-writes the array quirk. Switching a code cell to markdown drops execution_count/outputs; switching to code adds them. Atomic temp+rename write, JSON validated. - tools/index.ts: register notebook_edit. - notebookEdit.test.ts: 9 cases (replace by id/index, insert at position/append, delete, missing-id failure, insert-without-cell_type, replace-without- new_source, diff preview, code→markdown field drop, path containment). Prefer this over edit_file/write_file for .ipynb so the JSON structure stays valid (and avoids regenerating a whole notebook JSON in the model's output — a context-window win for local models). Verified: typecheck clean, build 266.91 KB, 284 tests pass (+9).
This commit is contained in:
@@ -8,6 +8,7 @@ import { multiEditTool } from "./multiEdit.js";
|
||||
import { gitCommitTool, gitStatusTool } from "./git.js";
|
||||
import { grepTool } from "./grep.js";
|
||||
import { listFilesTool } from "./listFiles.js";
|
||||
import { notebookEditTool } from "./notebookEdit.js";
|
||||
import { readFileTool } from "./readFile.js";
|
||||
import { todoWriteTool } from "./todoWrite.js";
|
||||
import { webFetchTool } from "./webFetch.js";
|
||||
@@ -28,6 +29,7 @@ export const TOOLS: ToolDef[] = [
|
||||
writeFileTool,
|
||||
editFileTool,
|
||||
multiEditTool,
|
||||
notebookEditTool,
|
||||
bashTool,
|
||||
bashOutputTool,
|
||||
bashKillTool,
|
||||
|
||||
@@ -0,0 +1,186 @@
|
||||
import { mkdtemp, readFile, rm, writeFile } from "node:fs/promises";
|
||||
import { tmpdir } from "node:os";
|
||||
import path from "node:path";
|
||||
import { describe, expect, it } from "vitest";
|
||||
import { notebookEditTool } from "./notebookEdit.js";
|
||||
|
||||
/** A minimal valid nbformat 4 notebook with two code cells. */
|
||||
function minimalNotebook(): string {
|
||||
return (
|
||||
JSON.stringify(
|
||||
{
|
||||
nbformat: 4,
|
||||
nbformat_minor: 5,
|
||||
metadata: {},
|
||||
cells: [
|
||||
{ cell_type: "code", id: "c1", source: ["print('a')\n"], metadata: {}, outputs: [], execution_count: null },
|
||||
{ cell_type: "code", id: "c2", source: ["print('b')\n"], metadata: {}, outputs: [], execution_count: null },
|
||||
],
|
||||
},
|
||||
null,
|
||||
2,
|
||||
) + "\n"
|
||||
);
|
||||
}
|
||||
|
||||
async function makeCwd(): Promise<string> {
|
||||
return mkdtemp(path.join(tmpdir(), "locode-notebook-"));
|
||||
}
|
||||
|
||||
function parseCells(content: string): { cell_type: string; id?: string; source: string[] }[] {
|
||||
return (JSON.parse(content) as { cells: { cell_type: string; id?: string; source: string[] }[] }).cells;
|
||||
}
|
||||
|
||||
describe("notebookEditTool", () => {
|
||||
it("rejects paths that escape the working directory", async () => {
|
||||
const cwd = await makeCwd();
|
||||
try {
|
||||
await expect(
|
||||
notebookEditTool.handler({ notebook_path: "../outside.ipynb", edit_mode: "delete" }, { cwd }),
|
||||
).rejects.toThrow(/outside the working directory/);
|
||||
} finally {
|
||||
await rm(cwd, { recursive: true, force: true });
|
||||
}
|
||||
});
|
||||
|
||||
it("replaces a cell source by cell_index", async () => {
|
||||
const cwd = await makeCwd();
|
||||
try {
|
||||
const file = path.join(cwd, "nb.ipynb");
|
||||
await writeFile(file, minimalNotebook());
|
||||
await notebookEditTool.handler(
|
||||
{ notebook_path: "nb.ipynb", cell_index: 0, edit_mode: "replace", new_source: "print('A')\n" },
|
||||
{ cwd },
|
||||
);
|
||||
const cells = parseCells(await readFile(file, "utf-8"));
|
||||
expect(cells[0]!.source).toEqual(["print('A')\n"]);
|
||||
expect(cells).toHaveLength(2);
|
||||
} finally {
|
||||
await rm(cwd, { recursive: true, force: true });
|
||||
}
|
||||
});
|
||||
|
||||
it("replaces a cell source by cell_id", async () => {
|
||||
const cwd = await makeCwd();
|
||||
try {
|
||||
const file = path.join(cwd, "nb.ipynb");
|
||||
await writeFile(file, minimalNotebook());
|
||||
await notebookEditTool.handler(
|
||||
{ notebook_path: "nb.ipynb", cell_id: "c2", edit_mode: "replace", new_source: "print('B2')" },
|
||||
{ cwd },
|
||||
);
|
||||
const cells = parseCells(await readFile(file, "utf-8"));
|
||||
expect(cells[1]!.source).toEqual(["print('B2')"]);
|
||||
} finally {
|
||||
await rm(cwd, { recursive: true, force: true });
|
||||
}
|
||||
});
|
||||
|
||||
it("inserts a new markdown cell at a position, shifting later cells down", async () => {
|
||||
const cwd = await makeCwd();
|
||||
try {
|
||||
const file = path.join(cwd, "nb.ipynb");
|
||||
await writeFile(file, minimalNotebook());
|
||||
await notebookEditTool.handler(
|
||||
{ notebook_path: "nb.ipynb", edit_mode: "insert", cell_index: 1, cell_type: "markdown", new_source: "# heading\n\ntext" },
|
||||
{ cwd },
|
||||
);
|
||||
const cells = parseCells(await readFile(file, "utf-8"));
|
||||
expect(cells).toHaveLength(3);
|
||||
expect(cells[1]!.cell_type).toBe("markdown");
|
||||
expect(cells[1]!.source).toEqual(["# heading\n", "\n", "text"]);
|
||||
// Original second cell (c2) is now at index 2.
|
||||
expect(cells[2]!.id).toBe("c2");
|
||||
} finally {
|
||||
await rm(cwd, { recursive: true, force: true });
|
||||
}
|
||||
});
|
||||
|
||||
it("appends a cell when insert omits cell_index", async () => {
|
||||
const cwd = await makeCwd();
|
||||
try {
|
||||
const file = path.join(cwd, "nb.ipynb");
|
||||
await writeFile(file, minimalNotebook());
|
||||
await notebookEditTool.handler(
|
||||
{ notebook_path: "nb.ipynb", edit_mode: "insert", cell_type: "code", new_source: "x = 1" },
|
||||
{ cwd },
|
||||
);
|
||||
const cells = parseCells(await readFile(file, "utf-8"));
|
||||
expect(cells).toHaveLength(3);
|
||||
expect(cells[2]!.source).toEqual(["x = 1"]);
|
||||
expect(cells[2]!.cell_type).toBe("code");
|
||||
} finally {
|
||||
await rm(cwd, { recursive: true, force: true });
|
||||
}
|
||||
});
|
||||
|
||||
it("deletes a cell by cell_id, and a missing id fails", async () => {
|
||||
const cwd = await makeCwd();
|
||||
try {
|
||||
const file = path.join(cwd, "nb.ipynb");
|
||||
await writeFile(file, minimalNotebook());
|
||||
await notebookEditTool.handler({ notebook_path: "nb.ipynb", cell_id: "c1", edit_mode: "delete" }, { cwd });
|
||||
const cells = parseCells(await readFile(file, "utf-8"));
|
||||
expect(cells).toHaveLength(1);
|
||||
expect(cells[0]!.id).toBe("c2");
|
||||
|
||||
// A missing id fails.
|
||||
await expect(
|
||||
notebookEditTool.handler({ notebook_path: "nb.ipynb", cell_id: "nope", edit_mode: "delete" }, { cwd }),
|
||||
).rejects.toThrow(/not found/);
|
||||
} finally {
|
||||
await rm(cwd, { recursive: true, force: true });
|
||||
}
|
||||
});
|
||||
|
||||
it("rejects insert without cell_type, and replace without new_source", async () => {
|
||||
const cwd = await makeCwd();
|
||||
try {
|
||||
const file = path.join(cwd, "nb.ipynb");
|
||||
await writeFile(file, minimalNotebook());
|
||||
await expect(
|
||||
notebookEditTool.handler({ notebook_path: "nb.ipynb", edit_mode: "insert", new_source: "x" }, { cwd }),
|
||||
).rejects.toThrow(/cell_type/);
|
||||
await expect(
|
||||
notebookEditTool.handler({ notebook_path: "nb.ipynb", cell_index: 0, edit_mode: "replace" }, { cwd }),
|
||||
).rejects.toThrow(/new_source/);
|
||||
} finally {
|
||||
await rm(cwd, { recursive: true, force: true });
|
||||
}
|
||||
});
|
||||
|
||||
it("produces a diff preview for an edit", async () => {
|
||||
const cwd = await makeCwd();
|
||||
try {
|
||||
const file = path.join(cwd, "nb.ipynb");
|
||||
await writeFile(file, minimalNotebook());
|
||||
const preview = await notebookEditTool.preview!(
|
||||
{ notebook_path: "nb.ipynb", cell_index: 0, edit_mode: "replace", new_source: "print('A')\n" },
|
||||
{ cwd },
|
||||
);
|
||||
expect(preview).toContain("@@");
|
||||
expect(preview).toContain("print('A')");
|
||||
expect(preview).toContain("print('a')");
|
||||
} finally {
|
||||
await rm(cwd, { recursive: true, force: true });
|
||||
}
|
||||
});
|
||||
|
||||
it("switching a code cell to markdown drops code-only fields", async () => {
|
||||
const cwd = await makeCwd();
|
||||
try {
|
||||
const file = path.join(cwd, "nb.ipynb");
|
||||
await writeFile(file, minimalNotebook());
|
||||
await notebookEditTool.handler(
|
||||
{ notebook_path: "nb.ipynb", cell_index: 0, edit_mode: "replace", cell_type: "markdown", new_source: "prose" },
|
||||
{ cwd },
|
||||
);
|
||||
const cell = parseCells(await readFile(file, "utf-8"))[0]!;
|
||||
expect(cell.cell_type).toBe("markdown");
|
||||
expect(cell).not.toHaveProperty("execution_count");
|
||||
expect(cell).not.toHaveProperty("outputs");
|
||||
} finally {
|
||||
await rm(cwd, { recursive: true, force: true });
|
||||
}
|
||||
});
|
||||
});
|
||||
@@ -0,0 +1,160 @@
|
||||
import { createPatch } from "diff";
|
||||
import { randomBytes } from "node:crypto";
|
||||
import { readFile as fsReadFile, rename as fsRename, unlink as fsUnlink, writeFile as fsWriteFile } from "node:fs/promises";
|
||||
import { z } from "zod";
|
||||
import { resolveWithinCwd } from "./pathGuard.js";
|
||||
import type { ToolDef } from "./types.js";
|
||||
|
||||
// .ipynb is a JSON document (nbformat 4): { nbformat, nbformat_minor, metadata, cells: Cell[] }.
|
||||
// Each cell is { cell_type: "code"|"markdown"|"raw", id?, source, metadata, outputs?, execution_count? }.
|
||||
// `source` is a list of strings where every line except the last carries a trailing "\n" (nbformat
|
||||
// convention). We convert the model's single-string new_source to/from that array form, so the
|
||||
// model never has to hand-write nbformat's line-array quirk — it just gives the full cell text.
|
||||
|
||||
type Notebook = { nbformat: number; nbformat_minor: number; metadata: Record<string, unknown>; cells: Cell[] };
|
||||
type Cell = { cell_type: string; id?: string; source: string[]; metadata: Record<string, unknown>; outputs?: unknown[]; execution_count?: unknown };
|
||||
|
||||
const schema = z.object({
|
||||
notebook_path: z.string().describe("Path to the .ipynb notebook to edit, relative to the working directory or absolute."),
|
||||
cell_id: z.string().optional().describe("The id of the cell to replace or delete. Ignored for insert."),
|
||||
cell_index: z
|
||||
.number()
|
||||
.int()
|
||||
.optional()
|
||||
.describe("0-based index of the cell to replace or delete; for insert, the position to insert at (defaults to append)."),
|
||||
cell_type: z.enum(["code", "markdown", "raw"]).optional().describe("Required for insert. For replace, overrides the existing cell's type if given."),
|
||||
edit_mode: z.enum(["replace", "insert", "delete"]).default("replace").describe("Whether to replace a cell, insert a new one, or delete."),
|
||||
new_source: z.string().optional().describe("The new cell source as a single string (required for replace and insert)."),
|
||||
});
|
||||
|
||||
/** Converts a plain multi-line string into nbformat's source array: each line carries a trailing
|
||||
* "\n" except the last, and a trailing newline in the input is preserved (so "a\n" => ["a\n"], not
|
||||
* ["a\n", ""]). An empty source becomes an empty array. Round-trips: array.join("") === input. */
|
||||
function toSourceArray(source: string): string[] {
|
||||
if (!source) return [];
|
||||
let lines = source.split("\n");
|
||||
// split("a\n") => ["a", ""] — the trailing "" is an artifact of the trailing newline, not a real
|
||||
// empty last line. Drop it and remember the input ended with \n so the now-last line keeps its \n.
|
||||
const endedWithNewline = lines.length > 1 && lines[lines.length - 1] === "";
|
||||
if (endedWithNewline) lines = lines.slice(0, -1);
|
||||
return lines.map((line, i) => (i < lines.length - 1 || endedWithNewline ? line + "\n" : line));
|
||||
}
|
||||
|
||||
/** Finds a cell's index by id (if present) falling back to the explicit index. Returns -1 if not
|
||||
* found. Used by replace/delete. */
|
||||
function findCellIndex(notebook: Notebook, cellId: string | undefined, cellIndex: number | undefined): number {
|
||||
if (cellId !== undefined) {
|
||||
return notebook.cells.findIndex((c) => c.id === cellId);
|
||||
}
|
||||
if (cellIndex !== undefined) {
|
||||
return cellIndex >= 0 && cellIndex < notebook.cells.length ? cellIndex : -1;
|
||||
}
|
||||
return -1;
|
||||
}
|
||||
|
||||
function applyNotebookEdit(notebook: Notebook, args: z.infer<typeof schema>, notebookPath: string): Notebook {
|
||||
const mode = args.edit_mode ?? "replace";
|
||||
if (mode === "insert") {
|
||||
if (!args.cell_type) throw new Error("insert requires 'cell_type'.");
|
||||
if (args.new_source === undefined) throw new Error("insert requires 'new_source'.");
|
||||
const cell: Cell = {
|
||||
cell_type: args.cell_type,
|
||||
source: toSourceArray(args.new_source),
|
||||
metadata: {},
|
||||
};
|
||||
if (args.cell_type === "code") {
|
||||
cell.execution_count = null;
|
||||
cell.outputs = [];
|
||||
}
|
||||
const insertAt = args.cell_index ?? notebook.cells.length;
|
||||
if (insertAt < 0 || insertAt > notebook.cells.length) {
|
||||
throw new Error(`insert cell_index ${insertAt} is out of range (0–${notebook.cells.length}).`);
|
||||
}
|
||||
notebook.cells.splice(insertAt, 0, cell);
|
||||
return notebook;
|
||||
}
|
||||
|
||||
const idx = findCellIndex(notebook, args.cell_id, args.cell_index);
|
||||
if (idx === -1) {
|
||||
const where = args.cell_id !== undefined ? `cell_id "${args.cell_id}"` : `cell_index ${args.cell_index}`;
|
||||
throw new Error(`${mode}: ${where} not found in ${notebookPath}.`);
|
||||
}
|
||||
if (mode === "delete") {
|
||||
notebook.cells.splice(idx, 1);
|
||||
return notebook;
|
||||
}
|
||||
// replace
|
||||
if (args.new_source === undefined) throw new Error("replace requires 'new_source'.");
|
||||
const cell = notebook.cells[idx]!;
|
||||
if (args.cell_type) cell.cell_type = args.cell_type;
|
||||
cell.source = toSourceArray(args.new_source);
|
||||
// Switching to a non-code cell type drops code-only fields; switching to code adds them.
|
||||
if (cell.cell_type === "code") {
|
||||
cell.execution_count ??= null;
|
||||
cell.outputs ??= [];
|
||||
} else {
|
||||
delete cell.execution_count;
|
||||
delete cell.outputs;
|
||||
}
|
||||
return notebook;
|
||||
}
|
||||
|
||||
export const notebookEditTool: ToolDef<z.infer<typeof schema>> = {
|
||||
name: "notebook_edit",
|
||||
description:
|
||||
"Edit a Jupyter (.ipynb) notebook cell-aware: replace, insert, or delete a cell by cell_id or cell_index. " +
|
||||
"new_source is the full new cell source as a single string. Read the notebook first (read_file shows the JSON). " +
|
||||
"Prefer this over edit_file/write_file for .ipynb so the JSON structure stays valid.",
|
||||
schema,
|
||||
mutating: true,
|
||||
preview: async (args, ctx) => {
|
||||
let resolved: string;
|
||||
try {
|
||||
resolved = resolveWithinCwd(ctx.cwd, args.notebook_path);
|
||||
} catch (err) {
|
||||
return (err as Error).message;
|
||||
}
|
||||
let original: string;
|
||||
try {
|
||||
original = await fsReadFile(resolved, "utf-8");
|
||||
} catch {
|
||||
return `Notebook ${resolved} does not exist.`;
|
||||
}
|
||||
let notebook: Notebook;
|
||||
try {
|
||||
notebook = JSON.parse(original) as Notebook;
|
||||
} catch {
|
||||
return `Warning: ${resolved} is not valid JSON — this edit will fail.`;
|
||||
}
|
||||
try {
|
||||
const updated = applyNotebookEdit(structuredClone(notebook), args, args.notebook_path);
|
||||
return createPatch(resolved, original, JSON.stringify(updated, null, 2) + "\n", "", "");
|
||||
} catch (err) {
|
||||
return `Warning: ${(err as Error).message} — this edit will fail.`;
|
||||
}
|
||||
},
|
||||
handler: async (args, ctx) => {
|
||||
const resolved = resolveWithinCwd(ctx.cwd, args.notebook_path);
|
||||
const original = await fsReadFile(resolved, "utf-8");
|
||||
let notebook: Notebook;
|
||||
try {
|
||||
notebook = JSON.parse(original) as Notebook;
|
||||
} catch {
|
||||
throw new Error(`${resolved} is not valid JSON — can't edit as a notebook.`);
|
||||
}
|
||||
if (!Array.isArray(notebook.cells)) throw new Error(`${resolved} has no cells array — not a valid .ipynb.`);
|
||||
applyNotebookEdit(notebook, args, args.notebook_path);
|
||||
const updated = JSON.stringify(notebook, null, 2) + "\n";
|
||||
// Atomic write via temp+rename (same rationale as edit_file/multi_edit): a crash mid-write can't
|
||||
// leave the notebook half-overwritten. Clean up the temp file if anything fails.
|
||||
const tmp = `${resolved}.locode-${randomBytes(4).toString("hex")}.tmp`;
|
||||
try {
|
||||
await fsWriteFile(tmp, updated, "utf-8");
|
||||
await fsRename(tmp, resolved);
|
||||
} catch (err) {
|
||||
await fsUnlink(tmp).catch(() => {});
|
||||
throw err;
|
||||
}
|
||||
return { path: resolved, edit_mode: args.edit_mode ?? "replace", cell_count: notebook.cells.length };
|
||||
},
|
||||
};
|
||||
Reference in New Issue
Block a user