shumengya/sproutclaw
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Add ReadonlySessionManager and refactor branch summarization

Browse Source 
- 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
This commit is contained in:
Mario Zechner
2025-12-29 21:22:50 +01:00
parent 5cbaf2be88
commit 08fab16e2d
5 changed files with 191 additions and 79 deletions
Download Patch File Download Diff File Expand all files Collapse all files

42
packages/coding-agent/src/core/agent-session.ts
View File

@@ -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<string>();
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 };

210
packages/coding-agent/src/core/compaction/branch-summarization.ts
View File

@@ -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<any>;
/** 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<any>,
apiKey: string,
signal: AbortSignal,
customInstructions?: string,
options: GenerateBranchSummaryOptions,
): Promise<BranchSummaryResult> {
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

1
packages/coding-agent/src/core/hooks/index.ts
View File

@@ -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";

14
packages/coding-agent/src/core/session-manager.ts
View File

@@ -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++) {