getcompanion-ai/co-mono
1
0
Fork 0
mirror of https://github.com/getcompanion-ai/co-mono.git synced 2026-04-15 17:00:59 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

Session tree structure with id/parentId linking

Browse source 
- 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 2025-12-25 23:46:44 +01:00
parent 04a764742e
commit c58d5f20a4
12 changed files with 6778 additions and 6297 deletions
Download patch file Download diff file Expand all files Collapse all files

12121
packages/ai/src/models.generated.ts
View file

File diff suppressed because it is too large Load diff

86
packages/coding-agent/docs/session-tree-plan.md Normal file
View file

@ -0,0 +1,86 @@
# Session Tree Implementation Plan
Reference: [session-tree.md](./session-tree.md)
## Phase 1: SessionManager Core
- [x] Update entry types with `id`, `parentId` fields (using TreeNode intersection)
- [x] Add `version` field to `SessionHeader`
- [x] Change `CompactionEntry.firstKeptEntryIndex` → `firstKeptEntryId`
- [x] Add `BranchSummaryEntry` type
- [x] Add `byId: Map<string, ConversationEntry>` index
- [x] Add `leafId: string` tracking
- [x] Implement `getPath(fromId?)` tree traversal
- [x] Implement `getEntry(id)` lookup
- [x] Implement `getLeafId()` helper
- [x] Update `_buildIndex()` to populate `byId` map
- [x] Update `saveMessage()` to include id/parentId (returns id)
- [x] Update `saveCompaction()` signature and fields (returns id)
- [x] Update `saveThinkingLevelChange()` to include id/parentId (returns id)
- [x] Update `saveModelChange()` to include id/parentId (returns id)
- [x] Update `buildSessionContext()` to use `getPath()` traversal
### Type Hierarchy
```typescript
// Tree fields (added by SessionManager)
interface TreeNode { id, parentId, timestamp }
// Content types (for input)
interface MessageContent { type: "message"; message: AppMessage }
interface CompactionContent { type: "compaction"; summary; firstKeptEntryId; tokensBefore }
// etc...
// Full entry types (TreeNode & Content)
type SessionMessageEntry = TreeNode & MessageContent;
type CompactionEntry = TreeNode & CompactionContent;
// etc...
```
## Phase 2: Migration
- [x] Add `CURRENT_SESSION_VERSION = 2` constant
- [x] Implement `_migrateToV2()` for v1→v2
- [x] Update `setSessionFile()` to detect version and migrate
- [x] Implement `_rewriteFile()` for post-migration persistence
- [x] Handle `firstKeptEntryIndex` → `firstKeptEntryId` conversion in migration
## Phase 3: Branching
- [x] Implement `branchInPlace(id)` - switch leaf pointer
- [x] Implement `branchWithSummary(id, summary)` - create summary entry
- [x] Update `branchToNewFile()` to use IDs (no remapping)
- [ ] Update `AgentSession.branch()` to use new API
## Phase 4: Compaction Integration
- [x] Update `compaction.ts` to work with IDs
- [x] Update `prepareCompaction()` to return `firstKeptEntryId`
- [x] Update `compact()` to return `CompactionResult` with `firstKeptEntryId`
- [x] Update `AgentSession` compaction methods
- [x] Add `firstKeptEntryId` to `before_compact` hook event
## Phase 5: Testing
- [ ] Add test fixtures from existing sessions
- [ ] Test migration of v1 sessions
- [ ] Test context building with tree structure
- [ ] Test branching operations
- [ ] Test compaction with IDs
- [x] Update existing tests for new types
## Phase 6: UI Integration
- [ ] Update `/branch` command for new API
- [ ] Add `/branch-here` command for in-place branching
- [ ] Add `/branches` command to list branches (future)
- [ ] Update session display to show tree info (future)
## Notes
- All save methods return the new entry's ID
- Migration rewrites file on first load if version < CURRENT_VERSION
- Existing sessions become linear chains after migration (parentId = previous entry)
- Tree features available immediately after migration
- SessionHeader does NOT have id/parentId (it's metadata, not part of tree)
- Content types allow clean input/output separation

10
packages/coding-agent/examples/hooks/custom-compaction.ts
View file

@ -94,14 +94,12 @@ Format the summary as structured markdown with clear sections.`,
return;
}
// Return a compaction entry that discards ALL messages
// firstKeptEntryIndex points past all current entries
// Return compaction content - SessionManager adds id/parentId
// Use firstKeptEntryId from event to keep recent messages
return {
compactionEntry: {
type: "compaction" as const,
timestamp: new Date().toISOString(),
compaction: {
summary,
firstKeptEntryIndex: entries.length,
firstKeptEntryId: event.firstKeptEntryId,
tokensBefore,
},
};

89
packages/coding-agent/src/core/agent-session.ts
View file

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

43
packages/coding-agent/src/core/compaction.ts
View file

@ -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,
};
}

10
packages/coding-agent/src/core/hooks/types.ts
View file

@ -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;
};
}
// ============================================================================

411
packages/coding-agent/src/core/session-manager.ts
View file

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

12
packages/coding-agent/src/index.ts
View file

@ -107,11 +107,20 @@ export {
readOnlyTools,
} from "./core/sdk.js";
export {
type BranchSummaryContent,
type BranchSummaryEntry,
buildSessionContext,
type CompactionContent,
type CompactionEntry,
type ConversationContent,
type ConversationEntry,
CURRENT_SESSION_VERSION,
createSummaryMessage,
getLatestCompactionEntry,
type MessageContent,
type ModelChangeContent,
type ModelChangeEntry,
migrateSessionEntries,
parseSessionEntries,
type SessionContext as LoadedSession,
type SessionEntry,
@ -122,6 +131,9 @@ export {
SUMMARY_PREFIX,
SUMMARY_SUFFIX,
type ThinkingLevelChangeEntry,
type ThinkingLevelContent,
// Tree types (v2)
type TreeNode,
} from "./core/session-manager.js";
export {
type CompactionSettings,

4
packages/coding-agent/test/agent-session-compaction.test.ts
View file

@ -156,9 +156,9 @@ describe.skipIf(!API_KEY)("AgentSession compaction e2e", () => {
expect(compaction.type).toBe("compaction");
if (compaction.type === "compaction") {
expect(compaction.summary.length).toBeGreaterThan(0);
// firstKeptEntryIndex can be 0 if all messages fit within keepRecentTokens
// firstKeptEntryId can be 0 if all messages fit within keepRecentTokens
// (which is the case for small conversations)
expect(compaction.firstKeptEntryIndex).toBeGreaterThanOrEqual(0);
expect(compaction.firstKeptEntryId).toBeGreaterThanOrEqual(0);
expect(compaction.tokensBefore).toBeGreaterThan(0);
}
}, 120000);

20
packages/coding-agent/test/compaction-hooks-example.test.ts
View file

@ -4,7 +4,6 @@
import { describe, expect, it } from "vitest";
import type { HookAPI } from "../src/core/hooks/index.js";
import type { CompactionEntry } from "../src/core/session-manager.js";
describe("Documentation example", () => {
it("custom compaction example should type-check correctly", () => {
@ -20,29 +19,30 @@ describe("Documentation example", () => {
const tokensBefore = event.tokensBefore;
const model = event.model;
const resolveApiKey = event.resolveApiKey;
const firstKeptEntryId = event.firstKeptEntryId;
// Verify types
expect(Array.isArray(messages)).toBe(true);
expect(Array.isArray(messagesToKeep)).toBe(true);
expect(typeof cutPoint.firstKeptEntryIndex).toBe("number");
expect(typeof cutPoint.firstKeptEntryIndex).toBe("number"); // cutPoint still uses index
expect(typeof tokensBefore).toBe("number");
expect(model).toBeDefined();
expect(typeof resolveApiKey).toBe("function");
expect(typeof firstKeptEntryId).toBe("string");
const summary = messages
.filter((m) => m.role === "user")
.map((m) => `- ${typeof m.content === "string" ? m.content.slice(0, 100) : "[complex]"}`)
.join("\n");
const compactionEntry: CompactionEntry = {
type: "compaction",
timestamp: new Date().toISOString(),
summary: `User requests:\n${summary}`,
firstKeptEntryIndex: event.cutPoint.firstKeptEntryIndex,
tokensBefore: event.tokensBefore,
// Hooks return compaction content - SessionManager adds id/parentId
return {
compaction: {
summary: `User requests:\n${summary}`,
firstKeptEntryId,
tokensBefore,
},
};
return { compactionEntry };
});
};

214
packages/coding-agent/test/compaction.test.ts
View file

@ -3,7 +3,7 @@ import type { AssistantMessage, Usage } from "@mariozechner/pi-ai";
import { getModel } from "@mariozechner/pi-ai";
import { readFileSync } from "fs";
import { join } from "path";
import { describe, expect, it } from "vitest";
import { beforeEach, describe, expect, it } from "vitest";
import {
type CompactionSettings,
calculateContextTokens,
@ -17,9 +17,12 @@ import {
buildSessionContext,
type CompactionEntry,
createSummaryMessage,
type ModelChangeEntry,
migrateSessionEntries,
parseSessionEntries,
type SessionEntry,
type SessionMessageEntry,
type ThinkingLevelChangeEntry,
} from "../src/core/session-manager.js";
// ============================================================================
@ -29,7 +32,9 @@ import {
function loadLargeSessionEntries(): SessionEntry[] {
const sessionPath = join(__dirname, "fixtures/large-session.jsonl");
const content = readFileSync(sessionPath, "utf-8");
return parseSessionEntries(content);
const entries = parseSessionEntries(content);
migrateSessionEntries(entries); // Add id/parentId for v1 fixtures
return entries;
}
function createMockUsage(input: number, output: number, cacheRead = 0, cacheWrite = 0): Usage {
@ -60,18 +65,82 @@ function createAssistantMessage(text: string, usage?: Usage): AssistantMessage {
};
}
function createMessageEntry(message: AppMessage): SessionMessageEntry {
return { type: "message", timestamp: new Date().toISOString(), message };
let entryCounter = 0;
let lastId: string | null = null;
function resetEntryCounter() {
entryCounter = 0;
lastId = null;
}
function createCompactionEntry(summary: string, firstKeptEntryIndex: number): CompactionEntry {
// Reset counter before each test to get predictable IDs
beforeEach(() => {
resetEntryCounter();
});
function createSessionHeader() {
return {
type: "session" as const,
version: 2,
id: "test-session",
timestamp: "",
cwd: "",
};
}
function createMessageEntry(message: AppMessage): SessionMessageEntry {
const id = `test-id-${entryCounter++}`;
const entry: SessionMessageEntry = {
type: "message",
id,
parentId: lastId,
timestamp: new Date().toISOString(),
message,
};
lastId = id;
return entry;
}
function createCompactionEntry(summary: string, firstKeptEntryId: string): CompactionEntry {
const id = `test-id-${entryCounter++}`;
const entry: CompactionEntry = {
type: "compaction",
id,
parentId: lastId,
timestamp: new Date().toISOString(),
summary,
firstKeptEntryIndex,
firstKeptEntryId,
tokensBefore: 10000,
};
lastId = id;
return entry;
}
function createModelChangeEntry(provider: string, modelId: string): ModelChangeEntry {
const id = `test-id-${entryCounter++}`;
const entry: ModelChangeEntry = {
type: "model_change",
id,
parentId: lastId,
timestamp: new Date().toISOString(),
provider,
modelId,
};
lastId = id;
return entry;
}
function createThinkingLevelEntry(thinkingLevel: string): ThinkingLevelChangeEntry {
const id = `test-id-${entryCounter++}`;
const entry: ThinkingLevelChangeEntry = {
type: "thinking_level_change",
id,
parentId: lastId,
timestamp: new Date().toISOString(),
thinkingLevel,
};
lastId = id;
return entry;
}
// ============================================================================
@ -248,78 +317,59 @@ describe("buildSessionContext", () => {
});
it("should handle single compaction", () => {
// indices: 0=session, 1=u1, 2=a1, 3=u2, 4=a2, 5=compaction, 6=u3, 7=a3
const entries: SessionEntry[] = [
{
type: "session",
id: "1",
timestamp: "",
cwd: "",
},
createMessageEntry(createUserMessage("1")),
createMessageEntry(createAssistantMessage("a")),
createMessageEntry(createUserMessage("2")),
createMessageEntry(createAssistantMessage("b")),
createCompactionEntry("Summary of 1,a,2,b", 3), // keep from index 3 (u2) onwards
createMessageEntry(createUserMessage("3")),
createMessageEntry(createAssistantMessage("c")),
];
// IDs: u1=test-id-0, a1=test-id-1, u2=test-id-2, a2=test-id-3, compaction=test-id-4, u3=test-id-5, a3=test-id-6
const u1 = createMessageEntry(createUserMessage("1"));
const a1 = createMessageEntry(createAssistantMessage("a"));
const u2 = createMessageEntry(createUserMessage("2"));
const a2 = createMessageEntry(createAssistantMessage("b"));
const compaction = createCompactionEntry("Summary of 1,a,2,b", u2.id); // keep from u2 onwards
const u3 = createMessageEntry(createUserMessage("3"));
const a3 = createMessageEntry(createAssistantMessage("c"));
const entries: SessionEntry[] = [createSessionHeader(), u1, a1, u2, a2, compaction, u3, a3];
const loaded = buildSessionContext(entries);
// summary + kept (u2,a2 from idx 3-4) + after (u3,a3 from idx 6-7) = 5
// summary + kept (u2, a2) + after (u3, a3) = 5
expect(loaded.messages.length).toBe(5);
expect(loaded.messages[0].role).toBe("user");
expect((loaded.messages[0] as any).content).toContain("Summary of 1,a,2,b");
});
it("should handle multiple compactions (only latest matters)", () => {
// indices: 0=session, 1=u1, 2=a1, 3=compact1, 4=u2, 5=b, 6=u3, 7=c, 8=compact2, 9=u4, 10=d
const entries: SessionEntry[] = [
{
type: "session",
id: "1",
timestamp: "",
cwd: "",
},
createMessageEntry(createUserMessage("1")),
createMessageEntry(createAssistantMessage("a")),
createCompactionEntry("First summary", 1), // keep from index 1
createMessageEntry(createUserMessage("2")),
createMessageEntry(createAssistantMessage("b")),
createMessageEntry(createUserMessage("3")),
createMessageEntry(createAssistantMessage("c")),
createCompactionEntry("Second summary", 6), // keep from index 6 (u3) onwards
createMessageEntry(createUserMessage("4")),
createMessageEntry(createAssistantMessage("d")),
];
// First batch
const u1 = createMessageEntry(createUserMessage("1"));
const a1 = createMessageEntry(createAssistantMessage("a"));
const compact1 = createCompactionEntry("First summary", u1.id);
// Second batch
const u2 = createMessageEntry(createUserMessage("2"));
const b = createMessageEntry(createAssistantMessage("b"));
const u3 = createMessageEntry(createUserMessage("3"));
const c = createMessageEntry(createAssistantMessage("c"));
const compact2 = createCompactionEntry("Second summary", u3.id); // keep from u3 onwards
// After second compaction
const u4 = createMessageEntry(createUserMessage("4"));
const d = createMessageEntry(createAssistantMessage("d"));
const entries: SessionEntry[] = [createSessionHeader(), u1, a1, compact1, u2, b, u3, c, compact2, u4, d];
const loaded = buildSessionContext(entries);
// summary + kept from idx 6 (u3,c) + after (u4,d) = 5
// summary + kept from u3 (u3, c) + after (u4, d) = 5
expect(loaded.messages.length).toBe(5);
expect((loaded.messages[0] as any).content).toContain("Second summary");
});
it("should clamp firstKeptEntryIndex to valid range", () => {
// indices: 0=session, 1=u1, 2=a1, 3=compact1, 4=u2, 5=b, 6=compact2
const entries: SessionEntry[] = [
{
type: "session",
id: "1",
timestamp: "",
cwd: "",
},
createMessageEntry(createUserMessage("1")),
createMessageEntry(createAssistantMessage("a")),
createCompactionEntry("First summary", 1),
createMessageEntry(createUserMessage("2")),
createMessageEntry(createAssistantMessage("b")),
createCompactionEntry("Second summary", 0), // index 0 is before compaction1, should still work
];
it("should keep all messages when firstKeptEntryId is first entry", () => {
const u1 = createMessageEntry(createUserMessage("1"));
const a1 = createMessageEntry(createAssistantMessage("a"));
const compact1 = createCompactionEntry("First summary", u1.id); // keep from first entry
const u2 = createMessageEntry(createUserMessage("2"));
const b = createMessageEntry(createAssistantMessage("b"));
const entries: SessionEntry[] = [createSessionHeader(), u1, a1, compact1, u2, b];
const loaded = buildSessionContext(entries);
// Keeps from index 0, but compaction entries are skipped, so u1,a1,u2,b = 4 + summary = 5
// Actually index 0 is session header, so messages are u1,a1,u2,b
expect(loaded.messages.length).toBe(5); // summary + 4 messages
// summary + all messages (u1, a1, u2, b) = 5
expect(loaded.messages.length).toBe(5);
});
it("should track model and thinking level changes", () => {
@ -331,9 +381,9 @@ describe("buildSessionContext", () => {
cwd: "",
},
createMessageEntry(createUserMessage("1")),
{ type: "model_change", timestamp: "", provider: "openai", modelId: "gpt-4" },
createModelChangeEntry("openai", "gpt-4"),
createMessageEntry(createAssistantMessage("a")),
{ type: "thinking_level_change", timestamp: "", thinkingLevel: "high" },
createThinkingLevelEntry("high"),
];
const loaded = buildSessionContext(entries);
@ -380,27 +430,26 @@ describe("Large session fixture", () => {
// ============================================================================
describe.skipIf(!process.env.ANTHROPIC_OAUTH_TOKEN)("LLM summarization", () => {
it("should generate a compaction event for the large session", async () => {
it("should generate a compaction result for the large session", async () => {
const entries = loadLargeSessionEntries();
const model = getModel("anthropic", "claude-sonnet-4-5")!;
const compactionEvent = await compact(
const compactionResult = await compact(
entries,
model,
DEFAULT_COMPACTION_SETTINGS,
process.env.ANTHROPIC_OAUTH_TOKEN!,
);
expect(compactionEvent.type).toBe("compaction");
expect(compactionEvent.summary.length).toBeGreaterThan(100);
expect(compactionEvent.firstKeptEntryIndex).toBeGreaterThan(0);
expect(compactionEvent.tokensBefore).toBeGreaterThan(0);
expect(compactionResult.summary.length).toBeGreaterThan(100);
expect(compactionResult.firstKeptEntryId).toBeTruthy();
expect(compactionResult.tokensBefore).toBeGreaterThan(0);
console.log("Summary length:", compactionEvent.summary.length);
console.log("First kept entry index:", compactionEvent.firstKeptEntryIndex);
console.log("Tokens before:", compactionEvent.tokensBefore);
console.log("Summary length:", compactionResult.summary.length);
console.log("First kept entry ID:", compactionResult.firstKeptEntryId);
console.log("Tokens before:", compactionResult.tokensBefore);
console.log("\n--- SUMMARY ---\n");
console.log(compactionEvent.summary);
console.log(compactionResult.summary);
}, 60000);
it("should produce valid session after compaction", async () => {
@ -408,21 +457,30 @@ describe.skipIf(!process.env.ANTHROPIC_OAUTH_TOKEN)("LLM summarization", () => {
const loaded = buildSessionContext(entries);
const model = getModel("anthropic", "claude-sonnet-4-5")!;
const compactionEvent = await compact(
const compactionResult = await compact(
entries,
model,
DEFAULT_COMPACTION_SETTINGS,
process.env.ANTHROPIC_OAUTH_TOKEN!,
);
// Simulate appending compaction to entries
const newEntries = [...entries, compactionEvent];
// Simulate appending compaction to entries by creating a proper entry
const lastEntry = entries[entries.length - 1];
const parentId = lastEntry.type === "session" ? null : lastEntry.id;
const compactionEntry: CompactionEntry = {
type: "compaction",
id: "compaction-test-id",
parentId,
timestamp: new Date().toISOString(),
...compactionResult,
};
const newEntries = [...entries, compactionEntry];
const reloaded = buildSessionContext(newEntries);
// Should have summary + kept messages
expect(reloaded.messages.length).toBeLessThan(loaded.messages.length);
expect(reloaded.messages[0].role).toBe("user");
expect((reloaded.messages[0] as any).content).toContain(compactionEvent.summary);
expect((reloaded.messages[0] as any).content).toContain(compactionResult.summary);
console.log("Original messages:", loaded.messages.length);
console.log("After compaction:", reloaded.messages.length);

55
packages/mom/src/context.ts
View file

@ -15,10 +15,14 @@ import {
buildSessionContext,
type CompactionEntry,
type LoadedSession,
type MessageContent,
type ModelChangeContent,
type ModelChangeEntry,
type SessionEntry,
type SessionMessageEntry,
type ThinkingLevelChangeEntry,
type ThinkingLevelContent,
type TreeNode,
} from "@mariozechner/pi-coding-agent";
import { randomBytes } from "crypto";
import { appendFileSync, existsSync, mkdirSync, readFileSync, writeFileSync } from "fs";
@ -49,6 +53,7 @@ export class MomSessionManager {
private channelDir: string;
private flushed: boolean = false;
private inMemoryEntries: SessionEntry[] = [];
private leafId: string | null = null;
constructor(channelDir: string) {
this.channelDir = channelDir;
@ -64,12 +69,14 @@ export class MomSessionManager {
if (existsSync(this.contextFile)) {
this.inMemoryEntries = this.loadEntriesFromFile();
this.sessionId = this.extractSessionId() || uuidv4();
this._updateLeafId();
this.flushed = true;
} else {
this.sessionId = uuidv4();
this.inMemoryEntries = [
{
type: "session",
version: 2,
id: this.sessionId,
timestamp: new Date().toISOString(),
cwd: this.channelDir,
@ -79,6 +86,28 @@ export class MomSessionManager {
// Note: syncFromLog() is called explicitly from agent.ts with excludeTimestamp
}
private _updateLeafId(): void {
for (let i = this.inMemoryEntries.length - 1; i >= 0; i--) {
const entry = this.inMemoryEntries[i];
if (entry.type !== "session") {
this.leafId = entry.id;
return;
}
}
this.leafId = null;
}
private _createTreeNode(): TreeNode {
const id = uuidv4();
const node: TreeNode = {
id,
parentId: this.leafId,
timestamp: new Date().toISOString(),
};
this.leafId = id;
return node;
}
private _persist(entry: SessionEntry): void {
const hasAssistant = this.inMemoryEntries.some((e) => e.type === "message" && e.message.role === "assistant");
if (!hasAssistant) return;
@ -206,11 +235,15 @@ export class MomSessionManager {
newMessages.sort((a, b) => new Date(a.timestamp).getTime() - new Date(b.timestamp).getTime());
for (const { timestamp, message } of newMessages) {
const id = uuidv4();
const entry: SessionMessageEntry = {
type: "message",
id,
parentId: this.leafId,
timestamp, // Use log date as entry timestamp for consistent deduplication
message,
};
this.leafId = id;
this.inMemoryEntries.push(entry);
appendFileSync(this.contextFile, `${JSON.stringify(entry)}\n`);
@ -247,32 +280,22 @@ export class MomSessionManager {
}
saveMessage(message: AppMessage): void {
const entry: SessionMessageEntry = {
type: "message",
timestamp: new Date().toISOString(),
message,
};
const content: MessageContent = { type: "message", message };
const entry: SessionMessageEntry = { ...this._createTreeNode(), ...content };
this.inMemoryEntries.push(entry);
this._persist(entry);
}
saveThinkingLevelChange(thinkingLevel: string): void {
const entry: ThinkingLevelChangeEntry = {
type: "thinking_level_change",
timestamp: new Date().toISOString(),
thinkingLevel,
};
const content: ThinkingLevelContent = { type: "thinking_level_change", thinkingLevel };
const entry: ThinkingLevelChangeEntry = { ...this._createTreeNode(), ...content };
this.inMemoryEntries.push(entry);
this._persist(entry);
}
saveModelChange(provider: string, modelId: string): void {
const entry: ModelChangeEntry = {
type: "model_change",
timestamp: new Date().toISOString(),
provider,
modelId,
};
const content: ModelChangeContent = { type: "model_change", provider, modelId };
const entry: ModelChangeEntry = { ...this._createTreeNode(), ...content };
this.inMemoryEntries.push(entry);
this._persist(entry);
}