From 08fab16e2df76a77336fb6245b382cc1b426660e Mon Sep 17 00:00:00 2001 From: Mario Zechner Date: Mon, 29 Dec 2025 21:22:50 +0100 Subject: [PATCH] Add ReadonlySessionManager and refactor branch summarization - Add ReadonlySessionManager interface to session-manager.ts - Re-export from hooks/index.ts - Add collectEntriesForBranchSummary() to extract entries for summarization - Don't stop at compaction boundaries (include their summaries as context) - Add token budget support to prepareBranchEntries() - Walk entries newest-to-oldest to prioritize recent context - Use options object for generateBranchSummary() - Handle compaction entries as context summaries - Export new types: CollectEntriesResult, GenerateBranchSummaryOptions --- .../coding-agent/src/core/agent-session.ts | 42 +--- .../core/compaction/branch-summarization.ts | 210 ++++++++++++++---- packages/coding-agent/src/core/hooks/index.ts | 1 + .../coding-agent/src/core/session-manager.ts | 14 ++ packages/coding-agent/src/index.ts | 3 + 5 files changed, 191 insertions(+), 79 deletions(-) diff --git a/packages/coding-agent/src/core/agent-session.ts b/packages/coding-agent/src/core/agent-session.ts index 978fb1b8..c7277c3c 100644 --- a/packages/coding-agent/src/core/agent-session.ts +++ b/packages/coding-agent/src/core/agent-session.ts @@ -21,6 +21,7 @@ import { type BashResult, executeBash as executeBashCommand } from "./bash-execu import { type CompactionResult, calculateContextTokens, + collectEntriesForBranchSummary, compact, generateBranchSummary, prepareCompaction, @@ -42,7 +43,7 @@ import type { } from "./hooks/index.js"; import type { BashExecutionMessage, HookMessage } from "./messages.js"; import type { ModelRegistry } from "./model-registry.js"; -import type { BranchSummaryEntry, CompactionEntry, SessionEntry, SessionManager } from "./session-manager.js"; +import type { BranchSummaryEntry, CompactionEntry, SessionManager } from "./session-manager.js"; import type { SettingsManager, SkillsSettings } from "./settings-manager.js"; import { expandSlashCommand, type FileSlashCommand } from "./slash-commands.js"; @@ -1601,30 +1602,12 @@ export class AgentSession { throw new Error(`Entry ${targetId} not found`); } - // Find common ancestor (if oldLeafId is null, there's no old path) - const oldPath = oldLeafId ? new Set(this.sessionManager.getPath(oldLeafId).map((e) => e.id)) : new Set(); - const targetPath = this.sessionManager.getPath(targetId); - let commonAncestorId: string | null = null; - for (const entry of targetPath) { - if (oldPath.has(entry.id)) { - commonAncestorId = entry.id; - break; - } - } - - // Collect entries to summarize (old leaf back to common ancestor, stop at compaction) - const entriesToSummarize: SessionEntry[] = []; - if (options.summarize && oldLeafId) { - let current: string | null = oldLeafId; - while (current && current !== commonAncestorId) { - const entry = this.sessionManager.getEntry(current); - if (!entry) break; - if (entry.type === "compaction") break; - entriesToSummarize.push(entry); - current = entry.parentId; - } - entriesToSummarize.reverse(); // Chronological order - } + // Collect entries to summarize (from old leaf to common ancestor) + const { entries: entriesToSummarize, commonAncestorId } = collectEntriesForBranchSummary( + this.sessionManager, + oldLeafId, + targetId, + ); // Prepare event data const preparation: TreePreparation = { @@ -1667,13 +1650,12 @@ export class AgentSession { if (!apiKey) { throw new Error(`No API key for ${model.provider}`); } - const result = await generateBranchSummary( - entriesToSummarize, + const result = await generateBranchSummary(entriesToSummarize, { model, apiKey, - this._branchSummaryAbortController.signal, - options.customInstructions, - ); + signal: this._branchSummaryAbortController.signal, + customInstructions: options.customInstructions, + }); this._branchSummaryAbortController = undefined; if (result.aborted) { return { cancelled: true, aborted: true }; diff --git a/packages/coding-agent/src/core/compaction/branch-summarization.ts b/packages/coding-agent/src/core/compaction/branch-summarization.ts index 9c549ba2..d79fd301 100644 --- a/packages/coding-agent/src/core/compaction/branch-summarization.ts +++ b/packages/coding-agent/src/core/compaction/branch-summarization.ts @@ -7,7 +7,7 @@ import type { Model } from "@mariozechner/pi-ai"; import { complete } from "@mariozechner/pi-ai"; -import type { SessionEntry } from "../session-manager.js"; +import type { ReadonlySessionManager, SessionEntry } from "../session-manager.js"; // ============================================================================ // Types @@ -26,18 +26,100 @@ export interface FileOperations { } export interface BranchPreparation { - /** Messages extracted for summarization */ - messages: Array<{ role: string; content: string }>; + /** Messages extracted for summarization, in chronological order */ + messages: Array<{ role: string; content: string; tokens: number }>; /** File operations extracted from tool calls */ fileOps: FileOperations; - /** Previous summaries found in entries */ - previousSummaries: string[]; + /** Total tokens in messages */ + totalTokens: number; +} + +export interface CollectEntriesResult { + /** Entries to summarize, in chronological order */ + entries: SessionEntry[]; + /** Common ancestor between old and new position, if any */ + commonAncestorId: string | null; +} + +export interface GenerateBranchSummaryOptions { + /** Model to use for summarization */ + model: Model; + /** API key for the model */ + apiKey: string; + /** Abort signal for cancellation */ + signal: AbortSignal; + /** Optional custom instructions for summarization */ + customInstructions?: string; + /** Reserve this fraction of context window for summary (default 0.2) */ + reserveFraction?: number; +} + +// ============================================================================ +// Entry Collection +// ============================================================================ + +/** + * Collect entries that should be summarized when navigating from one position to another. + * + * Walks from oldLeafId back to the common ancestor with targetId, collecting entries + * along the way. Does NOT stop at compaction boundaries - those are included and their + * summaries become context. + * + * @param session - Session manager (read-only access) + * @param oldLeafId - Current position (where we're navigating from) + * @param targetId - Target position (where we're navigating to) + * @returns Entries to summarize and the common ancestor + */ +export function collectEntriesForBranchSummary( + session: ReadonlySessionManager, + oldLeafId: string | null, + targetId: string, +): CollectEntriesResult { + // If no old position, nothing to summarize + if (!oldLeafId) { + return { entries: [], commonAncestorId: null }; + } + + // Find common ancestor + const oldPath = new Set(session.getPath(oldLeafId).map((e) => e.id)); + const targetPath = session.getPath(targetId); + + let commonAncestorId: string | null = null; + for (const entry of targetPath) { + if (oldPath.has(entry.id)) { + commonAncestorId = entry.id; + break; + } + } + + // Collect entries from old leaf back to common ancestor + const entries: SessionEntry[] = []; + let current: string | null = oldLeafId; + + while (current && current !== commonAncestorId) { + const entry = session.getEntry(current); + if (!entry) break; + entries.push(entry); + current = entry.parentId; + } + + // Reverse to get chronological order + entries.reverse(); + + return { entries, commonAncestorId }; } // ============================================================================ // Entry Parsing // ============================================================================ +/** + * Estimate token count for a string using chars/4 heuristic. + */ +function estimateStringTokens(text: string): number { + return Math.ceil(text.length / 4); +} + /** * Extract text content from any message type. */ @@ -84,44 +166,55 @@ function extractFileOpsFromToolCalls(message: any, fileOps: FileOperations): voi } /** - * Prepare entries for summarization. + * Prepare entries for summarization with token budget. * - * Extracts: - * - Messages (user, assistant text, custom_message) - * - File operations from tool calls - * - Previous branch summaries + * Walks entries from NEWEST to OLDEST, adding messages until we hit the token budget. + * This ensures we keep the most recent context when the branch is too long. + * + * Handles: + * - message (user, assistant) - extracts text, counts tokens + * - custom_message - treated as user message + * - branch_summary - included as context + * - compaction - includes summary as context * * Skips: - * - toolResult messages (context already in assistant message) + * - toolResult messages (context already in assistant's tool call) * - thinking_level_change, model_change, custom, label entries - * - compaction entries (these are boundaries, shouldn't be in the input) + * + * @param entries - Entries in chronological order + * @param tokenBudget - Maximum tokens to include (0 = no limit) */ -export function prepareBranchEntries(entries: SessionEntry[]): BranchPreparation { - const messages: Array<{ role: string; content: string }> = []; +export function prepareBranchEntries(entries: SessionEntry[], tokenBudget: number = 0): BranchPreparation { + const messages: Array<{ role: string; content: string; tokens: number }> = []; const fileOps: FileOperations = { read: new Set(), written: new Set(), edited: new Set(), }; - const previousSummaries: string[] = []; + let totalTokens = 0; + + // Walk from newest to oldest to prioritize recent context + for (let i = entries.length - 1; i >= 0; i--) { + const entry = entries[i]; + let role: string | undefined; + let content: string | undefined; - for (const entry of entries) { switch (entry.type) { case "message": { - const role = entry.message.role; + const msgRole = entry.message.role; - // Skip tool results - the context is in the assistant's tool call - if (role === "toolResult") continue; + // Skip tool results - context is in assistant's tool call + if (msgRole === "toolResult") continue; // Extract file ops from assistant tool calls - if (role === "assistant") { + if (msgRole === "assistant") { extractFileOpsFromToolCalls(entry.message, fileOps); } - // Extract text content const text = extractMessageText(entry.message); if (text) { - messages.push({ role, content: text }); + role = msgRole; + content = text; } break; } @@ -135,27 +228,56 @@ export function prepareBranchEntries(entries: SessionEntry[]): BranchPreparation .map((c) => c.text) .join(""); if (text) { - messages.push({ role: "user", content: text }); + role = "user"; + content = text; } break; } case "branch_summary": { - previousSummaries.push(entry.summary); + role = "context"; + content = `[Branch summary: ${entry.summary}]`; break; } - // Skip these entry types - they don't contribute to conversation content - case "compaction": + case "compaction": { + role = "context"; + content = `[Session summary: ${entry.summary}]`; + break; + } + + // Skip these - don't contribute to conversation content case "thinking_level_change": case "model_change": case "custom": case "label": + continue; + } + + if (role && content) { + const tokens = estimateStringTokens(content); + + // Check budget before adding + if (tokenBudget > 0 && totalTokens + tokens > tokenBudget) { + // If this is a summary entry, try to fit it anyway as it's important context + if (entry.type === "compaction" || entry.type === "branch_summary") { + // Add truncated version or skip + if (totalTokens < tokenBudget * 0.9) { + // Still have some room, add it + messages.unshift({ role, content, tokens }); + totalTokens += tokens; + } + } + // Stop - we've hit the budget break; + } + + messages.unshift({ role, content, tokens }); + totalTokens += tokens; } } - return { messages, fileOps, previousSummaries }; + return { messages, fileOps, totalTokens }; } // ============================================================================ @@ -202,37 +324,27 @@ function formatFileOperations(fileOps: FileOperations): string { /** * Generate a summary of abandoned branch entries. * - * @param entries - Session entries to summarize - * @param model - Model to use for summarization - * @param apiKey - API key for the model - * @param signal - Abort signal for cancellation - * @param customInstructions - Optional custom instructions for summarization + * @param entries - Session entries to summarize (chronological order) + * @param options - Generation options */ export async function generateBranchSummary( entries: SessionEntry[], - model: Model, - apiKey: string, - signal: AbortSignal, - customInstructions?: string, + options: GenerateBranchSummaryOptions, ): Promise { - const { messages, fileOps, previousSummaries } = prepareBranchEntries(entries); + const { model, apiKey, signal, customInstructions, reserveFraction = 0.2 } = options; + + // Calculate token budget (leave room for summary generation) + const contextWindow = model.contextWindow || 128000; + const tokenBudget = Math.floor(contextWindow * (1 - reserveFraction)); + + const { messages, fileOps } = prepareBranchEntries(entries, tokenBudget); if (messages.length === 0) { return { summary: "No content to summarize" }; } // Build conversation text - const parts: string[] = []; - - // Include previous summaries as context - if (previousSummaries.length > 0) { - parts.push(`[Previous context: ${previousSummaries.join(" | ")}]`); - } - - // Add conversation - parts.push(messages.map((m) => `${m.role}: ${m.content}`).join("\n\n")); - - const conversationText = parts.join("\n\n"); + const conversationText = messages.map((m) => `${m.role}: ${m.content}`).join("\n\n"); const instructions = customInstructions || BRANCH_SUMMARY_PROMPT; const prompt = `${instructions}\n\nConversation:\n${conversationText}`; @@ -248,7 +360,7 @@ export async function generateBranchSummary( }, ], }, - { apiKey, signal, maxTokens: 1024 }, + { apiKey, signal, maxTokens: 2048 }, ); // Check if aborted or errored diff --git a/packages/coding-agent/src/core/hooks/index.ts b/packages/coding-agent/src/core/hooks/index.ts index fc49f2ae..cb0806ed 100644 --- a/packages/coding-agent/src/core/hooks/index.ts +++ b/packages/coding-agent/src/core/hooks/index.ts @@ -10,3 +10,4 @@ export { export { execCommand, HookRunner, type HookErrorListener } from "./runner.js"; export { wrapToolsWithHooks, wrapToolWithHooks } from "./tool-wrapper.js"; export type * from "./types.js"; +export type { ReadonlySessionManager } from "../session-manager.js"; diff --git a/packages/coding-agent/src/core/session-manager.ts b/packages/coding-agent/src/core/session-manager.ts index ccf6e4ea..d8c2b420 100644 --- a/packages/coding-agent/src/core/session-manager.ts +++ b/packages/coding-agent/src/core/session-manager.ts @@ -155,6 +155,20 @@ export interface SessionInfo { allMessagesText: string; } +/** + * Read-only interface for SessionManager. + * Used by compaction/summarization utilities that only need to read session data. + */ +export interface ReadonlySessionManager { + getLeafId(): string | null; + getEntry(id: string): SessionEntry | undefined; + getPath(fromId?: string): SessionEntry[]; + getEntries(): SessionEntry[]; + getChildren(parentId: string): SessionEntry[]; + getTree(): SessionTreeNode[]; + getLabel(id: string): string | undefined; +} + /** Generate a unique short ID (8 hex chars, collision-checked) */ function generateId(byId: { has(id: string): boolean }): string { for (let i = 0; i < 100; i++) { diff --git a/packages/coding-agent/src/index.ts b/packages/coding-agent/src/index.ts index 44419113..1dae4688 100644 --- a/packages/coding-agent/src/index.ts +++ b/packages/coding-agent/src/index.ts @@ -14,15 +14,18 @@ export { type ApiKeyCredential, type AuthCredential, AuthStorage, type OAuthCred export { type BranchPreparation, type BranchSummaryResult, + type CollectEntriesResult, type CompactionResult, type CutPointResult, calculateContextTokens, + collectEntriesForBranchSummary, compact, DEFAULT_COMPACTION_SETTINGS, estimateTokens, type FileOperations, findCutPoint, findTurnStartIndex, + type GenerateBranchSummaryOptions, generateBranchSummary, generateSummary, getLastAssistantUsage,