lovey/pi/extensions/budget-compaction.ts

/**
 * Budget compaction: replaces pi's summary step with one that has an enforced size limit.
 *
 * Pi's default update prompt tells the model to "PRESERVE all existing information" from the previous
 * summary, so every compaction re-emits the whole old summary plus the new work. On a long session the
 * summary grows until rewriting it no longer fits in maxTokens, and auto-compaction fails forever with
 * "generation hit the token cap". This extension instead:
 *
 * - condenses an over-limit previous summary on its own first (a short call)
 * - writes a fresh checkpoint from that + the new messages, with a word allowance per section
 * - checks the word count and runs up to two condense passes on the result if it is still over
 * - always summarizes with thinking off, so the whole output budget goes to the text
 * - falls back to pi's default compaction if a call fails
 *
 * An instruction alone is not enough: given a long previous summary, Qwen3.8-Flash-Next copied it
 * forward at ~4x the limit. Condensing a lone document to a word count is something it does follow.
 *
 * Install: copy to ~/.pi/agent/extensions/ and restart pi. Override the limit with PI_COMPACT_WORDS.
 */

import type { ExtensionAPI, ExtensionContext, FileOperations } from "@earendil-works/pi-coding-agent";
import { convertToLlm, serializeConversation } from "@earendil-works/pi-coding-agent";

const WORD_BUDGET = Number(process.env.PI_COMPACT_WORDS) || 2500;
const SLACK = 1.25;         // accept up to 25% over before running a condense pass
const MAX_CONDENSE = 2;     // condense passes on the merged checkpoint
const MAX_READ_FILES = 25;  // read-only files listed after the summary; modified files are always listed

// Share of the word budget per section, in output order.
const SECTIONS: [string, number][] = [
	["## Goal", 0.06],
	["## Constraints & Preferences", 0.12],
	["## Progress\n### Done", 0.2],
	["### In Progress", 0.1],
	["### Blocked", 0.04],
	["## Key Decisions", 0.16],
	["## Next Steps", 0.1],
	["## Critical Context", 0.22],
];

const SYSTEM_PROMPT =
	"You write compact checkpoints of coding sessions so another assistant can continue the work. " +
	"You never continue the conversation, call tools or answer questions in it; you only output the checkpoint. " +
	"You always respect the word limits you are given.";

const RULES = `Rules:
- Describe the CURRENT state, not the history. Condense finished work into short outcome bullets; merge
  related items; drop finished items that no longer matter for what comes next.
- Keep constraints, preferences and decisions that are still in force. Drop superseded ones.
- Keep exact file paths, function names, commands and error messages only where they are still needed.
- Leave out screenshot or temp-file lists, logs, tool output and code listings.
- If a part of the session is still in progress, say exactly where it stands.`;

function format(words: number): string {
	return SECTIONS.map(([h, share]) => `${h}\n(at most ${Math.max(20, Math.round(words * share))} words)`).join("\n\n");
}

function countWords(text: string): number {
	return text.split(/\s+/).filter(Boolean).length;
}

function mergePrompt(words: number, hasPrevious: boolean, focus?: string): string {
	return `Write ONE fresh checkpoint of this session${hasPrevious ? ", merging the previous checkpoint with the new messages" : ""}.

HARD LIMIT: ${words} words in total. This replaces the old checkpoint, it does not append to it.

${RULES}

Use exactly these sections, each within its word limit:

${format(words)}${focus ? `\n\nAdditional focus from the user: ${focus}` : ""}`;
}

function condensePrompt(words: number, current: number, focus?: string): string {
	return `The checkpoint above is ${current} words. Rewrite it to at most ${words} words in total.

Keep what is needed to continue the work: current state, open tasks, decisions and constraints still in
force, exact paths and names that are still used. Cut history, merge bullets, drop detail on finished work.

${RULES}

Use exactly these sections, each within its word limit:

${format(words)}${focus ? `\n\nAdditional focus from the user: ${focus}` : ""}`;
}

function fileLists(ops: FileOperations): { readFiles: string[]; modifiedFiles: string[]; text: string } {
	const modified = new Set([...ops.edited, ...ops.written]);
	const modifiedFiles = [...modified].sort();
	const readFiles = [...ops.read].filter((f) => !modified.has(f)).sort();
	const shown = readFiles.slice(0, MAX_READ_FILES);
	const more = readFiles.length - shown.length;
	let text = "";
	if (shown.length) text += `\n\n<read-files>\n${shown.join("\n")}${more > 0 ? `\n(${more} more)` : ""}\n</read-files>`;
	if (modifiedFiles.length) text += `\n\n<modified-files>\n${modifiedFiles.join("\n")}\n</modified-files>`;
	return { readFiles, modifiedFiles, text };
}

