@andares/pi-toolkits 0.3.1 → 0.4.1

package/src/features/compact/index.ts

@@ -0,0 +1,371 @@
+/**
+ * compact-model feature — dedicated model for context compaction.
+ *
+ * pi's /compact, auto-compact (threshold), and overflow-recovery compaction
+ * all summarize with the CURRENT session model by default. This feature lets
+ * a cheaper/faster model take over every summarization while the session
+ * model stays focused on the conversation:
+ *
+ *  - takeover point: the `session_before_compact` extension event, which pi
+ *    emits for ALL three trigger reasons (manual / threshold / overflow)
+ *    before running its own compaction
+ *  - prompts, message assembly, split-turn handling, output budgets, and the
+ *    summary file-list tail are verbatim pi 0.87.1 built-in behavior — only
+ *    the executing model differs
+ *  - lossless fallback: unset config, unavailable model, empty/failed
+ *    summary, or user abort → return undefined and pi's default compaction
+ *    (session model) runs untouched. The extension can never make
+ *    compaction worse than stock.
+ *
+ * Config lives in `<agent-dir>/pi-toolkits-compact.json` ("provider/id",
+ * see store.ts). /compact-model views/changes/clears it; auto-compact paths
+ * never open dialogs (handler only uses notify).
+ */
+import type {
+	ExtensionAPI,
+	ExtensionContext,
+	ModelRegistry,
+	SessionBeforeCompactEvent,
+	SessionBeforeCompactResult,
+	SessionStartEvent,
+	Theme,
+} from "@earendil-works/pi-coding-agent";
+import {
+	convertToLlm,
+	serializeConversation,
+} from "@earendil-works/pi-coding-agent";
+import { uuidv7 } from "@earendil-works/pi-ai";
+import {
+	modelKey,
+	resolveConfiguredModel,
+	type AvailableModel,
+} from "../../lib/model-config.js";
+import { COMPACT_STATUS_KEY } from "./constants.js";
+import {
+	buildHistoryPromptText,
+	buildTurnPrefixPromptText,
+	combineUsage,
+	computeFileLists,
+	formatFileOperations,
+	SUMMARIZATION_SYSTEM_PROMPT,
+	type UsageLike,
+} from "./prompt.js";
+import { getCompactStore, type CompactStore } from "./store.js";
+
+/** Usage record type returned by ModelRegistry.complete(). */
+type CompleteUsage = NonNullable<
+	Awaited<ReturnType<ModelRegistry["complete"]>>["usage"]
+>;
+
+/** Thrown when a summarization response cannot safely become a checkpoint. */
+class SummarizationFailureError extends Error {}
+
+/**
+ * One-shot summarization call through the configured compact model,
+ * mirroring pi's completeSummarization options: signal honored, no prompt
+ * cache writes, fresh routing session id.
+ */
+async function completeSummary(
+	ctx: ExtensionContext,
+	model: AvailableModel,
+	promptText: string,
+	maxTokens: number,
+	signal: AbortSignal,
+): Promise<{ text: string; usage: CompleteUsage }> {
+	const response = await ctx.modelRegistry.complete(
+		model,
+		{
+			systemPrompt: SUMMARIZATION_SYSTEM_PROMPT,
+			messages: [
+				{
+					role: "user",
+					content: [{ type: "text", text: promptText }],
+					timestamp: Date.now(),
+				},
+			],
+		},
+		{
+			maxTokens,
+			signal,
+			cacheRetention: "none",
+			sessionId: uuidv7(),
+		},
+	);
+	// Same safety gates as the built-in compaction: an errored or
+	// length-capped response is not a usable checkpoint, and a summarizer
+	// that tries to call tools has gone off the rails.
+	if (response.stopReason === "error" || response.stopReason === "length") {
+		throw new SummarizationFailureError(
+			response.stopReason === "error"
+				? `Summarization failed: ${response.errorMessage ?? "unknown error"}`
+				: "Summarization failed: generation hit the token cap and the summary is incomplete",
+		);
+	}
+	if (response.content.some((block) => block.type === "toolCall")) {
+		throw new SummarizationFailureError(
+			"Summarization attempted to call a tool",
+		);
+	}
+	const text = response.content
+		.flatMap((block) => (block.type === "text" ? [block.text] : []))
+		.join("\n");
+	return { text, usage: response.usage };
+}
+
+/**
+ * pi's output-budget formula (0.87.1 compaction.js): a fraction of
+ * reserveTokens clamped by the model's own output cap.
+ */
+function summarizationBudget(
+	fraction: number,
+	reserveTokens: number,
+	modelMaxTokens: number,
+): number {
+	return Math.min(
+		Math.floor(fraction * reserveTokens),
+		modelMaxTokens > 0 ? modelMaxTokens : Number.POSITIVE_INFINITY,
+	);
+}
+
+/**
+ * Run the full compaction summarization with the configured model,
+ * replicating pi 0.87.1's compact(): history summary (first-connection or
+ * iterative-update prompt, custom focus appended), plus a separate
+ * turn-prefix checkpoint when the cut point splits a turn, merged into one
+ * summary with the file-list tail. Returns undefined on any failure after
+ * notifying (pi default compaction then takes over).
+ */
+async function runCompactModelCompaction(
+	ctx: ExtensionContext,
+	event: SessionBeforeCompactEvent,
+	model: AvailableModel,
+): Promise<SessionBeforeCompactResult | undefined> {
+	const { preparation, signal, customInstructions } = event;
+	const {
+		messagesToSummarize,
+		turnPrefixMessages,
+		isSplitTurn,
+		tokensBefore,
+		firstKeptEntryId,
+		previousSummary,
+		fileOps,
+		settings,
+	} = preparation;
+
+	try {
+		let summaryText: string;
+		let usage: CompleteUsage | undefined;
+
+		if (isSplitTurn && turnPrefixMessages.length > 0) {
+			// Split turn: summarize history (if any) and the turn prefix
+			// separately, then merge — exactly like the built-in compact().
+			let historyText = previousSummary ?? "No prior history.";
+			if (messagesToSummarize.length > 0) {
+				const history = await completeSummary(
+					ctx,
+					model,
+					buildHistoryPromptText(
+						serializeConversation(convertToLlm(messagesToSummarize)),
+						previousSummary,
+						customInstructions,
+					),
+					summarizationBudget(
+						0.8,
+						settings.reserveTokens,
+						model.maxTokens,
+					),
+					signal,
+				);
+				historyText = history.text;
+				usage = history.usage;
+			}
+			const turnPrefix = await completeSummary(
+				ctx,
+				model,
+				buildTurnPrefixPromptText(
+					serializeConversation(convertToLlm(turnPrefixMessages)),
+				),
+				summarizationBudget(0.5, settings.reserveTokens, model.maxTokens),
+				signal,
+			);
+			summaryText = `${historyText}\n\n---\n\n**Turn Context (split turn):**\n\n${turnPrefix.text}`;
+			usage = usage
+				? combineUsage(
+						usage as UsageLike,
+						turnPrefix.usage as UsageLike,
+					) as CompleteUsage
+				: turnPrefix.usage;
+		} else {
+			const history = await completeSummary(
+				ctx,
+				model,
+				buildHistoryPromptText(
+					serializeConversation(convertToLlm(messagesToSummarize)),
+					previousSummary,
+					customInstructions,
+				),
+				summarizationBudget(0.8, settings.reserveTokens, model.maxTokens),
+				signal,
+			);
+			summaryText = history.text;
+			usage = history.usage;
+		}
+
+		if (!summaryText.trim()) {
+			if (!signal.aborted) {
+				ctx.ui.notify(
+					"compact-model: 压缩摘要为空，已回落 pi 默认压缩。",
+					"warning",
+				);
+			}
+			return undefined;
+		}
+
+		// File-list tail + details carry (pi does not auto-carry for
+		// fromHook compactions — see plan), matching CompactionDetails shape.
+		const { readFiles, modifiedFiles } = computeFileLists(fileOps);
+		summaryText += formatFileOperations(readFiles, modifiedFiles);
+
+		return {
+			compaction: {
+				summary: summaryText,
+				firstKeptEntryId,
+				tokensBefore,
+				usage,
+				details: { readFiles, modifiedFiles },
+			},
+		};
+	} catch (error) {
+		// User abort is a clean cancel — pi handles the abort messaging.
+		if (signal.aborted) return undefined;
+		const message =
+			error instanceof Error ? error.message : String(error);
+		ctx.ui.notify(
+			`compact-model: 压缩摘要失败（${message}），已回落 pi 默认压缩。`,
+			"error",
+		);
+		return undefined;
+	}
+}
+
+/**
+ * Let the user pick a compact model (or clear the config) and persist the
+ * choice. Mirrors review's pickReviewModel, plus a leading clear option.
+ */
+export async function pickCompactModel(
+	ctx: ExtensionContext,
+	store: CompactStore,
+): Promise<AvailableModel | undefined> {
+	if (!ctx.hasUI) {
+		ctx.ui.notify(
+			"compact-model: 当前环境不支持模型选择对话框，请编辑 <agent-dir>/pi-toolkits-compact.json 手动配置。",
+			"error",
+		);
+		return undefined;
+	}
+	const available = ctx.modelRegistry.getAvailable();
+	if (available.length === 0) {
+		ctx.ui.notify(
+			"compact-model: 当前没有可用模型，请先在 pi 中配置模型。",
+			"error",
+		);
+		return undefined;
+	}
+	const configured = store.getModel();
+	const clearLabel = "✕ 清除配置（用 pi 默认压缩）";
+	const options = [
+		clearLabel,
+		...available.map(
+			(m: AvailableModel) => `${m.name} (${m.provider}/${m.id})`,
+		),
+	];
+	const choice = await ctx.ui.select(
+		configured
+			? `选择压缩模型 (当前: ${configured})`
+			: "选择压缩模型 (compact-model)",
+		options,
+	);
+	if (choice === undefined) {
+		ctx.ui.notify("compact-model: 已取消，压缩模型配置未变更。", "info");
+		return undefined;
+	}
+	if (choice === clearLabel) {
+		store.setModel(undefined);
+		ctx.ui.notify(
+			"compact-model: 已清除配置，压缩恢复 pi 默认（当前会话模型）。",
+			"info",
+		);
+		return undefined;
+	}
+	const model = available[options.indexOf(choice) - 1];
+	if (!model) return undefined;
+	store.setModel(modelKey(model));
+	ctx.ui.notify(
+		`compact-model: 已设置压缩模型「${model.name}」，后续压缩总结将由它执行。`,
+		"info",
+	);
+	return model;
+}
+
+export function registerCompact(pi: ExtensionAPI): void {
+	// Take over every compaction trigger (manual / threshold / overflow)
+	// when a compact model is configured; otherwise stay invisible.
+	pi.on(
+		"session_before_compact",
+		async (
+			event: SessionBeforeCompactEvent,
+			ctx: ExtensionContext,
+		): Promise<SessionBeforeCompactResult | undefined> => {
+			const configured = getCompactStore().getModel();
+			if (!configured) return undefined;
+			const model = resolveConfiguredModel(ctx, configured);
+			if (!model) {
+				ctx.ui.notify(
+					`compact-model: 配置的压缩模型 ${configured} 不存在或不可用，已回落 pi 默认压缩。`,
+					"warning",
+				);
+				return undefined;
+			}
+			return runCompactModelCompaction(ctx, event, model);
+		},
+	);
+
+	// Config entry point: view, change, or clear the compact model.
+	pi.registerCommand("compact-model", {
+		description:
+			"Set, change, or clear the dedicated model used for /compact summaries",
+		handler: async (_args: string, ctx: ExtensionContext) => {
+			await pickCompactModel(ctx, getCompactStore());
+		},
+	});
+
+	// Footer indicator: muted "compact" while a compact model is configured.
+	let theme: Theme | undefined;
+	let latestCtx: ExtensionContext | undefined;
+
+	const updateFooterStatus = (ctx: ExtensionContext): void => {
+		ctx.ui.setStatus(
+			COMPACT_STATUS_KEY,
+			getCompactStore().getModel() && theme
+				? theme.fg("muted", "compact")
+				: undefined,
+		);
+	};
+
+	pi.on("session_start", (_event: SessionStartEvent, ctx: ExtensionContext) => {
+		latestCtx = ctx;
+		try {
+			const t = ctx.ui.theme;
+			if (t) {
+				theme = t;
+			}
+		} catch {
+			// No theme in this mode — footer indicator stays off.
+		}
+		updateFooterStatus(ctx);
+	});
+
+	getCompactStore().subscribe(() => {
+		if (latestCtx) updateFooterStatus(latestCtx);
+	});
+}

