Files
oh-my-pi/packages/coding-agent/src/session/compaction/pruning.ts
T
can1357 9b90dabaa9 feat(coding-agent): added session compaction enhancements with auto-continue, tool pruning, and remote endpoint support
- Added session compaction auto-continue feature that automatically sends a synthetic continuation prompt after compaction completes when enabled in settings.
- Added tool output pruning system that intelligently removes verbose tool outputs from session history during compaction to reduce token usage while preserving critical information.
- Added short summary generation for compaction results, providing concise PR-style summaries of conversation changes alongside full summaries.
- Added session.compacting extension hook that allows extensions to customize compaction prompts, add context, and preserve metadata during summarization.
- Added synthetic message support to mark system-injected messages (like auto-continue prompts) and prevent them from being persisted to editor history.
2026-01-28 03:15:28 +01:00

92 lines
2.7 KiB
TypeScript

/**
* Tool output pruning utilities for compaction.
*/
import type { AgentMessage } from "@oh-my-pi/pi-agent-core";
import type { ToolResultMessage } from "@oh-my-pi/pi-ai";
import type { SessionEntry, SessionMessageEntry } from "../session-manager";
import { estimateTokens } from "./compaction";
export interface PruneConfig {
/** Keep the most recent tool output tokens intact. */
protectTokens: number;
/** Only prune if total savings meets this threshold. */
minimumSavings: number;
/** Tool names that should never be pruned. */
protectedTools: string[];
}
export const DEFAULT_PRUNE_CONFIG: PruneConfig = {
protectTokens: 40_000,
minimumSavings: 20_000,
protectedTools: ["skill", "read"],
};
export interface PruneResult {
prunedCount: number;
tokensSaved: number;
}
function createPrunedNotice(tokens: number): string {
return `[Output truncated - ${tokens} tokens]`;
}
function getToolResultMessage(entry: SessionEntry): ToolResultMessage | undefined {
if (entry.type !== "message") return undefined;
const message = entry.message as AgentMessage;
if (message.role !== "toolResult") return undefined;
return message as ToolResultMessage;
}
function estimatePrunedSavings(tokens: number): number {
const noticeTokens = Math.ceil(createPrunedNotice(tokens).length / 4);
return Math.max(0, tokens - noticeTokens);
}
export function pruneToolOutputs(entries: SessionEntry[], config: PruneConfig = DEFAULT_PRUNE_CONFIG): PruneResult {
let accumulatedTokens = 0;
let tokensSaved = 0;
let prunedCount = 0;
const candidates: Array<{ entry: SessionMessageEntry; tokens: number }> = [];
for (let i = entries.length - 1; i >= 0; i--) {
const entry = entries[i];
const message = getToolResultMessage(entry);
if (!message) continue;
const tokens = estimateTokens(message as AgentMessage);
const isProtected = config.protectedTools.includes(message.toolName);
if (message.prunedAt !== undefined) {
accumulatedTokens += tokens;
continue;
}
if (accumulatedTokens < config.protectTokens || isProtected) {
accumulatedTokens += tokens;
continue;
}
candidates.push({ entry: entry as SessionMessageEntry, tokens });
accumulatedTokens += tokens;
}
for (const candidate of candidates) {
tokensSaved += estimatePrunedSavings(candidate.tokens);
}
if (tokensSaved < config.minimumSavings || candidates.length === 0) {
return { prunedCount: 0, tokensSaved: 0 };
}
const prunedAt = Date.now();
for (const candidate of candidates) {
const message = candidate.entry.message as ToolResultMessage;
message.content = [{ type: "text", text: createPrunedNotice(candidate.tokens) }];
message.prunedAt = prunedAt;
prunedCount++;
}
return { prunedCount, tokensSaved };
}