one-code-extension 0.1.0

package/extensions/notebook/index.ts

@@ -0,0 +1,77 @@
+/**
+ * notebook extension — Claude Code's NotebookEdit, deferred behind tool_search
+ * (most sessions never touch a notebook, so its schema stays out of the prompt
+ * until needed).
+ */
+
+import { randomUUID } from "node:crypto";
+import { readFileSync, writeFileSync } from "node:fs";
+import { isAbsolute, resolve } from "node:path";
+import { StringEnum } from "@earendil-works/pi-ai";
+import type { ExtensionAPI } from "@earendil-works/pi-coding-agent";
+import { Type } from "typebox";
+import { DEFER_CHANNEL } from "../lib/deferred.ts";
+import { ccToolRenderers } from "../lib/tui-render.ts";
+import { applyEdit, type EditMode, parseNotebook } from "./notebook.ts";
+
+export default function notebookExtension(pi: ExtensionAPI) {
+	pi.registerTool({
+		name: "notebook_edit",
+		label: "Notebook Edit",
+		...ccToolRenderers("Notebook Edit"),
+		description:
+			"Edit a Jupyter notebook (.ipynb): replace a cell's source, insert a new cell, or delete a cell. Cells are addressed by their `id`; read the notebook first to get ids. Editing a code cell clears its outputs.",
+		parameters: Type.Object({
+			path: Type.String({ description: "Absolute or workspace-relative path to the .ipynb file" }),
+			edit_mode: StringEnum(["replace", "insert", "delete"] as const),
+			cell_id: Type.Optional(
+				Type.String({
+					description:
+						"Target cell id. Required for replace and delete; for insert the new cell goes after it (omit to insert at the top).",
+				}),
+			),
+			new_source: Type.Optional(Type.String({ description: "New cell source. Required unless deleting." })),
+			cell_type: Type.Optional(
+				StringEnum(["code", "markdown"] as const, {
+					description: "Cell type. Required when inserting; when replacing, changes the cell's type.",
+				}),
+			),
+		}),
+		async execute(_toolCallId, params, _signal, _onUpdate, ctx) {
+			const path = isAbsolute(params.path) ? params.path : resolve(ctx.cwd, params.path);
+			try {
+				const notebook = parseNotebook(readFileSync(path, "utf-8"));
+				const { notebook: updated, summary } = applyEdit(
+					notebook,
+					{
+						cellId: params.cell_id,
+						newSource: params.new_source,
+						cellType: params.cell_type as "code" | "markdown" | undefined,
+						editMode: params.edit_mode as EditMode,
+					},
+					() => randomUUID().slice(0, 8),
+				);
+				writeFileSync(path, `${JSON.stringify(updated, null, 1)}\n`);
+				return {
+					content: [{ type: "text", text: `${summary} in ${params.path} (${updated.cells.length} cells).` }],
+					details: { path, cellCount: updated.cells.length },
+				};
+			} catch (error) {
+				const code = (error as NodeJS.ErrnoException).code;
+				const hint =
+					code === "ENOENT"
+						? " — no such file; check the path points to an existing .ipynb"
+						: error instanceof SyntaxError
+							? " — the file is not valid notebook JSON"
+							: "";
+				return {
+					content: [{ type: "text", text: `Notebook edit failed: ${(error as Error).message}${hint}` }],
+					details: {},
+					isError: true,
+				};
+			}
+		},
+	});
+
+	pi.events.emit(DEFER_CHANNEL, { name: "notebook_edit", keywords: ["notebook", "jupyter", "ipynb", "cell"] });
+}

package/extensions/notebook/notebook.ts

