extensions!

pi/files/agent/extensions/truncated-tool.ts

@@ -0,0 +1,192 @@
+/**
+ * Truncated Tool Example - Demonstrates proper output truncation for custom tools
+ *
+ * Custom tools MUST truncate their output to avoid overwhelming the LLM context.
+ * The built-in limit is 50KB (~10k tokens) and 2000 lines, whichever is hit first.
+ *
+ * This example shows how to:
+ * 1. Use the built-in truncation utilities
+ * 2. Write full output to a temp file when truncated
+ * 3. Inform the LLM where to find the complete output
+ * 4. Custom rendering of tool calls and results
+ *
+ * The `rg` tool here wraps ripgrep with proper truncation. Compare this to the
+ * built-in `grep` tool in src/core/tools/grep.ts for a more complete implementation.
+ */
+
+import type { ExtensionAPI } from "@mariozechner/pi-coding-agent";
+import {
+	DEFAULT_MAX_BYTES,
+	DEFAULT_MAX_LINES,
+	formatSize,
+	type TruncationResult,
+	truncateHead,
+} from "@mariozechner/pi-coding-agent";
+import { Text } from "@mariozechner/pi-tui";
+import { Type } from "@sinclair/typebox";
+import { execSync } from "child_process";
+import { mkdtempSync, writeFileSync } from "fs";
+import { tmpdir } from "os";
+import { join } from "path";
+
+const RgParams = Type.Object({
+	pattern: Type.String({ description: "Search pattern (regex)" }),
+	path: Type.Optional(Type.String({ description: "Directory to search (default: current directory)" })),
+	glob: Type.Optional(Type.String({ description: "File glob pattern, e.g. '*.ts'" })),
+});
+
+interface RgDetails {
+	pattern: string;
+	path?: string;
+	glob?: string;
+	matchCount: number;
+	truncation?: TruncationResult;
+	fullOutputPath?: string;
+}
+
+export default function (pi: ExtensionAPI) {
+	pi.registerTool({
+		name: "rg",
+		label: "ripgrep",
+		// Document the truncation limits in the tool description so the LLM knows
+		description: `Search file contents using ripgrep. Output is truncated to ${DEFAULT_MAX_LINES} lines or ${formatSize(DEFAULT_MAX_BYTES)} (whichever is hit first). If truncated, full output is saved to a temp file.`,
+		parameters: RgParams,
+
+		async execute(_toolCallId, params, _signal, _onUpdate, ctx) {
+			const { pattern, path: searchPath, glob } = params;
+
+			// Build the ripgrep command
+			const args = ["rg", "--line-number", "--color=never"];
+			if (glob) args.push("--glob", glob);
+			args.push(pattern);
+			args.push(searchPath || ".");
+
+			let output: string;
+			try {
+				output = execSync(args.join(" "), {
+					cwd: ctx.cwd,
+					encoding: "utf-8",
+					maxBuffer: 100 * 1024 * 1024, // 100MB buffer to capture full output
+				});
+			} catch (err: any) {
+				// ripgrep exits with 1 when no matches found
+				if (err.status === 1) {
+					return {
+						content: [{ type: "text", text: "No matches found" }],
+						details: { pattern, path: searchPath, glob, matchCount: 0 } as RgDetails,
+					};
+				}
+				throw new Error(`ripgrep failed: ${err.message}`);
+			}
+
+			if (!output.trim()) {
+				return {
+					content: [{ type: "text", text: "No matches found" }],
+					details: { pattern, path: searchPath, glob, matchCount: 0 } as RgDetails,
+				};
+			}
+
+			// Apply truncation using built-in utilities
+			// truncateHead keeps the first N lines/bytes (good for search results)
+			// truncateTail keeps the last N lines/bytes (good for logs/command output)
+			const truncation = truncateHead(output, {
+				maxLines: DEFAULT_MAX_LINES,
+				maxBytes: DEFAULT_MAX_BYTES,
+			});
+
+			// Count matches (each non-empty line with a match)
+			const matchCount = output.split("\n").filter((line) => line.trim()).length;
+
+			const details: RgDetails = {
+				pattern,
+				path: searchPath,
+				glob,
+				matchCount,
+			};
+
+			let resultText = truncation.content;
+
+			if (truncation.truncated) {
+				// Save full output to a temp file so LLM can access it if needed
+				const tempDir = mkdtempSync(join(tmpdir(), "pi-rg-"));
+				const tempFile = join(tempDir, "output.txt");
+				writeFileSync(tempFile, output);
+
+				details.truncation = truncation;
+				details.fullOutputPath = tempFile;
+
+				// Add truncation notice - this helps the LLM understand the output is incomplete
+				const truncatedLines = truncation.totalLines - truncation.outputLines;
+				const truncatedBytes = truncation.totalBytes - truncation.outputBytes;
+
+				resultText += `\n\n[Output truncated: showing ${truncation.outputLines} of ${truncation.totalLines} lines`;
+				resultText += ` (${formatSize(truncation.outputBytes)} of ${formatSize(truncation.totalBytes)}).`;
+				resultText += ` ${truncatedLines} lines (${formatSize(truncatedBytes)}) omitted.`;
+				resultText += ` Full output saved to: ${tempFile}]`;
+			}
+
+			return {
+				content: [{ type: "text", text: resultText }],
+				details,
+			};
+		},
+
+		// Custom rendering of the tool call (shown before/during execution)
+		renderCall(args, theme) {
+			let text = theme.fg("toolTitle", theme.bold("rg "));
+			text += theme.fg("accent", `"${args.pattern}"`);
+			if (args.path) {
+				text += theme.fg("muted", ` in ${args.path}`);
+			}
+			if (args.glob) {
+				text += theme.fg("dim", ` --glob ${args.glob}`);
+			}
+			return new Text(text, 0, 0);
+		},
+
+		// Custom rendering of the tool result
+		renderResult(result, { expanded, isPartial }, theme) {
+			const details = result.details as RgDetails | undefined;
+
+			// Handle streaming/partial results
+			if (isPartial) {
+				return new Text(theme.fg("warning", "Searching..."), 0, 0);
+			}
+
+			// No matches
+			if (!details || details.matchCount === 0) {
+				return new Text(theme.fg("dim", "No matches found"), 0, 0);
+			}
+
+			// Build result display
+			let text = theme.fg("success", `${details.matchCount} matches`);
+
+			// Show truncation warning if applicable
+			if (details.truncation?.truncated) {
+				text += theme.fg("warning", " (truncated)");
+			}
+
+			// In expanded view, show the actual matches
+			if (expanded) {
+				const content = result.content[0];
+				if (content?.type === "text") {
+					// Show first 20 lines in expanded view, or all if fewer
+					const lines = content.text.split("\n").slice(0, 20);
+					for (const line of lines) {
+						text += `\n${theme.fg("dim", line)}`;
+					}
+					if (content.text.split("\n").length > 20) {
+						text += `\n${theme.fg("muted", "... (use read tool to see full output)")}`;
+					}
+				}
+
+				// Show temp file path if truncated
+				if (details.fullOutputPath) {
+					text += `\n${theme.fg("dim", `Full output: ${details.fullOutputPath}`)}`;
+				}
+			}
+
+			return new Text(text, 0, 0);
+		},
+	});
+}