@oh-my-pi/pi-coding-agent 6.1.0 → 6.7.0

package/src/core/tools/notebook.ts

@@ -1,8 +1,8 @@
-import type { AgentTool } from "@oh-my-pi/pi-agent-core";
+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 { Type } from "@sinclair/typebox";
+import { type Static, Type } from "@sinclair/typebox";
 import type { Theme } from "../../modes/interactive/theme/theme";
 import type { RenderResultOptions } from "../custom-tools/types";
 import type { ToolSession } from "../sdk";
@@ -63,133 +63,135 @@ function splitIntoLines(content: string): string[] {
 	return content.split("\n").map((line, i, arr) => (i < arr.length - 1 ? `${line}\n` : line));
 }
 
-export function createNotebookTool(session: ToolSession): AgentTool<typeof notebookSchema> {
-	return {
-		name: "notebook",
-		label: "Notebook",
-		description:
-			"Completely replaces the contents of a specific cell in a Jupyter notebook (.ipynb file) with new source. Jupyter notebooks are interactive documents that combine code, text, and visualizations, commonly used for data analysis and scientific computing. The notebook_path parameter must be an absolute path, not a relative path. The cell_number is 0-indexed. Use edit_mode=insert to add a new cell at the index specified by cell_number. Use edit_mode=delete to delete the cell at the index specified by cell_number.",
-		parameters: notebookSchema,
-		execute: async (
-			_toolCallId: string,
-			{
-				action,
-				notebook_path,
-				cell_index,
-				content,
-				cell_type,
-			}: { action: string; notebook_path: string; cell_index: number; content?: string; cell_type?: string },
-			signal?: AbortSignal,
-		) => {
-			const absolutePath = resolveToCwd(notebook_path, session.cwd);
-
-			return untilAborted(signal, async () => {
-				// Check if file exists
-				const file = Bun.file(absolutePath);
-				if (!(await file.exists())) {
-					throw new Error(`Notebook not found: ${notebook_path}`);
-				}
+type NotebookParams = Static<typeof notebookSchema>;
 
-				// Read and parse notebook
-				let notebook: Notebook;
-				try {
-					notebook = await file.json();
-				} catch {
-					throw new Error(`Invalid JSON in notebook: ${notebook_path}`);
-				}
+export class NotebookTool implements AgentTool<typeof notebookSchema, NotebookToolDetails> {
+	public readonly name = "notebook";
+	public readonly label = "Notebook";
+	public readonly description =
+		"Completely replaces the contents of a specific cell in a Jupyter notebook (.ipynb file) with new source. Jupyter notebooks are interactive documents that combine code, text, and visualizations, commonly used for data analysis and scientific computing. The notebook_path parameter must be an absolute path, not a relative path. The cell_number is 0-indexed. Use edit_mode=insert to add a new cell at the index specified by cell_number. Use edit_mode=delete to delete the cell at the index specified by cell_number.";
+	public readonly parameters = notebookSchema;
 
-				// Validate notebook structure
-				if (!notebook.cells || !Array.isArray(notebook.cells)) {
-					throw new Error(`Invalid notebook structure (missing cells array): ${notebook_path}`);
-				}
+	private readonly session: ToolSession;
 
-				const cellCount = notebook.cells.length;
+	constructor(session: ToolSession) {
+		this.session = session;
+	}
 
-				// 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}`);
-					}
+	public async execute(
+		_toolCallId: string,
+		params: NotebookParams,
+		signal?: AbortSignal,
+		_onUpdate?: AgentToolUpdateCallback<NotebookToolDetails>,
+		_context?: AgentToolContext,
+	): Promise<AgentToolResult<NotebookToolDetails>> {
+		const { action, notebook_path, cell_index, content, cell_type } = params;
+		const absolutePath = resolveToCwd(notebook_path, this.session.cwd);
+
+		return untilAborted(signal, async () => {
+			// Check if file exists
+			const file = Bun.file(absolutePath);
+			if (!(await file.exists())) {
+				throw new Error(`Notebook not found: ${notebook_path}`);
+			}
+
+			// Read and parse notebook
+			let notebook: Notebook;
+			try {
+				notebook = await file.json();
+			} catch {
+				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}`);
 				}
-
-				// Validate content for edit/insert
-				if ((action === "edit" || action === "insert") && content === undefined) {
-					throw new Error(`Content is required for ${action} action`);
+			} else {
+				if (cell_index < 0 || cell_index >= cellCount) {
+					throw new Error(`Cell index ${cell_index} out of range (0-${cellCount - 1}) in ${notebook_path}`);
 				}
-
-				// 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}`);
+			}
+
+			// 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,
+			// 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,
+				},
+			};
+		});
+	}
 }
 
 // =============================================================================

