diff --git a/packages/coding-agent/CHANGELOG.md b/packages/coding-agent/CHANGELOG.md index 43246cf7..987e4d5f 100644 --- a/packages/coding-agent/CHANGELOG.md +++ b/packages/coding-agent/CHANGELOG.md @@ -2,13 +2,41 @@ ## [Unreleased] +### Breaking Changes + +- **Session tree structure (v2)**: Sessions now store entries as a tree with `id`/`parentId` fields, enabling in-place branching without creating new files. Existing v1 sessions are auto-migrated on load. +- **SessionManager API**: + - `saveXXX()` renamed to `appendXXX()` (e.g., `appendMessage`, `appendCompaction`) + - `branchInPlace()` renamed to `branch()` + - `reset()` renamed to `newSession()` + - `createBranchedSessionFromEntries(entries, index)` replaced with `createBranchedSession(leafId)` + - `saveCompaction(entry)` replaced with `appendCompaction(summary, firstKeptEntryId, tokensBefore)` + - `getEntries()` now excludes the session header (use `getHeader()` separately) + - New methods: `getTree()`, `getPath()`, `getLeafUuid()`, `getLeafEntry()`, `getEntry()`, `branchWithSummary()` + - New `appendCustomEntry(customType, data)` for hooks to store custom data +- **Compaction API**: + - `compact()` now returns `CompactionResult` (`{ summary, firstKeptEntryId, tokensBefore }`) instead of `CompactionEntry` + - `CompactionEntry.firstKeptEntryIndex` replaced with `firstKeptEntryId` + - `prepareCompaction()` now returns `firstKeptEntryId` in its result +- **Hook types**: + - `SessionEventResult.compactionEntry` replaced with `SessionEventResult.compaction` (content only, SessionManager adds id/parentId) + - `before_compact` event now includes `firstKeptEntryId` field for hooks that return custom compaction + ### Added - **`enabledModels` setting**: Configure whitelisted models in `settings.json` (same format as `--models` CLI flag). CLI `--models` takes precedence over the setting. +### Changed + +- **Entry IDs**: Session entries now use short 8-character hex IDs instead of full UUIDs +- **API key priority**: `ANTHROPIC_OAUTH_TOKEN` now takes precedence over `ANTHROPIC_API_KEY` +- **New entry types**: `BranchSummaryEntry` for branch context, `CustomEntry` for hook data + ### Fixed - **Edit tool fails on Windows due to CRLF line endings**: Files with CRLF line endings now match correctly when LLMs send LF-only text. Line endings are normalized before matching and restored to original style on write. ([#355](https://github.com/badlogic/pi-mono/issues/355)) +- **Session file validation**: `findMostRecentSession()` now validates session headers before returning, preventing non-session JSONL files from being loaded +- **Compaction error handling**: `generateSummary()` and `generateTurnPrefixSummary()` now throw on LLM errors instead of returning empty strings ## [0.30.2] - 2025-12-26 diff --git a/packages/coding-agent/src/core/agent-session.ts b/packages/coding-agent/src/core/agent-session.ts index b6f45978..1547ba61 100644 --- a/packages/coding-agent/src/core/agent-session.ts +++ b/packages/coding-agent/src/core/agent-session.ts @@ -211,7 +211,7 @@ export class AgentSession { // Handle session persistence if (event.type === "message_end") { - this.sessionManager.saveMessage(event.message); + this.sessionManager.appendMessage(event.message); // Track assistant message for auto-compaction (checked on agent_end) if (event.message.role === "assistant") { @@ -535,7 +535,7 @@ export class AgentSession { this._disconnectFromAgent(); await this.abort(); this.agent.reset(); - this.sessionManager.reset(); + this.sessionManager.newSession(); this._queuedMessages = []; this._reconnectToAgent(); @@ -572,7 +572,7 @@ export class AgentSession { } this.agent.setModel(model); - this.sessionManager.saveModelChange(model.provider, model.id); + this.sessionManager.appendModelChange(model.provider, model.id); this.settingsManager.setDefaultModelAndProvider(model.provider, model.id); // Re-clamp thinking level for new model's capabilities @@ -611,7 +611,7 @@ export class AgentSession { // Apply model this.agent.setModel(next.model); - this.sessionManager.saveModelChange(next.model.provider, next.model.id); + this.sessionManager.appendModelChange(next.model.provider, next.model.id); this.settingsManager.setDefaultModelAndProvider(next.model.provider, next.model.id); // Apply thinking level (setThinkingLevel clamps to model capabilities) @@ -638,7 +638,7 @@ export class AgentSession { } this.agent.setModel(nextModel); - this.sessionManager.saveModelChange(nextModel.provider, nextModel.id); + this.sessionManager.appendModelChange(nextModel.provider, nextModel.id); this.settingsManager.setDefaultModelAndProvider(nextModel.provider, nextModel.id); // Re-clamp thinking level for new model's capabilities @@ -671,7 +671,7 @@ export class AgentSession { effectiveLevel = "high"; } this.agent.setThinkingLevel(effectiveLevel); - this.sessionManager.saveThinkingLevelChange(effectiveLevel); + this.sessionManager.appendThinkingLevelChange(effectiveLevel); this.settingsManager.setDefaultThinkingLevel(effectiveLevel); } @@ -831,7 +831,7 @@ export class AgentSession { throw new Error("Compaction cancelled"); } - this.sessionManager.saveCompaction(summary, firstKeptEntryId, tokensBefore); + this.sessionManager.appendCompaction(summary, firstKeptEntryId, tokensBefore); const newEntries = this.sessionManager.getEntries(); const sessionContext = this.sessionManager.buildSessionContext(); this.agent.replaceMessages(sessionContext.messages); @@ -1013,7 +1013,7 @@ export class AgentSession { return; } - this.sessionManager.saveCompaction(summary, firstKeptEntryId, tokensBefore); + this.sessionManager.appendCompaction(summary, firstKeptEntryId, tokensBefore); const newEntries = this.sessionManager.getEntries(); const sessionContext = this.sessionManager.buildSessionContext(); this.agent.replaceMessages(sessionContext.messages); @@ -1271,7 +1271,7 @@ export class AgentSession { this.agent.appendMessage(bashMessage); // Save to session - this.sessionManager.saveMessage(bashMessage); + this.sessionManager.appendMessage(bashMessage); } return result; @@ -1309,7 +1309,7 @@ export class AgentSession { this.agent.appendMessage(bashMessage); // Save to session - this.sessionManager.saveMessage(bashMessage); + this.sessionManager.appendMessage(bashMessage); } this._pendingBashMessages = []; @@ -1431,8 +1431,12 @@ export class AgentSession { skipConversationRestore = result?.skipConversationRestore ?? false; } - // Create branched session (returns null in --no-session mode) - const newSessionFile = this.sessionManager.createBranchedSessionFromEntries(entries, entryIndex); + // Create branched session ending before the selected message (returns null in --no-session mode) + // User will re-enter/edit the selected message + if (!selectedEntry.parentId) { + throw new Error("Cannot branch from first message"); + } + const newSessionFile = this.sessionManager.createBranchedSession(selectedEntry.parentId); // Update session file if we have one (file-based mode) if (newSessionFile !== null) { diff --git a/packages/coding-agent/src/core/session-manager.ts b/packages/coding-agent/src/core/session-manager.ts index 47a41c7c..148eec05 100644 --- a/packages/coding-agent/src/core/session-manager.ts +++ b/packages/coding-agent/src/core/session-manager.ts @@ -26,54 +26,47 @@ export interface SessionHeader { branchedFrom?: string; } -export interface MessageContent { - type: "message"; - message: AppMessage; -} - -export interface ThinkingLevelContent { - type: "thinking_level_change"; - thinkingLevel: string; -} - -export interface ModelChangeContent { - type: "model_change"; - provider: string; - modelId: string; -} - -export interface CompactionContent { - type: "compaction"; - summary: string; - firstKeptEntryId: string; - tokensBefore: number; -} - -export interface BranchSummaryContent { - type: "branch_summary"; - summary: string; -} - -/** Union of all content types (for "write" methods in SessionManager) */ -export type SessionContent = - | MessageContent - | ThinkingLevelContent - | ModelChangeContent - | CompactionContent - | BranchSummaryContent; - -export interface TreeNode { +export interface SessionEntryBase { type: string; id: string; parentId: string | null; timestamp: string; } -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; +export interface SessionMessageEntry extends SessionEntryBase { + type: "message"; + message: AppMessage; +} + +export interface ThinkingLevelChangeEntry extends SessionEntryBase { + type: "thinking_level_change"; + thinkingLevel: string; +} + +export interface ModelChangeEntry extends SessionEntryBase { + type: "model_change"; + provider: string; + modelId: string; +} + +export interface CompactionEntry extends SessionEntryBase { + type: "compaction"; + summary: string; + firstKeptEntryId: string; + tokensBefore: number; +} + +export interface BranchSummaryEntry extends SessionEntryBase { + type: "branch_summary"; + summary: string; +} + +/** Custom entry for hooks. Use customType to identify your hook's entries. */ +export interface CustomEntry extends SessionEntryBase { + type: "custom"; + customType: string; + data?: unknown; +} /** Session entry - has id/parentId for tree structure (returned by "read" methods in SessionManager) */ export type SessionEntry = @@ -81,11 +74,18 @@ export type SessionEntry = | ThinkingLevelChangeEntry | ModelChangeEntry | CompactionEntry - | BranchSummaryEntry; + | BranchSummaryEntry + | CustomEntry; /** Raw file entry (includes header) */ export type FileEntry = SessionHeader | SessionEntry; +/** Tree node for getTree() - defensive copy of session structure */ +export interface SessionTreeNode { + entry: SessionEntry; + children: SessionTreeNode[]; +} + export interface SessionContext { messages: AppMessage[]; thinkingLevel: string; @@ -387,6 +387,17 @@ export function findMostRecentSession(sessionDir: string): string | null { } } +/** + * Manages conversation sessions as append-only trees stored in JSONL files. + * + * Each session entry has an id and parentId forming a tree structure. The "leaf" + * pointer tracks the current position. Appending creates a child of the current leaf. + * Branching moves the leaf to an earlier entry, allowing new branches without + * modifying history. + * + * Use buildSessionContext() to get the resolved message list for the LLM, which + * handles compaction summaries and follows the path from root to current leaf. + */ export class SessionManager { private sessionId: string = ""; private sessionFile: string = ""; @@ -394,7 +405,7 @@ export class SessionManager { private cwd: string; private persist: boolean; private flushed: boolean = false; - private inMemoryEntries: FileEntry[] = []; + private fileEntries: FileEntry[] = []; private byId: Map = new Map(); private leafId: string = ""; @@ -409,7 +420,7 @@ export class SessionManager { if (sessionFile) { this.setSessionFile(sessionFile); } else { - this._initNewSession(); + this.newSession(); } } @@ -417,22 +428,22 @@ export class SessionManager { setSessionFile(sessionFile: string): void { this.sessionFile = resolve(sessionFile); if (existsSync(this.sessionFile)) { - this.inMemoryEntries = loadEntriesFromFile(this.sessionFile); - const header = this.inMemoryEntries.find((e) => e.type === "session") as SessionHeader | undefined; + this.fileEntries = loadEntriesFromFile(this.sessionFile); + const header = this.fileEntries.find((e) => e.type === "session") as SessionHeader | undefined; this.sessionId = header?.id ?? randomUUID(); - if (migrateToCurrentVersion(this.inMemoryEntries)) { + if (migrateToCurrentVersion(this.fileEntries)) { this._rewriteFile(); } this._buildIndex(); this.flushed = true; } else { - this._initNewSession(); + this.newSession(); } } - private _initNewSession(): void { + newSession(): void { this.sessionId = randomUUID(); const timestamp = new Date().toISOString(); const header: SessionHeader = { @@ -442,7 +453,7 @@ export class SessionManager { timestamp, cwd: this.cwd, }; - this.inMemoryEntries = [header]; + this.fileEntries = [header]; this.byId.clear(); this.leafId = ""; this.flushed = false; @@ -456,7 +467,7 @@ export class SessionManager { private _buildIndex(): void { this.byId.clear(); this.leafId = ""; - for (const entry of this.inMemoryEntries) { + for (const entry of this.fileEntries) { if (entry.type === "session") continue; this.byId.set(entry.id, entry); this.leafId = entry.id; @@ -465,7 +476,7 @@ export class SessionManager { private _rewriteFile(): void { if (!this.persist) return; - const content = `${this.inMemoryEntries.map((e) => JSON.stringify(e)).join("\n")}\n`; + const content = `${this.fileEntries.map((e) => JSON.stringify(e)).join("\n")}\n`; writeFileSync(this.sessionFile, content); } @@ -489,18 +500,14 @@ export class SessionManager { return this.sessionFile; } - reset(): void { - this._initNewSession(); - } - _persist(entry: SessionEntry): void { if (!this.persist) return; - const hasAssistant = this.inMemoryEntries.some((e) => e.type === "message" && e.message.role === "assistant"); + const hasAssistant = this.fileEntries.some((e) => e.type === "message" && e.message.role === "assistant"); if (!hasAssistant) return; if (!this.flushed) { - for (const e of this.inMemoryEntries) { + for (const e of this.fileEntries) { appendFileSync(this.sessionFile, `${JSON.stringify(e)}\n`); } this.flushed = true; @@ -510,13 +517,14 @@ export class SessionManager { } private _appendEntry(entry: SessionEntry): void { - this.inMemoryEntries.push(entry); + this.fileEntries.push(entry); this.byId.set(entry.id, entry); this.leafId = entry.id; this._persist(entry); } - saveMessage(message: AppMessage): string { + /** Append a message as child of current leaf, then advance leaf. Returns entry id. */ + appendMessage(message: AppMessage): string { const entry: SessionMessageEntry = { type: "message", id: generateId(this.byId), @@ -528,7 +536,8 @@ export class SessionManager { return entry.id; } - saveThinkingLevelChange(thinkingLevel: string): string { + /** Append a thinking level change as child of current leaf, then advance leaf. Returns entry id. */ + appendThinkingLevelChange(thinkingLevel: string): string { const entry: ThinkingLevelChangeEntry = { type: "thinking_level_change", id: generateId(this.byId), @@ -540,7 +549,8 @@ export class SessionManager { return entry.id; } - saveModelChange(provider: string, modelId: string): string { + /** Append a model change as child of current leaf, then advance leaf. Returns entry id. */ + appendModelChange(provider: string, modelId: string): string { const entry: ModelChangeEntry = { type: "model_change", id: generateId(this.byId), @@ -553,7 +563,8 @@ export class SessionManager { return entry.id; } - saveCompaction(summary: string, firstKeptEntryId: string, tokensBefore: number): string { + /** Append a compaction summary as child of current leaf, then advance leaf. Returns entry id. */ + appendCompaction(summary: string, firstKeptEntryId: string, tokensBefore: number): string { const entry: CompactionEntry = { type: "compaction", id: generateId(this.byId), @@ -567,6 +578,20 @@ export class SessionManager { return entry.id; } + /** Append a custom entry (for hooks) as child of current leaf, then advance leaf. Returns entry id. */ + appendCustomEntry(customType: string, data?: unknown): string { + const entry: CustomEntry = { + type: "custom", + customType, + data, + id: generateId(this.byId), + parentId: this.leafId || null, + timestamp: new Date().toISOString(), + }; + this._appendEntry(entry); + return entry.id; + } + // ========================================================================= // Tree Traversal // ========================================================================= @@ -575,11 +600,19 @@ export class SessionManager { return this.leafId; } + getLeafEntry(): SessionEntry | undefined { + return this.byId.get(this.leafId); + } + getEntry(id: string): SessionEntry | undefined { return this.byId.get(id); } - /** Walk from entry to root, returning path (conversation entries only) */ + /** + * Walk from entry to root, returning all entries in path order. + * Includes all entry types (messages, compaction, model changes, etc.). + * Use buildSessionContext() to get the resolved messages for the LLM. + */ getPath(fromId?: string): SessionEntry[] { const path: SessionEntry[] = []; let current = this.byId.get(fromId ?? this.leafId); @@ -602,31 +635,75 @@ export class SessionManager { * Get session header. */ getHeader(): SessionHeader | null { - const h = this.inMemoryEntries.find((e) => e.type === "session"); + const h = this.fileEntries.find((e) => e.type === "session"); return h ? (h as SessionHeader) : null; } /** - * Get all session entries (excludes header). Returns a defensive copy. - * Use buildSessionContext() if you need the messages for the LLM. + * Get all session entries (excludes header). Returns a shallow copy. + * The session is append-only: use appendXXX() to add entries, branch() to + * change the leaf pointer. Entries cannot be modified or deleted. */ getEntries(): SessionEntry[] { - return this.inMemoryEntries.filter((e): e is SessionEntry => e.type !== "session"); + return this.fileEntries.filter((e): e is SessionEntry => e.type !== "session"); + } + + /** + * Get the session as a tree structure. Returns a shallow defensive copy of all entries. + * A well-formed session has exactly one root (first entry with parentId === null). + * Orphaned entries (broken parent chain) are also returned as roots. + */ + getTree(): SessionTreeNode[] { + const entries = this.getEntries(); + const nodeMap = new Map(); + const roots: SessionTreeNode[] = []; + + // Create nodes + for (const entry of entries) { + nodeMap.set(entry.id, { entry, children: [] }); + } + + // Build tree + for (const entry of entries) { + const node = nodeMap.get(entry.id)!; + if (entry.parentId === null) { + roots.push(node); + } else { + const parent = nodeMap.get(entry.parentId); + if (parent) { + parent.children.push(node); + } else { + // Orphan - treat as root + roots.push(node); + } + } + } + + return roots; } // ========================================================================= // Branching // ========================================================================= - /** Branch in-place by changing the leaf pointer */ - branchInPlace(branchFromId: string): void { + /** + * Start a new branch from an earlier entry. + * Moves the leaf pointer to the specified entry. The next appendXXX() call + * will create a child of that entry, forming a new branch. Existing entries + * are not modified or deleted. + */ + branch(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 */ + /** + * Start a new branch with a summary of the abandoned path. + * Same as branch(), but also appends a branch_summary entry that captures + * context from the abandoned conversation path. + */ branchWithSummary(branchFromId: string, summary: string): string { if (!this.byId.has(branchFromId)) { throw new Error(`Entry ${branchFromId} not found`); @@ -643,35 +720,41 @@ export class SessionManager { return entry.id; } - createBranchedSessionFromEntries(entries: FileEntry[], branchBeforeIndex: number): string | null { - const newSessionId = randomUUID(); - const timestamp = new Date().toISOString().replace(/[:.]/g, "-"); - const newSessionFile = join(this.getSessionDir(), `${timestamp}_${newSessionId}.jsonl`); - - const newEntries: FileEntry[] = []; - for (let i = 0; i < branchBeforeIndex; i++) { - const entry = entries[i]; - - if (entry.type === "session") { - newEntries.push({ - ...entry, - version: CURRENT_SESSION_VERSION, - id: newSessionId, - timestamp: new Date().toISOString(), - branchedFrom: this.persist ? this.sessionFile : undefined, - }); - } else { - newEntries.push(entry); - } + /** + * Create a new session file containing only the path from root to the specified leaf. + * Useful for extracting a single conversation path from a branched session. + * Returns the new session file path, or null if not persisting. + */ + createBranchedSession(leafId: string): string | null { + const path = this.getPath(leafId); + if (path.length === 0) { + throw new Error(`Entry ${leafId} not found`); } + const newSessionId = randomUUID(); + const timestamp = new Date().toISOString(); + const fileTimestamp = timestamp.replace(/[:.]/g, "-"); + const newSessionFile = join(this.getSessionDir(), `${fileTimestamp}_${newSessionId}.jsonl`); + + const header: SessionHeader = { + type: "session", + version: CURRENT_SESSION_VERSION, + id: newSessionId, + timestamp, + cwd: this.cwd, + branchedFrom: this.persist ? this.sessionFile : undefined, + }; + if (this.persist) { - for (const entry of newEntries) { + appendFileSync(newSessionFile, `${JSON.stringify(header)}\n`); + for (const entry of path) { appendFileSync(newSessionFile, `${JSON.stringify(entry)}\n`); } return newSessionFile; } - this.inMemoryEntries = newEntries; + + // In-memory mode: replace current session with the path + this.fileEntries = [header, ...path]; this.sessionId = newSessionId; this._buildIndex(); return null; diff --git a/packages/coding-agent/src/index.ts b/packages/coding-agent/src/index.ts index 2eafe3c4..6ddaedb1 100644 --- a/packages/coding-agent/src/index.ts +++ b/packages/coding-agent/src/index.ts @@ -107,23 +107,19 @@ export { readOnlyTools, } from "./core/sdk.js"; export { - type BranchSummaryContent, type BranchSummaryEntry, buildSessionContext, - type CompactionContent, type CompactionEntry, CURRENT_SESSION_VERSION, createSummaryMessage, type FileEntry, getLatestCompactionEntry, - type MessageContent, - type ModelChangeContent, type ModelChangeEntry, migrateSessionEntries, parseSessionEntries, - type SessionContent as ConversationContent, - type SessionContext as LoadedSession, + type SessionContext, type SessionEntry, + type SessionEntryBase, type SessionHeader, type SessionInfo, SessionManager, @@ -131,9 +127,6 @@ export { SUMMARY_PREFIX, SUMMARY_SUFFIX, type ThinkingLevelChangeEntry, - type ThinkingLevelContent, - // Tree types (v2) - type TreeNode, } from "./core/session-manager.js"; export { type CompactionSettings, diff --git a/packages/coding-agent/test/session-manager/save-entry.test.ts b/packages/coding-agent/test/session-manager/save-entry.test.ts new file mode 100644 index 00000000..2a618986 --- /dev/null +++ b/packages/coding-agent/test/session-manager/save-entry.test.ts @@ -0,0 +1,55 @@ +import { describe, expect, it } from "vitest"; +import { type CustomEntry, SessionManager } from "../../src/core/session-manager.js"; + +describe("SessionManager.saveCustomEntry", () => { + it("saves custom entries and includes them in tree traversal", () => { + const session = SessionManager.inMemory(); + + // Save a message + const msgId = session.appendMessage({ role: "user", content: "hello", timestamp: 1 }); + + // Save a custom entry + const customId = session.appendCustomEntry("my_hook", { foo: "bar" }); + + // Save another message + const msg2Id = session.appendMessage({ + role: "assistant", + content: [{ type: "text", text: "hi" }], + api: "anthropic-messages", + provider: "anthropic", + model: "test", + usage: { + input: 1, + output: 1, + cacheRead: 0, + cacheWrite: 0, + totalTokens: 2, + cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 }, + }, + stopReason: "stop", + timestamp: 2, + }); + + // Custom entry should be in entries + const entries = session.getEntries(); + expect(entries).toHaveLength(3); + + const customEntry = entries.find((e) => e.type === "custom") as CustomEntry; + expect(customEntry).toBeDefined(); + expect(customEntry.customType).toBe("my_hook"); + expect(customEntry.data).toEqual({ foo: "bar" }); + expect(customEntry.id).toBe(customId); + expect(customEntry.parentId).toBe(msgId); + + // Tree structure should be correct + const path = session.getPath(); + expect(path).toHaveLength(3); + expect(path[0].id).toBe(msgId); + expect(path[1].id).toBe(customId); + expect(path[2].id).toBe(msg2Id); + + // buildSessionContext should work (custom entries skipped in messages) + const ctx = session.buildSessionContext(); + expect(ctx.messages).toHaveLength(2); // only message entries + }); +}); diff --git a/packages/coding-agent/test/session-manager/tree-traversal.test.ts b/packages/coding-agent/test/session-manager/tree-traversal.test.ts new file mode 100644 index 00000000..fa788164 --- /dev/null +++ b/packages/coding-agent/test/session-manager/tree-traversal.test.ts @@ -0,0 +1,483 @@ +import { describe, expect, it } from "vitest"; +import { type CustomEntry, SessionManager } from "../../src/core/session-manager.js"; + +function userMsg(text: string) { + return { role: "user" as const, content: text, timestamp: Date.now() }; +} + +function assistantMsg(text: string) { + return { + role: "assistant" as const, + content: [{ type: "text" as const, text }], + api: "anthropic-messages" as const, + provider: "anthropic", + model: "test", + usage: { + input: 1, + output: 1, + cacheRead: 0, + cacheWrite: 0, + totalTokens: 2, + cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 }, + }, + stopReason: "stop" as const, + timestamp: Date.now(), + }; +} + +describe("SessionManager append and tree traversal", () => { + describe("append operations", () => { + it("appendMessage creates entry with correct parentId chain", () => { + const session = SessionManager.inMemory(); + + const id1 = session.appendMessage(userMsg("first")); + const id2 = session.appendMessage(assistantMsg("second")); + const id3 = session.appendMessage(userMsg("third")); + + const entries = session.getEntries(); + expect(entries).toHaveLength(3); + + expect(entries[0].id).toBe(id1); + expect(entries[0].parentId).toBeNull(); + expect(entries[0].type).toBe("message"); + + expect(entries[1].id).toBe(id2); + expect(entries[1].parentId).toBe(id1); + + expect(entries[2].id).toBe(id3); + expect(entries[2].parentId).toBe(id2); + }); + + it("appendThinkingLevelChange integrates into tree", () => { + const session = SessionManager.inMemory(); + + const msgId = session.appendMessage(userMsg("hello")); + const thinkingId = session.appendThinkingLevelChange("high"); + const _msg2Id = session.appendMessage(assistantMsg("response")); + + const entries = session.getEntries(); + expect(entries).toHaveLength(3); + + const thinkingEntry = entries.find((e) => e.type === "thinking_level_change"); + expect(thinkingEntry).toBeDefined(); + expect(thinkingEntry!.id).toBe(thinkingId); + expect(thinkingEntry!.parentId).toBe(msgId); + + expect(entries[2].parentId).toBe(thinkingId); + }); + + it("appendModelChange integrates into tree", () => { + const session = SessionManager.inMemory(); + + const msgId = session.appendMessage(userMsg("hello")); + const modelId = session.appendModelChange("openai", "gpt-4"); + const _msg2Id = session.appendMessage(assistantMsg("response")); + + const entries = session.getEntries(); + const modelEntry = entries.find((e) => e.type === "model_change"); + expect(modelEntry).toBeDefined(); + expect(modelEntry?.id).toBe(modelId); + expect(modelEntry?.parentId).toBe(msgId); + if (modelEntry?.type === "model_change") { + expect(modelEntry.provider).toBe("openai"); + expect(modelEntry.modelId).toBe("gpt-4"); + } + + expect(entries[2].parentId).toBe(modelId); + }); + + it("appendCompaction integrates into tree", () => { + const session = SessionManager.inMemory(); + + const id1 = session.appendMessage(userMsg("1")); + const id2 = session.appendMessage(assistantMsg("2")); + const compactionId = session.appendCompaction("summary", id1, 1000); + const _id3 = session.appendMessage(userMsg("3")); + + const entries = session.getEntries(); + const compactionEntry = entries.find((e) => e.type === "compaction"); + expect(compactionEntry).toBeDefined(); + expect(compactionEntry?.id).toBe(compactionId); + expect(compactionEntry?.parentId).toBe(id2); + if (compactionEntry?.type === "compaction") { + expect(compactionEntry.summary).toBe("summary"); + expect(compactionEntry.firstKeptEntryId).toBe(id1); + expect(compactionEntry.tokensBefore).toBe(1000); + } + + expect(entries[3].parentId).toBe(compactionId); + }); + + it("appendCustomEntry integrates into tree", () => { + const session = SessionManager.inMemory(); + + const msgId = session.appendMessage(userMsg("hello")); + const customId = session.appendCustomEntry("my_hook", { key: "value" }); + const _msg2Id = session.appendMessage(assistantMsg("response")); + + const entries = session.getEntries(); + const customEntry = entries.find((e) => e.type === "custom") as CustomEntry; + expect(customEntry).toBeDefined(); + expect(customEntry.id).toBe(customId); + expect(customEntry.parentId).toBe(msgId); + expect(customEntry.customType).toBe("my_hook"); + expect(customEntry.data).toEqual({ key: "value" }); + + expect(entries[2].parentId).toBe(customId); + }); + + it("leaf pointer advances after each append", () => { + const session = SessionManager.inMemory(); + + expect(session.getLeafUuid()).toBe(""); + + const id1 = session.appendMessage(userMsg("1")); + expect(session.getLeafUuid()).toBe(id1); + + const id2 = session.appendMessage(assistantMsg("2")); + expect(session.getLeafUuid()).toBe(id2); + + const id3 = session.appendThinkingLevelChange("high"); + expect(session.getLeafUuid()).toBe(id3); + }); + }); + + describe("getPath", () => { + it("returns empty array for empty session", () => { + const session = SessionManager.inMemory(); + expect(session.getPath()).toEqual([]); + }); + + it("returns single entry path", () => { + const session = SessionManager.inMemory(); + const id = session.appendMessage(userMsg("hello")); + + const path = session.getPath(); + expect(path).toHaveLength(1); + expect(path[0].id).toBe(id); + }); + + it("returns full path from root to leaf", () => { + const session = SessionManager.inMemory(); + + const id1 = session.appendMessage(userMsg("1")); + const id2 = session.appendMessage(assistantMsg("2")); + const id3 = session.appendThinkingLevelChange("high"); + const id4 = session.appendMessage(userMsg("3")); + + const path = session.getPath(); + expect(path).toHaveLength(4); + expect(path.map((e) => e.id)).toEqual([id1, id2, id3, id4]); + }); + + it("returns path from specified entry to root", () => { + const session = SessionManager.inMemory(); + + const id1 = session.appendMessage(userMsg("1")); + const id2 = session.appendMessage(assistantMsg("2")); + const _id3 = session.appendMessage(userMsg("3")); + const _id4 = session.appendMessage(assistantMsg("4")); + + const path = session.getPath(id2); + expect(path).toHaveLength(2); + expect(path.map((e) => e.id)).toEqual([id1, id2]); + }); + }); + + describe("getTree", () => { + it("returns empty array for empty session", () => { + const session = SessionManager.inMemory(); + expect(session.getTree()).toEqual([]); + }); + + it("returns single root for linear session", () => { + const session = SessionManager.inMemory(); + + const id1 = session.appendMessage(userMsg("1")); + const id2 = session.appendMessage(assistantMsg("2")); + const id3 = session.appendMessage(userMsg("3")); + + const tree = session.getTree(); + expect(tree).toHaveLength(1); + + const root = tree[0]; + expect(root.entry.id).toBe(id1); + expect(root.children).toHaveLength(1); + expect(root.children[0].entry.id).toBe(id2); + expect(root.children[0].children).toHaveLength(1); + expect(root.children[0].children[0].entry.id).toBe(id3); + expect(root.children[0].children[0].children).toHaveLength(0); + }); + + it("returns tree with branches after branch", () => { + const session = SessionManager.inMemory(); + + // Build: 1 -> 2 -> 3 + const id1 = session.appendMessage(userMsg("1")); + const id2 = session.appendMessage(assistantMsg("2")); + const id3 = session.appendMessage(userMsg("3")); + + // Branch from id2, add new path: 2 -> 4 + session.branch(id2); + const id4 = session.appendMessage(userMsg("4-branch")); + + const tree = session.getTree(); + expect(tree).toHaveLength(1); + + const root = tree[0]; + expect(root.entry.id).toBe(id1); + expect(root.children).toHaveLength(1); + + const node2 = root.children[0]; + expect(node2.entry.id).toBe(id2); + expect(node2.children).toHaveLength(2); // id3 and id4 are siblings + + const childIds = node2.children.map((c) => c.entry.id).sort(); + expect(childIds).toEqual([id3, id4].sort()); + }); + + it("handles multiple branches at same point", () => { + const session = SessionManager.inMemory(); + + const _id1 = session.appendMessage(userMsg("root")); + const id2 = session.appendMessage(assistantMsg("response")); + + // Branch A + session.branch(id2); + const idA = session.appendMessage(userMsg("branch-A")); + + // Branch B + session.branch(id2); + const idB = session.appendMessage(userMsg("branch-B")); + + // Branch C + session.branch(id2); + const idC = session.appendMessage(userMsg("branch-C")); + + const tree = session.getTree(); + const node2 = tree[0].children[0]; + expect(node2.entry.id).toBe(id2); + expect(node2.children).toHaveLength(3); + + const branchIds = node2.children.map((c) => c.entry.id).sort(); + expect(branchIds).toEqual([idA, idB, idC].sort()); + }); + + it("handles deep branching", () => { + const session = SessionManager.inMemory(); + + // Main path: 1 -> 2 -> 3 -> 4 + const _id1 = session.appendMessage(userMsg("1")); + const id2 = session.appendMessage(assistantMsg("2")); + const id3 = session.appendMessage(userMsg("3")); + const _id4 = session.appendMessage(assistantMsg("4")); + + // Branch from 2: 2 -> 5 -> 6 + session.branch(id2); + const id5 = session.appendMessage(userMsg("5")); + const _id6 = session.appendMessage(assistantMsg("6")); + + // Branch from 5: 5 -> 7 + session.branch(id5); + const _id7 = session.appendMessage(userMsg("7")); + + const tree = session.getTree(); + + // Verify structure + const node2 = tree[0].children[0]; + expect(node2.children).toHaveLength(2); // id3 and id5 + + const node5 = node2.children.find((c) => c.entry.id === id5)!; + expect(node5.children).toHaveLength(2); // id6 and id7 + + const node3 = node2.children.find((c) => c.entry.id === id3)!; + expect(node3.children).toHaveLength(1); // id4 + }); + }); + + describe("branch", () => { + it("moves leaf pointer to specified entry", () => { + const session = SessionManager.inMemory(); + + const id1 = session.appendMessage(userMsg("1")); + const _id2 = session.appendMessage(assistantMsg("2")); + const id3 = session.appendMessage(userMsg("3")); + + expect(session.getLeafUuid()).toBe(id3); + + session.branch(id1); + expect(session.getLeafUuid()).toBe(id1); + }); + + it("throws for non-existent entry", () => { + const session = SessionManager.inMemory(); + session.appendMessage(userMsg("hello")); + + expect(() => session.branch("nonexistent")).toThrow("Entry nonexistent not found"); + }); + + it("new appends become children of branch point", () => { + const session = SessionManager.inMemory(); + + const id1 = session.appendMessage(userMsg("1")); + const _id2 = session.appendMessage(assistantMsg("2")); + + session.branch(id1); + const id3 = session.appendMessage(userMsg("branched")); + + const entries = session.getEntries(); + const branchedEntry = entries.find((e) => e.id === id3)!; + expect(branchedEntry.parentId).toBe(id1); // sibling of id2 + }); + }); + + describe("branchWithSummary", () => { + it("inserts branch summary and advances leaf", () => { + const session = SessionManager.inMemory(); + + const id1 = session.appendMessage(userMsg("1")); + const _id2 = session.appendMessage(assistantMsg("2")); + const _id3 = session.appendMessage(userMsg("3")); + + const summaryId = session.branchWithSummary(id1, "Summary of abandoned work"); + + expect(session.getLeafUuid()).toBe(summaryId); + + const entries = session.getEntries(); + const summaryEntry = entries.find((e) => e.type === "branch_summary"); + expect(summaryEntry).toBeDefined(); + expect(summaryEntry?.parentId).toBe(id1); + if (summaryEntry?.type === "branch_summary") { + expect(summaryEntry.summary).toBe("Summary of abandoned work"); + } + }); + + it("throws for non-existent entry", () => { + const session = SessionManager.inMemory(); + session.appendMessage(userMsg("hello")); + + expect(() => session.branchWithSummary("nonexistent", "summary")).toThrow("Entry nonexistent not found"); + }); + }); + + describe("getLeafEntry", () => { + it("returns undefined for empty session", () => { + const session = SessionManager.inMemory(); + expect(session.getLeafEntry()).toBeUndefined(); + }); + + it("returns current leaf entry", () => { + const session = SessionManager.inMemory(); + + session.appendMessage(userMsg("1")); + const id2 = session.appendMessage(assistantMsg("2")); + + const leaf = session.getLeafEntry(); + expect(leaf).toBeDefined(); + expect(leaf!.id).toBe(id2); + }); + }); + + describe("getEntry", () => { + it("returns undefined for non-existent id", () => { + const session = SessionManager.inMemory(); + expect(session.getEntry("nonexistent")).toBeUndefined(); + }); + + it("returns entry by id", () => { + const session = SessionManager.inMemory(); + + const id1 = session.appendMessage(userMsg("first")); + const id2 = session.appendMessage(assistantMsg("second")); + + const entry1 = session.getEntry(id1); + expect(entry1).toBeDefined(); + expect(entry1?.type).toBe("message"); + if (entry1?.type === "message" && entry1.message.role === "user") { + expect(entry1.message.content).toBe("first"); + } + + const entry2 = session.getEntry(id2); + expect(entry2).toBeDefined(); + if (entry2?.type === "message" && entry2.message.role === "assistant") { + expect((entry2.message.content as any)[0].text).toBe("second"); + } + }); + }); + + describe("buildSessionContext with branches", () => { + it("returns messages from current branch only", () => { + const session = SessionManager.inMemory(); + + // Main: 1 -> 2 -> 3 + session.appendMessage(userMsg("msg1")); + const id2 = session.appendMessage(assistantMsg("msg2")); + session.appendMessage(userMsg("msg3")); + + // Branch from 2: 2 -> 4 + session.branch(id2); + session.appendMessage(assistantMsg("msg4-branch")); + + const ctx = session.buildSessionContext(); + expect(ctx.messages).toHaveLength(3); // msg1, msg2, msg4-branch (not msg3) + + expect((ctx.messages[0] as any).content).toBe("msg1"); + expect((ctx.messages[1] as any).content[0].text).toBe("msg2"); + expect((ctx.messages[2] as any).content[0].text).toBe("msg4-branch"); + }); + }); +}); + +describe("createBranchedSession", () => { + it("throws for non-existent entry", () => { + const session = SessionManager.inMemory(); + session.appendMessage(userMsg("hello")); + + expect(() => session.createBranchedSession("nonexistent")).toThrow("Entry nonexistent not found"); + }); + + it("creates new session with path to specified leaf (in-memory)", () => { + const session = SessionManager.inMemory(); + + // Build: 1 -> 2 -> 3 -> 4 + const id1 = session.appendMessage(userMsg("1")); + const id2 = session.appendMessage(assistantMsg("2")); + const id3 = session.appendMessage(userMsg("3")); + session.appendMessage(assistantMsg("4")); + + // Branch from 3: 3 -> 5 + session.branch(id3); + const _id5 = session.appendMessage(userMsg("5")); + + // Create branched session from id2 (should only have 1 -> 2) + const result = session.createBranchedSession(id2); + expect(result).toBeNull(); // in-memory returns null + + // Session should now only have entries 1 and 2 + const entries = session.getEntries(); + expect(entries).toHaveLength(2); + expect(entries[0].id).toBe(id1); + expect(entries[1].id).toBe(id2); + }); + + it("extracts correct path from branched tree", () => { + const session = SessionManager.inMemory(); + + // Build: 1 -> 2 -> 3 + const id1 = session.appendMessage(userMsg("1")); + const id2 = session.appendMessage(assistantMsg("2")); + session.appendMessage(userMsg("3")); + + // Branch from 2: 2 -> 4 -> 5 + session.branch(id2); + const id4 = session.appendMessage(userMsg("4")); + const id5 = session.appendMessage(assistantMsg("5")); + + // Create branched session from id5 (should have 1 -> 2 -> 4 -> 5) + session.createBranchedSession(id5); + + const entries = session.getEntries(); + expect(entries).toHaveLength(4); + expect(entries.map((e) => e.id)).toEqual([id1, id2, id4, id5]); + }); +}); diff --git a/packages/mom/src/agent.ts b/packages/mom/src/agent.ts index b5cb6333..3e13e8b6 100644 --- a/packages/mom/src/agent.ts +++ b/packages/mom/src/agent.ts @@ -441,7 +441,7 @@ function createRunner(sandboxConfig: SandboxConfig, channelId: string, channelDi }); // Load existing messages - const loadedSession = sessionManager.loadSession(); + const loadedSession = sessionManager.buildSessionContex(); if (loadedSession.messages.length > 0) { agent.replaceMessages(loadedSession.messages); log.logInfo(`[${channelId}] Loaded ${loadedSession.messages.length} messages from context.jsonl`); @@ -628,7 +628,7 @@ function createRunner(sandboxConfig: SandboxConfig, channelId: string, channelDi // Reload messages from context.jsonl // This picks up any messages synced from log.jsonl before this run - const reloadedSession = sessionManager.loadSession(); + const reloadedSession = sessionManager.buildSessionContex(); if (reloadedSession.messages.length > 0) { agent.replaceMessages(reloadedSession.messages); log.logInfo(`[${channelId}] Reloaded ${reloadedSession.messages.length} messages from context`); diff --git a/packages/mom/src/context.ts b/packages/mom/src/context.ts index a5a8024c..36ae4c1d 100644 --- a/packages/mom/src/context.ts +++ b/packages/mom/src/context.ts @@ -15,15 +15,12 @@ import { buildSessionContext, type CompactionEntry, type FileEntry, - type LoadedSession, - type MessageContent, - type ModelChangeContent, type ModelChangeEntry, + type SessionContext, type SessionEntry, + type SessionEntryBase, 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"; @@ -98,15 +95,15 @@ export class MomSessionManager { this.leafId = null; } - private _createTreeNode(): Omit { + private _createEntryBase(): Omit { const id = uuidv4(); - const node = { + const base = { id, parentId: this.leafId, timestamp: new Date().toISOString(), }; this.leafId = id; - return node; + return base; } private _persist(entry: SessionEntry): void { @@ -281,22 +278,23 @@ export class MomSessionManager { } saveMessage(message: AppMessage): void { - const content: MessageContent = { type: "message", message }; - const entry: SessionMessageEntry = { ...this._createTreeNode(), ...content }; + const entry: SessionMessageEntry = { ...this._createEntryBase(), type: "message", message }; this.inMemoryEntries.push(entry); this._persist(entry); } saveThinkingLevelChange(thinkingLevel: string): void { - const content: ThinkingLevelContent = { type: "thinking_level_change", thinkingLevel }; - const entry: ThinkingLevelChangeEntry = { ...this._createTreeNode(), ...content }; + const entry: ThinkingLevelChangeEntry = { + ...this._createEntryBase(), + type: "thinking_level_change", + thinkingLevel, + }; this.inMemoryEntries.push(entry); this._persist(entry); } saveModelChange(provider: string, modelId: string): void { - const content: ModelChangeContent = { type: "model_change", provider, modelId }; - const entry: ModelChangeEntry = { ...this._createTreeNode(), ...content }; + const entry: ModelChangeEntry = { ...this._createEntryBase(), type: "model_change", provider, modelId }; this.inMemoryEntries.push(entry); this._persist(entry); } @@ -307,7 +305,7 @@ export class MomSessionManager { } /** Load session with compaction support */ - loadSession(): LoadedSession { + buildSessionContex(): SessionContext { const entries = this.loadEntries(); return buildSessionContext(entries); } @@ -354,15 +352,15 @@ export class MomSessionManager { } loadModel(): { provider: string; modelId: string } | null { - return this.loadSession().model; + return this.buildSessionContex().model; } loadThinkingLevel(): string { - return this.loadSession().thinkingLevel; + return this.buildSessionContex().thinkingLevel; } /** Not used by mom but required by AgentSession interface */ - createBranchedSessionFromEntries(_entries: SessionEntry[], _branchBeforeIndex: number): string | null { + createBranchedSession(_leafId: string): string | null { return null; // Mom doesn't support branching } }