@@ -0,0 +1,108 @@
+/**
+ * Jupyter notebook editing (pure) — Claude Code's NotebookEdit semantics.
+ */
+
+export interface NotebookCell {
+	cell_type: string;
+	id?: string;
+	source: string | string[];
+	metadata?: Record<string, unknown>;
+	outputs?: unknown[];
+	execution_count?: number | null;
+}
+
+export interface Notebook {
+	cells: NotebookCell[];
+	[key: string]: unknown;
+}
+
+export type EditMode = "replace" | "insert" | "delete";
+
+export interface EditRequest {
+	cellId?: string;
+	newSource?: string;
+	cellType?: "code" | "markdown";
+	editMode: EditMode;
+}
+
+export function parseNotebook(content: string): Notebook {
+	const parsed = JSON.parse(content) as Notebook;
+	if (!Array.isArray(parsed.cells)) throw new Error("Not a Jupyter notebook: missing 'cells' array");
+	return parsed;
+}
+
+/** Notebooks store source as a line array with trailing newlines; match that. */
+export function toSourceLines(source: string): string[] {
+	const lines = source.split("\n");
+	return lines.map((line, i) => (i === lines.length - 1 ? line : `${line}\n`)).filter((l, i) => l !== "" || i === 0);
+}
+
+export function findCellIndex(notebook: Notebook, cellId: string): number {
+	return notebook.cells.findIndex((cell) => cell.id === cellId);
+}
+
+function newCell(cellType: "code" | "markdown", source: string, id: string): NotebookCell {
+	const cell: NotebookCell = { cell_type: cellType, id, metadata: {}, source: toSourceLines(source) };
+	if (cellType === "code") {
+		cell.outputs = [];
+		cell.execution_count = null;
+	}
+	return cell;
+}
+
+export interface EditResult {
+	notebook: Notebook;
+	summary: string;
+}
+
+/**
+ * Applies one edit and returns a new notebook object. `makeId` supplies the id
+ * for inserted cells (injected so callers can keep this deterministic in tests).
+ */
+export function applyEdit(notebook: Notebook, request: EditRequest, makeId: () => string): EditResult {
+	const cells = [...notebook.cells];
+
+	if (request.editMode === "delete") {
+		if (!request.cellId) throw new Error("cell_id is required for edit_mode 'delete'");
+		const index = findCellIndex(notebook, request.cellId);
+		if (index === -1) throw new Error(`No cell with id "${request.cellId}"`);
+		cells.splice(index, 1);
+		return { notebook: { ...notebook, cells }, summary: `Deleted cell ${request.cellId}` };
+	}
+
+	if (request.newSource === undefined) throw new Error("new_source is required unless edit_mode is 'delete'");
+
+	if (request.editMode === "insert") {
+		if (!request.cellType) throw new Error("cell_type is required for edit_mode 'insert'");
+		const id = makeId();
+		const cell = newCell(request.cellType, request.newSource, id);
+		// Claude Code inserts AFTER the given cell; no cell_id means insert first.
+		const index = request.cellId ? findCellIndex(notebook, request.cellId) : -1;
+		if (request.cellId && index === -1) throw new Error(`No cell with id "${request.cellId}"`);
+		cells.splice(index + 1, 0, cell);
+		return {
+			notebook: { ...notebook, cells },
+			summary: `Inserted ${request.cellType} cell ${id}${request.cellId ? ` after ${request.cellId}` : " at the top"}`,
+		};
+	}
+
+	if (!request.cellId) throw new Error("cell_id is required for edit_mode 'replace'");
+	const index = findCellIndex(notebook, request.cellId);
+	if (index === -1) throw new Error(`No cell with id "${request.cellId}"`);
+	const existing = cells[index];
+	const cellType = (request.cellType ?? existing.cell_type) as "code" | "markdown";
+	const replacement: NotebookCell = {
+		...existing,
+		cell_type: cellType,
+		source: toSourceLines(request.newSource),
+	};
+	if (cellType === "code") {
+		replacement.outputs = [];
+		replacement.execution_count = null;
+	} else {
+		delete replacement.outputs;
+		delete replacement.execution_count;
+	}
+	cells[index] = replacement;
+	return { notebook: { ...notebook, cells }, summary: `Replaced cell ${request.cellId}` };
+}