Refactor session manager: migration chain, validation, tests

- Add migrateV1ToV2/migrateToCurrentVersion for extensible migrations
- createSummaryMessage now takes timestamp from entry
- loadEntriesFromFile validates session header
- findMostRecentSession only returns valid session files (reads first 512 bytes)
- Remove ConversationEntry alias
- Fix mom context.ts TreeNode type

Tests:
- migration.test.ts: v1 migration, idempotency
- build-context.test.ts: 14 tests covering trivial, compaction, branches
- file-operations.test.ts: loadEntriesFromFile, findMostRecentSession
This commit is contained in:
Mario Zechner 2025-12-26 01:33:13 +01:00
parent 95312e00bb
commit beb70f126d
7 changed files with 606 additions and 102 deletions

View file

@ -273,9 +273,11 @@ describe("findCutPoint", () => {
});
describe("createSummaryMessage", () => {
it("should create user message with prefix", () => {
const msg = createSummaryMessage("This is the summary");
it("should create user message with prefix and correct timestamp", () => {
const ts = "2025-01-01T12:00:00.000Z";
const msg = createSummaryMessage("This is the summary", ts);
expect(msg.role).toBe("user");
expect(msg.timestamp).toBe(new Date(ts).getTime());
if (msg.role === "user") {
expect(msg.content).toContain(
"The conversation history before this point was compacted into the following summary:",

View file

@ -0,0 +1,269 @@
import { describe, expect, it } from "vitest";
import {
type BranchSummaryEntry,
buildSessionContext,
type CompactionEntry,
type ModelChangeEntry,
type SessionEntry,
type SessionMessageEntry,
SUMMARY_PREFIX,
type ThinkingLevelChangeEntry,
} from "../../src/core/session-manager.js";
function msg(id: string, parentId: string | null, role: "user" | "assistant", text: string): SessionMessageEntry {
const base = { type: "message" as const, id, parentId, timestamp: "2025-01-01T00:00:00Z" };
if (role === "user") {
return { ...base, message: { role, content: text, timestamp: 1 } };
}
return {
...base,
message: {
role,
content: [{ type: "text", text }],
api: "anthropic-messages",
provider: "anthropic",
model: "claude-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: 1,
},
};
}
function compaction(id: string, parentId: string | null, summary: string, firstKeptEntryId: string): CompactionEntry {
return {
type: "compaction",
id,
parentId,
timestamp: "2025-01-01T00:00:00Z",
summary,
firstKeptEntryId,
tokensBefore: 1000,
};
}
function branchSummary(id: string, parentId: string | null, summary: string): BranchSummaryEntry {
return { type: "branch_summary", id, parentId, timestamp: "2025-01-01T00:00:00Z", summary };
}
function thinkingLevel(id: string, parentId: string | null, level: string): ThinkingLevelChangeEntry {
return { type: "thinking_level_change", id, parentId, timestamp: "2025-01-01T00:00:00Z", thinkingLevel: level };
}
function modelChange(id: string, parentId: string | null, provider: string, modelId: string): ModelChangeEntry {
return { type: "model_change", id, parentId, timestamp: "2025-01-01T00:00:00Z", provider, modelId };
}
describe("buildSessionContext", () => {
describe("trivial cases", () => {
it("empty entries returns empty context", () => {
const ctx = buildSessionContext([]);
expect(ctx.messages).toEqual([]);
expect(ctx.thinkingLevel).toBe("off");
expect(ctx.model).toBeNull();
});
it("single user message", () => {
const entries: SessionEntry[] = [msg("1", null, "user", "hello")];
const ctx = buildSessionContext(entries);
expect(ctx.messages).toHaveLength(1);
expect(ctx.messages[0].role).toBe("user");
});
it("simple conversation", () => {
const entries: SessionEntry[] = [
msg("1", null, "user", "hello"),
msg("2", "1", "assistant", "hi there"),
msg("3", "2", "user", "how are you"),
msg("4", "3", "assistant", "great"),
];
const ctx = buildSessionContext(entries);
expect(ctx.messages).toHaveLength(4);
expect(ctx.messages.map((m) => m.role)).toEqual(["user", "assistant", "user", "assistant"]);
});
it("tracks thinking level changes", () => {
const entries: SessionEntry[] = [
msg("1", null, "user", "hello"),
thinkingLevel("2", "1", "high"),
msg("3", "2", "assistant", "thinking hard"),
];
const ctx = buildSessionContext(entries);
expect(ctx.thinkingLevel).toBe("high");
expect(ctx.messages).toHaveLength(2);
});
it("tracks model from assistant message", () => {
const entries: SessionEntry[] = [msg("1", null, "user", "hello"), msg("2", "1", "assistant", "hi")];
const ctx = buildSessionContext(entries);
expect(ctx.model).toEqual({ provider: "anthropic", modelId: "claude-test" });
});
it("tracks model from model change entry", () => {
const entries: SessionEntry[] = [
msg("1", null, "user", "hello"),
modelChange("2", "1", "openai", "gpt-4"),
msg("3", "2", "assistant", "hi"),
];
const ctx = buildSessionContext(entries);
// Assistant message overwrites model change
expect(ctx.model).toEqual({ provider: "anthropic", modelId: "claude-test" });
});
});
describe("with compaction", () => {
it("includes summary before kept messages", () => {
const entries: SessionEntry[] = [
msg("1", null, "user", "first"),
msg("2", "1", "assistant", "response1"),
msg("3", "2", "user", "second"),
msg("4", "3", "assistant", "response2"),
compaction("5", "4", "Summary of first two turns", "3"),
msg("6", "5", "user", "third"),
msg("7", "6", "assistant", "response3"),
];
const ctx = buildSessionContext(entries);
// Should have: summary + kept (3,4) + after (6,7) = 5 messages
expect(ctx.messages).toHaveLength(5);
expect((ctx.messages[0] as any).content).toContain("Summary of first two turns");
expect((ctx.messages[1] as any).content).toBe("second");
expect((ctx.messages[2] as any).content[0].text).toBe("response2");
expect((ctx.messages[3] as any).content).toBe("third");
expect((ctx.messages[4] as any).content[0].text).toBe("response3");
});
it("handles compaction keeping from first message", () => {
const entries: SessionEntry[] = [
msg("1", null, "user", "first"),
msg("2", "1", "assistant", "response"),
compaction("3", "2", "Empty summary", "1"),
msg("4", "3", "user", "second"),
];
const ctx = buildSessionContext(entries);
// Summary + all messages (1,2,4)
expect(ctx.messages).toHaveLength(4);
expect((ctx.messages[0] as any).content).toContain(SUMMARY_PREFIX);
});
it("multiple compactions uses latest", () => {
const entries: SessionEntry[] = [
msg("1", null, "user", "a"),
msg("2", "1", "assistant", "b"),
compaction("3", "2", "First summary", "1"),
msg("4", "3", "user", "c"),
msg("5", "4", "assistant", "d"),
compaction("6", "5", "Second summary", "4"),
msg("7", "6", "user", "e"),
];
const ctx = buildSessionContext(entries);
// Should use second summary, keep from 4
expect(ctx.messages).toHaveLength(4);
expect((ctx.messages[0] as any).content).toContain("Second summary");
});
});
describe("with branches", () => {
it("follows path to specified leaf", () => {
// Tree:
// 1 -> 2 -> 3 (branch A)
// \-> 4 (branch B)
const entries: SessionEntry[] = [
msg("1", null, "user", "start"),
msg("2", "1", "assistant", "response"),
msg("3", "2", "user", "branch A"),
msg("4", "2", "user", "branch B"),
];
const ctxA = buildSessionContext(entries, "3");
expect(ctxA.messages).toHaveLength(3);
expect((ctxA.messages[2] as any).content).toBe("branch A");
const ctxB = buildSessionContext(entries, "4");
expect(ctxB.messages).toHaveLength(3);
expect((ctxB.messages[2] as any).content).toBe("branch B");
});
it("includes branch summary in path", () => {
const entries: SessionEntry[] = [
msg("1", null, "user", "start"),
msg("2", "1", "assistant", "response"),
msg("3", "2", "user", "abandoned path"),
branchSummary("4", "2", "Summary of abandoned work"),
msg("5", "4", "user", "new direction"),
];
const ctx = buildSessionContext(entries, "5");
expect(ctx.messages).toHaveLength(4);
expect((ctx.messages[2] as any).content).toContain("Summary of abandoned work");
expect((ctx.messages[3] as any).content).toBe("new direction");
});
it("complex tree with multiple branches and compaction", () => {
// Tree:
// 1 -> 2 -> 3 -> 4 -> compaction(5) -> 6 -> 7 (main path)
// \-> 8 -> 9 (abandoned branch)
// \-> branchSummary(10) -> 11 (resumed from 3)
const entries: SessionEntry[] = [
msg("1", null, "user", "start"),
msg("2", "1", "assistant", "r1"),
msg("3", "2", "user", "q2"),
msg("4", "3", "assistant", "r2"),
compaction("5", "4", "Compacted history", "3"),
msg("6", "5", "user", "q3"),
msg("7", "6", "assistant", "r3"),
// Abandoned branch from 3
msg("8", "3", "user", "wrong path"),
msg("9", "8", "assistant", "wrong response"),
// Branch summary resuming from 3
branchSummary("10", "3", "Tried wrong approach"),
msg("11", "10", "user", "better approach"),
];
// Main path to 7: summary + kept(3,4) + after(6,7)
const ctxMain = buildSessionContext(entries, "7");
expect(ctxMain.messages).toHaveLength(5);
expect((ctxMain.messages[0] as any).content).toContain("Compacted history");
expect((ctxMain.messages[1] as any).content).toBe("q2");
expect((ctxMain.messages[2] as any).content[0].text).toBe("r2");
expect((ctxMain.messages[3] as any).content).toBe("q3");
expect((ctxMain.messages[4] as any).content[0].text).toBe("r3");
// Branch path to 11: 1,2,3 + branch_summary + 11
const ctxBranch = buildSessionContext(entries, "11");
expect(ctxBranch.messages).toHaveLength(5);
expect((ctxBranch.messages[0] as any).content).toBe("start");
expect((ctxBranch.messages[1] as any).content[0].text).toBe("r1");
expect((ctxBranch.messages[2] as any).content).toBe("q2");
expect((ctxBranch.messages[3] as any).content).toContain("Tried wrong approach");
expect((ctxBranch.messages[4] as any).content).toBe("better approach");
});
});
describe("edge cases", () => {
it("uses last entry when leafId not found", () => {
const entries: SessionEntry[] = [msg("1", null, "user", "hello"), msg("2", "1", "assistant", "hi")];
const ctx = buildSessionContext(entries, "nonexistent");
expect(ctx.messages).toHaveLength(2);
});
it("handles orphaned entries gracefully", () => {
const entries: SessionEntry[] = [
msg("1", null, "user", "hello"),
msg("2", "missing", "assistant", "orphan"), // parent doesn't exist
];
const ctx = buildSessionContext(entries, "2");
// Should only get the orphan since parent chain is broken
expect(ctx.messages).toHaveLength(1);
});
});
});

View file

@ -0,0 +1,127 @@
import { mkdirSync, rmSync, writeFileSync } from "fs";
import { tmpdir } from "os";
import { join } from "path";
import { afterEach, beforeEach, describe, expect, it } from "vitest";
import { findMostRecentSession, loadEntriesFromFile } from "../../src/core/session-manager.js";
describe("loadEntriesFromFile", () => {
let tempDir: string;
beforeEach(() => {
tempDir = join(tmpdir(), `session-test-${Date.now()}`);
mkdirSync(tempDir, { recursive: true });
});
afterEach(() => {
rmSync(tempDir, { recursive: true, force: true });
});
it("returns empty array for non-existent file", () => {
const entries = loadEntriesFromFile(join(tempDir, "nonexistent.jsonl"));
expect(entries).toEqual([]);
});
it("returns empty array for empty file", () => {
const file = join(tempDir, "empty.jsonl");
writeFileSync(file, "");
expect(loadEntriesFromFile(file)).toEqual([]);
});
it("returns empty array for file without valid session header", () => {
const file = join(tempDir, "no-header.jsonl");
writeFileSync(file, '{"type":"message","id":"1"}\n');
expect(loadEntriesFromFile(file)).toEqual([]);
});
it("returns empty array for malformed JSON", () => {
const file = join(tempDir, "malformed.jsonl");
writeFileSync(file, "not json\n");
expect(loadEntriesFromFile(file)).toEqual([]);
});
it("loads valid session file", () => {
const file = join(tempDir, "valid.jsonl");
writeFileSync(
file,
'{"type":"session","id":"abc","timestamp":"2025-01-01T00:00:00Z","cwd":"/tmp"}\n' +
'{"type":"message","id":"1","parentId":null,"timestamp":"2025-01-01T00:00:01Z","message":{"role":"user","content":"hi","timestamp":1}}\n',
);
const entries = loadEntriesFromFile(file);
expect(entries).toHaveLength(2);
expect(entries[0].type).toBe("session");
expect(entries[1].type).toBe("message");
});
it("skips malformed lines but keeps valid ones", () => {
const file = join(tempDir, "mixed.jsonl");
writeFileSync(
file,
'{"type":"session","id":"abc","timestamp":"2025-01-01T00:00:00Z","cwd":"/tmp"}\n' +
"not valid json\n" +
'{"type":"message","id":"1","parentId":null,"timestamp":"2025-01-01T00:00:01Z","message":{"role":"user","content":"hi","timestamp":1}}\n',
);
const entries = loadEntriesFromFile(file);
expect(entries).toHaveLength(2);
});
});
describe("findMostRecentSession", () => {
let tempDir: string;
beforeEach(() => {
tempDir = join(tmpdir(), `session-test-${Date.now()}`);
mkdirSync(tempDir, { recursive: true });
});
afterEach(() => {
rmSync(tempDir, { recursive: true, force: true });
});
it("returns null for empty directory", () => {
expect(findMostRecentSession(tempDir)).toBeNull();
});
it("returns null for non-existent directory", () => {
expect(findMostRecentSession(join(tempDir, "nonexistent"))).toBeNull();
});
it("ignores non-jsonl files", () => {
writeFileSync(join(tempDir, "file.txt"), "hello");
writeFileSync(join(tempDir, "file.json"), "{}");
expect(findMostRecentSession(tempDir)).toBeNull();
});
it("ignores jsonl files without valid session header", () => {
writeFileSync(join(tempDir, "invalid.jsonl"), '{"type":"message"}\n');
expect(findMostRecentSession(tempDir)).toBeNull();
});
it("returns single valid session file", () => {
const file = join(tempDir, "session.jsonl");
writeFileSync(file, '{"type":"session","id":"abc","timestamp":"2025-01-01T00:00:00Z","cwd":"/tmp"}\n');
expect(findMostRecentSession(tempDir)).toBe(file);
});
it("returns most recently modified session", async () => {
const file1 = join(tempDir, "older.jsonl");
const file2 = join(tempDir, "newer.jsonl");
writeFileSync(file1, '{"type":"session","id":"old","timestamp":"2025-01-01T00:00:00Z","cwd":"/tmp"}\n');
// Small delay to ensure different mtime
await new Promise((r) => setTimeout(r, 10));
writeFileSync(file2, '{"type":"session","id":"new","timestamp":"2025-01-01T00:00:00Z","cwd":"/tmp"}\n');
expect(findMostRecentSession(tempDir)).toBe(file2);
});
it("skips invalid files and returns valid one", async () => {
const invalid = join(tempDir, "invalid.jsonl");
const valid = join(tempDir, "valid.jsonl");
writeFileSync(invalid, '{"type":"not-session"}\n');
await new Promise((r) => setTimeout(r, 10));
writeFileSync(valid, '{"type":"session","id":"abc","timestamp":"2025-01-01T00:00:00Z","cwd":"/tmp"}\n');
expect(findMostRecentSession(tempDir)).toBe(valid);
});
});

View file

@ -0,0 +1,78 @@
import { describe, expect, it } from "vitest";
import { type FileEntry, migrateSessionEntries } from "../../src/core/session-manager.js";
describe("migrateSessionEntries", () => {
it("should add id/parentId to v1 entries", () => {
const entries: FileEntry[] = [
{ type: "session", id: "sess-1", timestamp: "2025-01-01T00:00:00Z", cwd: "/tmp" },
{ type: "message", timestamp: "2025-01-01T00:00:01Z", message: { role: "user", content: "hi", timestamp: 1 } },
{
type: "message",
timestamp: "2025-01-01T00:00:02Z",
message: {
role: "assistant",
content: [{ type: "text", text: "hello" }],
api: "test",
provider: "test",
model: "test",
usage: { input: 1, output: 1, cacheRead: 0, cacheWrite: 0 },
stopReason: "stop",
timestamp: 2,
},
},
] as FileEntry[];
migrateSessionEntries(entries);
// Header should have version set
expect((entries[0] as any).version).toBe(2);
// Entries should have id/parentId
const msg1 = entries[1] as any;
const msg2 = entries[2] as any;
expect(msg1.id).toBeDefined();
expect(msg1.id.length).toBe(8);
expect(msg1.parentId).toBeNull();
expect(msg2.id).toBeDefined();
expect(msg2.id.length).toBe(8);
expect(msg2.parentId).toBe(msg1.id);
});
it("should be idempotent (skip already migrated)", () => {
const entries: FileEntry[] = [
{ type: "session", id: "sess-1", version: 2, timestamp: "2025-01-01T00:00:00Z", cwd: "/tmp" },
{
type: "message",
id: "abc12345",
parentId: null,
timestamp: "2025-01-01T00:00:01Z",
message: { role: "user", content: "hi", timestamp: 1 },
},
{
type: "message",
id: "def67890",
parentId: "abc12345",
timestamp: "2025-01-01T00:00:02Z",
message: {
role: "assistant",
content: [{ type: "text", text: "hello" }],
api: "test",
provider: "test",
model: "test",
usage: { input: 1, output: 1, cacheRead: 0, cacheWrite: 0 },
stopReason: "stop",
timestamp: 2,
},
},
] as FileEntry[];
migrateSessionEntries(entries);
// IDs should be unchanged
expect((entries[1] as any).id).toBe("abc12345");
expect((entries[2] as any).id).toBe("def67890");
expect((entries[2] as any).parentId).toBe("abc12345");
});
});