package/src/core/tools/output.ts

@@ -6,8 +6,8 @@
 
 import * as fs from "node:fs";
 import * as path from "node:path";
-import type { AgentTool } from "@oh-my-pi/pi-agent-core";
-import { StringEnum, type TextContent } from "@oh-my-pi/pi-ai";
+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 { Type } from "@sinclair/typebox";
@@ -232,166 +232,182 @@ function extractPreviewLines(content: string, maxLines: number): string[] {
 	return preview;
 }
 
-export function createOutputTool(session: ToolSession): AgentTool<typeof outputSchema, OutputToolDetails> {
-	return {
-		name: "output",
-		label: "Output",
-		description: renderPromptTemplate(outputDescription),
-		parameters: outputSchema,
-		execute: async (
-			_toolCallId: string,
-			params: {
-				ids: string[];
-				format?: "raw" | "json" | "stripped";
-				query?: string;
-				offset?: number;
-				limit?: number;
-			},
-		): Promise<{ content: TextContent[]; details: OutputToolDetails }> => {
-			const sessionFile = session.getSessionFile();
-
-			if (!sessionFile) {
-				return {
-					content: [{ type: "text", text: "No session - output artifacts unavailable" }],
-					details: { outputs: [], notFound: params.ids },
-				};
-			}
+type OutputParams = {
+	ids: string[];
+	format?: "raw" | "json" | "stripped";
+	query?: string;
+	offset?: number;
+	limit?: number;
+};
 
-			const artifactsDir = getArtifactsDir(sessionFile);
-			if (!artifactsDir || !fs.existsSync(artifactsDir)) {
-				return {
-					content: [{ type: "text", text: "No artifacts directory found" }],
-					details: { outputs: [], notFound: params.ids },
-				};
-			}
+/**
+ * Output tool for reading agent/task outputs by ID.
+ *
+ * Resolves IDs like "reviewer_0" to artifact paths in the current session.
+ */
+export class OutputTool implements AgentTool<typeof outputSchema, OutputToolDetails> {
+	public readonly name = "output";
+	public readonly label = "Output";
+	public readonly description: string;
+	public readonly parameters = outputSchema;
 
-			const outputs: OutputEntry[] = [];
-			const notFound: string[] = [];
-			const outputContentById = new Map<string, string>();
-			const query = params.query?.trim();
-			const wantsQuery = query !== undefined && query.length > 0;
-			const format = params.format ?? (wantsQuery ? "json" : "raw");
+	private readonly session: ToolSession;
 
-			if (wantsQuery && (params.offset !== undefined || params.limit !== undefined)) {
-				throw new Error("query cannot be combined with offset/limit");
-			}
+	constructor(session: ToolSession) {
+		this.session = session;
+		this.description = renderPromptTemplate(outputDescription);
+	}
 
-			const queryResults: Array<{ id: string; value: unknown }> = [];
+	public async execute(
+		_toolCallId: string,
+		params: OutputParams,
+		_signal?: AbortSignal,
+		_onUpdate?: AgentToolUpdateCallback<OutputToolDetails>,
+		_context?: AgentToolContext,
+	): Promise<AgentToolResult<OutputToolDetails>> {
+		const sessionFile = this.session.getSessionFile();
 
-			for (const id of params.ids) {
-				const outputPath = path.join(artifactsDir, `${id}.out.md`);
+		if (!sessionFile) {
+			return {
+				content: [{ type: "text", text: "No session - output artifacts unavailable" }],
+				details: { outputs: [], notFound: params.ids },
+			};
+		}
 
-				if (!fs.existsSync(outputPath)) {
-					notFound.push(id);
-					continue;
-				}
+		const artifactsDir = getArtifactsDir(sessionFile);
+		if (!artifactsDir || !fs.existsSync(artifactsDir)) {
+			return {
+				content: [{ type: "text", text: "No artifacts directory found" }],
+				details: { outputs: [], notFound: params.ids },
+			};
+		}
 
-				const rawContent = fs.readFileSync(outputPath, "utf-8");
-				const rawLines = rawContent.split("\n");
-				const totalLines = rawLines.length;
-				const totalChars = rawContent.length;
-
-				let selectedContent = rawContent;
-				let range: OutputRange | undefined;
-
-				if (wantsQuery && query) {
-					let jsonValue: unknown;
-					try {
-						jsonValue = JSON.parse(rawContent);
-					} catch (err) {
-						const message = err instanceof Error ? err.message : String(err);
-						throw new Error(`Output ${id} is not valid JSON: ${message}`);
-					}
-					const value = applyQuery(jsonValue, query);
-					queryResults.push({ id, value });
-					try {
-						selectedContent = JSON.stringify(value, null, 2) ?? "null";
-					} catch {
-						selectedContent = String(value);
-					}
-				} else if (params.offset !== undefined || params.limit !== undefined) {
-					const startLine = Math.max(1, params.offset ?? 1);
-					if (startLine > totalLines) {
-						throw new Error(
-							`Offset ${params.offset ?? startLine} is beyond end of output (${totalLines} lines) for ${id}`,
-						);
-					}
-					const effectiveLimit = params.limit ?? totalLines - startLine + 1;
-					const endLine = Math.min(totalLines, startLine + effectiveLimit - 1);
-					const selectedLines = rawLines.slice(startLine - 1, endLine);
-					selectedContent = selectedLines.join("\n");
-					range = { startLine, endLine, totalLines };
-				}
+		const outputs: OutputEntry[] = [];
+		const notFound: string[] = [];
+		const outputContentById = new Map<string, string>();
+		const query = params.query?.trim();
+		const wantsQuery = query !== undefined && query.length > 0;
+		const format = params.format ?? (wantsQuery ? "json" : "raw");
 
-				outputContentById.set(id, selectedContent);
-				outputs.push({
-					id,
-					path: outputPath,
-					lineCount: wantsQuery ? selectedContent.split("\n").length : totalLines,
-					charCount: wantsQuery ? selectedContent.length : totalChars,
-					provenance: parseOutputProvenance(id),
-					previewLines: extractPreviewLines(selectedContent, 4),
-					range,
-					query: query,
-				});
-			}
+		if (wantsQuery && (params.offset !== undefined || params.limit !== undefined)) {
+			throw new Error("query cannot be combined with offset/limit");
+		}
 
-			// Error case: some IDs not found
-			if (notFound.length > 0) {
-				const available = listAvailableOutputs(artifactsDir);
-				const errorMsg =
-					available.length > 0
-						? `Not found: ${notFound.join(", ")}\nAvailable: ${available.join(", ")}`
-						: `Not found: ${notFound.join(", ")}\nNo outputs available in current session`;
-
-				return {
-					content: [{ type: "text", text: errorMsg }],
-					details: { outputs, notFound, availableIds: available },
-				};
+		const queryResults: Array<{ id: string; value: unknown }> = [];
+
+		for (const id of params.ids) {
+			const outputPath = path.join(artifactsDir, `${id}.out.md`);
+
+			if (!fs.existsSync(outputPath)) {
+				notFound.push(id);
+				continue;
 			}
 
-			// Success: build response based on format
-			let contentText: string;
-
-			if (format === "json") {
-				const jsonData = wantsQuery
-					? queryResults
-					: outputs.map((o) => ({
-							id: o.id,
-							lineCount: o.lineCount,
-							charCount: o.charCount,
-							provenance: o.provenance,
-							previewLines: o.previewLines,
-							range: o.range,
-							content: outputContentById.get(o.id) ?? "",
-						}));
-				contentText = JSON.stringify(jsonData, null, 2);
-			} else {
-				// raw or stripped
-				const parts = outputs.map((o) => {
-					let content = outputContentById.get(o.id) ?? "";
-					if (format === "stripped") {
-						content = stripAnsi(content);
-					}
-					if (o.range && o.range.endLine < o.range.totalLines) {
-						const nextOffset = o.range.endLine + 1;
-						content += `\n\n[Showing lines ${o.range.startLine}-${o.range.endLine} of ${o.range.totalLines}. Use offset=${nextOffset} to continue]`;
-					}
-					// Add header for multiple outputs
-					if (outputs.length > 1) {
-						return `=== ${o.id} (${o.lineCount} lines, ${formatBytes(o.charCount)}) ===\n${content}`;
-					}
-					return content;
-				});
-				contentText = parts.join("\n\n");
+			const rawContent = fs.readFileSync(outputPath, "utf-8");
+			const rawLines = rawContent.split("\n");
+			const totalLines = rawLines.length;
+			const totalChars = rawContent.length;
+
+			let selectedContent = rawContent;
+			let range: OutputRange | undefined;
+
+			if (wantsQuery && query) {
+				let jsonValue: unknown;
+				try {
+					jsonValue = JSON.parse(rawContent);
+				} catch (err) {
+					const message = err instanceof Error ? err.message : String(err);
+					throw new Error(`Output ${id} is not valid JSON: ${message}`);
+				}
+				const value = applyQuery(jsonValue, query);
+				queryResults.push({ id, value });
+				try {
+					selectedContent = JSON.stringify(value, null, 2) ?? "null";
+				} catch {
+					selectedContent = String(value);
+				}
+			} else if (params.offset !== undefined || params.limit !== undefined) {
+				const startLine = Math.max(1, params.offset ?? 1);
+				if (startLine > totalLines) {
+					throw new Error(
+						`Offset ${params.offset ?? startLine} is beyond end of output (${totalLines} lines) for ${id}`,
+					);
+				}
+				const effectiveLimit = params.limit ?? totalLines - startLine + 1;
+				const endLine = Math.min(totalLines, startLine + effectiveLimit - 1);
+				const selectedLines = rawLines.slice(startLine - 1, endLine);
+				selectedContent = selectedLines.join("\n");
+				range = { startLine, endLine, totalLines };
 			}
 
+			outputContentById.set(id, selectedContent);
+			outputs.push({
+				id,
+				path: outputPath,
+				lineCount: wantsQuery ? selectedContent.split("\n").length : totalLines,
+				charCount: wantsQuery ? selectedContent.length : totalChars,
+				provenance: parseOutputProvenance(id),
+				previewLines: extractPreviewLines(selectedContent, 4),
+				range,
+				query: query,
+			});
+		}
+
+		// Error case: some IDs not found
+		if (notFound.length > 0) {
+			const available = listAvailableOutputs(artifactsDir);
+			const errorMsg =
+				available.length > 0
+					? `Not found: ${notFound.join(", ")}\nAvailable: ${available.join(", ")}`
+					: `Not found: ${notFound.join(", ")}\nNo outputs available in current session`;
+
 			return {
-				content: [{ type: "text", text: contentText }],
-				details: { outputs },
+				content: [{ type: "text", text: errorMsg }],
+				details: { outputs, notFound, availableIds: available },
 			};
-		},
-	};
+		}
+
+		// Success: build response based on format
+		let contentText: string;
+
+		if (format === "json") {
+			const jsonData = wantsQuery
+				? queryResults
+				: outputs.map((o) => ({
+						id: o.id,
+						lineCount: o.lineCount,
+						charCount: o.charCount,
+						provenance: o.provenance,
+						previewLines: o.previewLines,
+						range: o.range,
+						content: outputContentById.get(o.id) ?? "",
+					}));
+			contentText = JSON.stringify(jsonData, null, 2);
+		} else {
+			// raw or stripped
+			const parts = outputs.map((o) => {
+				let content = outputContentById.get(o.id) ?? "";
+				if (format === "stripped") {
+					content = stripAnsi(content);
+				}
+				if (o.range && o.range.endLine < o.range.totalLines) {
+					const nextOffset = o.range.endLine + 1;
+					content += `\n\n[Showing lines ${o.range.startLine}-${o.range.endLine} of ${o.range.totalLines}. Use offset=${nextOffset} to continue]`;
+				}
+				// Add header for multiple outputs
+				if (outputs.length > 1) {
+					return `=== ${o.id} (${o.lineCount} lines, ${formatBytes(o.charCount)}) ===\n${content}`;
+				}
+				return content;
+			});
+			contentText = parts.join("\n\n");
+		}
+
+		return {
+			content: [{ type: "text", text: contentText }],
+			details: { outputs },
+		};
+	}
 }
 
 // =============================================================================