@oh-my-pi/pi-coding-agent 8.0.20 → 8.2.0

package/src/tools/python.ts

@@ -1,23 +1,24 @@
-import { relative, resolve, sep } from "node:path";
+import * as path from "node:path";
 import type { AgentTool, AgentToolContext, AgentToolResult, AgentToolUpdateCallback } from "@oh-my-pi/pi-agent-core";
 import type { ImageContent } from "@oh-my-pi/pi-ai";
-import { renderPromptTemplate } from "@oh-my-pi/pi-coding-agent/config/prompt-templates";
-import type { RenderResultOptions } from "@oh-my-pi/pi-coding-agent/extensibility/custom-tools/types";
-import { executePython, getPreludeDocs, type PythonExecutorOptions } from "@oh-my-pi/pi-coding-agent/ipy/executor";
-import type { PreludeHelper, PythonStatusEvent } from "@oh-my-pi/pi-coding-agent/ipy/kernel";
-import { truncateToVisualLines } from "@oh-my-pi/pi-coding-agent/modes/components/visual-truncate";
-import { highlightCode, type Theme } from "@oh-my-pi/pi-coding-agent/modes/theme/theme";
-import pythonDescription from "@oh-my-pi/pi-coding-agent/prompts/tools/python.md" with { type: "text" };
-import { OutputSink, type OutputSummary } from "@oh-my-pi/pi-coding-agent/session/streaming-output";
-import type { OutputMeta } from "@oh-my-pi/pi-coding-agent/tools/output-meta";
-import { ToolAbortError, ToolError } from "@oh-my-pi/pi-coding-agent/tools/tool-errors";
 import type { Component } from "@oh-my-pi/pi-tui";
-import { Text, truncateToWidth, visibleWidth } from "@oh-my-pi/pi-tui";
+import { Text, truncateToWidth } from "@oh-my-pi/pi-tui";
 import { type Static, Type } from "@sinclair/typebox";
-import type { ToolSession } from "./index";
+import { renderPromptTemplate } from "../config/prompt-templates";
+import type { RenderResultOptions } from "../extensibility/custom-tools/types";
+import { executePython, getPreludeDocs, type PythonExecutorOptions } from "../ipy/executor";
+import type { PreludeHelper, PythonStatusEvent } from "../ipy/kernel";
+import { truncateToVisualLines } from "../modes/components/visual-truncate";
+import type { Theme } from "../modes/theme/theme";
+import pythonDescription from "../prompts/tools/python.md" with { type: "text" };
+import { OutputSink, type OutputSummary } from "../session/streaming-output";
+import { getTreeBranch, getTreeContinuePrefix, renderCodeCell } from "../tui";
+import type { ToolSession } from ".";
+import type { OutputMeta } from "./output-meta";
 import { allocateOutputArtifact, createTailBuffer } from "./output-utils";
 import { resolveToCwd } from "./path-utils";
-import { getTreeBranch, getTreeContinuePrefix, shortenPath, ToolUIKit, truncate } from "./render-utils";
+import { shortenPath, ToolUIKit, truncate } from "./render-utils";
+import { ToolAbortError, ToolError } from "./tool-errors";
 import { toolResult } from "./tool-result";
 import { DEFAULT_MAX_BYTES } from "./truncate";
 
