refactor(agent): simplify harness session repo layout

This commit is contained in:
Mario Zechner
2026-05-04 00:06:35 +02:00
parent cdde2e893d
commit 3d0f57187d
15 changed files with 616 additions and 581 deletions

View File

@@ -18,7 +18,7 @@ import {
serializeConversation,
shouldCompact,
} from "../../src/harness/compaction/compaction.js";
import { buildSessionContext } from "../../src/harness/session/session-tree.js";
import { buildSessionContext } from "../../src/harness/session/session.js";
import type {
CompactionEntry,
CompactionSettings,

View File

@@ -0,0 +1,65 @@
import { existsSync } from "node:fs";
import { describe, expect, it } from "vitest";
import { JsonlSessionRepo } from "../../src/harness/session/repo/jsonl.js";
import { InMemorySessionRepo } from "../../src/harness/session/repo/memory.js";
import { createAssistantMessage, createTempDir, createUserMessage } from "./session-test-utils.js";
describe("InMemorySessionRepo", () => {
it("opens, deletes, and forks by metadata", async () => {
const repo = new InMemorySessionRepo();
const session = await repo.create({ id: "session-1" });
const metadata = await session.getMetadata();
const user1 = await session.appendMessage(createUserMessage("one"));
const assistant1 = await session.appendMessage(createAssistantMessage("two"));
const user2 = await session.appendMessage(createUserMessage("three"));
expect(await repo.open(metadata)).toBe(session);
expect((await repo.list()).map((info) => info.id)).toEqual(["session-1"]);
const fork = await repo.fork(metadata, { entryId: user2, id: "session-2" });
expect((await fork.getEntries()).map((entry) => entry.id)).toEqual([user1, assistant1]);
const fullFork = await repo.fork(metadata, { id: "session-3" });
expect((await fullFork.getEntries()).map((entry) => entry.id)).toEqual([user1, assistant1, user2]);
await repo.delete(metadata);
await expect(repo.open(metadata)).rejects.toThrow("Session not found: session-1");
});
});
describe("JsonlSessionRepo", () => {
it("stores sessions below encoded cwd directories and lists by cwd", async () => {
const root = createTempDir();
const cwd = "/tmp/my-project";
const otherCwd = "/tmp/other-project";
const repo = new JsonlSessionRepo({ sessionsRoot: root });
const session = await repo.create({ cwd, id: "019de8c2-de29-73e9-ae0c-e134db34c447" });
const otherSession = await repo.create({ cwd: otherCwd, id: "other-session" });
const metadata = await session.getMetadata();
const otherMetadata = await otherSession.getMetadata();
expect(metadata.path).toContain("--tmp-my-project--");
expect(otherMetadata.path).toContain("--tmp-other-project--");
expect(existsSync(metadata.path)).toBe(true);
expect((await repo.list({ cwd })).map((sessionMetadata) => sessionMetadata.id)).toEqual([metadata.id]);
expect((await repo.list()).map((sessionMetadata) => sessionMetadata.id).sort()).toEqual(
[metadata.id, otherMetadata.id].sort(),
);
});
it("opens, deletes, and forks by metadata", async () => {
const root = createTempDir();
const repo = new JsonlSessionRepo({ sessionsRoot: root });
const source = await repo.create({ cwd: "/tmp/source", id: "source-session" });
const sourceMetadata = await source.getMetadata();
const user1 = await source.appendMessage(createUserMessage("one"));
const assistant1 = await source.appendMessage(createAssistantMessage("two"));
const user2 = await source.appendMessage(createUserMessage("three"));
await expect((await repo.open(sourceMetadata)).getMetadata()).resolves.toEqual(sourceMetadata);
const fork = await repo.fork(sourceMetadata, { cwd: "/tmp/target", id: "fork-session", entryId: user2 });
const forkMetadata = await fork.getMetadata();
expect(forkMetadata.cwd).toBe("/tmp/target");
expect(forkMetadata.parentSessionPath).toBe(sourceMetadata.path);
expect((await fork.getEntries()).map((entry) => entry.id)).toEqual([user1, assistant1]);
const fullFork = await repo.fork(sourceMetadata, { cwd: "/tmp/target", id: "full-fork-session" });
expect((await fullFork.getEntries()).map((entry) => entry.id)).toEqual([user1, assistant1, user2]);
await repo.delete(sourceMetadata);
expect(existsSync(sourceMetadata.path)).toBe(false);
await expect(repo.open(sourceMetadata)).rejects.toThrow("Session not found");
});
});

View File

@@ -0,0 +1,279 @@
import { existsSync, readFileSync, writeFileSync } from "node:fs";
import { join } from "node:path";
import { describe, expect, it } from "vitest";
import { JsonlSessionStorage, loadJsonlSessionMetadata } from "../../src/harness/session/storage/jsonl.js";
import { InMemorySessionStorage } from "../../src/harness/session/storage/memory.js";
import type { MessageEntry, SessionMetadata } from "../../src/harness/types.js";
import { createAssistantMessage, createTempDir, createUserMessage } from "./session-test-utils.js";
describe("InMemorySessionStorage", () => {
it("returns configured session metadata", async () => {
const metadata: SessionMetadata = { id: "session-1", createdAt: "2026-01-01T00:00:00.000Z" };
const storage = new InMemorySessionStorage({ metadata });
expect(await storage.getMetadata()).toEqual(metadata);
});
it("copies initial entries and tracks leaf independently", async () => {
const entry: MessageEntry = {
type: "message",
id: "entry-1",
parentId: null,
timestamp: "2026-01-01T00:00:00.000Z",
message: createUserMessage("one"),
};
const initialEntries = [entry];
const storage = new InMemorySessionStorage({ entries: initialEntries });
initialEntries.push({ ...entry, id: "entry-2" });
expect((await storage.getEntries()).map((storedEntry) => storedEntry.id)).toEqual(["entry-1"]);
expect(await storage.getLeafId()).toBe("entry-1");
await storage.setLeafId(null);
expect(await storage.getLeafId()).toBeNull();
});
it("rejects invalid leaf ids", async () => {
const storage = new InMemorySessionStorage();
await expect(storage.setLeafId("missing")).rejects.toThrow("Entry missing not found");
expect(() => new InMemorySessionStorage({ leafId: "missing" })).toThrow("Entry missing not found");
});
it("finds entries by type", async () => {
const entry: MessageEntry = {
type: "message",
id: "entry-1",
parentId: null,
timestamp: "2026-01-01T00:00:00.000Z",
message: createUserMessage("one"),
};
const storage = new InMemorySessionStorage({ entries: [entry] });
expect((await storage.findEntries("message")).map((found) => found.id)).toEqual(["entry-1"]);
expect(await storage.findEntries("session_info")).toEqual([]);
});
it("maintains label lookup", async () => {
const entry: MessageEntry = {
type: "message",
id: "entry-1",
parentId: null,
timestamp: "2026-01-01T00:00:00.000Z",
message: createUserMessage("one"),
};
const storage = new InMemorySessionStorage({ entries: [entry] });
expect(await storage.getLabel("entry-1")).toBeUndefined();
await storage.appendEntry({
type: "label",
id: "label-1",
parentId: "entry-1",
timestamp: "2026-01-01T00:00:01.000Z",
targetId: "entry-1",
label: "checkpoint",
});
expect(await storage.getLabel("entry-1")).toBe("checkpoint");
await storage.appendEntry({
type: "label",
id: "label-2",
parentId: "label-1",
timestamp: "2026-01-01T00:00:02.000Z",
targetId: "entry-1",
label: undefined,
});
expect(await storage.getLabel("entry-1")).toBeUndefined();
});
it("walks paths to root", async () => {
const root: MessageEntry = {
type: "message",
id: "root",
parentId: null,
timestamp: "2026-01-01T00:00:00.000Z",
message: createUserMessage("root"),
};
const child: MessageEntry = {
...root,
id: "child",
parentId: "root",
message: createAssistantMessage("child"),
};
const storage = new InMemorySessionStorage({ entries: [root, child] });
expect((await storage.getPathToRoot("child")).map((entry) => entry.id)).toEqual(["root", "child"]);
expect(await storage.getPathToRoot(null)).toEqual([]);
});
});
describe("JsonlSessionStorage", () => {
it("throws for missing files when opening", async () => {
const dir = createTempDir();
const filePath = join(dir, "session.jsonl");
await expect(JsonlSessionStorage.open(filePath)).rejects.toMatchObject({ code: "ENOENT" });
});
it("writes the header on create", async () => {
const dir = createTempDir();
const filePath = join(dir, "session.jsonl");
const storage = await JsonlSessionStorage.create(filePath, { cwd: dir, sessionId: "session-1" });
expect(existsSync(filePath)).toBe(true);
expect(readFileSync(filePath, "utf8").trim().split("\n")).toHaveLength(1);
expect(await storage.getLeafId()).toBeNull();
expect(await storage.getEntries()).toEqual([]);
await storage.appendEntry({
type: "message",
id: "user-1",
parentId: null,
timestamp: "2026-01-01T00:00:00.000Z",
message: createUserMessage("one"),
});
const lines = readFileSync(filePath, "utf8").trim().split("\n");
expect(JSON.parse(lines[0]!).type).toBe("session");
expect(JSON.parse(lines[1]!).id).toBe("user-1");
expect(lines).toHaveLength(2);
});
it("throws for malformed session headers", async () => {
const dir = createTempDir();
const filePath = join(dir, "session.jsonl");
writeFileSync(filePath, "not json\n");
await expect(JsonlSessionStorage.open(filePath)).rejects.toThrow("first line is not a valid session header");
});
it("ignores malformed entry lines", async () => {
const dir = createTempDir();
const filePath = join(dir, "session.jsonl");
const header = {
type: "session",
version: 3,
id: "session-1",
timestamp: "2026-01-01T00:00:00.000Z",
cwd: dir,
};
const entry: MessageEntry = {
type: "message",
id: "entry-1",
parentId: null,
timestamp: "2026-01-01T00:00:00.000Z",
message: createUserMessage("one"),
};
writeFileSync(filePath, `${JSON.stringify(header)}\nnot json\n${JSON.stringify(entry)}\n`);
const storage = await JsonlSessionStorage.open(filePath);
expect((await storage.getEntries()).map((loadedEntry) => loadedEntry.id)).toEqual(["entry-1"]);
expect(await storage.getLeafId()).toBe("entry-1");
});
it("creates and reads session metadata from the header", async () => {
const dir = createTempDir();
const filePath = join(dir, "session.jsonl");
const storage = await JsonlSessionStorage.create(filePath, {
cwd: dir,
sessionId: "session-1",
parentSessionPath: "/tmp/parent.jsonl",
});
const metadata = await storage.getMetadata();
expect(metadata).toMatchObject({
id: "session-1",
cwd: dir,
path: filePath,
parentSessionPath: "/tmp/parent.jsonl",
});
await storage.appendEntry({
type: "message",
id: "user-1",
parentId: null,
timestamp: "2026-01-01T00:00:00.000Z",
message: createUserMessage("one"),
});
expect(await loadJsonlSessionMetadata(filePath)).toEqual(metadata);
});
it("loads existing entries and reconstructs leaf", async () => {
const dir = createTempDir();
const filePath = join(dir, "session.jsonl");
const storage = await JsonlSessionStorage.create(filePath, { cwd: dir, sessionId: "session-1" });
const root: MessageEntry = {
type: "message",
id: "root",
parentId: null,
timestamp: "2026-01-01T00:00:00.000Z",
message: createUserMessage("root"),
};
const child: MessageEntry = {
...root,
id: "child",
parentId: "root",
message: createAssistantMessage("child"),
};
await storage.appendEntry(root);
await storage.appendEntry(child);
const loaded = await JsonlSessionStorage.open(filePath);
expect(await loaded.getLeafId()).toBe("child");
expect((await loaded.getEntries()).map((entry) => entry.id)).toEqual(["root", "child"]);
expect((await loaded.getPathToRoot("child")).map((entry) => entry.id)).toEqual(["root", "child"]);
});
it("finds entries by type", async () => {
const dir = createTempDir();
const filePath = join(dir, "session.jsonl");
const storage = await JsonlSessionStorage.create(filePath, { cwd: dir, sessionId: "session-1" });
await storage.appendEntry({
type: "message",
id: "entry-1",
parentId: null,
timestamp: "2026-01-01T00:00:00.000Z",
message: createUserMessage("one"),
});
expect((await storage.findEntries("message")).map((found) => found.id)).toEqual(["entry-1"]);
expect(await storage.findEntries("session_info")).toEqual([]);
});
it("maintains label lookup", async () => {
const dir = createTempDir();
const filePath = join(dir, "session.jsonl");
const storage = await JsonlSessionStorage.create(filePath, { cwd: dir, sessionId: "session-1" });
await storage.appendEntry({
type: "message",
id: "entry-1",
parentId: null,
timestamp: "2026-01-01T00:00:00.000Z",
message: createUserMessage("one"),
});
expect(await storage.getLabel("entry-1")).toBeUndefined();
await storage.appendEntry({
type: "label",
id: "label-1",
parentId: "entry-1",
timestamp: "2026-01-01T00:00:01.000Z",
targetId: "entry-1",
label: "checkpoint",
});
expect(await storage.getLabel("entry-1")).toBe("checkpoint");
await storage.appendEntry({
type: "label",
id: "label-2",
parentId: "label-1",
timestamp: "2026-01-01T00:00:02.000Z",
targetId: "entry-1",
label: undefined,
});
expect(await storage.getLabel("entry-1")).toBeUndefined();
const loaded = await JsonlSessionStorage.open(filePath);
expect(await loaded.getLabel("entry-1")).toBeUndefined();
});
it("reads session metadata from only the first JSONL line", async () => {
const dir = createTempDir();
const filePath = join(dir, "session.jsonl");
const header = {
type: "session",
version: 3,
id: "session-1",
timestamp: "2026-01-01T00:00:00.000Z",
cwd: dir,
};
const malformedSecondLine = "{".repeat(10000);
writeFileSync(filePath, `${JSON.stringify(header)}\n${malformedSecondLine}\n`);
expect(await loadJsonlSessionMetadata(filePath)).toEqual({
id: "session-1",
createdAt: "2026-01-01T00:00:00.000Z",
cwd: dir,
path: filePath,
parentSessionPath: undefined,
});
});
});

View File

@@ -0,0 +1,55 @@
import { existsSync, mkdirSync, rmSync } from "node:fs";
import { tmpdir } from "node:os";
import { join } from "node:path";
import type { AgentMessage } from "@mariozechner/pi-agent-core";
import { afterEach } from "vitest";
export function createUserMessage(text: string): AgentMessage {
return {
role: "user",
content: [{ type: "text", text }],
timestamp: Date.now(),
};
}
export function createAssistantMessage(text: string): AgentMessage {
return {
role: "assistant",
content: [{ type: "text", text }],
api: "anthropic-messages",
provider: "anthropic",
model: "claude-sonnet-4-5",
usage: {
input: 0,
output: 0,
cacheRead: 0,
cacheWrite: 0,
totalTokens: 0,
cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 },
},
stopReason: "stop",
timestamp: Date.now(),
};
}
const tempDirs: string[] = [];
export function createTempDir(): string {
const dir = join(tmpdir(), `pi-agent-session-${Date.now()}-${Math.random().toString(36).slice(2)}`);
mkdirSync(dir, { recursive: true });
tempDirs.push(dir);
return dir;
}
export function getLatestTempDir(): string {
return tempDirs[tempDirs.length - 1]!;
}
afterEach(() => {
while (tempDirs.length > 0) {
const dir = tempDirs.pop()!;
if (existsSync(dir)) {
rmSync(dir, { recursive: true, force: true });
}
}
});

