my-pi/pi-overflow/index.ts

import { getAgentDir, type ExtensionAPI } from "@earendil-works/pi-coding-agent";
import { Type } from "typebox";
import { stat } from "node:fs/promises";
import { join } from "node:path";
import { readOverflowSettings, type OverflowSettings } from "./settings.ts";
import { getStash } from "./src/stash.ts";
import { buildView, collectViewFromFile, collectViewFromText } from "./src/view.ts";
import { cleanStore, entryId, storeEntry } from "./src/store.ts";
import { runOverflowQuery } from "./src/query.ts";

function formatBytes(bytes: number): string {
	if (bytes < 1024) return `${bytes} B`;
	if (bytes < 1_048_576) return `${(bytes / 1024).toFixed(1)} KB`;
	return `${(bytes / 1_048_576).toFixed(1)} MB`;
}

export interface OverflowExtensionOptions {
	settings?: OverflowSettings;
	storeRoot?: string;
}

/**
 * Bound stage: loaded LAST among content extensions (before system-prompt).
 * By the time this handler runs, every other tool_result handler — RTK
 * included — has already transformed the result, so the bounded view is
 * built from the final visible text while the stored original comes from
 * capture.ts's stash or bash's fullOutputPath file.
 */
export default function overflowExtension(pi: ExtensionAPI, options: OverflowExtensionOptions = {}): void {
	const settings = options.settings ?? readOverflowSettings();
	const storeRoot = () => options.storeRoot ?? join(getAgentDir(), "my-pi", "pi-overflow");

	pi.on("session_start", () => {
		// Retention/size cleanup never throws (cleanStore swallows failures)
		// and the returned promise lets tests await it.
		return cleanStore(storeRoot(), settings.retentionDays, settings.maxStoreBytes);
	});

	pi.on("tool_result", async (event, ctx) => {
		const stash = getStash();
		const staged = stash.get(event.toolCallId);
		try {
			if (
				!settings.enabled ||
				settings.exemptTools.includes(event.toolName) ||
				event.parentToolCallId
			) {
				return undefined;
			}
			const current = event.content
				.filter((part): part is { type: "text"; text: string } => part.type === "text")
				.map((part) => part.text)
				.join("\n");
			const fullOutputPath =
				typeof event.details?.fullOutputPath === "string"
					? event.details.fullOutputPath
					: staged?.fullOutputPath;

			// Resolve the original: bash's spilled file wins (it is complete even
			// when the in-context text was already cut), then the capture stash,
			// then the current text itself.
			let source: { kind: "text"; text: string } | { kind: "file"; path: string };
			let originalSize: number;
			if (fullOutputPath) {
				try {
					const info = await stat(fullOutputPath);
					if (!info.isFile()) throw new Error("not a file");
					source = { kind: "file", path: fullOutputPath };
					originalSize = info.size;
				} catch {
					const fallback = staged?.rawText ?? current;
					source = { kind: "text", text: fallback };
					originalSize = fallback.length;
				}
			} else {
				const fallback = staged?.rawText ?? current;
				source = { kind: "text", text: fallback };
				originalSize = fallback.length;
			}

			if (current.length <= settings.thresholdChars && originalSize <= settings.thresholdChars) {
				return undefined;
			}

			const sessionId = ctx.sessionManager?.getSessionId?.() ?? "session";
			const id = entryId(sessionId, event.toolCallId);
			const stored = await storeEntry(
				storeRoot(),
				id,
				source,
				{ tool: event.toolName, args: event.input },
				settings.maxEntryBytes,
			);
			const overflow = {
				id,
				totalLines: stored.totalLines,
				totalBytes: stored.totalBytes,
				truncated: current.length > settings.thresholdChars,
			};

			const images = event.content.filter((part) => part.type !== "text");
			let text: string;
			if (current.length > settings.thresholdChars) {
				// The view always describes the stored file — its line numbers and
				// totals must match what overflow_query reads, which is not true
				// of `current` when RTK rewrote or bash pre-truncated the result.
				const viewInput =
					stored.storedText !== undefined
						? collectViewFromText(stored.storedText, settings)
						: await collectViewFromFile(join(storeRoot(), `${id}.txt`), settings);
				text = buildView(viewInput, settings, { id, tool: event.toolName });
			} else {
				text = `${current}\n[pi-overflow id=${id}: full original output stored (${stored.totalLines} lines, ${formatBytes(stored.totalBytes)}); use overflow_query to inspect it]`;
			}

			return {
				content: [{ type: "text", text }, ...images],
				structuredContent: event.structuredContent,
				details: { ...event.details, overflow },
			};
		} finally {
			stash.delete(event.toolCallId);
		}
	});

	pi.registerTool({
		name: "overflow_query",
		label: "Overflow Query",
		description:
			"Inspect the full original of a tool output that pi-overflow stored because it was too long to show " +
			"(search with regex, read line ranges, or extract a JSON path).",
		promptSnippet: "Inspect a stored oversized tool output by id: regex search, line ranges, or JSON path.",
		promptGuidelines: [
			"When a tool result carries a [pi-overflow id=…] marker, the model is seeing a bounded view; the complete original is stored on disk.",
			"Call overflow_query with the marker's id and a regex, line range, or jsonPath instead of re-running the tool or asking for the full dump.",
		],
		parameters: Type.Object({
			id: Type.String({ description: "12-hex storage id from a [pi-overflow id=…] marker." }),
			search: Type.Optional(
				Type.String({ description: "Case-insensitive regular expression matched per line." }),
			),
			context: Type.Optional(
				Type.Number({ minimum: 0, maximum: 10, description: "Context lines around each search hit (default 2)." }),
			),
			lines: Type.Optional(
				Type.Object({
					start: Type.Number({ minimum: 1, description: "First line to return (1-based, inclusive)." }),
					end: Type.Number({ minimum: 1, description: "Last line to return (inclusive)." }),
				}),
			),
			jsonPath: Type.Optional(
				Type.String({ description: "Dot/bracket path into the stored JSON, e.g. snapshot.nodes[3].name." }),
			),
			maxResults: Type.Optional(
				Type.Number({ minimum: 1, maximum: 200, description: "Maximum search hits to report (default 50)." }),
			),
		}),
		async execute(_toolCallId, params) {
			const result = await runOverflowQuery(storeRoot(), params, settings);
			return {
				content: [{ type: "text", text: result.text }],
				details: { error: result.error === true, truncated: result.truncated },
				isError: result.error === true,
			};
		},
	});
}