type Usage = { input: number; output: number; cacheRead: number; cacheWrite: number; totalTokens: number; [k: string]: unknown };

function addUsage(a: Usage | undefined, b: Usage | undefined): Usage | undefined {
	if (!a) return b;
	if (!b) return a;
	return {
		...a,
		input: a.input + b.input,
		output: a.output + b.output,
		cacheRead: a.cacheRead + b.cacheRead,
		cacheWrite: a.cacheWrite + b.cacheWrite,
		totalTokens: a.totalTokens + b.totalTokens,
	};
}

/** One thinking-off request. Throws on error or a reply cut off at the token cap. */
async function ask(ctx: ExtensionContext, text: string, signal: AbortSignal): Promise<{ text: string; usage?: Usage }> {
	const model = ctx.model!;
	// No `reasoning` option: thinking off. Strata gets enable_thinking=false via the chat-template compat.
	const response = await ctx.modelRegistry
		.streamSimple(
			model,
			{ systemPrompt: SYSTEM_PROMPT, messages: [{ role: "user", content: [{ type: "text", text }], timestamp: Date.now() }] },
			{ maxTokens: model.maxTokens, signal, cacheRetention: "none" },
		)
		.result();
	if (response.stopReason === "length") throw new Error("reply hit the token cap");
	if (response.stopReason === "error") throw new Error(response.errorMessage || "request failed");
	const out = response.content
		.filter((c): c is { type: "text"; text: string } => c.type === "text")
		.map((c) => c.text)
		.join("\n")
		.trim();
	if (!out) throw new Error("empty reply");
	return { text: out, usage: response.usage as Usage | undefined };
}

export default function (pi: ExtensionAPI) {
	pi.on("session_before_compact", async (event, ctx) => {
		const { preparation, customInstructions, signal } = event;
		const { messagesToSummarize, turnPrefixMessages, firstKeptEntryId, tokensBefore, fileOps } = preparation;
		if (!ctx.model) return;  // let pi handle it
		const words = WORD_BUDGET;
		const limit = Math.round(words * SLACK);
		let usage: Usage | undefined;

		try {
			let previous = preparation.previousSummary;
			ctx.ui.notify(`Compacting ${tokensBefore.toLocaleString()} tokens (limit ${words} words, thinking off)`, "info");

			// 1. A previous summary from pi's default compaction can be huge; shrink it alone first.
			if (previous && countWords(previous) > limit) {
				const before = countWords(previous);
				const r = await ask(ctx, `<checkpoint>\n${previous}\n</checkpoint>\n\n${condensePrompt(words, before)}`, signal);
				previous = r.text;
				usage = addUsage(usage, r.usage);
				ctx.ui.notify(`Condensed previous summary ${before} → ${countWords(previous)} words`, "info");
			}

			// 2. Merge the previous checkpoint with the new messages.
			let body = "";
			if (previous) body += `<previous-checkpoint>\n${previous}\n</previous-checkpoint>\n\n`;
			if (messagesToSummarize.length)
				body += `<conversation>\n${serializeConversation(convertToLlm(messagesToSummarize))}\n</conversation>\n\n`;
			if (turnPrefixMessages.length)
				body += `<current-turn-start>\n${serializeConversation(convertToLlm(turnPrefixMessages))}\n</current-turn-start>\n` +
					"(The rest of the current turn is kept verbatim after the checkpoint.)\n\n";
			const merged = await ask(ctx, body + mergePrompt(words, !!previous, customInstructions), signal);
			let summary = merged.text;
			usage = addUsage(usage, merged.usage);

			// 3. Enforce the limit on the result.
			for (let pass = 1; pass <= MAX_CONDENSE && countWords(summary) > limit; pass++) {
				const before = countWords(summary);
				const r = await ask(ctx, `<checkpoint>\n${summary}\n</checkpoint>\n\n${condensePrompt(words, before, customInstructions)}`, signal);
				summary = r.text;
				usage = addUsage(usage, r.usage);
				ctx.ui.notify(`Condensed checkpoint ${before} → ${countWords(summary)} words`, "info");
			}
			if (signal.aborted) return;
			if (countWords(summary) > limit)
				ctx.ui.notify(`Checkpoint is ${countWords(summary)} words, over the ${words}-word limit`, "warning");

			const files = fileLists(fileOps);
			return {
				compaction: {
					summary: summary + files.text,
					firstKeptEntryId,
					tokensBefore,
					usage,
					details: { readFiles: files.readFiles, modifiedFiles: files.modifiedFiles },
				},
			};
		} catch (error) {
			if (signal.aborted) return;
			ctx.ui.notify(`Budget compaction failed (${error instanceof Error ? error.message : error}); using pi's default`, "error");
			return;
		}
	});
}