From 9d3f8117a4259e979ba08ca04ce9ef05234935db Mon Sep 17 00:00:00 2001 From: Mario Zechner Date: Sat, 17 Jan 2026 11:39:46 +0100 Subject: [PATCH] feat(coding-agent): add extension compaction helpers --- .../examples/extensions/README.md | 1 + .../examples/extensions/trigger-compact.ts | 32 +++++++++ .../coding-agent/src/core/agent-session.ts | 22 ++++++ .../src/core/compaction/compaction.ts | 68 ++++++++++++++++++- .../coding-agent/src/core/extensions/index.ts | 4 +- .../src/core/extensions/runner.ts | 8 +++ .../coding-agent/src/core/extensions/types.ts | 21 ++++++ packages/coding-agent/src/index.ts | 2 + packages/coding-agent/src/modes/print-mode.ts | 12 ++++ .../coding-agent/src/modes/rpc/rpc-mode.ts | 12 ++++ .../test/compaction-extensions.test.ts | 12 ++++ 11 files changed, 190 insertions(+), 4 deletions(-) create mode 100644 packages/coding-agent/examples/extensions/trigger-compact.ts diff --git a/packages/coding-agent/examples/extensions/README.md b/packages/coding-agent/examples/extensions/README.md index f4154de9..5ccbc0ae 100644 --- a/packages/coding-agent/examples/extensions/README.md +++ b/packages/coding-agent/examples/extensions/README.md @@ -77,6 +77,7 @@ cp permission-gate.ts ~/.pi/agent/extensions/ | `pirate.ts` | Demonstrates `systemPromptAppend` to dynamically modify system prompt | | `claude-rules.ts` | Scans `.claude/rules/` folder and lists rules in system prompt | | `custom-compaction.ts` | Custom compaction that summarizes entire conversation | +| `trigger-compact.ts` | Triggers compaction when context usage exceeds 100k tokens and adds `/trigger-compact` command | ### System Integration diff --git a/packages/coding-agent/examples/extensions/trigger-compact.ts b/packages/coding-agent/examples/extensions/trigger-compact.ts new file mode 100644 index 00000000..9d12f9e5 --- /dev/null +++ b/packages/coding-agent/examples/extensions/trigger-compact.ts @@ -0,0 +1,32 @@ +import type { ExtensionAPI, ExtensionContext } from "@mariozechner/pi-coding-agent"; + +const COMPACT_THRESHOLD_TOKENS = 100_000; + +export default function (pi: ExtensionAPI) { + const triggerCompaction = (ctx: ExtensionContext, customInstructions?: string) => { + ctx.compact({ + customInstructions, + onError: (error) => { + if (ctx.hasUI) { + ctx.ui.notify(`Compaction failed: ${error.message}`, "error"); + } + }, + }); + }; + + pi.on("turn_end", (_event, ctx) => { + const usage = ctx.getContextUsage(); + if (!usage || usage.tokens <= COMPACT_THRESHOLD_TOKENS) { + return; + } + triggerCompaction(ctx); + }); + + pi.registerCommand("trigger-compact", { + description: "Trigger compaction immediately", + handler: async (args, ctx) => { + const instructions = args.trim() || undefined; + triggerCompaction(ctx, instructions); + }, + }); +} diff --git a/packages/coding-agent/src/core/agent-session.ts b/packages/coding-agent/src/core/agent-session.ts index caf01412..21c51c3e 100644 --- a/packages/coding-agent/src/core/agent-session.ts +++ b/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) diff --git a/packages/coding-agent/src/core/compaction/compaction.ts b/packages/coding-agent/src/core/compaction/compaction.ts index c0a88a3b..d59ff0ef 100644 --- a/packages/coding-agent/src/core/compaction/compaction.ts +++ b/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); diff --git a/packages/coding-agent/src/core/extensions/index.ts b/packages/coding-agent/src/core/extensions/index.ts index 9289a440..3d487e89 100644 --- a/packages/coding-agent/src/core/extensions/index.ts +++ b/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 diff --git a/packages/coding-agent/src/core/extensions/runner.ts b/packages/coding-agent/src/core/extensions/runner.ts index 4fa869ae..64ffd0b0 100644 --- a/packages/coding-agent/src/core/extensions/runner.ts +++ b/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 = 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), }; } diff --git a/packages/coding-agent/src/core/extensions/types.ts b/packages/coding-agent/src/core/extensions/types.ts index 1753576a..a9918344 100644 --- a/packages/coding-agent/src/core/extensions/types.ts +++ b/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; } /** diff --git a/packages/coding-agent/src/index.ts b/packages/coding-agent/src/index.ts index 6a159afd..c79f4f8b 100644 --- a/packages/coding-agent/src/index.ts +++ b/packages/coding-agent/src/index.ts @@ -45,7 +45,9 @@ export type { AgentToolUpdateCallback, AppAction, BeforeAgentStartEvent, + CompactOptions, ContextEvent, + ContextUsage, ExecOptions, ExecResult, Extension, diff --git a/packages/coding-agent/src/modes/print-mode.ts b/packages/coding-agent/src/modes/print-mode.ts index 2a9a5277..8af94fac 100644 --- a/packages/coding-agent/src/modes/print-mode.ts +++ b/packages/coding-agent/src/modes/print-mode.ts @@ -79,6 +79,18 @@ export async function runPrintMode(session: AgentSession, options: PrintModeOpti abort: () => session.abort(), hasPendingMessages: () => session.pendingMessageCount > 0, shutdown: () => {}, + getContextUsage: () => session.getContextUsage(), + compact: (options) => { + void (async () => { + try { + const result = await session.compact(options?.customInstructions); + options?.onComplete?.(result); + } catch (error) { + const err = error instanceof Error ? error : new Error(String(error)); + options?.onError?.(err); + } + })(); + }, }, // ExtensionCommandContextActions - commands invokable via prompt("/command") { diff --git a/packages/coding-agent/src/modes/rpc/rpc-mode.ts b/packages/coding-agent/src/modes/rpc/rpc-mode.ts index 50c4ed3c..11687b7b 100644 --- a/packages/coding-agent/src/modes/rpc/rpc-mode.ts +++ b/packages/coding-agent/src/modes/rpc/rpc-mode.ts @@ -294,6 +294,18 @@ export async function runRpcMode(session: AgentSession): Promise { shutdown: () => { shutdownRequested = true; }, + getContextUsage: () => session.getContextUsage(), + compact: (options) => { + void (async () => { + try { + const result = await session.compact(options?.customInstructions); + options?.onComplete?.(result); + } catch (error) { + const err = error instanceof Error ? error : new Error(String(error)); + options?.onError?.(err); + } + })(); + }, }, // ExtensionCommandContextActions - commands invokable via prompt("/command") { diff --git a/packages/coding-agent/test/compaction-extensions.test.ts b/packages/coding-agent/test/compaction-extensions.test.ts index 957985b5..62db5cce 100644 --- a/packages/coding-agent/test/compaction-extensions.test.ts +++ b/packages/coding-agent/test/compaction-extensions.test.ts @@ -124,6 +124,18 @@ describe.skipIf(!API_KEY)("Compaction extensions", () => { abort: () => session.abort(), hasPendingMessages: () => session.pendingMessageCount > 0, shutdown: () => {}, + getContextUsage: () => session.getContextUsage(), + compact: (options) => { + void (async () => { + try { + const result = await session.compact(options?.customInstructions); + options?.onComplete?.(result); + } catch (error) { + const err = error instanceof Error ? error : new Error(String(error)); + options?.onError?.(err); + } + })(); + }, }, );