mirror of
https://github.com/getcompanion-ai/co-mono.git
synced 2026-04-15 11:02:17 +00:00
feat(coding-agent): add extension compaction helpers
This commit is contained in:
parent
673916f63c
commit
9d3f8117a4
11 changed files with 190 additions and 4 deletions
|
|
@ -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
|
||||
|
||||
|
|
|
|||
32
packages/coding-agent/examples/extensions/trigger-compact.ts
Normal file
32
packages/coding-agent/examples/extensions/trigger-compact.ts
Normal 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);
|
||||
},
|
||||
});
|
||||
}
|
||||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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);
|
||||
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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),
|
||||
};
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -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;
|
||||
}
|
||||
|
||||
/**
|
||||
|
|
|
|||
|
|
@ -45,7 +45,9 @@ export type {
|
|||
AgentToolUpdateCallback,
|
||||
AppAction,
|
||||
BeforeAgentStartEvent,
|
||||
CompactOptions,
|
||||
ContextEvent,
|
||||
ContextUsage,
|
||||
ExecOptions,
|
||||
ExecResult,
|
||||
Extension,
|
||||
|
|
|
|||
|
|
@ -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")
|
||||
{
|
||||
|
|
|
|||
|
|
@ -294,6 +294,18 @@ export async function runRpcMode(session: AgentSession): Promise<never> {
|
|||
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")
|
||||
{
|
||||
|
|
|
|||
|
|
@ -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);
|
||||
}
|
||||
})();
|
||||
},
|
||||
},
|
||||
);
|
||||
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue