feat(coding-agent): add extension compaction helpers

This commit is contained in:
Mario Zechner 2026-01-17 11:39:46 +01:00
parent 673916f63c
commit 9d3f8117a4
11 changed files with 190 additions and 4 deletions

View file

@ -77,6 +77,7 @@ cp permission-gate.ts ~/.pi/agent/extensions/
| `pirate.ts` | Demonstrates `systemPromptAppend` to dynamically modify system prompt | | `pirate.ts` | Demonstrates `systemPromptAppend` to dynamically modify system prompt |
| `claude-rules.ts` | Scans `.claude/rules/` folder and lists rules in system prompt | | `claude-rules.ts` | Scans `.claude/rules/` folder and lists rules in system prompt |
| `custom-compaction.ts` | Custom compaction that summarizes entire conversation | | `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 ### System Integration

View file

@ -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);
},
});
}

View file

@ -33,6 +33,7 @@ import {
calculateContextTokens, calculateContextTokens,
collectEntriesForBranchSummary, collectEntriesForBranchSummary,
compact, compact,
estimateContextTokens,
generateBranchSummary, generateBranchSummary,
prepareCompaction, prepareCompaction,
shouldCompact, shouldCompact,
@ -40,6 +41,7 @@ import {
import { exportSessionToHtml, type ToolHtmlRenderer } from "./export-html/index.js"; import { exportSessionToHtml, type ToolHtmlRenderer } from "./export-html/index.js";
import { createToolHtmlRenderer } from "./export-html/tool-renderer.js"; import { createToolHtmlRenderer } from "./export-html/tool-renderer.js";
import type { import type {
ContextUsage,
ExtensionRunner, ExtensionRunner,
InputSource, InputSource,
SessionBeforeCompactResult, 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. * Export session to HTML.
* @param outputPath Optional output path (defaults to session directory) * @param outputPath Optional output path (defaults to session directory)

View file

@ -8,7 +8,12 @@
import type { AgentMessage } from "@mariozechner/pi-agent-core"; import type { AgentMessage } from "@mariozechner/pi-agent-core";
import type { AssistantMessage, Model, Usage } from "@mariozechner/pi-ai"; import type { AssistantMessage, Model, Usage } from "@mariozechner/pi-ai";
import { completeSimple } 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 type { CompactionEntry, SessionEntry } from "../session-manager.js";
import { import {
computeFileLists, computeFileLists,
@ -81,6 +86,9 @@ function getMessageFromEntry(entry: SessionEntry): AgentMessage | undefined {
if (entry.type === "branch_summary") { if (entry.type === "branch_summary") {
return createBranchSummaryMessage(entry.summary, entry.fromId, entry.timestamp); return createBranchSummaryMessage(entry.summary, entry.fromId, entry.timestamp);
} }
if (entry.type === "compaction") {
return createCompactionSummaryMessage(entry.summary, entry.tokensBefore, entry.timestamp);
}
return undefined; return undefined;
} }
@ -149,6 +157,55 @@ export function getLastAssistantUsage(entries: SessionEntry[]): Usage | undefine
return undefined; 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. * Check if compaction should trigger based on context usage.
*/ */
@ -555,8 +612,13 @@ export function prepareCompaction(
const boundaryStart = prevCompactionIndex + 1; const boundaryStart = prevCompactionIndex + 1;
const boundaryEnd = pathEntries.length; const boundaryEnd = pathEntries.length;
const lastUsage = getLastAssistantUsage(pathEntries); const usageStart = prevCompactionIndex >= 0 ? prevCompactionIndex : 0;
const tokensBefore = lastUsage ? calculateContextTokens(lastUsage) : 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); const cutPoint = findCutPoint(pathEntries, boundaryStart, boundaryEnd, settings.keepRecentTokens);

View file

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

View file

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

View file

@ -193,6 +193,21 @@ export interface ExtensionUIContext {
// Extension Context // 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. * Context passed to extension event handlers.
*/ */
@ -217,6 +232,10 @@ export interface ExtensionContext {
hasPendingMessages(): boolean; hasPendingMessages(): boolean;
/** Gracefully shutdown pi and exit. Available in all contexts. */ /** Gracefully shutdown pi and exit. Available in all contexts. */
shutdown(): void; 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; abort: () => void;
hasPendingMessages: () => boolean; hasPendingMessages: () => boolean;
shutdown: () => void; shutdown: () => void;
getContextUsage: () => ContextUsage | undefined;
compact: (options?: CompactOptions) => void;
} }
/** /**

View file

@ -45,7 +45,9 @@ export type {
AgentToolUpdateCallback, AgentToolUpdateCallback,
AppAction, AppAction,
BeforeAgentStartEvent, BeforeAgentStartEvent,
CompactOptions,
ContextEvent, ContextEvent,
ContextUsage,
ExecOptions, ExecOptions,
ExecResult, ExecResult,
Extension, Extension,

View file

@ -79,6 +79,18 @@ export async function runPrintMode(session: AgentSession, options: PrintModeOpti
abort: () => session.abort(), abort: () => session.abort(),
hasPendingMessages: () => session.pendingMessageCount > 0, hasPendingMessages: () => session.pendingMessageCount > 0,
shutdown: () => {}, 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") // ExtensionCommandContextActions - commands invokable via prompt("/command")
{ {

View file

@ -294,6 +294,18 @@ export async function runRpcMode(session: AgentSession): Promise<never> {
shutdown: () => { shutdown: () => {
shutdownRequested = true; 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") // ExtensionCommandContextActions - commands invokable via prompt("/command")
{ {

View file

@ -124,6 +124,18 @@ describe.skipIf(!API_KEY)("Compaction extensions", () => {
abort: () => session.abort(), abort: () => session.abort(),
hasPendingMessages: () => session.pendingMessageCount > 0, hasPendingMessages: () => session.pendingMessageCount > 0,
shutdown: () => {}, 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);
}
})();
},
}, },
); );