package/src/features/compact/prompt.ts

@@ -0,0 +1,240 @@
+/**
+ * compact-model feature — summarization prompts.
+ *
+ * Every prompt below is copied VERBATIM from pi 0.87.1's built-in
+ * compaction (dist/core/compaction/compaction.js and utils.js) so that
+ * summaries produced by the configured compact model are formatted exactly
+ * like pi's own compaction output. Only the executing model differs.
+ *
+ * Update these strings whenever pi changes its built-in prompts (diff
+ * dist/core/compaction/compaction.js between versions).
+ */
+
+/** System prompt for every summarization call (pi 0.87.1 utils.js). */
+export const SUMMARIZATION_SYSTEM_PROMPT =
+	"You are a context summarization assistant. Your task is to read a conversation between a user and an AI assistant, then produce a structured summary following the exact format specified.\n\nDo NOT continue the conversation. Do NOT respond to any questions in the conversation. ONLY output the structured summary.";
+
+/** First-connection history summary prompt (pi 0.87.1 compaction.js). */
+export const SUMMARIZATION_PROMPT = `The messages above are a conversation to summarize. Create a structured context checkpoint summary that another LLM will use to continue the work.
+
+Use this EXACT format:
+
+## Goal
+[What is the user trying to accomplish? Can be multiple items if the session covers different tasks.]
+
+## Constraints & Preferences
+- [Any constraints, preferences, or requirements mentioned by user]
+- [Or "(none)" if none were mentioned]
+
+## Progress
+### Done
+- [x] [Completed tasks/changes]
+
+### In Progress
+- [ ] [Current work]
+
+### Blocked
+- [Issues preventing progress, if any]
+
+## Key Decisions
+- **[Decision]**: [Brief rationale]
+
+## Next Steps
+1. [Ordered list of what should happen next]
+
+## Critical Context
+- [Any data, examples, or references needed to continue]
+- [Or "(none)" if not applicable]
+
+Keep each section concise. Preserve exact file paths, function names, and error messages.`;
+
+/**
+ * Iterative-update rules, used when a previous summary exists
+ * (pi 0.87.1 compaction.js).
+ */
+export const UPDATE_SUMMARIZATION_INSTRUCTIONS = `Update the existing structured summary with new information. RULES:
+- PRESERVE all existing information from the previous summary
+- ADD new progress, decisions, and context from the new messages
+- UPDATE the Progress section: move items from "In Progress" to "Done" when completed
+- UPDATE "Next Steps" based on what was accomplished
+- PRESERVE exact file paths, function names, and error messages
+- If something is no longer relevant, you may remove it
+
+Use this EXACT format:
+
+## Goal
+[Preserve existing goals, add new ones if the task expanded]
+
+## Constraints & Preferences
+- [Preserve existing, add new ones discovered]
+
+## Progress
+### Done
+- [x] [Include previously done items AND newly completed items]
+
+### In Progress
+- [ ] [Current work - update based on progress]
+
+### Blocked
+- [Current blockers - remove if resolved]
+
+## Key Decisions
+- **[Decision]**: [Brief rationale] (preserve all previous, add new)
+
+## Next Steps
+1. [Update based on current state]
+
+## Critical Context
+- [Preserve important context, add new if needed]
+
+Keep each section concise. Preserve exact file paths, function names, and error messages.`;
+
+/** Full iterative-update prompt (leading sentence + rules, pi 0.87.1). */
+export const UPDATE_SUMMARIZATION_PROMPT = `The messages above are NEW conversation messages to incorporate into the existing summary provided in <previous-summary> tags.
+
+${UPDATE_SUMMARIZATION_INSTRUCTIONS}`;
+
+/** Split-turn prefix checkpoint prompt (pi 0.87.1 compaction.js). */
+export const TURN_PREFIX_SUMMARIZATION_PROMPT = `The messages above are earlier context from an ongoing conversation. Later messages are stored separately and do not need to be reconstructed.
+
+Create a concise checkpoint of the user's request and the progress shown above. This checkpoint will be placed before the later messages so the conversation can continue with the necessary context.
+
+## Original Request
+[What did the user ask for?]
+
+## Progress So Far
+- [Key decisions and work completed in these messages]
+
+## Context Needed to Continue
+- [Information from these messages needed to understand the later work]
+
+Only summarize information explicitly present above. Do not infer or recreate later messages.`;
+
+/**
+ * Build the history-summary prompt text (pi 0.87.1 generateSummaryWithUsage):
+ * conversation wrapped in tags, optional previous summary, then the base
+ * prompt with the custom focus appended (`Additional focus: ...` — the exact
+ * phrase pi's built-in compaction uses for /compact <instructions>).
+ */
+export function buildHistoryPromptText(
+	conversationText: string,
+	previousSummary: string | undefined,
+	customInstructions: string | undefined,
+): string {
+	let basePrompt = previousSummary
+		? UPDATE_SUMMARIZATION_PROMPT
+		: SUMMARIZATION_PROMPT;
+	if (customInstructions) {
+		basePrompt = `${basePrompt}\n\nAdditional focus: ${customInstructions}`;
+	}
+	let promptText = `<conversation>\n${conversationText}\n</conversation>\n\n`;
+	if (previousSummary) {
+		promptText += `<previous-summary>\n${previousSummary}\n</previous-summary>\n\n`;
+	}
+	promptText += basePrompt;
+	return promptText;
+}
+
+/**
+ * Build the split-turn prefix prompt text (pi 0.87.1
+ * generateTurnPrefixSummary): `# Conversation` block followed by
+ * `# Instructions`.
+ */
+export function buildTurnPrefixPromptText(conversationText: string): string {
+	return `# Conversation\n${conversationText}\n\n# Instructions\n${TURN_PREFIX_SUMMARIZATION_PROMPT}`;
+}
+
+/** File-operation sets extracted by pi's prepareCompaction. */
+export interface FileOpsLike {
+	read: Set<string>;
+	written: Set<string>;
+	edited: Set<string>;
+}
+
+/**
+ * Compute read/modified file lists from compaction file operations
+ * (pi 0.87.1 utils.js computeFileLists): modified = written ∪ edited,
+ * read-only = read minus modified, both sorted.
+ */
+export function computeFileLists(fileOps: FileOpsLike): {
+	readFiles: string[];
+	modifiedFiles: string[];
+} {
+	const modified = new Set([...fileOps.edited, ...fileOps.written]);
+	// pi-lens-ignore: no-sort-without-comparator, — 逐字对齐 pi 0.87.1 的裸 .sort()：文件路径按码元序排序（加 locale 比较器反而会改变顺序）
+	const readOnly = [...fileOps.read]
+		.filter((f) => !modified.has(f))
+		.sort();
+	// pi-lens-ignore: no-sort-without-comparator, — 同上（pi 0.87.1 computeFileLists 原文）
+	const modifiedFiles = [...modified].sort();
+	return { readFiles: readOnly, modifiedFiles };
+}
+
+/**
+ * Format file lists as XML tags appended to the summary (pi 0.87.1 utils.js
+ * formatFileOperations): sections joined by a blank line, prefixed by one.
+ */
+export function formatFileOperations(
+	readFiles: string[],
+	modifiedFiles: string[],
+): string {
+	const sections: string[] = [];
+	if (readFiles.length > 0) {
+		sections.push(`<read-files>\n${readFiles.join("\n")}\n</read-files>`);
+	}
+	if (modifiedFiles.length > 0) {
+		sections.push(
+			`<modified-files>\n${modifiedFiles.join("\n")}\n</modified-files>`,
+		);
+	}
+	if (sections.length === 0) return "";
+	return `\n\n${sections.join("\n\n")}`;
+}
+
+/**
+ * Combine two provider usage records field-wise (pi 0.87.1 compaction.js
+ * combineUsage): used when a split turn produces two summarization calls.
+ */
+export function combineUsage<T extends UsageLike>(first: T, second: T): T {
+	return {
+		...first,
+		input: first.input + second.input,
+		output: first.output + second.output,
+		cacheRead: first.cacheRead + second.cacheRead,
+		cacheWrite: first.cacheWrite + second.cacheWrite,
+		...(first.cacheWrite1h !== undefined || second.cacheWrite1h !== undefined
+			? {
+					cacheWrite1h: (first.cacheWrite1h ?? 0) + (second.cacheWrite1h ?? 0),
+				}
+			: {}),
+		...(first.reasoning !== undefined || second.reasoning !== undefined
+			? { reasoning: (first.reasoning ?? 0) + (second.reasoning ?? 0) }
+			: {}),
+		totalTokens: first.totalTokens + second.totalTokens,
+		cost: {
+			input: first.cost.input + second.cost.input,
+			output: first.cost.output + second.cost.output,
+			cacheRead: first.cost.cacheRead + second.cost.cacheRead,
+			cacheWrite: first.cost.cacheWrite + second.cost.cacheWrite,
+			total: first.cost.total + second.cost.total,
+		},
+	} as T;
+}
+
+/** Structural view of pi's Usage record (the fields combineUsage touches). */
+export interface UsageLike {
+	input: number;
+	output: number;
+	cacheRead: number;
+	cacheWrite: number;
+	cacheWrite1h?: number;
+	reasoning?: number;
+	totalTokens: number;
+	cost: {
+		input: number;
+		output: number;
+		cacheRead: number;
+		cacheWrite: number;
+		total: number;
+	};
+}