dem-compact.index.ts

/**
 * demi-compact: a poor man's compaction.
 *
 * Instead of summarizing history into a single block, it truncates tool
 * call inputs (to 80 chars) and drops tool result outputs entirely from
 * messages older than the cutoff recorded when the user runs
 * `/demi-compact`. The fact that a tool was called is preserved, but the
 * verbose I/O that dominates token usage is stripped. User/assistant
 * text remains fully intact.
 *
 * The cutoff is persisted as a custom session entry, so it survives
 * restarts and re-applies on every LLM call via the `context` event.
 */

import type { ExtensionAPI } from "@mariozechner/pi-coding-agent";

type DemiCompactState = { cutoff: number };

const ENTRY_TYPE = "demi-compact-cutoff";

export default function (pi: ExtensionAPI) {
	let cutoff: number | undefined;

	const restoreCutoff = (entries: ReadonlyArray<any>) => {
		cutoff = undefined;
		// Walk the active branch newest-first, latest cutoff wins.
		for (let i = entries.length - 1; i >= 0; i--) {
			const e = entries[i];
			if (e?.type === "custom" && e.customType === ENTRY_TYPE) {
				const data = e.data as DemiCompactState | undefined;
				if (data && typeof data.cutoff === "number") {
					cutoff = data.cutoff;
					return;
				}
			}
		}
	};

	pi.on("session_start", async (_event, ctx) => {
		restoreCutoff(ctx.sessionManager.getBranch());
	});

	pi.on("context", async (event) => {
		if (cutoff === undefined) return;
		const cut = cutoff;

		const messages: any[] = [];
		for (const m of event.messages as any[]) {
			const ts: number | undefined =
				typeof m?.timestamp === "number" ? m.timestamp : undefined;
			const isOlder = ts !== undefined && ts < cut;
			if (!isOlder) {
				messages.push(m);
				continue;
			}

			// Replace tool result content with a tiny stub. We keep the
			// message itself so the assistant's toolCall has a matching
			// toolResult (providers require this pairing).
			if (m.role === "toolResult") {
				messages.push({
					...m,
					content: [{ type: "text", text: "[demi-compact: output elided]" }],
				});
				continue;
			}

			// For assistant messages, truncate toolCall arguments and drop thinking.
			if (m.role === "assistant" && Array.isArray(m.content)) {
				const transformed = m.content
					.filter((c: any) => c?.type !== "thinking")
					.map((c: any) => {
						if (c?.type !== "toolCall") return c;
						let argsStr: string;
						try {
							if (typeof c.arguments === "string") {
								argsStr = c.arguments;
							} else if (c.arguments === undefined || c.arguments === null) {
								argsStr = "";
							} else {
								argsStr = JSON.stringify(c.arguments) ?? "";
							}
						} catch {
							argsStr = String(c.arguments ?? "");
						}
						if (argsStr.length > 80) {
							argsStr = argsStr.slice(0, 80) + "…";
						}
						return { ...c, arguments: { _elided: argsStr } };
					});
				messages.push({ ...m, content: transformed });
				continue;
			}

			messages.push(m);
		}

		return { messages };
	});

	pi.registerCommand("demi-compact", {
		description:
			"Strip tool calls and tool results from history older than now (poor man's compaction)",
		handler: async (_args, ctx) => {
			const now = Date.now();
			cutoff = now;
			pi.appendEntry(ENTRY_TYPE, { cutoff: now } satisfies DemiCompactState);

			// Quick stats for user feedback.
			let stripped = 0;
			let droppedResults = 0;
			for (const e of ctx.sessionManager.getBranch()) {
				if (e?.type !== "message") continue;
				const m: any = e.message;
				const ts: number | undefined =
					typeof m?.timestamp === "number" ? m.timestamp : undefined;
				if (ts === undefined || ts >= now) continue;
				if (m.role === "toolResult") {
					droppedResults++;
				} else if (m.role === "assistant" && Array.isArray(m.content)) {
					if (m.content.some((c: any) => c?.type === "toolCall")) stripped++;
				}
			}

			ctx.ui.notify(
				`demi-compact: truncating ${stripped} assistant tool-call message(s) and eliding ${droppedResults} tool result(s) in future LLM calls.`,
				"info",
			);
		},
	});
}