feat(session): auto-compact long transcripts via transformContext
pi's `transformContext` hook fired on every LLM call but returned the transcript unchanged, so a long-running agent paid per-token for its entire history on every turn. The README promised "sessions persist + auto-compact" but the main loop had no compaction wired in at all — only rotation at the long-cycle boundary. `context-compactor` closes that gap. A pure, provider-free function estimates tokens (~4 chars/token), short-circuits below a 120K threshold, and on overflow keeps the last 20 messages verbatim while collapsing older ones into a single "context summary" user turn. The summary preserves the original task and the first twelve assistant highlights so the agent has continuity without forwarding every historical turn. Compaction is non-destructive: the Agent's in-memory transcript and the SQLite session store still hold the full history, feeding UI timelines and audit. Only the LLM payload shrinks. Rotation remains the long-cycle boundary (separate commit) — compaction is what keeps each individual LLM call within budget.
This commit is contained in:
1 parent
7f23f40e87
commit
e53ac36176
2 files changed
+166
-6
No files matched your search
@@ -0,0 +1,158 @@
|
||||
/**
|
||||
* Deterministic pi context compactor.
|
||||
*
|
||||
* pi's `transformContext` hook runs before each LLM call and returns
|
||||
* the message list that will actually be sent. A long-running agent
|
||||
* can accumulate hundreds of messages — if we forward every one on
|
||||
* every call we pay per-token for the whole history, and large models
|
||||
* throttle when the context approaches their window.
|
||||
*
|
||||
* This compactor keeps the hook pure and provider-free: no LLM
|
||||
* round-trip, no mutation, no IO. It estimates token count using a
|
||||
* 4-chars-per-token heuristic, keeps recent traffic verbatim, and
|
||||
* collapses older messages into a single "context summary" user turn
|
||||
* that preserves the assistant's narrative plus tool-call headlines.
|
||||
*
|
||||
* Triggering only when the estimate clears the threshold means short
|
||||
* runs pay nothing; long runs see roughly constant per-call cost once
|
||||
* they cross the threshold.
|
||||
*/
|
||||
|
||||
import type { AgentMessage } from "@mariozechner/pi-agent-core";
|
||||
|
||||
/** Token-budget threshold that triggers compaction. */
|
||||
const COMPACTION_TRIGGER_TOKENS = 120_000;
|
||||
/**
|
||||
* Number of recent messages preserved verbatim. The tail must include
|
||||
* the current turn plus enough context for the agent to answer without
|
||||
* losing mid-stream state (tool-call pairs, user intent).
|
||||
*/
|
||||
const RECENT_MESSAGE_TAIL = 20;
|
||||
/** Rough tokens-per-char estimator; good enough for gating. */
|
||||
const CHARS_PER_TOKEN = 4;
|
||||
/** Per-assistant-message summary cap to keep the compacted turn readable. */
|
||||
const MAX_SUMMARY_CHARS_PER_MESSAGE = 160;
|
||||
|
||||
interface MessageLike {
|
||||
role?: string;
|
||||
content?: unknown;
|
||||
timestamp?: number;
|
||||
}
|
||||
|
||||
export interface CompactionOutcome {
|
||||
changed: boolean;
|
||||
messages: AgentMessage[] | null;
|
||||
}
|
||||
|
||||
/**
|
||||
* Returns a compacted message list when the incoming transcript
|
||||
* exceeds the budget, or `{ changed: false }` when no compaction is
|
||||
* needed. The transformContext hook forwards the result verbatim.
|
||||
*/
|
||||
export function compactMessagesIfNeeded(messages: AgentMessage[]): CompactionOutcome {
|
||||
const estimatedTokens = estimateTokens(messages);
|
||||
if (estimatedTokens < COMPACTION_TRIGGER_TOKENS) {
|
||||
return { changed: false, messages: null };
|
||||
}
|
||||
if (messages.length <= RECENT_MESSAGE_TAIL + 1) {
|
||||
return { changed: false, messages: null };
|
||||
}
|
||||
|
||||
const cutoff = messages.length - RECENT_MESSAGE_TAIL;
|
||||
const older = messages.slice(0, cutoff);
|
||||
const recent = messages.slice(cutoff);
|
||||
|
||||
const summaryText = buildSummaryText(older);
|
||||
const summary: AgentMessage = {
|
||||
role: "user",
|
||||
content: [{ type: "text", text: summaryText }],
|
||||
timestamp: Date.now(),
|
||||
} as AgentMessage;
|
||||
|
||||
return { changed: true, messages: [summary, ...recent] };
|
||||
}
|
||||
|
||||
function estimateTokens(messages: AgentMessage[]): number {
|
||||
let total = 0;
|
||||
for (const message of messages) {
|
||||
total += estimateMessageTokens(message);
|
||||
}
|
||||
return total;
|
||||
}
|
||||
|
||||
function estimateMessageTokens(message: AgentMessage): number {
|
||||
const raw = extractTextForEstimate(message as MessageLike);
|
||||
return Math.ceil(raw.length / CHARS_PER_TOKEN);
|
||||
}
|
||||
|
||||
function extractTextForEstimate(message: MessageLike): string {
|
||||
const content = message.content;
|
||||
if (typeof content === "string") return content;
|
||||
if (!Array.isArray(content)) return "";
|
||||
const parts: string[] = [];
|
||||
for (const block of content) {
|
||||
if (!block || typeof block !== "object") continue;
|
||||
const obj = block as { type?: string; text?: string; thinking?: string; name?: string };
|
||||
if (typeof obj.text === "string") parts.push(obj.text);
|
||||
else if (typeof obj.thinking === "string") parts.push(obj.thinking);
|
||||
else if (obj.type === "toolCall" && typeof obj.name === "string") parts.push(`tool:${obj.name}`);
|
||||
}
|
||||
return parts.join(" ");
|
||||
}
|
||||
|
||||
function buildSummaryText(older: AgentMessage[]): string {
|
||||
const lines: string[] = [
|
||||
"[Context summary — older turns compacted to stay within the token budget]",
|
||||
"",
|
||||
];
|
||||
const firstUserText = findFirstUserPrompt(older);
|
||||
if (firstUserText) {
|
||||
lines.push(`Original task: ${truncate(firstUserText, 200)}`);
|
||||
lines.push("");
|
||||
}
|
||||
const assistantPoints = collectAssistantHighlights(older);
|
||||
if (assistantPoints.length > 0) {
|
||||
lines.push("Key turns:");
|
||||
for (const point of assistantPoints) {
|
||||
lines.push(`- ${point}`);
|
||||
}
|
||||
}
|
||||
lines.push("");
|
||||
lines.push(`(${older.length} earlier message${older.length === 1 ? "" : "s"} compacted.)`);
|
||||
return lines.join("\n");
|
||||
}
|
||||
|
||||
function findFirstUserPrompt(messages: AgentMessage[]): string | null {
|
||||
for (const message of messages) {
|
||||
const raw = message as MessageLike;
|
||||
if (raw.role !== "user") continue;
|
||||
const text = extractTextForEstimate(raw).trim();
|
||||
if (text.length > 0) return text;
|
||||
}
|
||||
return null;
|
||||
}
|
||||
|
||||
function collectAssistantHighlights(messages: AgentMessage[]): string[] {
|
||||
const highlights: string[] = [];
|
||||
for (const message of messages) {
|
||||
const raw = message as MessageLike;
|
||||
if (raw.role !== "assistant") continue;
|
||||
const text = extractTextForEstimate(raw).trim();
|
||||
if (text.length === 0) continue;
|
||||
highlights.push(truncate(firstSentence(text), MAX_SUMMARY_CHARS_PER_MESSAGE));
|
||||
if (highlights.length >= 12) break;
|
||||
}
|
||||
return highlights;
|
||||
}
|
||||
|
||||
function firstSentence(text: string): string {
|
||||
const collapsed = text.replace(/\s+/g, " ").trim();
|
||||
const sentenceEnd = collapsed.search(/[.?!](\s|$)/);
|
||||
if (sentenceEnd > 0) return collapsed.slice(0, sentenceEnd + 1);
|
||||
return collapsed;
|
||||
}
|
||||
|
||||
function truncate(text: string, maxLen: number): string {
|
||||
if (text.length <= maxLen) return text;
|
||||
return text.slice(0, maxLen - 1).trimEnd() + "…";
|
||||
}
|
||||
@@ -29,6 +29,7 @@ import {
|
||||
projectToolExecutionToWorkProduct,
|
||||
type ArtifactProjectorRunContext,
|
||||
} from "./artifact-projector";
|
||||
import { compactMessagesIfNeeded } from "./context-compactor";
|
||||
|
||||
export interface ActiveRunContext {
|
||||
runId: string;
|
||||
@@ -174,15 +175,16 @@ export function handleAfterToolCall(
|
||||
}
|
||||
|
||||
/**
|
||||
* Context-transform hook. Placeholder pass-through; auto-compaction
|
||||
* wiring lives in session-management-service and plugs in here once
|
||||
* the compactor is available (tracked separately from the contract
|
||||
* wiring).
|
||||
* Context-transform hook. Delegates to `compactMessagesIfNeeded`,
|
||||
* which is a pure, provider-free deterministic compactor. Runs before
|
||||
* every LLM call; short transcripts short-circuit on the token-budget
|
||||
* check so the steady-state cost is negligible.
|
||||
*/
|
||||
export function handleTransformContext(
|
||||
_deps: HookHandlerDeps,
|
||||
_ctx: ActiveRunContext,
|
||||
_payload: TransformContextRpcPayload,
|
||||
payload: TransformContextRpcPayload,
|
||||
): TransformContextRpcResult {
|
||||
return { changed: false, messages: null };
|
||||
const outcome = compactMessagesIfNeeded(payload.messages);
|
||||
return { changed: outcome.changed, messages: outcome.messages };
|
||||
}
|
||||
Reference in new issue
Block a user