mirror of
https://github.com/getcompanion-ai/co-mono.git
synced 2026-04-21 11:04:35 +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 |
|
| `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
|
||||||
|
|
||||||
|
|
|
||||||
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,
|
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)
|
||||||
|
|
|
||||||
|
|
@ -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);
|
||||||
|
|
||||||
|
|
|
||||||
|
|
@ -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
|
||||||
|
|
|
||||||
|
|
@ -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),
|
||||||
};
|
};
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|
|
||||||
|
|
@ -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;
|
||||||
}
|
}
|
||||||
|
|
||||||
/**
|
/**
|
||||||
|
|
|
||||||
|
|
@ -45,7 +45,9 @@ export type {
|
||||||
AgentToolUpdateCallback,
|
AgentToolUpdateCallback,
|
||||||
AppAction,
|
AppAction,
|
||||||
BeforeAgentStartEvent,
|
BeforeAgentStartEvent,
|
||||||
|
CompactOptions,
|
||||||
ContextEvent,
|
ContextEvent,
|
||||||
|
ContextUsage,
|
||||||
ExecOptions,
|
ExecOptions,
|
||||||
ExecResult,
|
ExecResult,
|
||||||
Extension,
|
Extension,
|
||||||
|
|
|
||||||
|
|
@ -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")
|
||||||
{
|
{
|
||||||
|
|
|
||||||
|
|
@ -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")
|
||||||
{
|
{
|
||||||
|
|
|
||||||
|
|
@ -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);
|
||||||
|
}
|
||||||
|
})();
|
||||||
|
},
|
||||||
},
|
},
|
||||||
);
|
);
|
||||||
|
|
||||||
|
|
|
||||||
Loading…
Add table
Add a link
Reference in a new issue