mirror of
https://github.com/getcompanion-ai/co-mono.git
synced 2026-04-15 20:03:05 +00:00
Session tree structure with id/parentId linking
- Add TreeNode base type with id, parentId, timestamp - Add *Content types for clean input/output separation - Entry types are now TreeNode & *Content intersections - SessionManager assigns id/parentId on save, tracks leafId - Add migrateSessionEntries() for v1 to v2 conversion - Migration runs on load, rewrites file - buildSessionContext() uses tree traversal from leaf - Compaction returns CompactionResult (content only) - Hooks return compaction content, not full entries - Add firstKeptEntryId to before_compact hook event - Update mom package for tree fields - Better error messages for compaction failures
This commit is contained in:
Mario Zechner
parent
04a764742e
commit
c58d5f20a4
12 changed files with 6778 additions and 6297 deletions
|
|
@ -754,7 +754,12 @@ export class AgentSession {
|
|||
|
||||
const preparation = prepareCompaction(entries, settings);
|
||||
if (!preparation) {
|
||||
throw new Error("Already compacted");
|
||||
// Check why we can't compact
|
||||
const lastEntry = entries[entries.length - 1];
|
||||
if (lastEntry?.type === "compaction") {
|
||||
throw new Error("Already compacted");
|
||||
}
|
||||
throw new Error("Nothing to compact (session too small or needs migration)");
|
||||
}
|
||||
|
||||
// Find previous compaction summary if any
|
||||
|
|
@ -766,7 +771,7 @@ export class AgentSession {
|
|||
}
|
||||
}
|
||||
|
||||
let compactionEntry: CompactionEntry | undefined;
|
||||
let hookCompaction: { summary: string; firstKeptEntryId: string; tokensBefore: number } | undefined;
|
||||
let fromHook = false;
|
||||
|
||||
if (this._hookRunner?.hasHandlers("session")) {
|
||||
|
|
@ -777,6 +782,7 @@ export class AgentSession {
|
|||
previousSessionFile: null,
|
||||
reason: "before_compact",
|
||||
cutPoint: preparation.cutPoint,
|
||||
firstKeptEntryId: preparation.firstKeptEntryId,
|
||||
previousSummary,
|
||||
messagesToSummarize: [...preparation.messagesToSummarize],
|
||||
messagesToKeep: [...preparation.messagesToKeep],
|
||||
|
|
@ -791,14 +797,24 @@ export class AgentSession {
|
|||
throw new Error("Compaction cancelled");
|
||||
}
|
||||
|
||||
if (result?.compactionEntry) {
|
||||
compactionEntry = result.compactionEntry;
|
||||
if (result?.compaction) {
|
||||
hookCompaction = result.compaction;
|
||||
fromHook = true;
|
||||
}
|
||||
}
|
||||
|
||||
if (!compactionEntry) {
|
||||
compactionEntry = await compact(
|
||||
let summary: string;
|
||||
let firstKeptEntryId: string;
|
||||
let tokensBefore: number;
|
||||
|
||||
if (hookCompaction) {
|
||||
// Hook provided compaction content
|
||||
summary = hookCompaction.summary;
|
||||
firstKeptEntryId = hookCompaction.firstKeptEntryId;
|
||||
tokensBefore = hookCompaction.tokensBefore;
|
||||
} else {
|
||||
// Generate compaction result
|
||||
const result = await compact(
|
||||
entries,
|
||||
this.model,
|
||||
settings,
|
||||
|
|
@ -806,33 +822,41 @@ export class AgentSession {
|
|||
this._compactionAbortController.signal,
|
||||
customInstructions,
|
||||
);
|
||||
summary = result.summary;
|
||||
firstKeptEntryId = result.firstKeptEntryId;
|
||||
tokensBefore = result.tokensBefore;
|
||||
}
|
||||
|
||||
if (this._compactionAbortController.signal.aborted) {
|
||||
throw new Error("Compaction cancelled");
|
||||
}
|
||||
|
||||
this.sessionManager.saveCompaction(compactionEntry);
|
||||
this.sessionManager.saveCompaction(summary, firstKeptEntryId, tokensBefore);
|
||||
const newEntries = this.sessionManager.getEntries();
|
||||
const sessionContext = this.sessionManager.buildSessionContext();
|
||||
this.agent.replaceMessages(sessionContext.messages);
|
||||
|
||||
if (this._hookRunner) {
|
||||
// Get the saved compaction entry for the hook
|
||||
const savedCompactionEntry = newEntries.find((e) => e.type === "compaction" && e.summary === summary) as
|
||||
| CompactionEntry
|
||||
| undefined;
|
||||
|
||||
if (this._hookRunner && savedCompactionEntry) {
|
||||
await this._hookRunner.emit({
|
||||
type: "session",
|
||||
entries: newEntries,
|
||||
sessionFile: this.sessionFile,
|
||||
previousSessionFile: null,
|
||||
reason: "compact",
|
||||
compactionEntry,
|
||||
tokensBefore: compactionEntry.tokensBefore,
|
||||
compactionEntry: savedCompactionEntry,
|
||||
tokensBefore,
|
||||
fromHook,
|
||||
});
|
||||
}
|
||||
|
||||
return {
|
||||
tokensBefore: compactionEntry.tokensBefore,
|
||||
summary: compactionEntry.summary,
|
||||
tokensBefore,
|
||||
summary,
|
||||
};
|
||||
} finally {
|
||||
this._compactionAbortController = null;
|
||||
|
|
@ -928,7 +952,7 @@ export class AgentSession {
|
|||
}
|
||||
}
|
||||
|
||||
let compactionEntry: CompactionEntry | undefined;
|
||||
let hookCompaction: { summary: string; firstKeptEntryId: string; tokensBefore: number } | undefined;
|
||||
let fromHook = false;
|
||||
|
||||
if (this._hookRunner?.hasHandlers("session")) {
|
||||
|
|
@ -939,6 +963,7 @@ export class AgentSession {
|
|||
previousSessionFile: null,
|
||||
reason: "before_compact",
|
||||
cutPoint: preparation.cutPoint,
|
||||
firstKeptEntryId: preparation.firstKeptEntryId,
|
||||
previousSummary,
|
||||
messagesToSummarize: [...preparation.messagesToSummarize],
|
||||
messagesToKeep: [...preparation.messagesToKeep],
|
||||
|
|
@ -954,20 +979,33 @@ export class AgentSession {
|
|||
return;
|
||||
}
|
||||
|
||||
if (hookResult?.compactionEntry) {
|
||||
compactionEntry = hookResult.compactionEntry;
|
||||
if (hookResult?.compaction) {
|
||||
hookCompaction = hookResult.compaction;
|
||||
fromHook = true;
|
||||
}
|
||||
}
|
||||
|
||||
if (!compactionEntry) {
|
||||
compactionEntry = await compact(
|
||||
let summary: string;
|
||||
let firstKeptEntryId: string;
|
||||
let tokensBefore: number;
|
||||
|
||||
if (hookCompaction) {
|
||||
// Hook provided compaction content
|
||||
summary = hookCompaction.summary;
|
||||
firstKeptEntryId = hookCompaction.firstKeptEntryId;
|
||||
tokensBefore = hookCompaction.tokensBefore;
|
||||
} else {
|
||||
// Generate compaction result
|
||||
const compactResult = await compact(
|
||||
entries,
|
||||
this.model,
|
||||
settings,
|
||||
apiKey,
|
||||
this._autoCompactionAbortController.signal,
|
||||
);
|
||||
summary = compactResult.summary;
|
||||
firstKeptEntryId = compactResult.firstKeptEntryId;
|
||||
tokensBefore = compactResult.tokensBefore;
|
||||
}
|
||||
|
||||
if (this._autoCompactionAbortController.signal.aborted) {
|
||||
|
|
@ -975,27 +1013,32 @@ export class AgentSession {
|
|||
return;
|
||||
}
|
||||
|
||||
this.sessionManager.saveCompaction(compactionEntry);
|
||||
this.sessionManager.saveCompaction(summary, firstKeptEntryId, tokensBefore);
|
||||
const newEntries = this.sessionManager.getEntries();
|
||||
const sessionContext = this.sessionManager.buildSessionContext();
|
||||
this.agent.replaceMessages(sessionContext.messages);
|
||||
|
||||
if (this._hookRunner) {
|
||||
// Get the saved compaction entry for the hook
|
||||
const savedCompactionEntry = newEntries.find((e) => e.type === "compaction" && e.summary === summary) as
|
||||
| CompactionEntry
|
||||
| undefined;
|
||||
|
||||
if (this._hookRunner && savedCompactionEntry) {
|
||||
await this._hookRunner.emit({
|
||||
type: "session",
|
||||
entries: newEntries,
|
||||
sessionFile: this.sessionFile,
|
||||
previousSessionFile: null,
|
||||
reason: "compact",
|
||||
compactionEntry,
|
||||
tokensBefore: compactionEntry.tokensBefore,
|
||||
compactionEntry: savedCompactionEntry,
|
||||
tokensBefore,
|
||||
fromHook,
|
||||
});
|
||||
}
|
||||
|
||||
const result: CompactionResult = {
|
||||
tokensBefore: compactionEntry.tokensBefore,
|
||||
summary: compactionEntry.summary,
|
||||
tokensBefore,
|
||||
summary,
|
||||
};
|
||||
this._emit({ type: "auto_compaction_end", result, aborted: false, willRetry });
|
||||
|
||||
|
|
|
|||
|
|
@ -9,7 +9,14 @@ import type { AppMessage } from "@mariozechner/pi-agent-core";
|
|||
import type { AssistantMessage, Model, Usage } from "@mariozechner/pi-ai";
|
||||
import { complete } from "@mariozechner/pi-ai";
|
||||
import { messageTransformer } from "./messages.js";
|
||||
import type { CompactionEntry, SessionEntry } from "./session-manager.js";
|
||||
import type { CompactionEntry, ConversationEntry, SessionEntry } from "./session-manager.js";
|
||||
|
||||
/** Result from compact() - SessionManager adds uuid/parentUuid when saving */
|
||||
export interface CompactionResult {
|
||||
summary: string;
|
||||
firstKeptEntryId: string;
|
||||
tokensBefore: number;
|
||||
}
|
||||
|
||||
// ============================================================================
|
||||
// Types
|
||||
|
|
@ -327,6 +334,8 @@ export async function generateSummary(
|
|||
|
||||
export interface CompactionPreparation {
|
||||
cutPoint: CutPointResult;
|
||||
/** UUID of first entry to keep */
|
||||
firstKeptEntryId: string;
|
||||
/** Messages that will be summarized and discarded */
|
||||
messagesToSummarize: AppMessage[];
|
||||
/** Messages that will be kept after the summary (recent turns) */
|
||||
|
|
@ -355,6 +364,16 @@ export function prepareCompaction(entries: SessionEntry[], settings: CompactionS
|
|||
|
||||
const cutPoint = findCutPoint(entries, boundaryStart, boundaryEnd, settings.keepRecentTokens);
|
||||
|
||||
// Get UUID of first kept entry
|
||||
const firstKeptEntry = entries[cutPoint.firstKeptEntryIndex];
|
||||
if (firstKeptEntry.type === "session") {
|
||||
return null; // Can't compact if first kept is header
|
||||
}
|
||||
const firstKeptEntryId = (firstKeptEntry as ConversationEntry).id;
|
||||
if (!firstKeptEntryId) {
|
||||
return null; // Session needs migration
|
||||
}
|
||||
|
||||
const historyEnd = cutPoint.isSplitTurn ? cutPoint.turnStartIndex : cutPoint.firstKeptEntryIndex;
|
||||
|
||||
// Messages to summarize (will be discarded after summary)
|
||||
|
|
@ -375,7 +394,7 @@ export function prepareCompaction(entries: SessionEntry[], settings: CompactionS
|
|||
}
|
||||
}
|
||||
|
||||
return { cutPoint, messagesToSummarize, messagesToKeep, tokensBefore, boundaryStart };
|
||||
return { cutPoint, firstKeptEntryId, messagesToSummarize, messagesToKeep, tokensBefore, boundaryStart };
|
||||
}
|
||||
|
||||
// ============================================================================
|
||||
|
|
@ -394,9 +413,9 @@ Be concise. Focus on information needed to understand the retained recent work.`
|
|||
|
||||
/**
|
||||
* Calculate compaction and generate summary.
|
||||
* Returns the CompactionEntry to append to the session file.
|
||||
* Returns CompactionResult - SessionManager adds uuid/parentUuid when saving.
|
||||
*
|
||||
* @param entries - All session entries
|
||||
* @param entries - All session entries (must have uuid fields for v2)
|
||||
* @param model - Model to use for summarization
|
||||
* @param settings - Compaction settings
|
||||
* @param apiKey - API key for LLM
|
||||
|
|
@ -410,7 +429,7 @@ export async function compact(
|
|||
apiKey: string,
|
||||
signal?: AbortSignal,
|
||||
customInstructions?: string,
|
||||
): Promise<CompactionEntry> {
|
||||
): Promise<CompactionResult> {
|
||||
// Don't compact if the last entry is already a compaction
|
||||
if (entries.length > 0 && entries[entries.length - 1].type === "compaction") {
|
||||
throw new Error("Already compacted");
|
||||
|
|
@ -490,11 +509,19 @@ export async function compact(
|
|||
);
|
||||
}
|
||||
|
||||
// Get UUID of first kept entry
|
||||
const firstKeptEntry = entries[cutResult.firstKeptEntryIndex];
|
||||
if (firstKeptEntry.type === "session") {
|
||||
throw new Error("Cannot compact: first kept entry is session header");
|
||||
}
|
||||
const firstKeptEntryId = (firstKeptEntry as ConversationEntry).id;
|
||||
if (!firstKeptEntryId) {
|
||||
throw new Error("First kept entry has no UUID - session may need migration");
|
||||
}
|
||||
|
||||
return {
|
||||
type: "compaction",
|
||||
timestamp: new Date().toISOString(),
|
||||
summary,
|
||||
firstKeptEntryIndex: cutResult.firstKeptEntryIndex,
|
||||
firstKeptEntryId,
|
||||
tokensBefore,
|
||||
};
|
||||
}
|
||||
|
|
|
|||
|
|
@ -130,6 +130,8 @@ export type SessionEvent =
|
|||
| (SessionEventBase & {
|
||||
reason: "before_compact";
|
||||
cutPoint: CutPointResult;
|
||||
/** ID of first entry to keep (for hooks that return CompactionEntry) */
|
||||
firstKeptEntryId: string;
|
||||
/** Summary from previous compaction, if any. Include this in your summary to preserve context. */
|
||||
previousSummary?: string;
|
||||
/** Messages that will be summarized and discarded */
|
||||
|
|
@ -351,8 +353,12 @@ export interface SessionEventResult {
|
|||
cancel?: boolean;
|
||||
/** If true (for before_branch only), skip restoring conversation to branch point while still creating the branched session file */
|
||||
skipConversationRestore?: boolean;
|
||||
/** Custom compaction entry (for before_compact event) */
|
||||
compactionEntry?: CompactionEntry;
|
||||
/** Custom compaction result (for before_compact event) - SessionManager adds id/parentId */
|
||||
compaction?: {
|
||||
summary: string;
|
||||
firstKeptEntryId: string;
|
||||
tokensBefore: number;
|
||||
};
|
||||
}
|
||||
|
||||
// ============================================================================
|
||||
|
|
|
|||
|
|
@ -1,9 +1,11 @@
|
|||
import type { AppMessage } from "@mariozechner/pi-agent-core";
|
||||
import { randomBytes } from "crypto";
|
||||
import { appendFileSync, existsSync, mkdirSync, readdirSync, readFileSync, statSync } from "fs";
|
||||
import { appendFileSync, existsSync, mkdirSync, readdirSync, readFileSync, statSync, writeFileSync } from "fs";
|
||||
import { join, resolve } from "path";
|
||||
import { getAgentDir as getDefaultAgentDir } from "../config.js";
|
||||
|
||||
export const CURRENT_SESSION_VERSION = 2;
|
||||
|
||||
function uuidv4(): string {
|
||||
const bytes = randomBytes(16);
|
||||
bytes[6] = (bytes[6] & 0x0f) | 0x40;
|
||||
|
|
@ -12,47 +14,89 @@ function uuidv4(): string {
|
|||
return `${hex.slice(0, 8)}-${hex.slice(8, 12)}-${hex.slice(12, 16)}-${hex.slice(16, 20)}-${hex.slice(20, 32)}`;
|
||||
}
|
||||
|
||||
// ============================================================================
|
||||
// Session Header (metadata, not part of conversation tree)
|
||||
// ============================================================================
|
||||
|
||||
export interface SessionHeader {
|
||||
type: "session";
|
||||
version?: number; // v1 sessions don't have this
|
||||
id: string;
|
||||
timestamp: string;
|
||||
cwd: string;
|
||||
branchedFrom?: string;
|
||||
}
|
||||
|
||||
export interface SessionMessageEntry {
|
||||
type: "message";
|
||||
// ============================================================================
|
||||
// Tree Node (added by SessionManager to all conversation entries)
|
||||
// ============================================================================
|
||||
|
||||
export interface TreeNode {
|
||||
id: string;
|
||||
parentId: string | null;
|
||||
timestamp: string;
|
||||
}
|
||||
|
||||
// ============================================================================
|
||||
// Content Types (what distinguishes entries - used for input)
|
||||
// ============================================================================
|
||||
|
||||
export interface MessageContent {
|
||||
type: "message";
|
||||
message: AppMessage;
|
||||
}
|
||||
|
||||
export interface ThinkingLevelChangeEntry {
|
||||
export interface ThinkingLevelContent {
|
||||
type: "thinking_level_change";
|
||||
timestamp: string;
|
||||
thinkingLevel: string;
|
||||
}
|
||||
|
||||
export interface ModelChangeEntry {
|
||||
export interface ModelChangeContent {
|
||||
type: "model_change";
|
||||
timestamp: string;
|
||||
provider: string;
|
||||
modelId: string;
|
||||
}
|
||||
|
||||
export interface CompactionEntry {
|
||||
export interface CompactionContent {
|
||||
type: "compaction";
|
||||
timestamp: string;
|
||||
summary: string;
|
||||
firstKeptEntryIndex: number;
|
||||
firstKeptEntryId: string;
|
||||
tokensBefore: number;
|
||||
}
|
||||
|
||||
export type SessionEntry =
|
||||
| SessionHeader
|
||||
export interface BranchSummaryContent {
|
||||
type: "branch_summary";
|
||||
summary: string;
|
||||
}
|
||||
|
||||
/** Union of all content types (for input) */
|
||||
export type ConversationContent =
|
||||
| MessageContent
|
||||
| ThinkingLevelContent
|
||||
| ModelChangeContent
|
||||
| CompactionContent
|
||||
| BranchSummaryContent;
|
||||
|
||||
// ============================================================================
|
||||
// Full Entry Types (TreeNode + Content - returned from SessionManager)
|
||||
// ============================================================================
|
||||
|
||||
export type SessionMessageEntry = TreeNode & MessageContent;
|
||||
export type ThinkingLevelChangeEntry = TreeNode & ThinkingLevelContent;
|
||||
export type ModelChangeEntry = TreeNode & ModelChangeContent;
|
||||
export type CompactionEntry = TreeNode & CompactionContent;
|
||||
export type BranchSummaryEntry = TreeNode & BranchSummaryContent;
|
||||
|
||||
/** Conversation entry - has id/parentId for tree structure */
|
||||
export type ConversationEntry =
|
||||
| SessionMessageEntry
|
||||
| ThinkingLevelChangeEntry
|
||||
| ModelChangeEntry
|
||||
| CompactionEntry;
|
||||
| CompactionEntry
|
||||
| BranchSummaryEntry;
|
||||
|
||||
/** Any session entry (header or conversation) */
|
||||
export type SessionEntry = SessionHeader | ConversationEntry;
|
||||
|
||||
export interface SessionContext {
|
||||
messages: AppMessage[];
|
||||
|
|
@ -87,6 +131,45 @@ export function createSummaryMessage(summary: string): AppMessage {
|
|||
};
|
||||
}
|
||||
|
||||
/**
|
||||
* Migrate v1 entries to v2 format by adding id/parentId fields.
|
||||
* Mutates entries in place. Safe to call on already-migrated entries.
|
||||
*/
|
||||
export function migrateSessionEntries(entries: SessionEntry[]): void {
|
||||
// Check if already migrated
|
||||
const firstConv = entries.find((e) => e.type !== "session");
|
||||
if (firstConv && "id" in firstConv && firstConv.id) {
|
||||
return; // Already migrated
|
||||
}
|
||||
|
||||
let prevId: string | null = null;
|
||||
for (const entry of entries) {
|
||||
if (entry.type === "session") {
|
||||
entry.version = CURRENT_SESSION_VERSION;
|
||||
continue;
|
||||
}
|
||||
|
||||
// Add id/parentId to conversation entries
|
||||
const convEntry = entry as ConversationEntry;
|
||||
convEntry.id = uuidv4();
|
||||
convEntry.parentId = prevId;
|
||||
prevId = convEntry.id;
|
||||
|
||||
// Convert firstKeptEntryIndex to firstKeptEntryId for compaction
|
||||
if (entry.type === "compaction") {
|
||||
const comp = entry as CompactionEntry & { firstKeptEntryIndex?: number };
|
||||
if (typeof comp.firstKeptEntryIndex === "number") {
|
||||
// Find the entry at that index and get its id
|
||||
const targetEntry = entries[comp.firstKeptEntryIndex];
|
||||
if (targetEntry && targetEntry.type !== "session") {
|
||||
comp.firstKeptEntryId = (targetEntry as ConversationEntry).id;
|
||||
}
|
||||
delete comp.firstKeptEntryIndex;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/** Exported for compaction.test.ts */
|
||||
export function parseSessionEntries(content: string): SessionEntry[] {
|
||||
const entries: SessionEntry[] = [];
|
||||
|
|
@ -115,59 +198,108 @@ export function getLatestCompactionEntry(entries: SessionEntry[]): CompactionEnt
|
|||
}
|
||||
|
||||
/**
|
||||
* Build the session context from entries. This is what gets sent to the LLM.
|
||||
*
|
||||
* If there's a compaction entry, returns the summary message plus messages
|
||||
* from `firstKeptEntryIndex` onwards. Otherwise returns all messages.
|
||||
*
|
||||
* Also extracts the current thinking level and model from the entries.
|
||||
* Build the session context from entries using tree traversal.
|
||||
* If leafId is provided, walks from that entry to root.
|
||||
* Handles compaction and branch summaries along the path.
|
||||
*/
|
||||
export function buildSessionContext(entries: SessionEntry[]): SessionContext {
|
||||
export function buildSessionContext(entries: SessionEntry[], leafId?: string): SessionContext {
|
||||
// Build uuid index for conversation entries
|
||||
const byId = new Map<string, ConversationEntry>();
|
||||
for (const entry of entries) {
|
||||
if (entry.type !== "session") {
|
||||
byId.set(entry.id, entry);
|
||||
}
|
||||
}
|
||||
|
||||
// Find leaf
|
||||
let leaf: ConversationEntry | undefined;
|
||||
if (leafId) {
|
||||
leaf = byId.get(leafId);
|
||||
} else {
|
||||
// Find last conversation entry
|
||||
for (let i = entries.length - 1; i >= 0; i--) {
|
||||
if (entries[i].type !== "session") {
|
||||
leaf = entries[i] as ConversationEntry;
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if (!leaf) {
|
||||
return { messages: [], thinkingLevel: "off", model: null };
|
||||
}
|
||||
|
||||
// Walk from leaf to root, collecting path
|
||||
const path: ConversationEntry[] = [];
|
||||
let current: ConversationEntry | undefined = leaf;
|
||||
while (current) {
|
||||
path.unshift(current);
|
||||
current = current.parentId ? byId.get(current.parentId) : undefined;
|
||||
}
|
||||
|
||||
// Extract settings and find compaction
|
||||
let thinkingLevel = "off";
|
||||
let model: { provider: string; modelId: string } | null = null;
|
||||
let compaction: CompactionEntry | null = null;
|
||||
|
||||
for (const entry of entries) {
|
||||
for (const entry of path) {
|
||||
if (entry.type === "thinking_level_change") {
|
||||
thinkingLevel = entry.thinkingLevel;
|
||||
} else if (entry.type === "model_change") {
|
||||
model = { provider: entry.provider, modelId: entry.modelId };
|
||||
} else if (entry.type === "message" && entry.message.role === "assistant") {
|
||||
model = { provider: entry.message.provider, modelId: entry.message.model };
|
||||
} else if (entry.type === "compaction") {
|
||||
compaction = entry;
|
||||
}
|
||||
}
|
||||
|
||||
let latestCompactionIndex = -1;
|
||||
for (let i = entries.length - 1; i >= 0; i--) {
|
||||
if (entries[i].type === "compaction") {
|
||||
latestCompactionIndex = i;
|
||||
break;
|
||||
}
|
||||
}
|
||||
// Build messages - handle compaction ordering correctly
|
||||
// When there's a compaction, we need to:
|
||||
// 1. Emit summary first
|
||||
// 2. Emit kept messages (from firstKeptEntryId up to compaction)
|
||||
// 3. Emit messages after compaction
|
||||
const messages: AppMessage[] = [];
|
||||
|
||||
if (latestCompactionIndex === -1) {
|
||||
const messages: AppMessage[] = [];
|
||||
for (const entry of entries) {
|
||||
if (entry.type === "message") {
|
||||
if (compaction) {
|
||||
// Emit summary first
|
||||
messages.push(createSummaryMessage(compaction.summary));
|
||||
|
||||
// Find compaction index in path
|
||||
const compactionIdx = path.findIndex((e) => e.type === "compaction" && e.id === compaction.id);
|
||||
|
||||
// Emit kept messages (before compaction, starting from firstKeptEntryId)
|
||||
let foundFirstKept = false;
|
||||
for (let i = 0; i < compactionIdx; i++) {
|
||||
const entry = path[i];
|
||||
if (entry.id === compaction.firstKeptEntryId) {
|
||||
foundFirstKept = true;
|
||||
}
|
||||
if (foundFirstKept && entry.type === "message") {
|
||||
messages.push(entry.message);
|
||||
}
|
||||
}
|
||||
return { messages, thinkingLevel, model };
|
||||
}
|
||||
|
||||
const compactionEvent = entries[latestCompactionIndex] as CompactionEntry;
|
||||
|
||||
const keptMessages: AppMessage[] = [];
|
||||
for (let i = compactionEvent.firstKeptEntryIndex; i < entries.length; i++) {
|
||||
const entry = entries[i];
|
||||
if (entry.type === "message") {
|
||||
keptMessages.push(entry.message);
|
||||
// Emit messages after compaction
|
||||
for (let i = compactionIdx + 1; i < path.length; i++) {
|
||||
const entry = path[i];
|
||||
if (entry.type === "message") {
|
||||
messages.push(entry.message);
|
||||
} else if (entry.type === "branch_summary") {
|
||||
messages.push(createSummaryMessage(entry.summary));
|
||||
}
|
||||
}
|
||||
} else {
|
||||
// No compaction - emit all messages, handle branch summaries
|
||||
for (const entry of path) {
|
||||
if (entry.type === "message") {
|
||||
messages.push(entry.message);
|
||||
} else if (entry.type === "branch_summary") {
|
||||
messages.push(createSummaryMessage(entry.summary));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
const messages: AppMessage[] = [];
|
||||
messages.push(createSummaryMessage(compactionEvent.summary));
|
||||
messages.push(...keptMessages);
|
||||
|
||||
return { messages, thinkingLevel, model };
|
||||
}
|
||||
|
||||
|
|
@ -229,6 +361,10 @@ export class SessionManager {
|
|||
private flushed: boolean = false;
|
||||
private inMemoryEntries: SessionEntry[] = [];
|
||||
|
||||
// Tree structure (v2)
|
||||
private byId: Map<string, ConversationEntry> = new Map();
|
||||
private leafId: string = "";
|
||||
|
||||
private constructor(cwd: string, sessionDir: string, sessionFile: string | null, persist: boolean) {
|
||||
this.cwd = cwd;
|
||||
this.sessionDir = sessionDir;
|
||||
|
|
@ -240,10 +376,7 @@ export class SessionManager {
|
|||
if (sessionFile) {
|
||||
this.setSessionFile(sessionFile);
|
||||
} else {
|
||||
this.sessionId = uuidv4();
|
||||
const timestamp = new Date().toISOString().replace(/[:.]/g, "-");
|
||||
const sessionFile = join(this.getSessionDir(), `${timestamp}_${this.sessionId}.jsonl`);
|
||||
this.setSessionFile(sessionFile);
|
||||
this._initNewSession();
|
||||
}
|
||||
}
|
||||
|
||||
|
|
@ -252,23 +385,61 @@ export class SessionManager {
|
|||
this.sessionFile = resolve(sessionFile);
|
||||
if (existsSync(this.sessionFile)) {
|
||||
this.inMemoryEntries = loadEntriesFromFile(this.sessionFile);
|
||||
const header = this.inMemoryEntries.find((e) => e.type === "session");
|
||||
this.sessionId = header ? (header as SessionHeader).id : uuidv4();
|
||||
const header = this.inMemoryEntries.find((e) => e.type === "session") as SessionHeader | undefined;
|
||||
this.sessionId = header?.id ?? uuidv4();
|
||||
|
||||
// Migrate v1 to v2 if needed
|
||||
const version = header?.version ?? 1;
|
||||
if (version < CURRENT_SESSION_VERSION) {
|
||||
this._migrateToV2();
|
||||
this._rewriteFile();
|
||||
}
|
||||
|
||||
this._buildIndex();
|
||||
this.flushed = true;
|
||||
} else {
|
||||
this.sessionId = uuidv4();
|
||||
this.inMemoryEntries = [];
|
||||
this.flushed = false;
|
||||
const entry: SessionHeader = {
|
||||
type: "session",
|
||||
id: this.sessionId,
|
||||
timestamp: new Date().toISOString(),
|
||||
cwd: this.cwd,
|
||||
};
|
||||
this.inMemoryEntries.push(entry);
|
||||
this._initNewSession();
|
||||
}
|
||||
}
|
||||
|
||||
private _initNewSession(): void {
|
||||
this.sessionId = uuidv4();
|
||||
const timestamp = new Date().toISOString();
|
||||
const header: SessionHeader = {
|
||||
type: "session",
|
||||
version: CURRENT_SESSION_VERSION,
|
||||
id: this.sessionId,
|
||||
timestamp,
|
||||
cwd: this.cwd,
|
||||
};
|
||||
this.inMemoryEntries = [header];
|
||||
this.byId.clear();
|
||||
this.leafId = "";
|
||||
this.flushed = false;
|
||||
const fileTimestamp = timestamp.replace(/[:.]/g, "-");
|
||||
this.sessionFile = join(this.getSessionDir(), `${fileTimestamp}_${this.sessionId}.jsonl`);
|
||||
}
|
||||
|
||||
private _migrateToV2(): void {
|
||||
migrateSessionEntries(this.inMemoryEntries);
|
||||
}
|
||||
|
||||
private _buildIndex(): void {
|
||||
this.byId.clear();
|
||||
this.leafId = "";
|
||||
for (const entry of this.inMemoryEntries) {
|
||||
if (entry.type === "session") continue;
|
||||
this.byId.set(entry.id, entry);
|
||||
this.leafId = entry.id;
|
||||
}
|
||||
}
|
||||
|
||||
private _rewriteFile(): void {
|
||||
if (!this.persist) return;
|
||||
const content = `${this.inMemoryEntries.map((e) => JSON.stringify(e)).join("\n")}\n`;
|
||||
writeFileSync(this.sessionFile, content);
|
||||
}
|
||||
|
||||
isPersisted(): boolean {
|
||||
return this.persist;
|
||||
}
|
||||
|
|
@ -290,18 +461,7 @@ export class SessionManager {
|
|||
}
|
||||
|
||||
reset(): void {
|
||||
this.sessionId = uuidv4();
|
||||
this.flushed = false;
|
||||
const timestamp = new Date().toISOString().replace(/[:.]/g, "-");
|
||||
this.sessionFile = join(this.getSessionDir(), `${timestamp}_${this.sessionId}.jsonl`);
|
||||
this.inMemoryEntries = [
|
||||
{
|
||||
type: "session",
|
||||
id: this.sessionId,
|
||||
timestamp: new Date().toISOString(),
|
||||
cwd: this.cwd,
|
||||
},
|
||||
];
|
||||
this._initNewSession();
|
||||
}
|
||||
|
||||
_persist(entry: SessionEntry): void {
|
||||
|
|
@ -320,49 +480,93 @@ export class SessionManager {
|
|||
}
|
||||
}
|
||||
|
||||
saveMessage(message: AppMessage): void {
|
||||
private _appendEntry(entry: ConversationEntry): void {
|
||||
this.inMemoryEntries.push(entry);
|
||||
this.byId.set(entry.id, entry);
|
||||
this.leafId = entry.id;
|
||||
this._persist(entry);
|
||||
}
|
||||
|
||||
saveMessage(message: AppMessage): string {
|
||||
const entry: SessionMessageEntry = {
|
||||
type: "message",
|
||||
id: uuidv4(),
|
||||
parentId: this.leafId || null,
|
||||
timestamp: new Date().toISOString(),
|
||||
message,
|
||||
};
|
||||
this.inMemoryEntries.push(entry);
|
||||
this._persist(entry);
|
||||
this._appendEntry(entry);
|
||||
return entry.id;
|
||||
}
|
||||
|
||||
saveThinkingLevelChange(thinkingLevel: string): void {
|
||||
saveThinkingLevelChange(thinkingLevel: string): string {
|
||||
const entry: ThinkingLevelChangeEntry = {
|
||||
type: "thinking_level_change",
|
||||
id: uuidv4(),
|
||||
parentId: this.leafId || null,
|
||||
timestamp: new Date().toISOString(),
|
||||
thinkingLevel,
|
||||
};
|
||||
this.inMemoryEntries.push(entry);
|
||||
this._persist(entry);
|
||||
this._appendEntry(entry);
|
||||
return entry.id;
|
||||
}
|
||||
|
||||
saveModelChange(provider: string, modelId: string): void {
|
||||
saveModelChange(provider: string, modelId: string): string {
|
||||
const entry: ModelChangeEntry = {
|
||||
type: "model_change",
|
||||
id: uuidv4(),
|
||||
parentId: this.leafId || null,
|
||||
timestamp: new Date().toISOString(),
|
||||
provider,
|
||||
modelId,
|
||||
};
|
||||
this.inMemoryEntries.push(entry);
|
||||
this._persist(entry);
|
||||
this._appendEntry(entry);
|
||||
return entry.id;
|
||||
}
|
||||
|
||||
saveCompaction(entry: CompactionEntry): void {
|
||||
this.inMemoryEntries.push(entry);
|
||||
this._persist(entry);
|
||||
saveCompaction(summary: string, firstKeptEntryId: string, tokensBefore: number): string {
|
||||
const entry: CompactionEntry = {
|
||||
type: "compaction",
|
||||
id: uuidv4(),
|
||||
parentId: this.leafId || null,
|
||||
timestamp: new Date().toISOString(),
|
||||
summary,
|
||||
firstKeptEntryId,
|
||||
tokensBefore,
|
||||
};
|
||||
this._appendEntry(entry);
|
||||
return entry.id;
|
||||
}
|
||||
|
||||
// =========================================================================
|
||||
// Tree Traversal
|
||||
// =========================================================================
|
||||
|
||||
getLeafUuid(): string {
|
||||
return this.leafId;
|
||||
}
|
||||
|
||||
getEntry(id: string): ConversationEntry | undefined {
|
||||
return this.byId.get(id);
|
||||
}
|
||||
|
||||
/** Walk from entry to root, returning path (conversation entries only) */
|
||||
getPath(fromId?: string): ConversationEntry[] {
|
||||
const path: ConversationEntry[] = [];
|
||||
let current = this.byId.get(fromId ?? this.leafId);
|
||||
while (current) {
|
||||
path.unshift(current);
|
||||
current = current.parentId ? this.byId.get(current.parentId) : undefined;
|
||||
}
|
||||
return path;
|
||||
}
|
||||
|
||||
/**
|
||||
* Build the session context (what gets sent to the LLM).
|
||||
* If compacted, returns summary + kept messages. Otherwise all messages.
|
||||
* Includes thinking level and model.
|
||||
* Uses tree traversal from current leaf.
|
||||
*/
|
||||
buildSessionContext(): SessionContext {
|
||||
return buildSessionContext(this.getEntries());
|
||||
return buildSessionContext(this.getEntries(), this.leafId);
|
||||
}
|
||||
|
||||
/**
|
||||
|
|
@ -373,6 +577,35 @@ export class SessionManager {
|
|||
return [...this.inMemoryEntries];
|
||||
}
|
||||
|
||||
// =========================================================================
|
||||
// Branching
|
||||
// =========================================================================
|
||||
|
||||
/** Branch in-place by changing the leaf pointer */
|
||||
branchInPlace(branchFromId: string): void {
|
||||
if (!this.byId.has(branchFromId)) {
|
||||
throw new Error(`Entry ${branchFromId} not found`);
|
||||
}
|
||||
this.leafId = branchFromId;
|
||||
}
|
||||
|
||||
/** Branch with a summary of the abandoned path */
|
||||
branchWithSummary(branchFromId: string, summary: string): string {
|
||||
if (!this.byId.has(branchFromId)) {
|
||||
throw new Error(`Entry ${branchFromId} not found`);
|
||||
}
|
||||
this.leafId = branchFromId;
|
||||
const entry: BranchSummaryEntry = {
|
||||
type: "branch_summary",
|
||||
id: uuidv4(),
|
||||
parentId: branchFromId,
|
||||
timestamp: new Date().toISOString(),
|
||||
summary,
|
||||
};
|
||||
this._appendEntry(entry);
|
||||
return entry.id;
|
||||
}
|
||||
|
||||
createBranchedSessionFromEntries(entries: SessionEntry[], branchBeforeIndex: number): string | null {
|
||||
const newSessionId = uuidv4();
|
||||
const timestamp = new Date().toISOString().replace(/[:.]/g, "-");
|
||||
|
|
@ -385,6 +618,7 @@ export class SessionManager {
|
|||
if (entry.type === "session") {
|
||||
newEntries.push({
|
||||
...entry,
|
||||
version: CURRENT_SESSION_VERSION,
|
||||
id: newSessionId,
|
||||
timestamp: new Date().toISOString(),
|
||||
branchedFrom: this.persist ? this.sessionFile : undefined,
|
||||
|
|
@ -402,6 +636,7 @@ export class SessionManager {
|
|||
}
|
||||
this.inMemoryEntries = newEntries;
|
||||
this.sessionId = newSessionId;
|
||||
this._buildIndex();
|
||||
return null;
|
||||
}
|
||||
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue