/** * 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); }, }); }