/** * Context compaction for long sessions. * * Pure functions for compaction logic. The session manager handles I/O, * and after compaction the session is reloaded. */ import type { AgentMessage } from "@oh-my-pi/pi-agent-core"; import type { AssistantMessage, Model, Usage } from "@oh-my-pi/pi-ai"; import { completeSimple } from "@oh-my-pi/pi-ai"; import { logger } from "@oh-my-pi/pi-utils"; import { renderPromptTemplate } from "../../config/prompt-templates"; import compactionShortSummaryPrompt from "../../prompts/compaction/compaction-short-summary.md" with { type: "text" }; import compactionSummaryPrompt from "../../prompts/compaction/compaction-summary.md" with { type: "text" }; import compactionTurnPrefixPrompt from "../../prompts/compaction/compaction-turn-prefix.md" with { type: "text" }; import compactionUpdateSummaryPrompt from "../../prompts/compaction/compaction-update-summary.md" with { type: "text" }; import { convertToLlm, createBranchSummaryMessage, createCustomMessage } from "../../session/messages"; import type { CompactionEntry, SessionEntry } from "../../session/session-manager"; import { computeFileLists, createFileOps, extractFileOpsFromMessage, type FileOperations, formatFileOperations, SUMMARIZATION_SYSTEM_PROMPT, serializeConversation, } from "./utils"; // ============================================================================ // File Operation Tracking // ============================================================================ /** Details stored in CompactionEntry.details for file tracking */ export interface CompactionDetails { readFiles: string[]; modifiedFiles: string[]; } /** * Extract file operations from messages and previous compaction entries. */ function extractFileOperations( messages: AgentMessage[], entries: SessionEntry[], prevCompactionIndex: number, ): FileOperations { const fileOps = createFileOps(); // Collect from previous compaction's details (if pi-generated) if (prevCompactionIndex >= 0) { const prevCompaction = entries[prevCompactionIndex] as CompactionEntry; if (!prevCompaction.fromExtension && prevCompaction.details) { const details = prevCompaction.details as CompactionDetails; if (Array.isArray(details.readFiles)) { for (const f of details.readFiles) fileOps.read.add(f); } if (Array.isArray(details.modifiedFiles)) { for (const f of details.modifiedFiles) fileOps.edited.add(f); } } } // Extract from tool calls in messages for (const msg of messages) { extractFileOpsFromMessage(msg, fileOps); } return fileOps; } // ============================================================================ // Message Extraction // ============================================================================ /** * Extract AgentMessage from an entry if it produces one. * Returns undefined for entries that don't contribute to LLM context. */ function getMessageFromEntry(entry: SessionEntry): AgentMessage | undefined { if (entry.type === "message") { return entry.message; } if (entry.type === "custom_message") { return createCustomMessage(entry.customType, entry.content, entry.display, entry.details, entry.timestamp); } if (entry.type === "branch_summary") { return createBranchSummaryMessage(entry.summary, entry.fromId, entry.timestamp); } return undefined; } /** Result from compact() - SessionManager adds uuid/parentUuid when saving */ export interface CompactionResult { summary: string; /** Short PR-style summary for display purposes. */ shortSummary?: string; firstKeptEntryId: string; tokensBefore: number; /** Hook-specific data (e.g., ArtifactIndex, version markers for structured compaction) */ details?: T; /** Hook-provided data to persist alongside compaction entry. */ preserveData?: Record; } // ============================================================================ // Types // ============================================================================ export interface CompactionSettings { enabled: boolean; reserveTokens: number; keepRecentTokens: number; autoContinue?: boolean; remoteEndpoint?: string; } export const DEFAULT_COMPACTION_SETTINGS: CompactionSettings = { enabled: true, reserveTokens: 16384, keepRecentTokens: 20000, autoContinue: true, }; // ============================================================================ // Token calculation // ============================================================================ /** * Calculate total context tokens from usage. * Uses the native totalTokens field when available, falls back to computing from components. */ export function calculateContextTokens(usage: Usage): number { return usage.totalTokens || usage.input + usage.output + usage.cacheRead + usage.cacheWrite; } export function calculatePromptTokens(usage: Usage): number { const promptTokens = usage.input + usage.cacheRead + usage.cacheWrite; if (promptTokens > 0) { return promptTokens; } return calculateContextTokens(usage); } /** * Get usage from an assistant message if available. * Skips aborted and error messages as they don't have valid usage data. */ function getAssistantUsage(msg: AgentMessage): Usage | undefined { if (msg.role === "assistant" && "usage" in msg) { const assistantMsg = msg as AssistantMessage; if (assistantMsg.stopReason !== "aborted" && assistantMsg.stopReason !== "error" && assistantMsg.usage) { return assistantMsg.usage; } } return undefined; } /** * Find the last non-aborted assistant message usage from session entries. */ export function getLastAssistantUsage(entries: SessionEntry[]): Usage | undefined { for (let i = entries.length - 1; i >= 0; i--) { const entry = entries[i]; if (entry.type === "message") { const usage = getAssistantUsage(entry.message); if (usage) return usage; } } return undefined; } /** * Check if compaction should trigger based on context usage. */ export function shouldCompact(contextTokens: number, contextWindow: number, settings: CompactionSettings): boolean { if (!settings.enabled) return false; return contextTokens > contextWindow - settings.reserveTokens; } // ============================================================================ // Cut point detection // ============================================================================ /** * Estimate token count for a message using chars/4 heuristic. * This is conservative (overestimates tokens). */ export function estimateTokens(message: AgentMessage): number { let chars = 0; switch (message.role) { case "user": { const content = (message as { content: string | Array<{ type: string; text?: string }> }).content; if (typeof content === "string") { chars = content.length; } else if (Array.isArray(content)) { for (const block of content) { if (block.type === "text" && block.text) { chars += block.text.length; } } } return Math.ceil(chars / 4); } case "assistant": { const assistant = message as AssistantMessage; for (const block of assistant.content) { if (block.type === "text") { chars += block.text.length; } else if (block.type === "thinking") { chars += block.thinking.length; } else if (block.type === "toolCall") { chars += block.name.length + JSON.stringify(block.arguments).length; } } return Math.ceil(chars / 4); } case "hookMessage": case "toolResult": { if (typeof message.content === "string") { chars = message.content.length; } else { for (const block of message.content) { if (block.type === "text" && block.text) { chars += block.text.length; } if (block.type === "image") { chars += 4800; // Estimate images as 4000 chars, or 1200 tokens } } } return Math.ceil(chars / 4); } case "bashExecution": { chars = message.command.length + message.output.length; return Math.ceil(chars / 4); } case "branchSummary": case "compactionSummary": { chars = message.summary.length; return Math.ceil(chars / 4); } } return 0; } function estimateEntriesTokens(entries: SessionEntry[], startIndex: number, endIndex: number): number { let total = 0; for (let i = startIndex; i < endIndex; i++) { const msg = getMessageFromEntry(entries[i]); if (msg) { total += estimateTokens(msg); } } return total; } /** * Find valid cut points: indices of user, assistant, custom, or bashExecution messages. * Never cut at tool results (they must follow their tool call). * When we cut at an assistant message with tool calls, its tool results follow it * and will be kept. * BashExecutionMessage is treated like a user message (user-initiated context). */ function findValidCutPoints(entries: SessionEntry[], startIndex: number, endIndex: number): number[] { const cutPoints: number[] = []; for (let i = startIndex; i < endIndex; i++) { const entry = entries[i]; switch (entry.type) { case "message": { const role = entry.message.role; switch (role) { case "bashExecution": case "hookMessage": case "branchSummary": case "compactionSummary": case "user": case "assistant": cutPoints.push(i); break; case "toolResult": break; } break; } case "thinking_level_change": case "model_change": case "compaction": case "branch_summary": case "custom": case "custom_message": case "label": } // branch_summary and custom_message are user-role messages, valid cut points if (entry.type === "branch_summary" || entry.type === "custom_message") { cutPoints.push(i); } } return cutPoints; } /** * Find the user message (or bashExecution) that starts the turn containing the given entry index. * Returns -1 if no turn start found before the index. * BashExecutionMessage is treated like a user message for turn boundaries. */ export function findTurnStartIndex(entries: SessionEntry[], entryIndex: number, startIndex: number): number { for (let i = entryIndex; i >= startIndex; i--) { const entry = entries[i]; // branch_summary and custom_message are user-role messages, can start a turn if (entry.type === "branch_summary" || entry.type === "custom_message") { return i; } if (entry.type === "message") { const role = entry.message.role; if (role === "user" || role === "bashExecution") { return i; } } } return -1; } export interface CutPointResult { /** Index of first entry to keep */ firstKeptEntryIndex: number; /** Index of user message that starts the turn being split, or -1 if not splitting */ turnStartIndex: number; /** Whether this cut splits a turn (cut point is not a user message) */ isSplitTurn: boolean; } /** * Find the cut point in session entries that keeps approximately `keepRecentTokens`. * * Algorithm: Walk backwards from newest, accumulating estimated message sizes. * Stop when we've accumulated >= keepRecentTokens. Cut at that point. * * Can cut at user OR assistant messages (never tool results). When cutting at an * assistant message with tool calls, its tool results come after and will be kept. * * Returns CutPointResult with: * - firstKeptEntryIndex: the entry index to start keeping from * - turnStartIndex: if cutting mid-turn, the user message that started that turn * - isSplitTurn: whether we're cutting in the middle of a turn * * Only considers entries between `startIndex` and `endIndex` (exclusive). */ export function findCutPoint( entries: SessionEntry[], startIndex: number, endIndex: number, keepRecentTokens: number, ): CutPointResult { const cutPoints = findValidCutPoints(entries, startIndex, endIndex); if (cutPoints.length === 0) { return { firstKeptEntryIndex: startIndex, turnStartIndex: -1, isSplitTurn: false }; } // Walk backwards from newest, accumulating estimated message sizes let accumulatedTokens = 0; let cutIndex = cutPoints[0]; // Default: keep from first message (not header) for (let i = endIndex - 1; i >= startIndex; i--) { const entry = entries[i]; if (entry.type !== "message") continue; // Estimate this message's size const messageTokens = estimateTokens(entry.message); accumulatedTokens += messageTokens; // Check if we've exceeded the budget if (accumulatedTokens >= keepRecentTokens) { // Find the closest valid cut point at or after this entry for (let c = 0; c < cutPoints.length; c++) { if (cutPoints[c] >= i) { cutIndex = cutPoints[c]; break; } } break; } } // Scan backwards from cutIndex to include any non-message entries (bash, settings, etc.) while (cutIndex > startIndex) { const prevEntry = entries[cutIndex - 1]; // Stop at session header or compaction boundaries if (prevEntry.type === "compaction") { break; } if (prevEntry.type === "message") { // Stop if we hit any message break; } // Include this non-message entry (bash, settings change, etc.) cutIndex--; } // Determine if this is a split turn const cutEntry = entries[cutIndex]; const isUserMessage = cutEntry.type === "message" && cutEntry.message.role === "user"; const turnStartIndex = isUserMessage ? -1 : findTurnStartIndex(entries, cutIndex, startIndex); return { firstKeptEntryIndex: cutIndex, turnStartIndex, isSplitTurn: !isUserMessage && turnStartIndex !== -1, }; } // ============================================================================ // Summarization // ============================================================================ const SUMMARIZATION_PROMPT = renderPromptTemplate(compactionSummaryPrompt); const UPDATE_SUMMARIZATION_PROMPT = renderPromptTemplate(compactionUpdateSummaryPrompt); const SHORT_SUMMARY_PROMPT = renderPromptTemplate(compactionShortSummaryPrompt); function formatAdditionalContext(context: string[] | undefined): string { if (!context || context.length === 0) return ""; const lines = context.map(line => `- ${line}`).join("\n"); return `\n${lines}\n\n\n`; } interface RemoteCompactionRequest { systemPrompt: string; prompt: string; } interface RemoteCompactionResponse { summary: string; shortSummary?: string; } async function requestRemoteCompaction( endpoint: string, request: RemoteCompactionRequest, ): Promise { const response = await fetch(endpoint, { method: "POST", headers: { "content-type": "application/json" }, body: JSON.stringify(request), }); if (!response.ok) { const errorText = await response.text().catch(() => ""); logger.warn("Remote compaction failed", { endpoint, status: response.status, statusText: response.statusText, errorText, }); throw new Error(`Remote compaction failed (${response.status} ${response.statusText})`); } const data = (await response.json()) as RemoteCompactionResponse | undefined; if (!data || typeof data.summary !== "string") { throw new Error("Remote compaction response missing summary"); } return data; } /** * Generate a summary of the conversation using the LLM. * If previousSummary is provided, uses the update prompt to merge. */ export interface SummaryOptions { promptOverride?: string; extraContext?: string[]; remoteEndpoint?: string; } export async function generateSummary( currentMessages: AgentMessage[], model: Model, reserveTokens: number, apiKey: string, signal?: AbortSignal, customInstructions?: string, previousSummary?: string, options?: SummaryOptions, ): Promise { const maxTokens = Math.floor(0.8 * reserveTokens); // Use update prompt if we have a previous summary, otherwise initial prompt let basePrompt = previousSummary ? UPDATE_SUMMARIZATION_PROMPT : SUMMARIZATION_PROMPT; if (options?.promptOverride) { basePrompt = options.promptOverride; } if (customInstructions) { basePrompt = `${basePrompt}\n\nAdditional focus: ${customInstructions}`; } // Serialize conversation to text so model doesn't try to continue it // Convert to LLM messages first (handles custom types like bashExecution, hookMessage, etc.) const llmMessages = convertToLlm(currentMessages); const conversationText = serializeConversation(llmMessages); // Build the prompt with conversation wrapped in tags let promptText = `\n${conversationText}\n\n\n`; if (previousSummary) { promptText += `\n${previousSummary}\n\n\n`; } promptText += formatAdditionalContext(options?.extraContext); promptText += basePrompt; const summarizationMessages = [ { role: "user" as const, content: [{ type: "text" as const, text: promptText }], timestamp: Date.now(), }, ]; if (options?.remoteEndpoint) { const remote = await requestRemoteCompaction(options.remoteEndpoint, { systemPrompt: SUMMARIZATION_SYSTEM_PROMPT, prompt: promptText, }); return remote.summary; } const response = await completeSimple( model, { systemPrompt: SUMMARIZATION_SYSTEM_PROMPT, messages: summarizationMessages }, { maxTokens, signal, apiKey, reasoning: "high" }, ); if (response.stopReason === "error") { throw new Error(`Summarization failed: ${response.errorMessage || "Unknown error"}`); } const textContent = response.content .filter((c): c is { type: "text"; text: string } => c.type === "text") .map(c => c.text) .join("\n"); return textContent; } async function generateShortSummary( recentMessages: AgentMessage[], historySummary: string | undefined, model: Model, reserveTokens: number, apiKey: string, signal?: AbortSignal, options?: SummaryOptions, ): Promise { const maxTokens = Math.min(512, Math.floor(0.2 * reserveTokens)); const llmMessages = convertToLlm(recentMessages); const conversationText = serializeConversation(llmMessages); let promptText = `\n${conversationText}\n\n\n`; if (historySummary) { promptText += `\n${historySummary}\n\n\n`; } promptText += formatAdditionalContext(options?.extraContext); promptText += SHORT_SUMMARY_PROMPT; if (options?.remoteEndpoint) { const remote = await requestRemoteCompaction(options.remoteEndpoint, { systemPrompt: SUMMARIZATION_SYSTEM_PROMPT, prompt: promptText, }); return remote.summary; } const response = await completeSimple( model, { systemPrompt: SUMMARIZATION_SYSTEM_PROMPT, messages: [{ role: "user", content: [{ type: "text", text: promptText }], timestamp: Date.now() }], }, { maxTokens, signal, apiKey, reasoning: "high" }, ); if (response.stopReason === "error") { throw new Error(`Short summary failed: ${response.errorMessage || "Unknown error"}`); } return response.content .filter((c): c is { type: "text"; text: string } => c.type === "text") .map(c => c.text) .join("\n"); } // ============================================================================ // Compaction Preparation (for hooks) // ============================================================================ export interface CompactionPreparation { /** UUID of first entry to keep */ firstKeptEntryId: string; /** Messages that will be summarized and discarded */ messagesToSummarize: AgentMessage[]; /** Messages that will be turned into turn prefix summary (if splitting) */ turnPrefixMessages: AgentMessage[]; /** Messages kept in full after compaction (recent history) */ recentMessages: AgentMessage[]; /** Whether this is a split turn (cut point in middle of turn) */ isSplitTurn: boolean; tokensBefore: number; /** Summary from previous compaction, for iterative update */ previousSummary?: string; /** File operations extracted from messagesToSummarize */ fileOps: FileOperations; /** Compaction settions from settings.jsonl */ settings: CompactionSettings; } export function prepareCompaction( pathEntries: SessionEntry[], settings: CompactionSettings, ): CompactionPreparation | undefined { if (pathEntries.length > 0 && pathEntries[pathEntries.length - 1].type === "compaction") { return undefined; } let prevCompactionIndex = -1; for (let i = pathEntries.length - 1; i >= 0; i--) { if (pathEntries[i].type === "compaction") { prevCompactionIndex = i; break; } } const boundaryStart = prevCompactionIndex + 1; const boundaryEnd = pathEntries.length; const lastUsage = getLastAssistantUsage(pathEntries); const tokensBefore = lastUsage ? calculateContextTokens(lastUsage) : 0; let keepRecentTokens = settings.keepRecentTokens; if (lastUsage) { const estimatedTokens = estimateEntriesTokens(pathEntries, boundaryStart, boundaryEnd); const promptTokens = calculatePromptTokens(lastUsage); const ratio = estimatedTokens > 0 ? promptTokens / estimatedTokens : 0; if (Number.isFinite(ratio) && ratio > 1) { keepRecentTokens = Math.max(1, Math.floor(keepRecentTokens / ratio)); } } const cutPoint = findCutPoint(pathEntries, boundaryStart, boundaryEnd, keepRecentTokens); // Get ID of first kept entry const firstKeptEntry = pathEntries[cutPoint.firstKeptEntryIndex]; if (!firstKeptEntry?.id) { return undefined; // Session needs migration } const firstKeptEntryId = firstKeptEntry.id; const historyEnd = cutPoint.isSplitTurn ? cutPoint.turnStartIndex : cutPoint.firstKeptEntryIndex; // Messages to summarize (will be discarded after summary) const messagesToSummarize: AgentMessage[] = []; for (let i = boundaryStart; i < historyEnd; i++) { const msg = getMessageFromEntry(pathEntries[i]); if (msg) messagesToSummarize.push(msg); } // Messages for turn prefix summary (if splitting a turn) const turnPrefixMessages: AgentMessage[] = []; if (cutPoint.isSplitTurn) { for (let i = cutPoint.turnStartIndex; i < cutPoint.firstKeptEntryIndex; i++) { const msg = getMessageFromEntry(pathEntries[i]); if (msg) turnPrefixMessages.push(msg); } } // Messages kept after compaction (recent history) const recentMessages: AgentMessage[] = []; for (let i = cutPoint.firstKeptEntryIndex; i < boundaryEnd; i++) { const msg = getMessageFromEntry(pathEntries[i]); if (msg) recentMessages.push(msg); } // Get previous summary for iterative update let previousSummary: string | undefined; if (prevCompactionIndex >= 0) { const prevCompaction = pathEntries[prevCompactionIndex] as CompactionEntry; previousSummary = prevCompaction.summary; } // Extract file operations from messages and previous compaction const fileOps = extractFileOperations(messagesToSummarize, pathEntries, prevCompactionIndex); // Also extract file ops from turn prefix if splitting if (cutPoint.isSplitTurn) { for (const msg of turnPrefixMessages) { extractFileOpsFromMessage(msg, fileOps); } } return { firstKeptEntryId, messagesToSummarize, turnPrefixMessages, recentMessages, isSplitTurn: cutPoint.isSplitTurn, tokensBefore, previousSummary, fileOps, settings, }; } // ============================================================================ // Main compaction function // ============================================================================ const TURN_PREFIX_SUMMARIZATION_PROMPT = renderPromptTemplate(compactionTurnPrefixPrompt); /** * Generate summaries for compaction using prepared data. * Returns CompactionResult - SessionManager adds id/parentId when saving. * * @param preparation - Pre-calculated preparation from prepareCompaction() * @param customInstructions - Optional custom focus for the summary */ export async function compact( preparation: CompactionPreparation, model: Model, apiKey: string, customInstructions?: string, signal?: AbortSignal, options?: SummaryOptions, ): Promise { const { firstKeptEntryId, messagesToSummarize, turnPrefixMessages, recentMessages, isSplitTurn, tokensBefore, previousSummary, fileOps, settings, } = preparation; const summaryOptions: SummaryOptions = { promptOverride: options?.promptOverride, extraContext: options?.extraContext, remoteEndpoint: settings.remoteEndpoint, }; // Generate summaries (can be parallel if both needed) and merge into one let summary: string; if (isSplitTurn && turnPrefixMessages.length > 0) { // Generate both summaries in parallel const [historyResult, turnPrefixResult] = await Promise.all([ messagesToSummarize.length > 0 ? generateSummary( messagesToSummarize, model, settings.reserveTokens, apiKey, signal, customInstructions, previousSummary, summaryOptions, ) : Promise.resolve("No prior history."), generateTurnPrefixSummary(turnPrefixMessages, model, settings.reserveTokens, apiKey, signal), ]); // Merge into single summary summary = `${historyResult}\n\n---\n\n**Turn Context (split turn):**\n\n${turnPrefixResult}`; } else if (messagesToSummarize.length > 0) { // Generate history summary from messages to summarize summary = await generateSummary( messagesToSummarize, model, settings.reserveTokens, apiKey, signal, customInstructions, previousSummary, summaryOptions, ); } else if (previousSummary) { // No new messages to summarize, preserve previous summary summary = previousSummary; } else { // No messages and no previous summary summary = "No prior history."; } const shortSummary = await generateShortSummary( recentMessages, summary, model, settings.reserveTokens, apiKey, signal, { extraContext: options?.extraContext, remoteEndpoint: settings.remoteEndpoint }, ); // Compute file lists and append to summary const { readFiles, modifiedFiles } = computeFileLists(fileOps); summary += formatFileOperations(readFiles, modifiedFiles); if (!firstKeptEntryId) { throw new Error("First kept entry has no ID - session may need migration"); } return { summary, shortSummary, firstKeptEntryId, tokensBefore, details: { readFiles, modifiedFiles } as CompactionDetails, }; } /** * Generate a summary for a turn prefix (when splitting a turn). */ async function generateTurnPrefixSummary( messages: AgentMessage[], model: Model, reserveTokens: number, apiKey: string, signal?: AbortSignal, ): Promise { const maxTokens = Math.floor(0.5 * reserveTokens); // Smaller budget for turn prefix const llmMessages = convertToLlm(messages); const conversationText = serializeConversation(llmMessages); const promptText = `\n${conversationText}\n\n\n${TURN_PREFIX_SUMMARIZATION_PROMPT}`; const summarizationMessages = [ { role: "user" as const, content: [{ type: "text" as const, text: promptText }], timestamp: Date.now(), }, ]; const response = await completeSimple( model, { systemPrompt: SUMMARIZATION_SYSTEM_PROMPT, messages: summarizationMessages }, { maxTokens, signal, apiKey, reasoning: "high" }, ); if (response.stopReason === "error") { throw new Error(`Turn prefix summarization failed: ${response.errorMessage || "Unknown error"}`); } return response.content .filter((c): c is { type: "text"; text: string } => c.type === "text") .map(c => c.text) .join("\n"); }