From 9478a3c1f55ef1a4dd522d9b635807d3298dea5e Mon Sep 17 00:00:00 2001 From: Mario Zechner Date: Fri, 26 Dec 2025 00:31:53 +0100 Subject: [PATCH] Fix SessionEntry type to exclude SessionHeader - SessionEntry now only contains conversation entries (messages, compaction, etc.) - SessionHeader is separate, not part of SessionEntry - FileEntry = SessionHeader | SessionEntry (for file storage) - getEntries() filters out header, returns SessionEntry[] - Added getHeader() for accessing session metadata - Updated compaction and tests to not expect header in entries - Updated mom package to use FileEntry for internal storage --- .../coding-agent/src/core/agent-session.ts | 2 +- packages/coding-agent/src/core/compaction.ts | 18 ++---- .../coding-agent/src/core/session-manager.ts | 57 ++++++++++--------- packages/coding-agent/src/index.ts | 1 + packages/coding-agent/test/compaction.test.ts | 32 ++--------- packages/mom/src/context.ts | 15 +++-- 6 files changed, 50 insertions(+), 75 deletions(-) diff --git a/packages/coding-agent/src/core/agent-session.ts b/packages/coding-agent/src/core/agent-session.ts index fe4804a1..b6f45978 100644 --- a/packages/coding-agent/src/core/agent-session.ts +++ b/packages/coding-agent/src/core/agent-session.ts @@ -759,7 +759,7 @@ export class AgentSession { if (lastEntry?.type === "compaction") { throw new Error("Already compacted"); } - throw new Error("Nothing to compact (session too small or needs migration)"); + throw new Error("Nothing to compact (session too small)"); } // Find previous compaction summary if any diff --git a/packages/coding-agent/src/core/compaction.ts b/packages/coding-agent/src/core/compaction.ts index 6a56a3ac..ea1e44a3 100644 --- a/packages/coding-agent/src/core/compaction.ts +++ b/packages/coding-agent/src/core/compaction.ts @@ -9,7 +9,7 @@ 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, ConversationEntry, SessionEntry } from "./session-manager.js"; +import type { CompactionEntry, SessionEntry } from "./session-manager.js"; /** Result from compact() - SessionManager adds uuid/parentUuid when saving */ export interface CompactionResult { @@ -251,7 +251,7 @@ export function findCutPoint( while (cutIndex > startIndex) { const prevEntry = entries[cutIndex - 1]; // Stop at session header or compaction boundaries - if (prevEntry.type === "session" || prevEntry.type === "compaction") { + if (prevEntry.type === "compaction") { break; } if (prevEntry.type === "message") { @@ -370,13 +370,10 @@ export function prepareCompaction(entries: SessionEntry[], settings: CompactionS // 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) { + if (!firstKeptEntry?.id) { return null; // Session needs migration } + const firstKeptEntryId = firstKeptEntry.id; const historyEnd = cutPoint.isSplitTurn ? cutPoint.turnStartIndex : cutPoint.firstKeptEntryIndex; @@ -405,7 +402,7 @@ export function prepareCompaction(entries: SessionEntry[], settings: CompactionS // Main compaction function // ============================================================================ -const TURN_PREFIX_SUMMARIZATION_PROMPT = `You are performing a CONTEXT CHECKPOINT COMPACTION for a split turn. +const TURN_PREFIX_SUMMARIZATION_PROMPT = `You are performing a CONTEXT CHECKPOINT COMPACTION for a split turn. This is the PREFIX of a turn that was too large to keep in full. The SUFFIX (recent work) is being kept. Create a handoff summary that captures: @@ -515,10 +512,7 @@ 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; + const firstKeptEntryId = firstKeptEntry.id; if (!firstKeptEntryId) { throw new Error("First kept entry has no UUID - session may need migration"); } diff --git a/packages/coding-agent/src/core/session-manager.ts b/packages/coding-agent/src/core/session-manager.ts index 3b9eda0b..e81791b9 100644 --- a/packages/coding-agent/src/core/session-manager.ts +++ b/packages/coding-agent/src/core/session-manager.ts @@ -87,16 +87,19 @@ 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 = +/** Session entry - has id/parentId for tree structure */ +export type SessionEntry = | SessionMessageEntry | ThinkingLevelChangeEntry | ModelChangeEntry | CompactionEntry | BranchSummaryEntry; -/** Any session entry (header or conversation) */ -export type SessionEntry = SessionHeader | ConversationEntry; +/** @deprecated Use SessionEntry */ +export type ConversationEntry = SessionEntry; + +/** Raw file entry (includes header) */ +export type FileEntry = SessionHeader | SessionEntry; export interface SessionContext { messages: AppMessage[]; @@ -135,7 +138,7 @@ 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 { +export function migrateSessionEntries(entries: FileEntry[]): void { // Check if already migrated const firstConv = entries.find((e) => e.type !== "session"); if (firstConv && "id" in firstConv && firstConv.id) { @@ -171,7 +174,7 @@ export function migrateSessionEntries(entries: SessionEntry[]): void { } /** Exported for compaction.test.ts */ -export function parseSessionEntries(content: string): SessionEntry[] { +export function parseSessionEntries(content: string): FileEntry[] { const entries: SessionEntry[] = []; const lines = content.trim().split("\n"); @@ -203,26 +206,18 @@ export function getLatestCompactionEntry(entries: SessionEntry[]): CompactionEnt * Handles compaction and branch summaries along the path. */ export function buildSessionContext(entries: SessionEntry[], leafId?: string): SessionContext { - // Build uuid index for conversation entries - const byId = new Map(); + // Build uuid index + const byId = new Map(); for (const entry of entries) { - if (entry.type !== "session") { - byId.set(entry.id, entry); - } + byId.set(entry.id, entry); } // Find leaf - let leaf: ConversationEntry | undefined; + let leaf: SessionEntry | 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; - } - } + leaf = entries[entries.length - 1]; } if (!leaf) { @@ -230,8 +225,8 @@ export function buildSessionContext(entries: SessionEntry[], leafId?: string): S } // Walk from leaf to root, collecting path - const path: ConversationEntry[] = []; - let current: ConversationEntry | undefined = leaf; + const path: SessionEntry[] = []; + let current: SessionEntry | undefined = leaf; while (current) { path.unshift(current); current = current.parentId ? byId.get(current.parentId) : undefined; @@ -316,7 +311,7 @@ function getDefaultSessionDir(cwd: string): string { return sessionDir; } -function loadEntriesFromFile(filePath: string): SessionEntry[] { +function loadEntriesFromFile(filePath: string): FileEntry[] { if (!existsSync(filePath)) return []; const content = readFileSync(filePath, "utf8"); @@ -359,7 +354,7 @@ export class SessionManager { private cwd: string; private persist: boolean; private flushed: boolean = false; - private inMemoryEntries: SessionEntry[] = []; + private inMemoryEntries: FileEntry[] = []; // Tree structure (v2) private byId: Map = new Map(); @@ -570,11 +565,19 @@ export class SessionManager { } /** - * Get all session entries. Returns a defensive copy. + * Get session header. + */ + getHeader(): SessionHeader | null { + const h = this.inMemoryEntries.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. */ getEntries(): SessionEntry[] { - return [...this.inMemoryEntries]; + return this.inMemoryEntries.filter((e): e is SessionEntry => e.type !== "session"); } // ========================================================================= @@ -606,12 +609,12 @@ export class SessionManager { return entry.id; } - createBranchedSessionFromEntries(entries: SessionEntry[], branchBeforeIndex: number): string | null { + createBranchedSessionFromEntries(entries: FileEntry[], branchBeforeIndex: number): string | null { const newSessionId = uuidv4(); const timestamp = new Date().toISOString().replace(/[:.]/g, "-"); const newSessionFile = join(this.getSessionDir(), `${timestamp}_${newSessionId}.jsonl`); - const newEntries: SessionEntry[] = []; + const newEntries: FileEntry[] = []; for (let i = 0; i < branchBeforeIndex; i++) { const entry = entries[i]; diff --git a/packages/coding-agent/src/index.ts b/packages/coding-agent/src/index.ts index 04151f0b..f9673384 100644 --- a/packages/coding-agent/src/index.ts +++ b/packages/coding-agent/src/index.ts @@ -116,6 +116,7 @@ export { type ConversationEntry, CURRENT_SESSION_VERSION, createSummaryMessage, + type FileEntry, getLatestCompactionEntry, type MessageContent, type ModelChangeContent, diff --git a/packages/coding-agent/test/compaction.test.ts b/packages/coding-agent/test/compaction.test.ts index 787e863e..c12748e4 100644 --- a/packages/coding-agent/test/compaction.test.ts +++ b/packages/coding-agent/test/compaction.test.ts @@ -34,7 +34,7 @@ function loadLargeSessionEntries(): SessionEntry[] { const content = readFileSync(sessionPath, "utf-8"); const entries = parseSessionEntries(content); migrateSessionEntries(entries); // Add id/parentId for v1 fixtures - return entries; + return entries.filter((e): e is SessionEntry => e.type !== "session"); } function createMockUsage(input: number, output: number, cacheRead = 0, cacheWrite = 0): Usage { @@ -78,16 +78,6 @@ 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 = { @@ -298,12 +288,6 @@ describe("createSummaryMessage", () => { describe("buildSessionContext", () => { it("should load all messages when no compaction", () => { const entries: SessionEntry[] = [ - { - type: "session", - id: "1", - timestamp: "", - cwd: "", - }, createMessageEntry(createUserMessage("1")), createMessageEntry(createAssistantMessage("a")), createMessageEntry(createUserMessage("2")), @@ -326,7 +310,7 @@ describe("buildSessionContext", () => { const u3 = createMessageEntry(createUserMessage("3")); const a3 = createMessageEntry(createAssistantMessage("c")); - const entries: SessionEntry[] = [createSessionHeader(), u1, a1, u2, a2, compaction, u3, a3]; + const entries: SessionEntry[] = [u1, a1, u2, a2, compaction, u3, a3]; const loaded = buildSessionContext(entries); // summary + kept (u2, a2) + after (u3, a3) = 5 @@ -350,7 +334,7 @@ describe("buildSessionContext", () => { 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 entries: SessionEntry[] = [u1, a1, compact1, u2, b, u3, c, compact2, u4, d]; const loaded = buildSessionContext(entries); // summary + kept from u3 (u3, c) + after (u4, d) = 5 @@ -365,7 +349,7 @@ describe("buildSessionContext", () => { const u2 = createMessageEntry(createUserMessage("2")); const b = createMessageEntry(createAssistantMessage("b")); - const entries: SessionEntry[] = [createSessionHeader(), u1, a1, compact1, u2, b]; + const entries: SessionEntry[] = [u1, a1, compact1, u2, b]; const loaded = buildSessionContext(entries); // summary + all messages (u1, a1, u2, b) = 5 @@ -374,12 +358,6 @@ describe("buildSessionContext", () => { it("should track model and thinking level changes", () => { const entries: SessionEntry[] = [ - { - type: "session", - id: "1", - timestamp: "", - cwd: "", - }, createMessageEntry(createUserMessage("1")), createModelChangeEntry("openai", "gpt-4"), createMessageEntry(createAssistantMessage("a")), @@ -466,7 +444,7 @@ describe.skipIf(!process.env.ANTHROPIC_OAUTH_TOKEN)("LLM summarization", () => { // 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 parentId = lastEntry.id; const compactionEntry: CompactionEntry = { type: "compaction", id: "compaction-test-id", diff --git a/packages/mom/src/context.ts b/packages/mom/src/context.ts index 2e24bf5a..58a43b24 100644 --- a/packages/mom/src/context.ts +++ b/packages/mom/src/context.ts @@ -14,6 +14,7 @@ import type { AppMessage } from "@mariozechner/pi-agent-core"; import { buildSessionContext, type CompactionEntry, + type FileEntry, type LoadedSession, type MessageContent, type ModelChangeContent, @@ -52,7 +53,7 @@ export class MomSessionManager { private logFile: string; private channelDir: string; private flushed: boolean = false; - private inMemoryEntries: SessionEntry[] = []; + private inMemoryEntries: FileEntry[] = []; private leafId: string | null = null; constructor(channelDir: string) { @@ -259,17 +260,17 @@ export class MomSessionManager { return null; } - private loadEntriesFromFile(): SessionEntry[] { + private loadEntriesFromFile(): FileEntry[] { if (!existsSync(this.contextFile)) return []; const content = readFileSync(this.contextFile, "utf8"); - const entries: SessionEntry[] = []; + const entries: FileEntry[] = []; const lines = content.trim().split("\n"); for (const line of lines) { if (!line.trim()) continue; try { - const entry = JSON.parse(line) as SessionEntry; + const entry = JSON.parse(line) as FileEntry; entries.push(entry); } catch { // Skip malformed lines @@ -313,10 +314,8 @@ export class MomSessionManager { loadEntries(): SessionEntry[] { // Re-read from file to get latest state - if (existsSync(this.contextFile)) { - return this.loadEntriesFromFile(); - } - return [...this.inMemoryEntries]; + const entries = existsSync(this.contextFile) ? this.loadEntriesFromFile() : this.inMemoryEntries; + return entries.filter((e): e is SessionEntry => e.type !== "session"); } getSessionId(): string {