diff --git a/packages/agent/src/harness/session/jsonl-session-storage.ts b/packages/agent/src/harness/session/jsonl-session-storage.ts index 2e2541b2..ffc7a437 100644 --- a/packages/agent/src/harness/session/jsonl-session-storage.ts +++ b/packages/agent/src/harness/session/jsonl-session-storage.ts @@ -1,7 +1,8 @@ -import { randomUUID } from "node:crypto"; +import { createReadStream } from "node:fs"; import { appendFile, mkdir, readFile, writeFile } from "node:fs/promises"; import { dirname, resolve } from "node:path"; -import type { CodingAgentSessionInfo, SessionInfo, SessionTreeEntry, SessionTreeStorage } from "../types.js"; +import { createInterface } from "node:readline"; +import type { JsonlSessionInfo, SessionTreeEntry, SessionTreeStorage } from "../types.js"; interface SessionHeader { type: "session"; @@ -12,116 +13,140 @@ interface SessionHeader { parentSession?: string; } -function headerToSessionInfo(header: SessionHeader, filePath?: string): CodingAgentSessionInfo { +function headerToSessionInfo(header: SessionHeader, path: string): JsonlSessionInfo { return { id: header.id, createdAt: header.timestamp, - parentSession: header.parentSession, - projectCwd: header.cwd, - filePath, + cwd: header.cwd, + path, + parentSessionPath: header.parentSession, }; } -async function loadJsonlStorage( - filePath: string, -): Promise<{ header?: SessionHeader; entries: SessionTreeEntry[]; leafId: string | null }> { +export async function loadJsonlSessionInfo(filePath: string): Promise { + const stream = createReadStream(filePath, { encoding: "utf8" }); + const lines = createInterface({ input: stream, crlfDelay: Infinity }); try { - const content = await readFile(filePath, "utf8"); - const entries: SessionTreeEntry[] = []; - let header: SessionHeader | undefined; - let leafId: string | null = null; - for (const line of content.split("\n")) { - if (!line.trim()) continue; + for await (const line of lines) { + if (!line.trim()) break; try { - const record = JSON.parse(line) as SessionHeader | SessionTreeEntry; - if (record.type === "session") { - header = record as SessionHeader; - continue; - } - entries.push(record as SessionTreeEntry); - leafId = (record as SessionTreeEntry).id; + const header = JSON.parse(line) as SessionHeader; + return headerToSessionInfo(header, resolve(filePath)); } catch { - // ignore malformed lines + throw new Error(`Invalid JSONL session file ${filePath}: first line is not a valid session header`); } } - return { header, entries, leafId }; - } catch { - return { entries: [], leafId: null }; + throw new Error(`Invalid JSONL session file ${filePath}: missing session header`); + } finally { + lines.close(); + stream.destroy(); } } -export class JsonlSessionTreeStorage implements SessionTreeStorage { - private filePath: string; - private cwd: string; - private headerInitialized = false; - private cacheLoaded = false; - private sessionInfo?: CodingAgentSessionInfo; - private entries: SessionTreeEntry[] = []; - private byId = new Map(); - private currentLeafId: string | null = null; - private requestedSessionId?: string; - private parentSession?: string; +async function loadJsonlStorage(filePath: string): Promise<{ + header: SessionHeader; + entries: SessionTreeEntry[]; + leafId: string | null; +}> { + const content = await readFile(filePath, "utf8"); + const lines = content.split("\n").filter((line) => line.trim()); + if (lines.length === 0) { + throw new Error(`Invalid JSONL session file ${filePath}: missing session header`); + } - constructor(filePath: string, options: { cwd: string; sessionId?: string; parentSession?: string }) { + let header: SessionHeader; + try { + header = JSON.parse(lines[0]!) as SessionHeader; + } catch { + throw new Error(`Invalid JSONL session file ${filePath}: first line is not a valid session header`); + } + + const entries: SessionTreeEntry[] = []; + let leafId: string | null = null; + for (const line of lines.slice(1)) { + try { + const entry = JSON.parse(line) as SessionTreeEntry; + entries.push(entry); + leafId = entry.id; + } catch { + // ignore malformed entry lines + } + } + return { header, entries, leafId }; +} + +export class JsonlSessionTreeStorage implements SessionTreeStorage { + private readonly filePath: string; + private readonly header: SessionHeader; + private readonly sessionInfo: JsonlSessionInfo; + private entries: SessionTreeEntry[]; + private byId: Map; + private currentLeafId: string | null; + private headerWritten: boolean; + + private constructor( + filePath: string, + header: SessionHeader, + entries: SessionTreeEntry[], + leafId: string | null, + headerWritten: boolean, + ) { this.filePath = resolve(filePath); - this.cwd = options.cwd; - this.requestedSessionId = options.sessionId; - this.parentSession = options.parentSession; + this.header = header; + this.sessionInfo = headerToSessionInfo(header, this.filePath); + this.entries = entries; + this.byId = new Map(entries.map((entry) => [entry.id, entry])); + this.currentLeafId = leafId; + this.headerWritten = headerWritten; } - private async ensureParentDir(): Promise { - await mkdir(dirname(this.filePath), { recursive: true }); + static async open(filePath: string): Promise { + const resolvedPath = resolve(filePath); + const loaded = await loadJsonlStorage(resolvedPath); + return new JsonlSessionTreeStorage(resolvedPath, loaded.header, loaded.entries, loaded.leafId, true); } - private async ensureLoaded(): Promise { - if (this.cacheLoaded) { - return; - } - const loaded = await loadJsonlStorage(this.filePath); - this.entries = loaded.entries; - this.byId = new Map(loaded.entries.map((entry) => [entry.id, entry])); - this.currentLeafId = loaded.leafId; - this.headerInitialized = loaded.header !== undefined; - if (loaded.header) { - this.sessionInfo = headerToSessionInfo(loaded.header, this.filePath); - } - this.cacheLoaded = true; - } - - private async ensureHeader(): Promise { - await this.ensureLoaded(); - if (this.headerInitialized) return; - await this.ensureParentDir(); + static async create( + filePath: string, + options: { + cwd: string; + sessionId: string; + parentSessionPath?: string; + }, + ): Promise { + const resolvedPath = resolve(filePath); const header: SessionHeader = { type: "session", version: 3, - id: this.requestedSessionId ?? randomUUID(), + id: options.sessionId, timestamp: new Date().toISOString(), - cwd: this.cwd, - parentSession: this.parentSession, + cwd: options.cwd, + parentSession: options.parentSessionPath, }; - await writeFile(this.filePath, `${JSON.stringify(header)}\n`); - this.sessionInfo = headerToSessionInfo(header, this.filePath); - this.headerInitialized = true; + return new JsonlSessionTreeStorage(resolvedPath, header, [], null, false); } - async getSessionInfo(): Promise { - await this.ensureHeader(); - return this.sessionInfo!; + async getSessionInfo(): Promise { + return this.sessionInfo; } async getLeafId(): Promise { - await this.ensureLoaded(); return this.currentLeafId; } async setLeafId(leafId: string | null): Promise { - await this.ensureLoaded(); + if (leafId !== null && !this.byId.has(leafId)) { + throw new Error(`Entry ${leafId} not found`); + } this.currentLeafId = leafId; } async appendEntry(entry: SessionTreeEntry): Promise { - await this.ensureHeader(); + if (!this.headerWritten) { + await mkdir(dirname(this.filePath), { recursive: true }); + await writeFile(this.filePath, `${JSON.stringify(this.header)}\n`); + this.headerWritten = true; + } await appendFile(this.filePath, `${JSON.stringify(entry)}\n`); this.entries.push(entry); this.byId.set(entry.id, entry); @@ -129,12 +154,10 @@ export class JsonlSessionTreeStorage implements SessionTreeStorage { } async getEntry(id: string): Promise { - await this.ensureLoaded(); return this.byId.get(id); } async getPathToRoot(leafId: string | null): Promise { - await this.ensureLoaded(); if (leafId === null) return []; const path: SessionTreeEntry[] = []; let current = this.byId.get(leafId); @@ -146,7 +169,6 @@ export class JsonlSessionTreeStorage implements SessionTreeStorage { } async getEntries(): Promise { - await this.ensureLoaded(); return [...this.entries]; } } diff --git a/packages/agent/src/harness/session/memory-session-storage.ts b/packages/agent/src/harness/session/memory-session-storage.ts index b1fe743f..4485d87b 100644 --- a/packages/agent/src/harness/session/memory-session-storage.ts +++ b/packages/agent/src/harness/session/memory-session-storage.ts @@ -1,15 +1,20 @@ -import { randomUUID } from "crypto"; +import { v7 as uuidv7 } from "uuid"; import type { SessionInfo, SessionTreeEntry, SessionTreeStorage } from "../types.js"; export class InMemorySessionTreeStorage implements SessionTreeStorage { + private readonly sessionInfo: SessionInfo; private entries: SessionTreeEntry[]; + private byId: Map; private leafId: string | null; - private sessionInfo: SessionInfo; constructor(options?: { entries?: SessionTreeEntry[]; leafId?: string | null; sessionInfo?: SessionInfo }) { this.entries = options?.entries ? [...options.entries] : []; + this.byId = new Map(this.entries.map((entry) => [entry.id, entry])); this.leafId = options?.leafId ?? this.entries[this.entries.length - 1]?.id ?? null; - this.sessionInfo = options?.sessionInfo ?? { id: randomUUID(), createdAt: new Date().toISOString() }; + if (this.leafId !== null && !this.byId.has(this.leafId)) { + throw new Error(`Entry ${this.leafId} not found`); + } + this.sessionInfo = options?.sessionInfo ?? { id: uuidv7(), createdAt: new Date().toISOString() }; } async getSessionInfo(): Promise { @@ -21,26 +26,29 @@ export class InMemorySessionTreeStorage implements SessionTreeStorage { } async setLeafId(leafId: string | null): Promise { + if (leafId !== null && !this.byId.has(leafId)) { + throw new Error(`Entry ${leafId} not found`); + } this.leafId = leafId; } async appendEntry(entry: SessionTreeEntry): Promise { this.entries.push(entry); + this.byId.set(entry.id, entry); this.leafId = entry.id; } async getEntry(id: string): Promise { - return this.entries.find((entry) => entry.id === id); + return this.byId.get(id); } async getPathToRoot(leafId: string | null): Promise { if (leafId === null) return []; - const byId = new Map(this.entries.map((entry) => [entry.id, entry])); const path: SessionTreeEntry[] = []; - let current = byId.get(leafId); + let current = this.byId.get(leafId); while (current) { path.unshift(current); - current = current.parentId ? byId.get(current.parentId) : undefined; + current = current.parentId ? this.byId.get(current.parentId) : undefined; } return path; } diff --git a/packages/agent/src/harness/session/session-repo.ts b/packages/agent/src/harness/session/session-repo.ts index 587b0f9d..bde9c291 100644 --- a/packages/agent/src/harness/session/session-repo.ts +++ b/packages/agent/src/harness/session/session-repo.ts @@ -1,13 +1,14 @@ -import { existsSync, mkdirSync, readdirSync, readFileSync, rmSync } from "node:fs"; +import { existsSync, mkdirSync, readdirSync, rmSync } from "node:fs"; import { join, resolve } from "node:path"; import { v7 as uuidv7 } from "uuid"; import type { - CodingAgentSessionInfo, - CodingAgentSessionRepo, + JsonlSessionInfo, + JsonlSessionRepo, Session, SessionInfo, SessionRepo, SessionTreeEntry, + SessionTreeStorage, } from "../types.js"; import { JsonlSessionTreeStorage } from "./jsonl-session-storage.js"; import { InMemorySessionTreeStorage } from "./memory-session-storage.js"; @@ -21,8 +22,11 @@ function createTimestamp(): string { return new Date().toISOString(); } -function toSession(info: TInfo, tree: DefaultSessionTree): Session { - return { info, tree }; +function toSession( + storage: SessionTreeStorage, + tree: DefaultSessionTree, +): Session { + return { storage, tree }; } function getPathEntriesToFork( @@ -59,14 +63,13 @@ function getPathEntriesToFork( export class InMemorySessionRepo implements SessionRepo { private sessions = new Map>(); - async create(options?: { id?: string; parentSession?: string }): Promise> { + async create(options?: { id?: string }): Promise> { const info: SessionInfo = { id: options?.id ?? createSessionId(), createdAt: createTimestamp(), - parentSession: options?.parentSession, }; const storage = new InMemorySessionTreeStorage({ sessionInfo: info }); - const session = toSession(info, new DefaultSessionTree(storage)); + const session = toSession(storage, new DefaultSessionTree(storage)); this.sessions.set(info.id, session); return session; } @@ -97,42 +100,16 @@ export class InMemorySessionRepo implements SessionRepo { const info: SessionInfo = { id: options.id ?? createSessionId(), createdAt: createTimestamp(), - parentSession: source.info.id, }; const leafId = forkedEntries[forkedEntries.length - 1]?.id ?? null; const storage = new InMemorySessionTreeStorage({ sessionInfo: info, entries: forkedEntries, leafId }); - const session = toSession(info, new DefaultSessionTree(storage)); + const session = toSession(storage, new DefaultSessionTree(storage)); this.sessions.set(info.id, session); return session; } } -function readJsonlHeader(filePath: string): CodingAgentSessionInfo | undefined { - try { - const content = readFileSync(filePath, "utf8"); - const firstLine = content.split("\n")[0]; - if (!firstLine) return undefined; - const header = JSON.parse(firstLine) as { - type: string; - id: string; - timestamp: string; - cwd: string; - parentSession?: string; - }; - if (header.type !== "session") return undefined; - return { - id: header.id, - createdAt: header.timestamp, - parentSession: header.parentSession, - projectCwd: header.cwd, - filePath, - }; - } catch { - return undefined; - } -} - -export class JsonlCodingAgentSessionRepo implements CodingAgentSessionRepo { +export class JsonlSessionFileRepo implements JsonlSessionRepo { private sessionDir: string; private cwd: string; @@ -146,55 +123,66 @@ export class JsonlCodingAgentSessionRepo implements CodingAgentSessionRepo> { + async create(options?: { id?: string; parentSessionPath?: string }): Promise> { const id = options?.id ?? createSessionId(); const createdAt = createTimestamp(); const filePath = this.createSessionFilePath(id, createdAt); - const storage = new JsonlSessionTreeStorage(filePath, { + const storage = await JsonlSessionTreeStorage.create(filePath, { cwd: this.cwd, sessionId: id, - parentSession: options?.parentSession, + parentSessionPath: options?.parentSessionPath, }); - const info = (await storage.getSessionInfo()) as CodingAgentSessionInfo; - return toSession(info, new DefaultSessionTree(storage)); + return toSession(storage, new DefaultSessionTree(storage)); } - async open(ref: string): Promise> { + async open(ref: string): Promise> { const filePath = ref.includes("/") || ref.endsWith(".jsonl") ? resolve(ref) : join(this.sessionDir, ref); if (!existsSync(filePath)) { throw new Error(`Session not found: ${ref}`); } - const storage = new JsonlSessionTreeStorage(filePath, { cwd: this.cwd }); - const info = (await storage.getSessionInfo()) as CodingAgentSessionInfo; - return toSession(info, new DefaultSessionTree(storage)); + const storage = await JsonlSessionTreeStorage.open(filePath); + return toSession(storage, new DefaultSessionTree(storage)); } - async list(): Promise>> { + async list(): Promise>> { if (!existsSync(this.sessionDir)) { return []; } const files = readdirSync(this.sessionDir) .filter((file) => file.endsWith(".jsonl")) .map((file) => join(this.sessionDir, file)); - const sessions: Array> = []; + const sessions: Array> = []; for (const filePath of files) { - const info = readJsonlHeader(filePath); - if (!info) continue; - sessions.push( - toSession(info, new DefaultSessionTree(new JsonlSessionTreeStorage(filePath, { cwd: info.projectCwd }))), - ); + try { + const storage = await JsonlSessionTreeStorage.open(filePath); + sessions.push(toSession(storage, new DefaultSessionTree(storage))); + } catch { + // Ignore invalid session files when listing a directory. + } } return sessions; } - async listByCwd(cwd: string): Promise>> { - return (await this.list()).filter((session) => session.info.projectCwd === cwd); + async listByCwd(cwd: string): Promise>> { + const sessions = await this.list(); + const result: Array> = []; + for (const session of sessions) { + if ((await session.storage.getSessionInfo()).cwd === cwd) { + result.push(session); + } + } + return result; } - async getMostRecentByCwd(cwd: string): Promise | undefined> { - const sessions = await this.listByCwd(cwd); - sessions.sort((a, b) => new Date(b.info.createdAt).getTime() - new Date(a.info.createdAt).getTime()); - return sessions[0]; + async getMostRecentByCwd(cwd: string): Promise | undefined> { + const sessionsWithInfo = await Promise.all( + (await this.listByCwd(cwd)).map(async (session) => ({ + session, + info: await session.storage.getSessionInfo(), + })), + ); + sessionsWithInfo.sort((a, b) => new Date(b.info.createdAt).getTime() - new Date(a.info.createdAt).getTime()); + return sessionsWithInfo[0]?.session; } async delete(ref: string): Promise { @@ -207,17 +195,18 @@ export class JsonlCodingAgentSessionRepo implements CodingAgentSessionRepo> { + ): Promise> { const source = await this.open(ref); const entries = await source.tree.getEntries(); const forkedEntries = getPathEntriesToFork(entries, options.entryId, options.position ?? "before"); const id = options.id ?? createSessionId(); const createdAt = createTimestamp(); const filePath = this.createSessionFilePath(id, createdAt); - const storage = new JsonlSessionTreeStorage(filePath, { - cwd: source.info.projectCwd, + const sourceInfo = await source.storage.getSessionInfo(); + const storage = await JsonlSessionTreeStorage.create(filePath, { + cwd: sourceInfo.cwd, sessionId: id, - parentSession: source.info.filePath ?? source.info.id, + parentSessionPath: sourceInfo.path, }); for (const entry of forkedEntries) { await storage.appendEntry(entry); @@ -225,7 +214,6 @@ export class JsonlCodingAgentSessionRepo implements CodingAgentSessionRepo implements SessionTree { + private storage: SessionTreeStorage; - constructor(storage?: SessionTreeStorage) { - this.storage = storage ?? new InMemorySessionTreeStorage(); + constructor(storage: SessionTreeStorage) { + this.storage = storage; } getLeafId(): Promise { @@ -107,7 +106,7 @@ export class DefaultSessionTree implements SessionTree { return buildSessionContext(await this.getBranch()); } - getSessionInfo(): Promise { + getSessionInfo(): Promise { return this.storage.getSessionInfo(); } diff --git a/packages/agent/src/harness/types.ts b/packages/agent/src/harness/types.ts index fa75a953..15f23853 100644 --- a/packages/agent/src/harness/types.ts +++ b/packages/agent/src/harness/types.ts @@ -160,16 +160,16 @@ export interface SessionContext { export interface SessionInfo { id: string; createdAt: string; - parentSession?: string; } -export interface CodingAgentSessionInfo extends SessionInfo { - projectCwd: string; - filePath?: string; +export interface JsonlSessionInfo extends SessionInfo { + cwd: string; + path: string; + parentSessionPath?: string; } -export interface SessionTreeStorage { - getSessionInfo(): Promise; +export interface SessionTreeStorage { + getSessionInfo(): Promise; getLeafId(): Promise; setLeafId(leafId: string | null): Promise; appendEntry(entry: SessionTreeEntry): Promise; @@ -179,7 +179,6 @@ export interface SessionTreeStorage { } export interface SessionTree { - getSessionInfo(): Promise; getLeafId(): Promise; getEntry(id: string): Promise; getEntries(): Promise; @@ -213,21 +212,40 @@ export interface SessionTree { } export interface Session { - info: TInfo; + storage: SessionTreeStorage; tree: SessionTree; } -export interface SessionRepo { - create(options?: { id?: string; parentSession?: string }): Promise>; +export interface SessionCreateOptions { + id?: string; +} + +export interface SessionForkOptions { + entryId: string; + position?: "before" | "at"; + id?: string; +} + +export interface SessionRepo< + TRef = string, + TInfo extends SessionInfo = SessionInfo, + TCreateOptions extends SessionCreateOptions = SessionCreateOptions, +> { + create(options?: TCreateOptions): Promise>; open(ref: TRef): Promise>; list(): Promise>>; delete(ref: TRef): Promise; - fork(ref: TRef, options: { entryId: string; position?: "before" | "at"; id?: string }): Promise>; + fork(ref: TRef, options: SessionForkOptions): Promise>; } -export interface CodingAgentSessionRepo extends SessionRepo { - listByCwd(cwd: string): Promise>>; - getMostRecentByCwd(cwd: string): Promise | undefined>; +export interface JsonlSessionCreateOptions extends SessionCreateOptions { + parentSessionPath?: string; +} + +export interface JsonlSessionRepo + extends SessionRepo { + listByCwd(cwd: string): Promise>>; + getMostRecentByCwd(cwd: string): Promise | undefined>; } export interface AgentHarnessPendingMutations { diff --git a/packages/agent/test/harness/session-tree.test.ts b/packages/agent/test/harness/session-tree.test.ts index eb1171c8..ffcb01d4 100644 --- a/packages/agent/test/harness/session-tree.test.ts +++ b/packages/agent/test/harness/session-tree.test.ts @@ -1,12 +1,12 @@ -import { existsSync, mkdirSync, readFileSync, rmSync } from "node:fs"; +import { existsSync, mkdirSync, readFileSync, rmSync, writeFileSync } from "node:fs"; import { tmpdir } from "node:os"; import { join } from "node:path"; import type { AgentMessage } from "@mariozechner/pi-agent-core"; import { afterEach, describe, expect, it } from "vitest"; -import { JsonlSessionTreeStorage } from "../../src/harness/session/jsonl-session-storage.js"; +import { JsonlSessionTreeStorage, loadJsonlSessionInfo } from "../../src/harness/session/jsonl-session-storage.js"; import { InMemorySessionTreeStorage } from "../../src/harness/session/memory-session-storage.js"; import { DefaultSessionTree } from "../../src/harness/session/session-tree.js"; -import type { SessionTreeStorage } from "../../src/harness/types.js"; +import type { MessageEntry, SessionInfo, SessionTreeStorage } from "../../src/harness/types.js"; function createUserMessage(text: string): AgentMessage { return { @@ -38,6 +38,13 @@ function createAssistantMessage(text: string): AgentMessage { 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()!; @@ -47,10 +54,14 @@ afterEach(() => { } }); -async function runSessionTreeSuite(name: string, createStorage: () => SessionTreeStorage, inspect?: () => void) { +async function runSessionTreeSuite( + name: string, + createStorage: () => SessionTreeStorage | Promise, + inspect?: () => void, +) { describe(name, () => { it("appends messages and builds context in order", async () => { - const tree = new DefaultSessionTree(createStorage()); + const tree = new DefaultSessionTree(await createStorage()); await tree.appendMessage(createUserMessage("one")); await tree.appendMessage(createAssistantMessage("two")); const context = await tree.buildContext(); @@ -58,7 +69,7 @@ async function runSessionTreeSuite(name: string, createStorage: () => SessionTre }); it("tracks model and thinking level changes", async () => { - const tree = new DefaultSessionTree(createStorage()); + const tree = new DefaultSessionTree(await createStorage()); await tree.appendMessage(createUserMessage("one")); await tree.appendModelChange("openai", "gpt-4.1"); await tree.appendThinkingLevelChange("high"); @@ -68,7 +79,7 @@ async function runSessionTreeSuite(name: string, createStorage: () => SessionTre }); it("supports branching by moving the leaf and appending a new branch", async () => { - const tree = new DefaultSessionTree(createStorage()); + const tree = new DefaultSessionTree(await createStorage()); const user1 = await tree.appendMessage(createUserMessage("one")); const assistant1 = await tree.appendMessage(createAssistantMessage("two")); await tree.appendMessage(createUserMessage("three")); @@ -82,7 +93,7 @@ async function runSessionTreeSuite(name: string, createStorage: () => SessionTre }); it("supports moving the leaf to root", async () => { - const tree = new DefaultSessionTree(createStorage()); + const tree = new DefaultSessionTree(await createStorage()); await tree.appendMessage(createUserMessage("one")); await tree.moveTo(null); expect(await tree.getLeafId()).toBeNull(); @@ -90,7 +101,7 @@ async function runSessionTreeSuite(name: string, createStorage: () => SessionTre }); it("reconstructs compaction summaries in context", async () => { - const tree = new DefaultSessionTree(createStorage()); + const tree = new DefaultSessionTree(await createStorage()); await tree.appendMessage(createUserMessage("one")); await tree.appendMessage(createAssistantMessage("two")); const user2 = await tree.appendMessage(createUserMessage("three")); @@ -103,7 +114,7 @@ async function runSessionTreeSuite(name: string, createStorage: () => SessionTre }); it("supports branch summary entries in context", async () => { - const tree = new DefaultSessionTree(createStorage()); + const tree = new DefaultSessionTree(await createStorage()); const user1 = await tree.appendMessage(createUserMessage("one")); await tree.appendBranchSummary(user1, "summary text"); const context = await tree.buildContext(); @@ -111,7 +122,7 @@ async function runSessionTreeSuite(name: string, createStorage: () => SessionTre }); it("supports custom message entries in context", async () => { - const tree = new DefaultSessionTree(createStorage()); + const tree = new DefaultSessionTree(await createStorage()); await tree.appendMessage(createUserMessage("one")); await tree.appendCustomMessageEntry("custom", "hello", true, { ok: true }); const context = await tree.buildContext(); @@ -119,7 +130,7 @@ async function runSessionTreeSuite(name: string, createStorage: () => SessionTre }); it("supports labels and session info entries without affecting context", async () => { - const tree = new DefaultSessionTree(createStorage()); + const tree = new DefaultSessionTree(await createStorage()); const user1 = await tree.appendMessage(createUserMessage("one")); await tree.appendLabelChange(user1, "checkpoint"); await tree.appendSessionInfo("name"); @@ -132,7 +143,7 @@ async function runSessionTreeSuite(name: string, createStorage: () => SessionTre }); it("persists leaf changes and appended entries via storage", async () => { - const storage = createStorage(); + const storage = await createStorage(); const tree = new DefaultSessionTree(storage); const user1 = await tree.appendMessage(createUserMessage("one")); await tree.appendMessage(createAssistantMessage("two")); @@ -150,15 +161,194 @@ async function runSessionTreeSuite(name: string, createStorage: () => SessionTre }); } +describe("InMemorySessionTreeStorage", () => { + it("returns configured session info", async () => { + const sessionInfo: SessionInfo = { id: "session-1", createdAt: "2026-01-01T00:00:00.000Z" }; + const storage = new InMemorySessionTreeStorage({ sessionInfo }); + expect(await storage.getSessionInfo()).toEqual(sessionInfo); + }); + + 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 InMemorySessionTreeStorage({ 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 InMemorySessionTreeStorage(); + await expect(storage.setLeafId("missing")).rejects.toThrow("Entry missing not found"); + expect(() => new InMemorySessionTreeStorage({ leafId: "missing" })).toThrow("Entry missing not found"); + }); + + 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 InMemorySessionTreeStorage({ entries: [root, child] }); + expect((await storage.getPathToRoot("child")).map((entry) => entry.id)).toEqual(["root", "child"]); + expect(await storage.getPathToRoot(null)).toEqual([]); + }); +}); + runSessionTreeSuite("SessionTree with in-memory storage", () => new InMemorySessionTreeStorage()); +describe("JsonlSessionTreeStorage", () => { + it("throws for missing files when opening", async () => { + const dir = createTempDir(); + const filePath = join(dir, "session.jsonl"); + await expect(JsonlSessionTreeStorage.open(filePath)).rejects.toMatchObject({ code: "ENOENT" }); + }); + + it("writes the header before the first appended entry", async () => { + const dir = createTempDir(); + const filePath = join(dir, "session.jsonl"); + const storage = await JsonlSessionTreeStorage.create(filePath, { cwd: dir, sessionId: "session-1" }); + expect(existsSync(filePath)).toBe(false); + 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(JsonlSessionTreeStorage.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 JsonlSessionTreeStorage.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 JsonlSessionTreeStorage.create(filePath, { + cwd: dir, + sessionId: "session-1", + parentSessionPath: "/tmp/parent.jsonl", + }); + const info = await storage.getSessionInfo(); + expect(info).toMatchObject({ + id: "session-1", + cwd: dir, + path: filePath, + parentSessionPath: "/tmp/parent.jsonl", + }); + expect(existsSync(filePath)).toBe(false); + await storage.appendEntry({ + type: "message", + id: "user-1", + parentId: null, + timestamp: "2026-01-01T00:00:00.000Z", + message: createUserMessage("one"), + }); + expect(await loadJsonlSessionInfo(filePath)).toEqual(info); + }); + + it("loads existing entries and reconstructs leaf", async () => { + const dir = createTempDir(); + const filePath = join(dir, "session.jsonl"); + const storage = await JsonlSessionTreeStorage.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 JsonlSessionTreeStorage.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("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 loadJsonlSessionInfo(filePath)).toEqual({ + id: "session-1", + createdAt: "2026-01-01T00:00:00.000Z", + cwd: dir, + path: filePath, + parentSessionPath: undefined, + }); + }); +}); + runSessionTreeSuite( "SessionTree with JSONL storage", - () => { - const dir = join(tmpdir(), `pi-agent-session-tree-${Date.now()}-${Math.random().toString(36).slice(2)}`); - mkdirSync(dir, { recursive: true }); - tempDirs.push(dir); - return new JsonlSessionTreeStorage(join(dir, "session.jsonl"), { cwd: dir }); + async () => { + const dir = createTempDir(); + return await JsonlSessionTreeStorage.create(join(dir, "session.jsonl"), { cwd: dir, sessionId: "session-1" }); }, () => { const dir = tempDirs[tempDirs.length - 1]!;