feat(coding-agent): add extension compaction helpers

packages/coding-agent/src/core/agent-session.ts

@@ -33,6 +33,7 @@ import {
 	calculateContextTokens,
 	collectEntriesForBranchSummary,
 	compact,
+	estimateContextTokens,
 	generateBranchSummary,
 	prepareCompaction,
 	shouldCompact,
@@ -40,6 +41,7 @@ import {
 import { exportSessionToHtml, type ToolHtmlRenderer } from "./export-html/index.js";
 import { createToolHtmlRenderer } from "./export-html/tool-renderer.js";
 import type {
+	ContextUsage,
 	ExtensionRunner,
 	InputSource,
 	SessionBeforeCompactResult,
@@ -2250,6 +2252,26 @@ export class AgentSession {
 		};
 	}
 
+	getContextUsage(): ContextUsage | undefined {
+		const model = this.model;
+		if (!model) return undefined;
+
+		const contextWindow = model.contextWindow ?? 0;
+		if (contextWindow <= 0) return undefined;
+
+		const estimate = estimateContextTokens(this.messages);
+		const percent = (estimate.tokens / contextWindow) * 100;
+
+		return {
+			tokens: estimate.tokens,
+			contextWindow,
+			percent,
+			usageTokens: estimate.usageTokens,
+			trailingTokens: estimate.trailingTokens,
+			lastUsageIndex: estimate.lastUsageIndex,
+		};
+	}
+
 	/**
 	 * Export session to HTML.
 	 * @param outputPath Optional output path (defaults to session directory)

packages/coding-agent/src/core/compaction/compaction.ts

@@ -8,7 +8,12 @@
 import type { AgentMessage } from "@mariozechner/pi-agent-core";
 import type { AssistantMessage, Model, Usage } from "@mariozechner/pi-ai";
 import { completeSimple } from "@mariozechner/pi-ai";
-import { convertToLlm, createBranchSummaryMessage, createCustomMessage } from "../messages.js";
+import {
+	convertToLlm,
+	createBranchSummaryMessage,
+	createCompactionSummaryMessage,
+	createCustomMessage,
+} from "../messages.js";
 import type { CompactionEntry, SessionEntry } from "../session-manager.js";
 import {
 	computeFileLists,
@@ -81,6 +86,9 @@ function getMessageFromEntry(entry: SessionEntry): AgentMessage | undefined {
 	if (entry.type === "branch_summary") {
 		return createBranchSummaryMessage(entry.summary, entry.fromId, entry.timestamp);
 	}
+	if (entry.type === "compaction") {
+		return createCompactionSummaryMessage(entry.summary, entry.tokensBefore, entry.timestamp);
+	}
 	return undefined;
 }
 
@@ -149,6 +157,55 @@ export function getLastAssistantUsage(entries: SessionEntry[]): Usage | undefine
 	return undefined;
 }
 
+export interface ContextUsageEstimate {
+	tokens: number;
+	usageTokens: number;
+	trailingTokens: number;
+	lastUsageIndex: number | null;
+}
+
+function getLastAssistantUsageInfo(messages: AgentMessage[]): { usage: Usage; index: number } | undefined {
+	for (let i = messages.length - 1; i >= 0; i--) {
+		const usage = getAssistantUsage(messages[i]);
+		if (usage) return { usage, index: i };
+	}
+	return undefined;
+}
+
+/**
+ * Estimate context tokens from messages, using the last assistant usage when available.
+ * If there are messages after the last usage, estimate their tokens with estimateTokens.
+ */
+export function estimateContextTokens(messages: AgentMessage[]): ContextUsageEstimate {
+	const usageInfo = getLastAssistantUsageInfo(messages);
+
+	if (!usageInfo) {
+		let estimated = 0;
+		for (const message of messages) {
+			estimated += estimateTokens(message);
+		}
+		return {
+			tokens: estimated,
+			usageTokens: 0,
+			trailingTokens: estimated,
+			lastUsageIndex: null,
+		};
+	}
+
+	const usageTokens = calculateContextTokens(usageInfo.usage);
+	let trailingTokens = 0;
+	for (let i = usageInfo.index + 1; i < messages.length; i++) {
+		trailingTokens += estimateTokens(messages[i]);
+	}
+
+	return {
+		tokens: usageTokens + trailingTokens,
+		usageTokens,
+		trailingTokens,
+		lastUsageIndex: usageInfo.index,
+	};
+}
+
 /**
  * Check if compaction should trigger based on context usage.
  */
@@ -555,8 +612,13 @@ export function prepareCompaction(
 	const boundaryStart = prevCompactionIndex + 1;
 	const boundaryEnd = pathEntries.length;
 
-	const lastUsage = getLastAssistantUsage(pathEntries);
-	const tokensBefore = lastUsage ? calculateContextTokens(lastUsage) : 0;
+	const usageStart = prevCompactionIndex >= 0 ? prevCompactionIndex : 0;
+	const usageMessages: AgentMessage[] = [];
+	for (let i = usageStart; i < boundaryEnd; i++) {
+		const msg = getMessageFromEntry(pathEntries[i]);
+		if (msg) usageMessages.push(msg);
+	}
+	const tokensBefore = estimateContextTokens(usageMessages).tokens;
 
 	const cutPoint = findCutPoint(pathEntries, boundaryStart, boundaryEnd, settings.keepRecentTokens);
 

packages/coding-agent/src/core/extensions/index.ts

@@ -28,10 +28,13 @@ export type {
 	BashToolResultEvent,
 	BeforeAgentStartEvent,
 	BeforeAgentStartEventResult,
+	// Context
+	CompactOptions,
 	// Events - Agent
 	ContextEvent,
 	// Event Results
 	ContextEventResult,
+	ContextUsage,
 	CustomToolResultEvent,
 	EditToolResultEvent,
 	ExecOptions,
@@ -42,7 +45,6 @@ export type {
 	ExtensionAPI,
 	ExtensionCommandContext,
 	ExtensionCommandContextActions,
-	// Context
 	ExtensionContext,
 	ExtensionContextActions,
 	// Errors

packages/coding-agent/src/core/extensions/runner.ts

@@ -11,8 +11,10 @@ import type { SessionManager } from "../session-manager.js";
 import type {
 	BeforeAgentStartEvent,
 	BeforeAgentStartEventResult,
+	CompactOptions,
 	ContextEvent,
 	ContextEventResult,
+	ContextUsage,
 	Extension,
 	ExtensionActions,
 	ExtensionCommandContext,
@@ -113,6 +115,8 @@ export class ExtensionRunner {
 	private waitForIdleFn: () => Promise<void> = async () => {};
 	private abortFn: () => void = () => {};
 	private hasPendingMessagesFn: () => boolean = () => false;
+	private getContextUsageFn: () => ContextUsage | undefined = () => undefined;
+	private compactFn: (options?: CompactOptions) => void = () => {};
 	private newSessionHandler: NewSessionHandler = async () => ({ cancelled: false });
 	private forkHandler: ForkHandler = async () => ({ cancelled: false });
 	private navigateTreeHandler: NavigateTreeHandler = async () => ({ cancelled: false });
@@ -158,6 +162,8 @@ export class ExtensionRunner {
 		this.abortFn = contextActions.abort;
 		this.hasPendingMessagesFn = contextActions.hasPendingMessages;
 		this.shutdownHandler = contextActions.shutdown;
+		this.getContextUsageFn = contextActions.getContextUsage;
+		this.compactFn = contextActions.compact;
 
 		// Command context actions (optional, only for interactive mode)
 		if (commandContextActions) {
@@ -337,6 +343,8 @@ export class ExtensionRunner {
 			abort: () => this.abortFn(),
 			hasPendingMessages: () => this.hasPendingMessagesFn(),
 			shutdown: () => this.shutdownHandler(),
+			getContextUsage: () => this.getContextUsageFn(),
+			compact: (options) => this.compactFn(options),
 		};
 	}
 

packages/coding-agent/src/core/extensions/types.ts

@@ -193,6 +193,21 @@ export interface ExtensionUIContext {
 // Extension Context
 // ============================================================================
 
+export interface ContextUsage {
+	tokens: number;
+	contextWindow: number;
+	percent: number;
+	usageTokens: number;
+	trailingTokens: number;
+	lastUsageIndex: number | null;
+}
+
+export interface CompactOptions {
+	customInstructions?: string;
+	onComplete?: (result: CompactionResult) => void;
+	onError?: (error: Error) => void;
+}
+
 /**
  * Context passed to extension event handlers.
  */
@@ -217,6 +232,10 @@ export interface ExtensionContext {
 	hasPendingMessages(): boolean;
 	/** Gracefully shutdown pi and exit. Available in all contexts. */
 	shutdown(): void;
+	/** Get current context usage for the active model. */
+	getContextUsage(): ContextUsage | undefined;
+	/** Trigger compaction without awaiting completion. */
+	compact(options?: CompactOptions): void;
 }
 
 /**
@@ -919,6 +938,8 @@ export interface ExtensionContextActions {
 	abort: () => void;
 	hasPendingMessages: () => boolean;
 	shutdown: () => void;
+	getContextUsage: () => ContextUsage | undefined;
+	compact: (options?: CompactOptions) => void;
 }
 
 /**