diff --git a/packages/agent/src/harness/compaction/compaction.ts b/packages/agent/src/harness/compaction/compaction.ts index da28abc3..9ba9ac01 100644 --- a/packages/agent/src/harness/compaction/compaction.ts +++ b/packages/agent/src/harness/compaction/compaction.ts @@ -14,7 +14,7 @@ import { createCompactionSummaryMessage, createCustomMessage, } from "../messages.js"; -import { buildSessionContext } from "../session/session-tree.js"; +import { buildSessionContext } from "../session/session.js"; import type { CompactionEntry, SessionTreeEntry } from "../types.js"; import { computeFileLists, diff --git a/packages/agent/src/harness/session/memory-session-repo.ts b/packages/agent/src/harness/session/memory-session-repo.ts deleted file mode 100644 index 318f6ed2..00000000 --- a/packages/agent/src/harness/session/memory-session-repo.ts +++ /dev/null @@ -1,55 +0,0 @@ -import type { Session, SessionMetadata, SessionRepo } from "../types.js"; -import { InMemorySessionStorage } from "./memory-session-storage.js"; -import { createSessionId, createTimestamp, getPathEntriesToFork, toSession } from "./session-repo.js"; - -export class InMemorySessionRepo implements SessionRepo { - private sessions = new Map>(); - - async create(options: { id?: string } = {}): Promise> { - const info: SessionMetadata = { - id: options.id ?? createSessionId(), - createdAt: createTimestamp(), - }; - const storage = new InMemorySessionStorage({ metadata: info }); - const session = toSession(storage); - this.sessions.set(info.id, session); - return session; - } - - async open(ref: string): Promise> { - const session = this.sessions.get(ref); - if (!session) { - throw new Error(`Session not found: ${ref}`); - } - return session; - } - - async list(): Promise { - return Promise.all([...this.sessions.values()].map((session) => session.getMetadata())); - } - - async delete(ref: string): Promise { - this.sessions.delete(ref); - } - - async fork( - ref: string, - options: { entryId: string; position?: "before" | "at"; id?: string }, - ): Promise> { - const source = await this.open(ref); - const forkedEntries = await getPathEntriesToFork( - source.getStorage(), - options.entryId, - options.position ?? "before", - ); - const info: SessionMetadata = { - id: options.id ?? createSessionId(), - createdAt: createTimestamp(), - }; - const leafId = forkedEntries[forkedEntries.length - 1]?.id ?? null; - const storage = new InMemorySessionStorage({ metadata: info, entries: forkedEntries, leafId }); - const session = toSession(storage); - this.sessions.set(info.id, session); - return session; - } -} diff --git a/packages/agent/src/harness/session/jsonl-session-repo.ts b/packages/agent/src/harness/session/repo/jsonl.ts similarity index 58% rename from packages/agent/src/harness/session/jsonl-session-repo.ts rename to packages/agent/src/harness/session/repo/jsonl.ts index 39a31ce0..a4c5ec5b 100644 --- a/packages/agent/src/harness/session/jsonl-session-repo.ts +++ b/packages/agent/src/harness/session/repo/jsonl.ts @@ -3,15 +3,13 @@ import { access, mkdir, readdir, rm } from "node:fs/promises"; import { join, resolve } from "node:path"; import type { JsonlSessionCreateOptions, - JsonlSessionListQuery, + JsonlSessionListOptions, JsonlSessionMetadata, - JsonlSessionRef, JsonlSessionRepoApi, - JsonlSessionResolveOptions, Session, -} from "../types.js"; -import { JsonlSessionStorage, loadJsonlSessionMetadata } from "./jsonl-session-storage.js"; -import { createSessionId, createTimestamp, getPathEntriesToFork, toSession } from "./session-repo.js"; +} from "../../types.js"; +import { JsonlSessionStorage, loadJsonlSessionMetadata } from "../storage/jsonl.js"; +import { createSessionId, createTimestamp, getEntriesToFork, toSession } from "./shared.js"; async function exists(path: string): Promise { try { @@ -41,10 +39,6 @@ export class JsonlSessionRepo implements JsonlSessionRepoApi { return join(this.getSessionDir(cwd), `${timestamp.replace(/[:.]/g, "-")}_${sessionId}.jsonl`); } - private refPath(ref: JsonlSessionRef): string { - return resolve(ref.path); - } - async create(options: JsonlSessionCreateOptions): Promise> { await mkdir(this.sessionsRoot, { recursive: true }); const id = options.id ?? createSessionId(); @@ -58,17 +52,16 @@ export class JsonlSessionRepo implements JsonlSessionRepoApi { return toSession(storage); } - async open(ref: JsonlSessionRef): Promise> { - const filePath = this.refPath(ref); - if (!(await exists(filePath))) { - throw new Error(`Session not found: ${filePath}`); + async open(metadata: JsonlSessionMetadata): Promise> { + if (!(await exists(metadata.path))) { + throw new Error(`Session not found: ${metadata.path}`); } - const storage = await JsonlSessionStorage.open(filePath); + const storage = await JsonlSessionStorage.open(metadata.path); return toSession(storage); } - async list(query: JsonlSessionListQuery = {}): Promise { - const dirs = query.cwd ? [this.getSessionDir(query.cwd)] : await this.listSessionDirs(); + async list(options: JsonlSessionListOptions = {}): Promise { + const dirs = options.cwd ? [this.getSessionDir(options.cwd)] : await this.listSessionDirs(); const sessions: JsonlSessionMetadata[] = []; for (const dir of dirs) { if (!(await exists(dir))) continue; @@ -85,48 +78,22 @@ export class JsonlSessionRepo implements JsonlSessionRepoApi { return sessions; } - async resolve(ref: string, options: JsonlSessionResolveOptions = {}): Promise { - if (ref.includes("/") || ref.includes("\\") || ref.endsWith(".jsonl")) { - try { - return [await loadJsonlSessionMetadata(resolve(ref))]; - } catch { - return []; - } - } - - const local = options.cwd - ? (await this.list({ cwd: options.cwd })).filter((session) => session.id.startsWith(ref)) - : []; - if (local.length > 0 || !options.searchAll) return local; - return (await this.list()).filter((session) => session.id.startsWith(ref)); - } - - async getMostRecent(query: JsonlSessionListQuery = {}): Promise { - return (await this.list(query))[0]; - } - - async delete(ref: JsonlSessionRef): Promise { - const filePath = this.refPath(ref); - await rm(filePath, { force: true }); + async delete(metadata: JsonlSessionMetadata): Promise { + await rm(metadata.path, { force: true }); } async fork( - ref: JsonlSessionRef, - options: JsonlSessionCreateOptions & { entryId: string; position?: "before" | "at"; id?: string }, + sourceMetadata: JsonlSessionMetadata, + options: JsonlSessionCreateOptions & { entryId?: string; position?: "before" | "at"; id?: string }, ): Promise> { - const source = await this.open(ref); - const forkedEntries = await getPathEntriesToFork( - source.getStorage(), - options.entryId, - options.position ?? "before", - ); - const sourceInfo = await source.getMetadata(); + const source = await this.open(sourceMetadata); + const forkedEntries = await getEntriesToFork(source.getStorage(), options); const id = options.id ?? createSessionId(); const createdAt = createTimestamp(); const storage = await JsonlSessionStorage.create(this.createSessionFilePath(options.cwd, id, createdAt), { cwd: options.cwd, sessionId: id, - parentSessionPath: options.parentSessionPath ?? sourceInfo.path, + parentSessionPath: options.parentSessionPath ?? sourceMetadata.path, }); for (const entry of forkedEntries) { await storage.appendEntry(entry); diff --git a/packages/agent/src/harness/session/repo/memory.ts b/packages/agent/src/harness/session/repo/memory.ts new file mode 100644 index 00000000..3846ae2a --- /dev/null +++ b/packages/agent/src/harness/session/repo/memory.ts @@ -0,0 +1,51 @@ +import type { Session, SessionMetadata, SessionRepo } from "../../types.js"; +import { InMemorySessionStorage } from "../storage/memory.js"; +import { createSessionId, createTimestamp, getEntriesToFork, toSession } from "./shared.js"; + +export class InMemorySessionRepo implements SessionRepo { + private sessions = new Map>(); + + async create(options: { id?: string } = {}): Promise> { + const metadata: SessionMetadata = { + id: options.id ?? createSessionId(), + createdAt: createTimestamp(), + }; + const storage = new InMemorySessionStorage({ metadata }); + const session = toSession(storage); + this.sessions.set(metadata.id, session); + return session; + } + + async open(metadata: SessionMetadata): Promise> { + const session = this.sessions.get(metadata.id); + if (!session) { + throw new Error(`Session not found: ${metadata.id}`); + } + return session; + } + + async list(): Promise { + return Promise.all([...this.sessions.values()].map((session) => session.getMetadata())); + } + + async delete(metadata: SessionMetadata): Promise { + this.sessions.delete(metadata.id); + } + + async fork( + sourceMetadata: SessionMetadata, + options: { entryId?: string; position?: "before" | "at"; id?: string }, + ): Promise> { + const source = await this.open(sourceMetadata); + const forkedEntries = await getEntriesToFork(source.getStorage(), options); + const metadata: SessionMetadata = { + id: options.id ?? createSessionId(), + createdAt: createTimestamp(), + }; + const leafId = forkedEntries[forkedEntries.length - 1]?.id ?? null; + const storage = new InMemorySessionStorage({ metadata, entries: forkedEntries, leafId }); + const session = toSession(storage); + this.sessions.set(metadata.id, session); + return session; + } +} diff --git a/packages/agent/src/harness/session/session-repo.ts b/packages/agent/src/harness/session/repo/shared.ts similarity index 60% rename from packages/agent/src/harness/session/session-repo.ts rename to packages/agent/src/harness/session/repo/shared.ts index 3362d4a0..f6395614 100644 --- a/packages/agent/src/harness/session/session-repo.ts +++ b/packages/agent/src/harness/session/repo/shared.ts @@ -1,6 +1,6 @@ import { v7 as uuidv7 } from "uuid"; -import type { Session, SessionMetadata, SessionStorage, SessionTreeEntry } from "../types.js"; -import { DefaultSession } from "./session-tree.js"; +import type { Session, SessionMetadata, SessionStorage, SessionTreeEntry } from "../../types.js"; +import { DefaultSession } from "../session.js"; export function createSessionId(): string { return uuidv7(); @@ -14,21 +14,21 @@ export function toSession(storage: SessionSto return new DefaultSession(storage); } -export async function getPathEntriesToFork( +export async function getEntriesToFork( storage: SessionStorage, - entryId: string, - position: "before" | "at", + options: { entryId?: string; position?: "before" | "at" }, ): Promise { - const target = await storage.getEntry(entryId); + if (!options.entryId) return storage.getEntries(); + const target = await storage.getEntry(options.entryId); if (!target) { - throw new Error(`Entry ${entryId} not found`); + throw new Error(`Entry ${options.entryId} not found`); } let effectiveLeafId: string | null; - if (position === "at") { + if ((options.position ?? "before") === "at") { effectiveLeafId = target.id; } else { if (target.type !== "message" || target.message.role !== "user") { - throw new Error(`Entry ${entryId} is not a user message`); + throw new Error(`Entry ${options.entryId} is not a user message`); } effectiveLeafId = target.parentId; } diff --git a/packages/agent/src/harness/session/session-tree.ts b/packages/agent/src/harness/session/session.ts similarity index 88% rename from packages/agent/src/harness/session/session-tree.ts rename to packages/agent/src/harness/session/session.ts index 3ed24975..62f59c38 100644 --- a/packages/agent/src/harness/session/session-tree.ts +++ b/packages/agent/src/harness/session/session.ts @@ -1,4 +1,3 @@ -import { randomUUID } from "node:crypto"; import type { ImageContent, TextContent } from "@mariozechner/pi-ai"; import type { AgentMessage } from "../../types.js"; import { createBranchSummaryMessage, createCompactionSummaryMessage, createCustomMessage } from "../messages.js"; @@ -19,14 +18,6 @@ import type { ThinkingLevelChangeEntry, } from "../types.js"; -function generateId(byId: { has(id: string): boolean }): string { - for (let i = 0; i < 100; i++) { - const id = randomUUID().slice(0, 8); - if (!byId.has(id)) return id; - } - return randomUUID(); -} - export function buildSessionContext(pathEntries: SessionTreeEntry[]): SessionContext { let thinkingLevel = "off"; let model: { provider: string; modelId: string } | null = null; @@ -125,19 +116,8 @@ export class DefaultSession } async getSessionName(): Promise { - const entries = await this.storage.getEntries(); - for (let i = entries.length - 1; i >= 0; i--) { - const entry = entries[i]!; - if (entry.type === "session_info") { - return entry.name?.trim() || undefined; - } - } - return undefined; - } - - private async makeEntryId(): Promise { - const entries = await this.storage.getEntries(); - return generateId(new Set(entries.map((entry) => entry.id))); + const entries = await this.storage.findEntries("session_info"); + return entries[entries.length - 1]?.name?.trim() || undefined; } private async appendTypedEntry(entry: TEntry): Promise { @@ -148,7 +128,7 @@ export class DefaultSession async appendMessage(message: AgentMessage): Promise { return this.appendTypedEntry({ type: "message", - id: await this.makeEntryId(), + id: await this.storage.createEntryId(), parentId: await this.storage.getLeafId(), timestamp: new Date().toISOString(), message, @@ -158,7 +138,7 @@ export class DefaultSession async appendThinkingLevelChange(thinkingLevel: string): Promise { return this.appendTypedEntry({ type: "thinking_level_change", - id: await this.makeEntryId(), + id: await this.storage.createEntryId(), parentId: await this.storage.getLeafId(), timestamp: new Date().toISOString(), thinkingLevel, @@ -168,7 +148,7 @@ export class DefaultSession async appendModelChange(provider: string, modelId: string): Promise { return this.appendTypedEntry({ type: "model_change", - id: await this.makeEntryId(), + id: await this.storage.createEntryId(), parentId: await this.storage.getLeafId(), timestamp: new Date().toISOString(), provider, @@ -185,7 +165,7 @@ export class DefaultSession ): Promise { return this.appendTypedEntry({ type: "compaction", - id: await this.makeEntryId(), + id: await this.storage.createEntryId(), parentId: await this.storage.getLeafId(), timestamp: new Date().toISOString(), summary, @@ -199,7 +179,7 @@ export class DefaultSession async appendCustomEntry(customType: string, data?: unknown): Promise { return this.appendTypedEntry({ type: "custom", - id: await this.makeEntryId(), + id: await this.storage.createEntryId(), parentId: await this.storage.getLeafId(), timestamp: new Date().toISOString(), customType, @@ -215,7 +195,7 @@ export class DefaultSession ): Promise { return this.appendTypedEntry({ type: "custom_message", - id: await this.makeEntryId(), + id: await this.storage.createEntryId(), parentId: await this.storage.getLeafId(), timestamp: new Date().toISOString(), customType, @@ -231,7 +211,7 @@ export class DefaultSession } return this.appendTypedEntry({ type: "label", - id: await this.makeEntryId(), + id: await this.storage.createEntryId(), parentId: await this.storage.getLeafId(), timestamp: new Date().toISOString(), targetId, @@ -242,7 +222,7 @@ export class DefaultSession async appendSessionName(name: string): Promise { return this.appendTypedEntry({ type: "session_info", - id: await this.makeEntryId(), + id: await this.storage.createEntryId(), parentId: await this.storage.getLeafId(), timestamp: new Date().toISOString(), name: name.trim(), @@ -260,7 +240,7 @@ export class DefaultSession if (!summary) return undefined; return this.appendTypedEntry({ type: "branch_summary", - id: await this.makeEntryId(), + id: await this.storage.createEntryId(), parentId: entryId, timestamp: new Date().toISOString(), fromId: entryId ?? "root", diff --git a/packages/agent/src/harness/session/jsonl-session-storage.ts b/packages/agent/src/harness/session/storage/jsonl.ts similarity index 90% rename from packages/agent/src/harness/session/jsonl-session-storage.ts rename to packages/agent/src/harness/session/storage/jsonl.ts index 0ec40ed1..72ce6ce3 100644 --- a/packages/agent/src/harness/session/jsonl-session-storage.ts +++ b/packages/agent/src/harness/session/storage/jsonl.ts @@ -1,8 +1,9 @@ +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 { createInterface } from "node:readline"; -import type { JsonlSessionMetadata, SessionStorage, SessionTreeEntry } from "../types.js"; +import type { JsonlSessionMetadata, SessionStorage, SessionTreeEntry } from "../../types.js"; interface SessionHeader { type: "session"; @@ -31,6 +32,14 @@ function buildLabelsById(entries: SessionTreeEntry[]): Map { return labelsById; } +function generateEntryId(byId: { has(id: string): boolean }): string { + for (let i = 0; i < 100; i++) { + const id = randomUUID().slice(0, 8); + if (!byId.has(id)) return id; + } + return randomUUID(); +} + function headerToSessionMetadata(header: SessionHeader, path: string): JsonlSessionMetadata { return { id: header.id, @@ -153,6 +162,10 @@ export class JsonlSessionStorage implements SessionStorage this.currentLeafId = leafId; } + async createEntryId(): Promise { + return generateEntryId(this.byId); + } + async appendEntry(entry: SessionTreeEntry): Promise { await appendFile(this.filePath, `${JSON.stringify(entry)}\n`); this.entries.push(entry); @@ -165,6 +178,12 @@ export class JsonlSessionStorage implements SessionStorage return this.byId.get(id); } + async findEntries( + type: TType, + ): Promise>> { + return this.entries.filter((entry): entry is Extract => entry.type === type); + } + async getLabel(id: string): Promise { return this.labelsById.get(id); } diff --git a/packages/agent/src/harness/session/memory-session-storage.ts b/packages/agent/src/harness/session/storage/memory.ts similarity index 81% rename from packages/agent/src/harness/session/memory-session-storage.ts rename to packages/agent/src/harness/session/storage/memory.ts index cdf3e791..648ae6b0 100644 --- a/packages/agent/src/harness/session/memory-session-storage.ts +++ b/packages/agent/src/harness/session/storage/memory.ts @@ -1,5 +1,6 @@ +import { randomUUID } from "node:crypto"; import { v7 as uuidv7 } from "uuid"; -import type { SessionMetadata, SessionStorage, SessionTreeEntry } from "../types.js"; +import type { SessionMetadata, SessionStorage, SessionTreeEntry } from "../../types.js"; function updateLabelCache(labelsById: Map, entry: SessionTreeEntry): void { if (entry.type !== "label") return; @@ -19,6 +20,14 @@ function buildLabelsById(entries: SessionTreeEntry[]): Map { return labelsById; } +function generateEntryId(byId: { has(id: string): boolean }): string { + for (let i = 0; i < 100; i++) { + const id = randomUUID().slice(0, 8); + if (!byId.has(id)) return id; + } + return randomUUID(); +} + export class InMemorySessionStorage implements SessionStorage { private readonly metadata: SessionMetadata; private entries: SessionTreeEntry[]; @@ -52,6 +61,10 @@ export class InMemorySessionStorage implements SessionStorage { this.leafId = leafId; } + async createEntryId(): Promise { + return generateEntryId(this.byId); + } + async appendEntry(entry: SessionTreeEntry): Promise { this.entries.push(entry); this.byId.set(entry.id, entry); @@ -63,6 +76,12 @@ export class InMemorySessionStorage implements SessionStorage { return this.byId.get(id); } + async findEntries( + type: TType, + ): Promise>> { + return this.entries.filter((entry): entry is Extract => entry.type === type); + } + async getLabel(id: string): Promise { return this.labelsById.get(id); } diff --git a/packages/agent/src/harness/types.ts b/packages/agent/src/harness/types.ts index 7975690b..605fe447 100644 --- a/packages/agent/src/harness/types.ts +++ b/packages/agent/src/harness/types.ts @@ -172,8 +172,12 @@ export interface SessionStorage; getLeafId(): Promise; setLeafId(leafId: string | null): Promise; + createEntryId(): Promise; appendEntry(entry: SessionTreeEntry): Promise; getEntry(id: string): Promise; + findEntries( + type: TType, + ): Promise>>; getLabel(id: string): Promise; getPathToRoot(leafId: string | null): Promise; getEntries(): Promise; @@ -222,7 +226,7 @@ export interface SessionCreateOptions { } export interface SessionForkOptions { - entryId: string; + entryId?: string; position?: "before" | "at"; id?: string; } @@ -230,14 +234,13 @@ export interface SessionForkOptions { export interface SessionRepo< TMetadata extends SessionMetadata = SessionMetadata, TCreateOptions extends SessionCreateOptions = SessionCreateOptions, - TRef = string, - TListQuery = void, + TListOptions = void, > { create(options: TCreateOptions): Promise>; - open(ref: TRef): Promise>; - list(query?: TListQuery): Promise; - delete(ref: TRef): Promise; - fork(ref: TRef, options: SessionForkOptions & TCreateOptions): Promise>; + open(metadata: TMetadata): Promise>; + list(options?: TListOptions): Promise; + delete(metadata: TMetadata): Promise; + fork(source: TMetadata, options: SessionForkOptions & TCreateOptions): Promise>; } export interface JsonlSessionCreateOptions extends SessionCreateOptions { @@ -245,22 +248,12 @@ export interface JsonlSessionCreateOptions extends SessionCreateOptions { parentSessionPath?: string; } -export type JsonlSessionRef = { path: string } | JsonlSessionMetadata; - -export interface JsonlSessionListQuery { +export interface JsonlSessionListOptions { cwd?: string; } -export interface JsonlSessionResolveOptions { - cwd?: string; - searchAll?: boolean; -} - export interface JsonlSessionRepoApi - extends SessionRepo { - resolve(ref: string, options?: JsonlSessionResolveOptions): Promise; - getMostRecent(query?: JsonlSessionListQuery): Promise; -} + extends SessionRepo {} export interface AgentHarnessPendingMutations { appendMessages: AgentMessage[]; diff --git a/packages/agent/src/index.ts b/packages/agent/src/index.ts index 43557ef0..d48629a4 100644 --- a/packages/agent/src/index.ts +++ b/packages/agent/src/index.ts @@ -25,10 +25,10 @@ export { export * from "./harness/execution-env.js"; export * from "./harness/messages.js"; export * from "./harness/prompt-templates.js"; -export * from "./harness/session/jsonl-session-repo.js"; -export * from "./harness/session/memory-session-repo.js"; -export * from "./harness/session/session-repo.js"; -export * from "./harness/session/session-tree.js"; +export * from "./harness/session/repo/jsonl.js"; +export * from "./harness/session/repo/memory.js"; +export * from "./harness/session/repo/shared.js"; +export * from "./harness/session/session.js"; // Harness export * from "./harness/types.js"; export * from "./harness/utils/shell-output.js"; diff --git a/packages/agent/test/harness/compaction.test.ts b/packages/agent/test/harness/compaction.test.ts index 0c260ab2..40585766 100644 --- a/packages/agent/test/harness/compaction.test.ts +++ b/packages/agent/test/harness/compaction.test.ts @@ -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, diff --git a/packages/agent/test/harness/session-repo.test.ts b/packages/agent/test/harness/session-repo.test.ts new file mode 100644 index 00000000..310d8898 --- /dev/null +++ b/packages/agent/test/harness/session-repo.test.ts @@ -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"); + }); +}); diff --git a/packages/agent/test/harness/session-storage.test.ts b/packages/agent/test/harness/session-storage.test.ts new file mode 100644 index 00000000..8c05702b --- /dev/null +++ b/packages/agent/test/harness/session-storage.test.ts @@ -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, + }); + }); +}); diff --git a/packages/agent/test/harness/session-test-utils.ts b/packages/agent/test/harness/session-test-utils.ts new file mode 100644 index 00000000..b1531918 --- /dev/null +++ b/packages/agent/test/harness/session-test-utils.ts @@ -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 }); + } + } +}); diff --git a/packages/agent/test/harness/session-tree.test.ts b/packages/agent/test/harness/session-tree.test.ts index 17d578fb..6fa49594 100644 --- a/packages/agent/test/harness/session-tree.test.ts +++ b/packages/agent/test/harness/session-tree.test.ts @@ -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, 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);