View File

@@ -1,475 +1,137 @@
import { existsSync, mkdirSync, readFileSync, rmSync, writeFileSync } from "node:fs";
import { tmpdir } from "node:os";
import { readFileSync } from "node:fs";
import { join } from "node:path";
import type { AgentMessage } from "@mariozechner/pi-agent-core";
import { afterEach, describe, expect, it } from "vitest";
import { JsonlSessionRepo } from "../../src/harness/session/jsonl-session-repo.js";
import { JsonlSessionStorage, loadJsonlSessionMetadata } from "../../src/harness/session/jsonl-session-storage.js";
import { InMemorySessionRepo } from "../../src/harness/session/memory-session-repo.js";
import { InMemorySessionStorage } from "../../src/harness/session/memory-session-storage.js";
import { DefaultSession } from "../../src/harness/session/session-tree.js";
import type { MessageEntry, SessionMetadata, SessionStorage } from "../../src/harness/types.js";
import { describe, expect, it } from "vitest";
import { DefaultSession } from "../../src/harness/session/session.js";
import { JsonlSessionStorage } from "../../src/harness/session/storage/jsonl.js";
import { InMemorySessionStorage } from "../../src/harness/session/storage/memory.js";
import type { SessionStorage } from "../../src/harness/types.js";
import { createAssistantMessage, createTempDir, createUserMessage, getLatestTempDir } from "./session-test-utils.js";
function createUserMessage(text: string): AgentMessage {
return {
role: "user",
content: [{ type: "text", text }],
timestamp: Date.now(),
};
}
function createAssistantMessage(text: string): AgentMessage {
return {
role: "assistant",
content: [{ type: "text", text }],
api: "anthropic-messages",
provider: "anthropic",
model: "claude-sonnet-4-5",
usage: {
input: 0,
output: 0,
cacheRead: 0,
cacheWrite: 0,
totalTokens: 0,
cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 },
},
stopReason: "stop",
timestamp: Date.now(),
};
}
const tempDirs: string[] = [];
function createTempDir(): string {
const dir = join(tmpdir(), `pi-agent-session-tree-${Date.now()}-${Math.random().toString(36).slice(2)}`);
mkdirSync(dir, { recursive: true });
tempDirs.push(dir);
return dir;
}
afterEach(() => {
while (tempDirs.length > 0) {
const dir = tempDirs.pop()!;
if (existsSync(dir)) {
rmSync(dir, { recursive: true, force: true });
}
}
});
async function runSessionTreeSuite(
async function runSessionSuite(
name: string,
createStorage: () => SessionStorage | Promise<SessionStorage>,
inspect?: () => void,
) {
describe(name, () => {
it("appends messages and builds context in order", async () => {
const tree = new DefaultSession(await createStorage());
await tree.appendMessage(createUserMessage("one"));
await tree.appendMessage(createAssistantMessage("two"));
const context = await tree.buildContext();
const session = new DefaultSession(await createStorage());
await session.appendMessage(createUserMessage("one"));
await session.appendMessage(createAssistantMessage("two"));
const context = await session.buildContext();
expect(context.messages.map((message) => message.role)).toEqual(["user", "assistant"]);
});
it("tracks model and thinking level changes", async () => {
const tree = new DefaultSession(await createStorage());
await tree.appendMessage(createUserMessage("one"));
await tree.appendModelChange("openai", "gpt-4.1");
await tree.appendThinkingLevelChange("high");
const context = await tree.buildContext();
const session = new DefaultSession(await createStorage());
await session.appendMessage(createUserMessage("one"));
await session.appendModelChange("openai", "gpt-4.1");
await session.appendThinkingLevelChange("high");
const context = await session.buildContext();
expect(context.thinkingLevel).toBe("high");
expect(context.model).toEqual({ provider: "openai", modelId: "gpt-4.1" });
});
it("supports branching by moving the leaf and appending a new branch", async () => {
const tree = new DefaultSession(await createStorage());
const user1 = await tree.appendMessage(createUserMessage("one"));
const assistant1 = await tree.appendMessage(createAssistantMessage("two"));
await tree.appendMessage(createUserMessage("three"));
await tree.moveTo(user1);
await tree.appendMessage(createAssistantMessage("branched"));
const branch = await tree.getBranch();
const session = new DefaultSession(await createStorage());
const user1 = await session.appendMessage(createUserMessage("one"));
const assistant1 = await session.appendMessage(createAssistantMessage("two"));
await session.appendMessage(createUserMessage("three"));
await session.moveTo(user1);
await session.appendMessage(createAssistantMessage("branched"));
const branch = await session.getBranch();
expect(branch.map((entry) => entry.id)).toContain(user1);
expect(branch.map((entry) => entry.id)).not.toContain(assistant1);
const context = await tree.buildContext();
const context = await session.buildContext();
expect(context.messages.map((message) => message.role)).toEqual(["user", "assistant"]);
});
it("supports moving the leaf to root", async () => {
const tree = new DefaultSession(await createStorage());
await tree.appendMessage(createUserMessage("one"));
await tree.moveTo(null);
expect(await tree.getLeafId()).toBeNull();
expect((await tree.buildContext()).messages).toEqual([]);
const session = new DefaultSession(await createStorage());
await session.appendMessage(createUserMessage("one"));
await session.moveTo(null);
expect(await session.getLeafId()).toBeNull();
expect((await session.buildContext()).messages).toEqual([]);
});
it("reconstructs compaction summaries in context", async () => {
const tree = new DefaultSession(await createStorage());
await tree.appendMessage(createUserMessage("one"));
await tree.appendMessage(createAssistantMessage("two"));
const user2 = await tree.appendMessage(createUserMessage("three"));
await tree.appendMessage(createAssistantMessage("four"));
await tree.appendCompaction("summary", user2, 1234);
await tree.appendMessage(createUserMessage("five"));
const context = await tree.buildContext();
const session = new DefaultSession(await createStorage());
await session.appendMessage(createUserMessage("one"));
await session.appendMessage(createAssistantMessage("two"));
const user2 = await session.appendMessage(createUserMessage("three"));
await session.appendMessage(createAssistantMessage("four"));
await session.appendCompaction("summary", user2, 1234);
await session.appendMessage(createUserMessage("five"));
const context = await session.buildContext();
expect(context.messages[0]?.role).toBe("compactionSummary");
expect(context.messages).toHaveLength(4);
});
it("supports moving with branch summary entries in context", async () => {
const tree = new DefaultSession(await createStorage());
const user1 = await tree.appendMessage(createUserMessage("one"));
const summaryId = await tree.moveTo(user1, { summary: "summary text" });
const session = new DefaultSession(await createStorage());
const user1 = await session.appendMessage(createUserMessage("one"));
const summaryId = await session.moveTo(user1, { summary: "summary text" });
expect(summaryId).toBeTruthy();
const summaryEntry = await tree.getEntry(summaryId!);
const summaryEntry = await session.getEntry(summaryId!);
expect(summaryEntry).toMatchObject({ type: "branch_summary", parentId: user1, fromId: user1 });
const context = await tree.buildContext();
const context = await session.buildContext();
expect(context.messages[1]?.role).toBe("branchSummary");
});
it("supports custom message entries in context", async () => {
const tree = new DefaultSession(await createStorage());
await tree.appendMessage(createUserMessage("one"));
await tree.appendCustomMessageEntry("custom", "hello", true, { ok: true });
const context = await tree.buildContext();
const session = new DefaultSession(await createStorage());
await session.appendMessage(createUserMessage("one"));
await session.appendCustomMessageEntry("custom", "hello", true, { ok: true });
const context = await session.buildContext();
expect(context.messages[1]?.role).toBe("custom");
});
it("supports labels and session info entries without affecting context", async () => {
const tree = new DefaultSession(await createStorage());
const user1 = await tree.appendMessage(createUserMessage("one"));
await tree.appendLabel(user1, "checkpoint");
await tree.appendSessionName("name");
const entries = await tree.getEntries();
const session = new DefaultSession(await createStorage());
const user1 = await session.appendMessage(createUserMessage("one"));
await session.appendLabel(user1, "checkpoint");
await session.appendSessionName("name");
const entries = await session.getEntries();
expect(entries.some((entry) => entry.type === "label")).toBe(true);
expect(entries.some((entry) => entry.type === "session_info")).toBe(true);
expect(await tree.getLabel(user1)).toBe("checkpoint");
expect(await tree.getSessionName()).toBe("name");
expect((await tree.buildContext()).messages).toHaveLength(1);
expect(await session.getLabel(user1)).toBe("checkpoint");
expect(await session.getSessionName()).toBe("name");
expect((await session.buildContext()).messages).toHaveLength(1);
});
it("rejects labels for missing entries", async () => {
const tree = new DefaultSession(await createStorage());
await expect(tree.appendLabel("missing", "checkpoint")).rejects.toThrow("Entry missing not found");
const session = new DefaultSession(await createStorage());
await expect(session.appendLabel("missing", "checkpoint")).rejects.toThrow("Entry missing not found");
});
it("persists leaf changes and appended entries via storage", async () => {
const storage = await createStorage();
const tree = new DefaultSession(storage);
const user1 = await tree.appendMessage(createUserMessage("one"));
await tree.appendMessage(createAssistantMessage("two"));
await tree.appendLabel(user1, "checkpoint");
await tree.appendSessionName("name");
await tree.moveTo(user1);
await tree.appendMessage(createAssistantMessage("branched"));
const tree2 = new DefaultSession(storage);
const context = await tree2.buildContext();
const session = new DefaultSession(storage);
const user1 = await session.appendMessage(createUserMessage("one"));
await session.appendMessage(createAssistantMessage("two"));
await session.appendLabel(user1, "checkpoint");
await session.appendSessionName("name");
await session.moveTo(user1);
await session.appendMessage(createAssistantMessage("branched"));
const session2 = new DefaultSession(storage);
const context = await session2.buildContext();
expect(context.messages.map((message) => message.role)).toEqual(["user", "assistant"]);
expect(await tree2.getLabel(user1)).toBe("checkpoint");
expect(await tree2.getSessionName()).toBe("name");
expect(await session2.getLabel(user1)).toBe("checkpoint");
expect(await session2.getSessionName()).toBe("name");
inspect?.();
});
});
}
describe("InMemorySessionStorage", () => {
it("returns configured session info", async () => {
const metadata: SessionMetadata = { id: "session-1", createdAt: "2026-01-01T00:00:00.000Z" };
const storage = new InMemorySessionStorage({ metadata });
expect(await storage.getMetadata()).toEqual(metadata);
});
runSessionSuite("Session with in-memory storage", () => new InMemorySessionStorage());
it("copies initial entries and tracks leaf independently", async () => {
const entry: MessageEntry = {
type: "message",
id: "entry-1",
parentId: null,
timestamp: "2026-01-01T00:00:00.000Z",
message: createUserMessage("one"),
};
const initialEntries = [entry];
const storage = new InMemorySessionStorage({ entries: initialEntries });
initialEntries.push({ ...entry, id: "entry-2" });
expect((await storage.getEntries()).map((storedEntry) => storedEntry.id)).toEqual(["entry-1"]);
expect(await storage.getLeafId()).toBe("entry-1");
await storage.setLeafId(null);
expect(await storage.getLeafId()).toBeNull();
});
it("rejects invalid leaf ids", async () => {
const storage = new InMemorySessionStorage();
await expect(storage.setLeafId("missing")).rejects.toThrow("Entry missing not found");
expect(() => new InMemorySessionStorage({ leafId: "missing" })).toThrow("Entry missing not found");
});
it("maintains label lookup", async () => {
const entry: MessageEntry = {
type: "message",
id: "entry-1",
parentId: null,
timestamp: "2026-01-01T00:00:00.000Z",
message: createUserMessage("one"),
};
const storage = new InMemorySessionStorage({ entries: [entry] });
expect(await storage.getLabel("entry-1")).toBeUndefined();
await storage.appendEntry({
type: "label",
id: "label-1",
parentId: "entry-1",
timestamp: "2026-01-01T00:00:01.000Z",
targetId: "entry-1",
label: "checkpoint",
});
expect(await storage.getLabel("entry-1")).toBe("checkpoint");
await storage.appendEntry({
type: "label",
id: "label-2",
parentId: "label-1",
timestamp: "2026-01-01T00:00:02.000Z",
targetId: "entry-1",
label: undefined,
});
expect(await storage.getLabel("entry-1")).toBeUndefined();
});
it("walks paths to root", async () => {
const root: MessageEntry = {
type: "message",
id: "root",
parentId: null,
timestamp: "2026-01-01T00:00:00.000Z",
message: createUserMessage("root"),
};
const child: MessageEntry = {
...root,
id: "child",
parentId: "root",
message: createAssistantMessage("child"),
};
const storage = new InMemorySessionStorage({ entries: [root, child] });
expect((await storage.getPathToRoot("child")).map((entry) => entry.id)).toEqual(["root", "child"]);
expect(await storage.getPathToRoot(null)).toEqual([]);
});
});
runSessionTreeSuite("Session with in-memory storage", () => new InMemorySessionStorage());
describe("InMemorySessionRepo", () => {
it("lists session infos and forks via storage path traversal", async () => {
const repo = new InMemorySessionRepo();
const session = await repo.create({ id: "session-1" });
const user1 = await session.appendMessage(createUserMessage("one"));
const assistant1 = await session.appendMessage(createAssistantMessage("two"));
const user2 = await session.appendMessage(createUserMessage("three"));
const infos = await repo.list();
expect(infos.map((info) => info.id)).toEqual(["session-1"]);
const fork = await repo.fork("session-1", { entryId: user2, id: "session-2" });
expect((await fork.getEntries()).map((entry) => entry.id)).toEqual([user1, assistant1]);
});
});
describe("JsonlSessionRepo", () => {
it("stores sessions below encoded cwd directories and resolves prefixes", async () => {
const root = createTempDir();
const cwd = "/tmp/my-project";
const repo = new JsonlSessionRepo({ sessionsRoot: root });
const session = await repo.create({ cwd, id: "019de8c2-de29-73e9-ae0c-e134db34c447" });
const info = await session.getMetadata();
expect(info.path).toContain("--tmp-my-project--");
expect(existsSync(info.path)).toBe(true);
expect((await repo.list({ cwd })).map((sessionInfo) => sessionInfo.id)).toEqual([info.id]);
expect((await repo.resolve("019de8c2", { cwd })).map((sessionInfo) => sessionInfo.path)).toEqual([info.path]);
});
it("forks sessions and records the parent path", async () => {
const root = createTempDir();
const repo = new JsonlSessionRepo({ sessionsRoot: root });
const source = await repo.create({ cwd: "/tmp/source", id: "source-session" });
const sourceInfo = await source.getMetadata();
const user1 = await source.appendMessage(createUserMessage("one"));
const assistant1 = await source.appendMessage(createAssistantMessage("two"));
const user2 = await source.appendMessage(createUserMessage("three"));
const fork = await repo.fork(sourceInfo, { cwd: "/tmp/target", id: "fork-session", entryId: user2 });
const forkInfo = await fork.getMetadata();
expect(forkInfo.cwd).toBe("/tmp/target");
expect(forkInfo.parentSessionPath).toBe(sourceInfo.path);
expect((await fork.getEntries()).map((entry) => entry.id)).toEqual([user1, assistant1]);
});
});
describe("JsonlSessionStorage", () => {
it("throws for missing files when opening", async () => {
const dir = createTempDir();
const filePath = join(dir, "session.jsonl");
await expect(JsonlSessionStorage.open(filePath)).rejects.toMatchObject({ code: "ENOENT" });
});
it("writes the header on create", async () => {
const dir = createTempDir();
const filePath = join(dir, "session.jsonl");
const storage = await JsonlSessionStorage.create(filePath, { cwd: dir, sessionId: "session-1" });
expect(existsSync(filePath)).toBe(true);
expect(readFileSync(filePath, "utf8").trim().split("\n")).toHaveLength(1);
expect(await storage.getLeafId()).toBeNull();
expect(await storage.getEntries()).toEqual([]);
await storage.appendEntry({
type: "message",
id: "user-1",
parentId: null,
timestamp: "2026-01-01T00:00:00.000Z",
message: createUserMessage("one"),
});
expect(existsSync(filePath)).toBe(true);
const lines = readFileSync(filePath, "utf8").trim().split("\n");
expect(JSON.parse(lines[0]!).type).toBe("session");
expect(JSON.parse(lines[1]!).id).toBe("user-1");
expect(lines).toHaveLength(2);
});
it("throws for malformed session headers", async () => {
const dir = createTempDir();
const filePath = join(dir, "session.jsonl");
writeFileSync(filePath, "not json\n");
await expect(JsonlSessionStorage.open(filePath)).rejects.toThrow("first line is not a valid session header");
});
it("ignores malformed entry lines", async () => {
const dir = createTempDir();
const filePath = join(dir, "session.jsonl");
const header = {
type: "session",
version: 3,
id: "session-1",
timestamp: "2026-01-01T00:00:00.000Z",
cwd: dir,
};
const entry: MessageEntry = {
type: "message",
id: "entry-1",
parentId: null,
timestamp: "2026-01-01T00:00:00.000Z",
message: createUserMessage("one"),
};
writeFileSync(filePath, `${JSON.stringify(header)}\nnot json\n${JSON.stringify(entry)}\n`);
const storage = await JsonlSessionStorage.open(filePath);
expect((await storage.getEntries()).map((loadedEntry) => loadedEntry.id)).toEqual(["entry-1"]);
expect(await storage.getLeafId()).toBe("entry-1");
});
it("creates and reads session info from the header", async () => {
const dir = createTempDir();
const filePath = join(dir, "session.jsonl");
const storage = await JsonlSessionStorage.create(filePath, {
cwd: dir,
sessionId: "session-1",
parentSessionPath: "/tmp/parent.jsonl",
});
const info = await storage.getMetadata();
expect(info).toMatchObject({
id: "session-1",
cwd: dir,
path: filePath,
parentSessionPath: "/tmp/parent.jsonl",
});
expect(existsSync(filePath)).toBe(true);
await storage.appendEntry({
type: "message",
id: "user-1",
parentId: null,
timestamp: "2026-01-01T00:00:00.000Z",
message: createUserMessage("one"),
});
expect(await loadJsonlSessionMetadata(filePath)).toEqual(info);
});
it("loads existing entries and reconstructs leaf", async () => {
const dir = createTempDir();
const filePath = join(dir, "session.jsonl");
const storage = await JsonlSessionStorage.create(filePath, { cwd: dir, sessionId: "session-1" });
const root: MessageEntry = {
type: "message",
id: "root",
parentId: null,
timestamp: "2026-01-01T00:00:00.000Z",
message: createUserMessage("root"),
};
const child: MessageEntry = {
...root,
id: "child",
parentId: "root",
message: createAssistantMessage("child"),
};
await storage.appendEntry(root);
await storage.appendEntry(child);
const loaded = await JsonlSessionStorage.open(filePath);
expect(await loaded.getLeafId()).toBe("child");
expect((await loaded.getEntries()).map((entry) => entry.id)).toEqual(["root", "child"]);
expect((await loaded.getPathToRoot("child")).map((entry) => entry.id)).toEqual(["root", "child"]);
});
it("maintains label lookup", async () => {
const dir = createTempDir();
const filePath = join(dir, "session.jsonl");
const storage = await JsonlSessionStorage.create(filePath, { cwd: dir, sessionId: "session-1" });
await storage.appendEntry({
type: "message",
id: "entry-1",
parentId: null,
timestamp: "2026-01-01T00:00:00.000Z",
message: createUserMessage("one"),
});
expect(await storage.getLabel("entry-1")).toBeUndefined();
await storage.appendEntry({
type: "label",
id: "label-1",
parentId: "entry-1",
timestamp: "2026-01-01T00:00:01.000Z",
targetId: "entry-1",
label: "checkpoint",
});
expect(await storage.getLabel("entry-1")).toBe("checkpoint");
await storage.appendEntry({
type: "label",
id: "label-2",
parentId: "label-1",
timestamp: "2026-01-01T00:00:02.000Z",
targetId: "entry-1",
label: undefined,
});
expect(await storage.getLabel("entry-1")).toBeUndefined();
const loaded = await JsonlSessionStorage.open(filePath);
expect(await loaded.getLabel("entry-1")).toBeUndefined();
});
it("reads session info from only the first JSONL line", async () => {
const dir = createTempDir();
const filePath = join(dir, "session.jsonl");
const header = {
type: "session",
version: 3,
id: "session-1",
timestamp: "2026-01-01T00:00:00.000Z",
cwd: dir,
};
const malformedSecondLine = "{".repeat(10000);
writeFileSync(filePath, `${JSON.stringify(header)}\n${malformedSecondLine}\n`);
expect(await loadJsonlSessionMetadata(filePath)).toEqual({
id: "session-1",
createdAt: "2026-01-01T00:00:00.000Z",
cwd: dir,
path: filePath,
parentSessionPath: undefined,
});
});
});
runSessionTreeSuite(
runSessionSuite(
"Session with JSONL storage",
async () => {
const dir = createTempDir();
return await JsonlSessionStorage.create(join(dir, "session.jsonl"), { cwd: dir, sessionId: "session-1" });
},
() => {
const dir = tempDirs[tempDirs.length - 1]!;
const dir = getLatestTempDir();
const filePath = join(dir, "session.jsonl");
const lines = readFileSync(filePath, "utf8").trim().split("\n");
expect(lines.length).toBeGreaterThan(1);