@@ -176,9 +177,8 @@ export class PythonTool implements AgentTool<typeof pythonSchema> {
 		// Clamp to reasonable range: 1s - 600s (10 min)
 		timeoutSec = Math.max(1, Math.min(600, timeoutSec));
 		const timeoutMs = timeoutSec * 1000;
-		const controller = new AbortController();
-		const onAbort = () => controller.abort();
-		signal?.addEventListener("abort", onAbort, { once: true });
+		const timeoutSignal = AbortSignal.timeout(timeoutMs);
+		const combinedSignal = signal ? AbortSignal.any([signal, timeoutSignal]) : timeoutSignal;
 		let outputSink: OutputSink | undefined;
 		let outputSummary: OutputSummary | undefined;
 		let outputDumped = false;
@@ -225,7 +225,7 @@ export class PythonTool implements AgentTool<typeof pythonSchema> {
 
 			const buildUpdateDetails = (): PythonToolDetails => {
 				const details: PythonToolDetails = {
-					cells: cellResults.map((cell) => ({
+					cells: cellResults.map(cell => ({
 						...cell,
 						statusEvents: cell.statusEvents ? [...cell.statusEvents] : undefined,
 					})),
@@ -257,7 +257,7 @@ export class PythonTool implements AgentTool<typeof pythonSchema> {
 			outputSink = new OutputSink({
 				artifactPath,
 				artifactId,
-				onChunk: (chunk) => {
+				onChunk: chunk => {
 					appendTail(chunk);
 					pushUpdate();
 				},
@@ -266,7 +266,7 @@ export class PythonTool implements AgentTool<typeof pythonSchema> {
 			const baseExecutorOptions: Omit<PythonExecutorOptions, "reset"> = {
 				cwd: commandCwd,
 				timeoutMs,
-				signal: controller.signal,
+				signal: combinedSignal,
 				sessionId,
 				kernelMode: this.session.settings?.getPythonKernelMode?.() ?? "session",
 				useSharedGateway: this.session.settings?.getPythonSharedGateway?.() ?? true,
@@ -288,7 +288,7 @@ export class PythonTool implements AgentTool<typeof pythonSchema> {
 				const executorOptions: PythonExecutorOptions = {
 					...baseExecutorOptions,
 					reset: isFirstCell ? reset : false,
-					onChunk: async (chunk) => {
+					onChunk: async chunk => {
 						await outputSink!.push(chunk);
 					},
 				};
@@ -397,7 +397,6 @@ export class PythonTool implements AgentTool<typeof pythonSchema> {
 
 			return resultBuilder.done();
 		} finally {
-			signal?.removeEventListener("abort", onAbort);
 			if (!outputDumped) {
 				try {
 					await finalizeOutput();
@@ -639,14 +638,14 @@ function formatStatusEventExpanded(event: PythonStatusEvent, theme: Theme): stri
 	switch (op) {
 		case "find":
 		case "glob":
-			if (data.matches) addItems(data.matches as unknown[], (m) => String(m));
+			if (data.matches) addItems(data.matches as unknown[], m => String(m));
 			break;
 		case "ls":
-			if (data.items) addItems(data.items as unknown[], (m) => String(m));
+			if (data.items) addItems(data.items as unknown[], m => String(m));
 			break;
 		case "grep":
 			if (data.hits) {
-				addItems(data.hits as unknown[], (h) => {
+				addItems(data.hits as unknown[], h => {
 					const hit = h as { line: number; text: string };
 					return `${hit.line}: ${truncate(hit.text, 60, theme.format.ellipsis)}`;
 				});
@@ -654,7 +653,7 @@ function formatStatusEventExpanded(event: PythonStatusEvent, theme: Theme): stri
 			break;
 		case "rgrep":
 			if (data.hits) {
-				addItems(data.hits as unknown[], (h) => {
+				addItems(data.hits as unknown[], h => {
 					const hit = h as { file: string; line: number; text: string };
 					return `${shortenPath(hit.file)}:${hit.line}: ${truncate(hit.text, 50, theme.format.ellipsis)}`;
 				});
@@ -662,28 +661,28 @@ function formatStatusEventExpanded(event: PythonStatusEvent, theme: Theme): stri
 			break;
 		case "rsed":
 			if (data.changed) {
-				addItems(data.changed as unknown[], (c) => {
+				addItems(data.changed as unknown[], c => {
 					const change = c as { file: string; count: number };
 					return `${shortenPath(change.file)}: ${change.count} replacement${change.count !== 1 ? "s" : ""}`;
 				});
 			}
 			break;
 		case "env":
-			if (data.keys) addItems(data.keys as unknown[], (k) => String(k), 10);
+			if (data.keys) addItems(data.keys as unknown[], k => String(k), 10);
 			break;
 		case "git_log":
 			if (data.entries) {
-				addItems(data.entries as unknown[], (e) => {
+				addItems(data.entries as unknown[], e => {
 					const entry = e as { sha: string; subject: string };
 					return `${entry.sha} ${truncate(entry.subject, 50, theme.format.ellipsis)}`;
 				});
 			}
 			break;
 		case "git_status":
-			if (data.files) addItems(data.files as unknown[], (f) => String(f));
+			if (data.files) addItems(data.files as unknown[], f => String(f));
 			break;
 		case "git_branch":
-			if (data.branches) addItems(data.branches as unknown[], (b) => String(b));
+			if (data.branches) addItems(data.branches as unknown[], b => String(b));
 			break;
 		case "read":
 		case "cat":
@@ -740,56 +739,6 @@ function renderStatusEvents(events: PythonStatusEvent[], theme: Theme, expanded:
 	return lines;
 }
 
-function applyCellBackground(line: string, width: number, bgFn?: (text: string) => string): string {
-	if (!bgFn) return line;
-	if (width <= 0) return bgFn(line);
-	const paddingNeeded = Math.max(0, width - visibleWidth(line));
-	const padded = line + " ".repeat(paddingNeeded);
-	return bgFn(padded);
-}
-
-function highlightPythonCode(code?: string): string[] {
-	return highlightCode(code ?? "", "python");
-}
-
-function formatCellStatus(cell: PythonCellResult, ui: ToolUIKit, spinnerFrame?: number): string | undefined {
-	switch (cell.status) {
-		case "pending":
-			return `${ui.statusIcon("pending")} ${ui.theme.fg("muted", "pending")}`;
-		case "running":
-			return `${ui.statusIcon("running", spinnerFrame)} ${ui.theme.fg("muted", "running")}`;
-		case "complete":
-			return ui.statusIcon("success");
-		case "error":
-			return ui.statusIcon("error");
-	}
-}
-
-function formatCellHeader(
-	cell: PythonCellResult,
-	index: number,
-	total: number,
-	ui: ToolUIKit,
-	spinnerFrame?: number,
-	workdirLabel?: string,
-): string {
-	const indexLabel = ui.theme.fg("accent", `[${index + 1}/${total}]`);
-	const title = cell.title ? ` ${cell.title}` : "";
-	const metaParts: string[] = [];
-	if (workdirLabel) {
-		metaParts.push(ui.theme.fg("dim", workdirLabel));
-	}
-	if (cell.durationMs !== undefined) {
-		metaParts.push(ui.theme.fg("dim", `(${ui.formatDuration(cell.durationMs)})`));
-	}
-	const statusLabel = formatCellStatus(cell, ui, spinnerFrame);
-	if (statusLabel) {
-		metaParts.push(statusLabel);
-	}
-	const meta = metaParts.length > 0 ? ` ${metaParts.join(ui.theme.fg("dim", ui.theme.sep.dot))}` : "";
-	return `${indexLabel}${title}${meta}`;
-}
-
 function formatCellOutputLines(
 	cell: PythonCellResult,
 	expanded: boolean,
@@ -799,7 +748,7 @@ function formatCellOutputLines(
 	const rawLines = cell.output ? cell.output.split("\n") : [];
 	const displayLines = expanded ? rawLines : rawLines.slice(-previewLines);
 	const hiddenCount = rawLines.length - displayLines.length;
-	const outputLines = displayLines.map((line) => theme.fg("toolOutput", line));
+	const outputLines = displayLines.map(line => theme.fg("toolOutput", line));
 
 	if (outputLines.length === 0) {
 		return { lines: [], hiddenCount: 0 };
@@ -808,101 +757,6 @@ function formatCellOutputLines(
 	return { lines: outputLines, hiddenCount };
 }
 
-function renderCellBlock(
-	cell: PythonCellResult,
-	index: number,
-	total: number,
-	ui: ToolUIKit,
-	options: {
-		expanded: boolean;
-		previewLines: number;
-		spinnerFrame?: number;
-		showOutput: boolean;
-		workdirLabel?: string;
-		width: number;
-		bgFn?: (text: string) => string;
-	},
-): string[] {
-	const { expanded, previewLines, spinnerFrame, showOutput, workdirLabel, width, bgFn } = options;
-	const h = ui.theme.boxSharp.horizontal;
-	const v = ui.theme.boxSharp.vertical;
-	const cap = h.repeat(3);
-	const border = (text: string) => ui.theme.fg("dim", text);
-	const lineWidth = Math.max(0, width);
-
-	const buildBarLine = (leftChar: string, label?: string): string => {
-		const left = border(`${leftChar}${cap}`);
-		if (lineWidth <= 0) return left;
-		const rawLabel = label ? ` ${label} ` : " ";
-		const maxLabelWidth = Math.max(0, lineWidth - visibleWidth(left));
-		const trimmedLabel = truncateToWidth(rawLabel, maxLabelWidth, ui.theme.format.ellipsis);
-		const fillCount = Math.max(0, lineWidth - visibleWidth(left + trimmedLabel));
-		return `${left}${trimmedLabel}${border(h.repeat(fillCount))}`;
-	};
-
-	const lines: string[] = [];
-	lines.push(
-		applyCellBackground(
-			buildBarLine(ui.theme.boxSharp.topLeft, formatCellHeader(cell, index, total, ui, spinnerFrame, workdirLabel)),
-			lineWidth,
-			bgFn,
-		),
-	);
-
-	const codePrefix = border(`${v} `);
-	const codeWidth = Math.max(0, lineWidth - visibleWidth(codePrefix));
-	const codeLines = highlightPythonCode(cell.code);
-	for (const line of codeLines) {
-		const text = truncateToWidth(line, codeWidth, ui.theme.format.ellipsis);
-		lines.push(applyCellBackground(`${codePrefix}${text}`, lineWidth, bgFn));
-	}
-
-	const statusLines = renderStatusEvents(cell.statusEvents ?? [], ui.theme, expanded);
-	const outputContent = formatCellOutputLines(cell, expanded, previewLines, ui.theme);
-	const hasOutput = outputContent.lines.length > 0;
-	const hasStatus = statusLines.length > 0;
-	const showOutputSection = showOutput && (hasOutput || hasStatus);
-
-	if (showOutputSection) {
-		lines.push(
-			applyCellBackground(
-				buildBarLine(ui.theme.boxSharp.teeRight, ui.theme.fg("toolTitle", "Output")),
-				lineWidth,
-				bgFn,
-			),
-		);
-
-		for (const line of outputContent.lines) {
-			const text = truncateToWidth(line, codeWidth, ui.theme.format.ellipsis);
-			lines.push(applyCellBackground(`${codePrefix}${text}`, lineWidth, bgFn));
-		}
-		if (!expanded && outputContent.hiddenCount > 0) {
-			const hint = ui.theme.fg(
-				"dim",
-				`${ui.theme.format.ellipsis} ${outputContent.hiddenCount} more lines (ctrl+o to expand)`,
-			);
-			lines.push(
-				applyCellBackground(
-					`${codePrefix}${truncateToWidth(hint, codeWidth, ui.theme.format.ellipsis)}`,
-					lineWidth,
-					bgFn,
-				),
-			);
-		}
-
-		for (const line of statusLines) {
-			const text = truncateToWidth(line, codeWidth, ui.theme.format.ellipsis);
-			lines.push(applyCellBackground(`${codePrefix}${text}`, lineWidth, bgFn));
-		}
-	}
-
-	const bottomLeft = border(`${ui.theme.boxSharp.bottomLeft}${cap}`);
-	const bottomFillCount = Math.max(0, lineWidth - visibleWidth(bottomLeft));
-	const bottomLine = `${bottomLeft}${border(h.repeat(bottomFillCount))}`;
-	lines.push(applyCellBackground(bottomLine, lineWidth, bgFn));
-	return lines;
-}
-
 export const pythonToolRenderer = {
 	renderCall(args: PythonRenderArgs, uiTheme: Theme): Component {
 		const ui = new ToolUIKit(uiTheme);
@@ -911,13 +765,14 @@ export const pythonToolRenderer = {
 		let displayWorkdir = args.cwd;
 
 		if (displayWorkdir) {
-			const resolvedCwd = resolve(cwd);
-			const resolvedWorkdir = resolve(displayWorkdir);
+			const resolvedCwd = path.resolve(cwd);
+			const resolvedWorkdir = path.resolve(displayWorkdir);
 			if (resolvedWorkdir === resolvedCwd) {
 				displayWorkdir = undefined;
 			} else {
-				const relativePath = relative(resolvedCwd, resolvedWorkdir);
-				const isWithinCwd = relativePath && !relativePath.startsWith("..") && !relativePath.startsWith(`..${sep}`);
+				const relativePath = path.relative(resolvedCwd, resolvedWorkdir);
+				const isWithinCwd =
+					relativePath && !relativePath.startsWith("..") && !relativePath.startsWith(`..${path.sep}`);
 				if (isWithinCwd) {
 					displayWorkdir = relativePath;
 				}
@@ -937,23 +792,24 @@ export const pythonToolRenderer = {
 				const lines: string[] = [];
 				for (let i = 0; i < cells.length; i++) {
 					const cell = cells[i];
-					const cellResult: PythonCellResult = {
-						index: i,
-						title: cell.title,
-						code: cell.code,
-						output: "",
-						status: "pending",
-					};
-					lines.push(
-						...renderCellBlock(cellResult, i, cells.length, ui, {
-							expanded: true,
-							previewLines: PYTHON_DEFAULT_PREVIEW_LINES,
-							showOutput: false,
-							workdirLabel: i === 0 ? workdirLabel : undefined,
+					const cellTitle = cell.title;
+					const combinedTitle =
+						cellTitle && workdirLabel ? `${workdirLabel} · ${cellTitle}` : (cellTitle ?? workdirLabel);
+					const cellLines = renderCodeCell(
+						{
+							code: cell.code,
+							language: "python",
+							index: i,
+							total: cells.length,
+							title: combinedTitle,
+							status: "pending",
 							width,
-							bgFn: (text: string) => uiTheme.bg("toolPendingBg", text),
-						}),
+							codeMaxLines: PYTHON_DEFAULT_PREVIEW_LINES,
+							expanded: true,
+						},
+						uiTheme,
 					);
+					lines.push(...cellLines);
 					if (i < cells.length - 1) {
 						lines.push("");
 					}
@@ -975,7 +831,7 @@ export const pythonToolRenderer = {
 
 		const expanded = renderContext?.expanded ?? options.expanded;
 		const previewLines = renderContext?.previewLines ?? PYTHON_DEFAULT_PREVIEW_LINES;
-		const output = renderContext?.output ?? (result.content?.find((c) => c.type === "text")?.text ?? "").trim();
+		const output = renderContext?.output ?? (result.content?.find(c => c.type === "text")?.text ?? "").trimEnd();
 
 		const jsonOutputs = details?.jsonOutputs ?? [];
 		const jsonLines = jsonOutputs.flatMap((value, index) => {
@@ -1015,23 +871,42 @@ export const pythonToolRenderer = {
 					const lines: string[] = [];
 					for (let i = 0; i < cellResults.length; i++) {
 						const cell = cellResults[i];
-						const showOutput = cell.status !== "pending";
-						const bgColor =
-							cell.status === "error"
-								? "toolErrorBg"
-								: cell.status === "complete"
-									? "toolSuccessBg"
-									: "toolPendingBg";
-						lines.push(
-							...renderCellBlock(cell, i, cellResults.length, ui, {
-								expanded,
-								previewLines,
+						const statusLines = renderStatusEvents(cell.statusEvents ?? [], uiTheme, expanded);
+						const outputContent = formatCellOutputLines(cell, expanded, previewLines, uiTheme);
+						const outputLines = [...outputContent.lines];
+						if (!expanded && outputContent.hiddenCount > 0) {
+							outputLines.push(
+								uiTheme.fg(
+									"dim",
+									`${uiTheme.format.ellipsis} ${outputContent.hiddenCount} more lines (ctrl+o to expand)`,
+								),
+							);
+						}
+						if (statusLines.length > 0) {
+							if (outputLines.length > 0) {
+								outputLines.push(uiTheme.fg("dim", "Status"));
+							}
+							outputLines.push(...statusLines);
+						}
+						const cellLines = renderCodeCell(
+							{
+								code: cell.code,
+								language: "python",
+								index: i,
+								total: cellResults.length,
+								title: cell.title,
+								status: cell.status,
 								spinnerFrame: options.spinnerFrame,
-								showOutput,
+								duration: cell.durationMs,
+								output: outputLines.length > 0 ? outputLines.join("\n") : undefined,
+								outputMaxLines: outputLines.length,
+								codeMaxLines: expanded ? Number.POSITIVE_INFINITY : PYTHON_DEFAULT_PREVIEW_LINES,
+								expanded,
 								width,
-								bgFn: (text: string) => uiTheme.bg(bgColor, text),
-							}),
+							},
+							uiTheme,
 						);
+						lines.push(...cellLines);
 						if (i < cellResults.length - 1) {
 							lines.push("");
 						}
@@ -1066,22 +941,29 @@ export const pythonToolRenderer = {
 		}
 
 		if (!combinedOutput && statusLines.length > 0) {
-			const lines = [...statusLines, timeoutLine, warningLine].filter(Boolean) as string[];
+			const lines = [uiTheme.fg("dim", "Status"), ...statusLines, timeoutLine, warningLine].filter(
+				Boolean,
+			) as string[];
 			return new Text(lines.join("\n"), 0, 0);
 		}
 
 		if (expanded) {
 			const styledOutput = combinedOutput
 				.split("\n")
-				.map((line) => uiTheme.fg("toolOutput", line))
+				.map(line => uiTheme.fg("toolOutput", line))
 				.join("\n");
-			const lines = [styledOutput, ...statusLines, timeoutLine, warningLine].filter(Boolean) as string[];
+			const lines = [
+				styledOutput,
+				...(statusLines.length > 0 ? [uiTheme.fg("dim", "Status"), ...statusLines] : []),
+				timeoutLine,
+				warningLine,
+			].filter(Boolean) as string[];
 			return new Text(lines.join("\n"), 0, 0);
 		}
 
 		const styledOutput = combinedOutput
 			.split("\n")
-			.map((line) => uiTheme.fg("toolOutput", line))
+			.map(line => uiTheme.fg("toolOutput", line))
 			.join("\n");
 		const textContent = `\n${styledOutput}`;
 
@@ -1107,8 +989,13 @@ export const pythonToolRenderer = {
 					outputLines.push(truncateToWidth(skippedLine, width, uiTheme.fg("dim", uiTheme.format.ellipsis)));
 				}
 				outputLines.push(...cachedLines);
-				for (const statusLine of statusLines) {
-					outputLines.push(truncateToWidth(statusLine, width, uiTheme.fg("dim", uiTheme.format.ellipsis)));
+				if (statusLines.length > 0) {
+					outputLines.push(
+						truncateToWidth(uiTheme.fg("dim", "Status"), width, uiTheme.fg("dim", uiTheme.format.ellipsis)),
+					);
+					for (const statusLine of statusLines) {
+						outputLines.push(truncateToWidth(statusLine, width, uiTheme.fg("dim", uiTheme.format.ellipsis)));
+					}
 				}
 				if (timeoutLine) {
 					outputLines.push(truncateToWidth(timeoutLine, width, uiTheme.fg("dim", uiTheme.format.ellipsis)));