mirror of
https://github.com/getcompanion-ai/co-mono.git
synced 2026-04-17 08:00:59 +00:00
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
This commit is contained in:
parent
5cbaf2be88
commit
08fab16e2d
5 changed files with 191 additions and 79 deletions
|
|
@ -21,6 +21,7 @@ import { type BashResult, executeBash as executeBashCommand } from "./bash-execu
|
||||||
import {
|
import {
|
||||||
type CompactionResult,
|
type CompactionResult,
|
||||||
calculateContextTokens,
|
calculateContextTokens,
|
||||||
|
collectEntriesForBranchSummary,
|
||||||
compact,
|
compact,
|
||||||
generateBranchSummary,
|
generateBranchSummary,
|
||||||
prepareCompaction,
|
prepareCompaction,
|
||||||
|
|
@ -42,7 +43,7 @@ import type {
|
||||||
} from "./hooks/index.js";
|
} from "./hooks/index.js";
|
||||||
import type { BashExecutionMessage, HookMessage } from "./messages.js";
|
import type { BashExecutionMessage, HookMessage } from "./messages.js";
|
||||||
import type { ModelRegistry } from "./model-registry.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 type { SettingsManager, SkillsSettings } from "./settings-manager.js";
|
||||||
import { expandSlashCommand, type FileSlashCommand } from "./slash-commands.js";
|
import { expandSlashCommand, type FileSlashCommand } from "./slash-commands.js";
|
||||||
|
|
||||||
|
|
@ -1601,30 +1602,12 @@ export class AgentSession {
|
||||||
throw new Error(`Entry ${targetId} not found`);
|
throw new Error(`Entry ${targetId} not found`);
|
||||||
}
|
}
|
||||||
|
|
||||||
// Find common ancestor (if oldLeafId is null, there's no old path)
|
// Collect entries to summarize (from old leaf to common ancestor)
|
||||||
const oldPath = oldLeafId ? new Set(this.sessionManager.getPath(oldLeafId).map((e) => e.id)) : new Set<string>();
|
const { entries: entriesToSummarize, commonAncestorId } = collectEntriesForBranchSummary(
|
||||||
const targetPath = this.sessionManager.getPath(targetId);
|
this.sessionManager,
|
||||||
let commonAncestorId: string | null = null;
|
oldLeafId,
|
||||||
for (const entry of targetPath) {
|
targetId,
|
||||||
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
|
|
||||||
}
|
|
||||||
|
|
||||||
// Prepare event data
|
// Prepare event data
|
||||||
const preparation: TreePreparation = {
|
const preparation: TreePreparation = {
|
||||||
|
|
@ -1667,13 +1650,12 @@ export class AgentSession {
|
||||||
if (!apiKey) {
|
if (!apiKey) {
|
||||||
throw new Error(`No API key for ${model.provider}`);
|
throw new Error(`No API key for ${model.provider}`);
|
||||||
}
|
}
|
||||||
const result = await generateBranchSummary(
|
const result = await generateBranchSummary(entriesToSummarize, {
|
||||||
entriesToSummarize,
|
|
||||||
model,
|
model,
|
||||||
apiKey,
|
apiKey,
|
||||||
this._branchSummaryAbortController.signal,
|
signal: this._branchSummaryAbortController.signal,
|
||||||
options.customInstructions,
|
customInstructions: options.customInstructions,
|
||||||
);
|
});
|
||||||
this._branchSummaryAbortController = undefined;
|
this._branchSummaryAbortController = undefined;
|
||||||
if (result.aborted) {
|
if (result.aborted) {
|
||||||
return { cancelled: true, aborted: true };
|
return { cancelled: true, aborted: true };
|
||||||
|
|
|
||||||
|
|
@ -7,7 +7,7 @@
|
||||||
|
|
||||||
import type { Model } from "@mariozechner/pi-ai";
|
import type { Model } from "@mariozechner/pi-ai";
|
||||||
import { complete } 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
|
// Types
|
||||||
|
|
@ -26,18 +26,100 @@ export interface FileOperations {
|
||||||
}
|
}
|
||||||
|
|
||||||
export interface BranchPreparation {
|
export interface BranchPreparation {
|
||||||
/** Messages extracted for summarization */
|
/** Messages extracted for summarization, in chronological order */
|
||||||
messages: Array<{ role: string; content: string }>;
|
messages: Array<{ role: string; content: string; tokens: number }>;
|
||||||
/** File operations extracted from tool calls */
|
/** File operations extracted from tool calls */
|
||||||
fileOps: FileOperations;
|
fileOps: FileOperations;
|
||||||
/** Previous summaries found in entries */
|
/** Total tokens in messages */
|
||||||
previousSummaries: string[];
|
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
|
// 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.
|
* 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:
|
* Walks entries from NEWEST to OLDEST, adding messages until we hit the token budget.
|
||||||
* - Messages (user, assistant text, custom_message)
|
* This ensures we keep the most recent context when the branch is too long.
|
||||||
* - File operations from tool calls
|
*
|
||||||
* - Previous branch summaries
|
* 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:
|
* 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
|
* - 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 {
|
export function prepareBranchEntries(entries: SessionEntry[], tokenBudget: number = 0): BranchPreparation {
|
||||||
const messages: Array<{ role: string; content: string }> = [];
|
const messages: Array<{ role: string; content: string; tokens: number }> = [];
|
||||||
const fileOps: FileOperations = {
|
const fileOps: FileOperations = {
|
||||||
read: new Set(),
|
read: new Set(),
|
||||||
written: new Set(),
|
written: new Set(),
|
||||||
edited: 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) {
|
switch (entry.type) {
|
||||||
case "message": {
|
case "message": {
|
||||||
const role = entry.message.role;
|
const msgRole = entry.message.role;
|
||||||
|
|
||||||
// Skip tool results - the context is in the assistant's tool call
|
// Skip tool results - context is in assistant's tool call
|
||||||
if (role === "toolResult") continue;
|
if (msgRole === "toolResult") continue;
|
||||||
|
|
||||||
// Extract file ops from assistant tool calls
|
// Extract file ops from assistant tool calls
|
||||||
if (role === "assistant") {
|
if (msgRole === "assistant") {
|
||||||
extractFileOpsFromToolCalls(entry.message, fileOps);
|
extractFileOpsFromToolCalls(entry.message, fileOps);
|
||||||
}
|
}
|
||||||
|
|
||||||
// Extract text content
|
|
||||||
const text = extractMessageText(entry.message);
|
const text = extractMessageText(entry.message);
|
||||||
if (text) {
|
if (text) {
|
||||||
messages.push({ role, content: text });
|
role = msgRole;
|
||||||
|
content = text;
|
||||||
}
|
}
|
||||||
break;
|
break;
|
||||||
}
|
}
|
||||||
|
|
@ -135,27 +228,56 @@ export function prepareBranchEntries(entries: SessionEntry[]): BranchPreparation
|
||||||
.map((c) => c.text)
|
.map((c) => c.text)
|
||||||
.join("");
|
.join("");
|
||||||
if (text) {
|
if (text) {
|
||||||
messages.push({ role: "user", content: text });
|
role = "user";
|
||||||
|
content = text;
|
||||||
}
|
}
|
||||||
break;
|
break;
|
||||||
}
|
}
|
||||||
|
|
||||||
case "branch_summary": {
|
case "branch_summary": {
|
||||||
previousSummaries.push(entry.summary);
|
role = "context";
|
||||||
|
content = `[Branch summary: ${entry.summary}]`;
|
||||||
break;
|
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 "thinking_level_change":
|
||||||
case "model_change":
|
case "model_change":
|
||||||
case "custom":
|
case "custom":
|
||||||
case "label":
|
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;
|
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.
|
* Generate a summary of abandoned branch entries.
|
||||||
*
|
*
|
||||||
* @param entries - Session entries to summarize
|
* @param entries - Session entries to summarize (chronological order)
|
||||||
* @param model - Model to use for summarization
|
* @param options - Generation options
|
||||||
* @param apiKey - API key for the model
|
|
||||||
* @param signal - Abort signal for cancellation
|
|
||||||
* @param customInstructions - Optional custom instructions for summarization
|
|
||||||
*/
|
*/
|
||||||
export async function generateBranchSummary(
|
export async function generateBranchSummary(
|
||||||
entries: SessionEntry[],
|
entries: SessionEntry[],
|
||||||
model: Model<any>,
|
options: GenerateBranchSummaryOptions,
|
||||||
apiKey: string,
|
|
||||||
signal: AbortSignal,
|
|
||||||
customInstructions?: string,
|
|
||||||
): Promise<BranchSummaryResult> {
|
): 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) {
|
if (messages.length === 0) {
|
||||||
return { summary: "No content to summarize" };
|
return { summary: "No content to summarize" };
|
||||||
}
|
}
|
||||||
|
|
||||||
// Build conversation text
|
// Build conversation text
|
||||||
const parts: string[] = [];
|
const conversationText = messages.map((m) => `${m.role}: ${m.content}`).join("\n\n");
|
||||||
|
|
||||||
// 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 instructions = customInstructions || BRANCH_SUMMARY_PROMPT;
|
const instructions = customInstructions || BRANCH_SUMMARY_PROMPT;
|
||||||
const prompt = `${instructions}\n\nConversation:\n${conversationText}`;
|
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
|
// Check if aborted or errored
|
||||||
|
|
|
||||||
|
|
@ -10,3 +10,4 @@ export {
|
||||||
export { execCommand, HookRunner, type HookErrorListener } from "./runner.js";
|
export { execCommand, HookRunner, type HookErrorListener } from "./runner.js";
|
||||||
export { wrapToolsWithHooks, wrapToolWithHooks } from "./tool-wrapper.js";
|
export { wrapToolsWithHooks, wrapToolWithHooks } from "./tool-wrapper.js";
|
||||||
export type * from "./types.js";
|
export type * from "./types.js";
|
||||||
|
export type { ReadonlySessionManager } from "../session-manager.js";
|
||||||
|
|
|
||||||
|
|
@ -155,6 +155,20 @@ export interface SessionInfo {
|
||||||
allMessagesText: string;
|
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) */
|
/** Generate a unique short ID (8 hex chars, collision-checked) */
|
||||||
function generateId(byId: { has(id: string): boolean }): string {
|
function generateId(byId: { has(id: string): boolean }): string {
|
||||||
for (let i = 0; i < 100; i++) {
|
for (let i = 0; i < 100; i++) {
|
||||||
|
|
|
||||||
|
|
@ -14,15 +14,18 @@ export { type ApiKeyCredential, type AuthCredential, AuthStorage, type OAuthCred
|
||||||
export {
|
export {
|
||||||
type BranchPreparation,
|
type BranchPreparation,
|
||||||
type BranchSummaryResult,
|
type BranchSummaryResult,
|
||||||
|
type CollectEntriesResult,
|
||||||
type CompactionResult,
|
type CompactionResult,
|
||||||
type CutPointResult,
|
type CutPointResult,
|
||||||
calculateContextTokens,
|
calculateContextTokens,
|
||||||
|
collectEntriesForBranchSummary,
|
||||||
compact,
|
compact,
|
||||||
DEFAULT_COMPACTION_SETTINGS,
|
DEFAULT_COMPACTION_SETTINGS,
|
||||||
estimateTokens,
|
estimateTokens,
|
||||||
type FileOperations,
|
type FileOperations,
|
||||||
findCutPoint,
|
findCutPoint,
|
||||||
findTurnStartIndex,
|
findTurnStartIndex,
|
||||||
|
type GenerateBranchSummaryOptions,
|
||||||
generateBranchSummary,
|
generateBranchSummary,
|
||||||
generateSummary,
|
generateSummary,
|
||||||
getLastAssistantUsage,
|
getLastAssistantUsage,
|
||||||
|
|
|
||||||
Loading…
Add table
Add a link
Reference in a new issue