diff --git a/packages/agent/src/harness/agent-harness.ts b/packages/agent/src/harness/agent-harness.ts index b03372d2..03518d12 100644 --- a/packages/agent/src/harness/agent-harness.ts +++ b/packages/agent/src/harness/agent-harness.ts @@ -19,7 +19,7 @@ import type { ExecutionEnv, NavigateTreeResult, PromptTemplate, - SessionTree, + Session, Skill, SystemPromptInputs, } from "./types.js"; @@ -56,7 +56,7 @@ export class DefaultAgentHarness implements AgentHarness { readonly conversation: AgentHarnessConversationState; readonly operation: AgentHarnessOperationState; - private sessionTree: SessionTree; + private session: Session; private promptTemplates: PromptTemplate[]; private skills: Skill[]; private requestAuth?: AgentHarnessOptions["requestAuth"]; @@ -70,7 +70,7 @@ export class DefaultAgentHarness implements AgentHarness { constructor(options: AgentHarnessOptions) { this.agent = options.agent; this.env = options.env; - this.sessionTree = options.sessionTree; + this.session = options.session; this.promptTemplates = options.promptTemplates ?? []; this.skills = options.skills ?? []; this.requestAuth = options.requestAuth; @@ -78,7 +78,7 @@ export class DefaultAgentHarness implements AgentHarness { this.toolRegistry.set(tool.name, tool); } this.conversation = { - sessionTree: options.sessionTree, + session: options.session, model: options.initialModel ?? this.agent.state.model, thinkingLevel: options.initialThinkingLevel ?? this.agent.state.thinkingLevel, activeToolNames: options.initialActiveToolNames ?? this.agent.state.tools.map((tool) => tool.name), @@ -191,7 +191,7 @@ export class DefaultAgentHarness implements AgentHarness { } private async syncFromTree(): Promise { - const context = await this.sessionTree.buildContext(); + const context = await this.session.buildContext(); this.agent.state.messages = context.messages; if (context.model && this.conversation.model) { // leave active model untouched; harness-level model is source of truth @@ -201,7 +201,7 @@ export class DefaultAgentHarness implements AgentHarness { private async applyPendingMutations(): Promise { for (const message of this.operation.pendingMutations.appendMessages) { - await this.sessionTree.appendMessage(message); + await this.session.appendMessage(message); } this.operation.pendingMutations.appendMessages = []; @@ -210,7 +210,7 @@ export class DefaultAgentHarness implements AgentHarness { const previousModel = this.conversation.model; this.conversation.model = model; this.agent.state.model = model; - await this.sessionTree.appendModelChange(model.provider, model.id); + await this.session.appendModelChange(model.provider, model.id); await this.emitOwn({ type: "model_select", model, previousModel, source: "set" }); this.operation.pendingMutations.model = undefined; } @@ -220,7 +220,7 @@ export class DefaultAgentHarness implements AgentHarness { const previousLevel = this.conversation.thinkingLevel; this.conversation.thinkingLevel = level; this.agent.state.thinkingLevel = level; - await this.sessionTree.appendThinkingLevelChange(level); + await this.session.appendThinkingLevelChange(level); await this.emitOwn({ type: "thinking_level_select", level, previousLevel }); this.operation.pendingMutations.thinkingLevel = undefined; } @@ -258,7 +258,7 @@ export class DefaultAgentHarness implements AgentHarness { } } if (event.type === "message_end") { - await this.sessionTree.appendMessage(event.message); + await this.session.appendMessage(event.message); } if (event.type === "turn_end") { const hadPendingMutations = @@ -342,7 +342,7 @@ export class DefaultAgentHarness implements AgentHarness { async appendMessage(message: AgentMessage): Promise { if (this.operation.idle) { - await this.sessionTree.appendMessage(message); + await this.session.appendMessage(message); await this.syncFromTree(); } else { this.operation.pendingMutations.appendMessages.push(message); @@ -371,7 +371,7 @@ export class DefaultAgentHarness implements AgentHarness { if (!model) throw new Error("No model set for compaction"); const auth = await this.requestAuth?.(model); if (!auth) throw new Error("No auth available for compaction"); - const branchEntries = await this.sessionTree.getBranch(); + const branchEntries = await this.session.getBranch(); const preparation = prepareCompaction(branchEntries, DEFAULT_COMPACTION_SETTINGS); if (!preparation) throw new Error("Nothing to compact"); const hookResult = await this.emitHook("session_before_compact", { @@ -394,14 +394,14 @@ export class DefaultAgentHarness implements AgentHarness { undefined, this.conversation.thinkingLevel, )); - const entryId = await this.sessionTree.appendCompaction( + const entryId = await this.session.appendCompaction( result.summary, result.firstKeptEntryId, result.tokensBefore, result.details, provided !== undefined, ); - const entry = await this.sessionTree.getEntry(entryId); + const entry = await this.session.getEntry(entryId); await this.syncFromTree(); if (entry?.type === "compaction") { await this.emitOwn({ type: "session_compact", compactionEntry: entry, fromHook: provided !== undefined }); @@ -414,11 +414,11 @@ export class DefaultAgentHarness implements AgentHarness { options?: { summarize?: boolean; customInstructions?: string; replaceInstructions?: boolean; label?: string }, ): Promise { if (!this.operation.idle) throw new Error("navigateTree() requires idle harness"); - const oldLeafId = await this.sessionTree.getLeafId(); + const oldLeafId = await this.session.getLeafId(); if (oldLeafId === targetId) return { cancelled: false }; - const targetEntry = await this.sessionTree.getEntry(targetId); + const targetEntry = await this.session.getEntry(targetId); if (!targetEntry) throw new Error(`Entry ${targetId} not found`); - const { entries, commonAncestorId } = await collectEntriesForBranchSummary(this.sessionTree, oldLeafId, targetId); + const { entries, commonAncestorId } = await collectEntriesForBranchSummary(this.session, oldLeafId, targetId); const preparation = { targetId, oldLeafId, @@ -469,7 +469,7 @@ export class DefaultAgentHarness implements AgentHarness { typeof content === "string" ? content : content - .filter((c): c is { type: "text"; text: string } => c.type === "text") + .filter((c): c is { readonly type: "text"; readonly text: string } => c.type === "text") .map((c) => c.text) .join(""); } else if (targetEntry.type === "custom_message") { @@ -478,26 +478,29 @@ export class DefaultAgentHarness implements AgentHarness { typeof targetEntry.content === "string" ? targetEntry.content : targetEntry.content - .filter((c): c is { type: "text"; text: string } => c.type === "text") + .filter((c): c is { readonly type: "text"; readonly text: string } => c.type === "text") .map((c) => c.text) .join(""); } else { newLeafId = targetId; } - await this.sessionTree.moveTo(newLeafId); - if (summaryText) { - const summaryId = await this.sessionTree.appendBranchSummary( - newLeafId ?? "root", - summaryText, - summaryDetails, - hookResult?.summary !== undefined, - ); - summaryEntry = await this.sessionTree.getEntry(summaryId); + const summaryId = await this.session.moveTo( + newLeafId, + summaryText + ? { + summary: summaryText, + details: summaryDetails, + fromHook: hookResult?.summary !== undefined, + } + : undefined, + ); + if (summaryId) { + summaryEntry = await this.session.getEntry(summaryId); } await this.syncFromTree(); await this.emitOwn({ type: "session_tree", - newLeafId: await this.sessionTree.getLeafId(), + newLeafId: await this.session.getLeafId(), oldLeafId, summaryEntry, fromHook: hookResult?.summary !== undefined, @@ -510,7 +513,7 @@ export class DefaultAgentHarness implements AgentHarness { const previousModel = this.conversation.model; this.conversation.model = model; this.agent.state.model = model; - await this.sessionTree.appendModelChange(model.provider, model.id); + await this.session.appendModelChange(model.provider, model.id); await this.emitOwn({ type: "model_select", model, previousModel, source: "set" }); } else { this.operation.pendingMutations.model = model; @@ -522,7 +525,7 @@ export class DefaultAgentHarness implements AgentHarness { const previousLevel = this.conversation.thinkingLevel; this.conversation.thinkingLevel = level; this.agent.state.thinkingLevel = level; - await this.sessionTree.appendThinkingLevelChange(level); + await this.session.appendThinkingLevelChange(level); await this.emitOwn({ type: "thinking_level_select", level, previousLevel }); } else { this.operation.pendingMutations.thinkingLevel = level; diff --git a/packages/agent/src/harness/compaction/branch-summarization.ts b/packages/agent/src/harness/compaction/branch-summarization.ts index 9bef8799..b1ff6fdf 100644 --- a/packages/agent/src/harness/compaction/branch-summarization.ts +++ b/packages/agent/src/harness/compaction/branch-summarization.ts @@ -5,7 +5,7 @@ * a summary of the branch being left so context isn't lost. */ -import type { Model } from "@mariozechner/pi-ai"; +import type { ImageContent, Model, TextContent } from "@mariozechner/pi-ai"; import { completeSimple } from "@mariozechner/pi-ai"; import type { AgentMessage } from "../../types.js"; import { @@ -14,7 +14,7 @@ import { createCompactionSummaryMessage, createCustomMessage, } from "../messages.js"; -import type { SessionTree, SessionTreeEntry } from "../types.js"; +import type { Session, SessionTreeEntry } from "../types.js"; import { estimateTokens } from "./compaction.js"; import { computeFileLists, @@ -96,7 +96,7 @@ export interface GenerateBranchSummaryOptions { * @returns Entries to summarize and the common ancestor */ export async function collectEntriesForBranchSummary( - session: SessionTree, + session: Session, oldLeafId: string | null, targetId: string, ): Promise { @@ -125,7 +125,7 @@ export async function collectEntriesForBranchSummary( while (current && current !== commonAncestorId) { const entry = await session.getEntry(current); if (!entry) break; - entries.push(entry); + entries.push(entry as SessionTreeEntry); current = entry.parentId; } @@ -148,10 +148,16 @@ function getMessageFromEntry(entry: SessionTreeEntry): AgentMessage | undefined case "message": // Skip tool results - context is in assistant's tool call if (entry.message.role === "toolResult") return undefined; - return entry.message; + return entry.message as AgentMessage; case "custom_message": - return createCustomMessage(entry.customType, entry.content, entry.display, entry.details, entry.timestamp); + return createCustomMessage( + entry.customType, + entry.content as string | (TextContent | ImageContent)[], + entry.display, + entry.details, + entry.timestamp, + ); case "branch_summary": return createBranchSummaryMessage(entry.summary, entry.fromId, entry.timestamp); diff --git a/packages/agent/src/harness/compaction/compaction.ts b/packages/agent/src/harness/compaction/compaction.ts index 2c4f58a2..da28abc3 100644 --- a/packages/agent/src/harness/compaction/compaction.ts +++ b/packages/agent/src/harness/compaction/compaction.ts @@ -5,7 +5,7 @@ * and after compaction the session is reloaded. */ -import type { AssistantMessage, Model, Usage } from "@mariozechner/pi-ai"; +import type { AssistantMessage, ImageContent, Model, TextContent, Usage } from "@mariozechner/pi-ai"; import { completeSimple } from "@mariozechner/pi-ai"; import type { AgentMessage, ThinkingLevel } from "../../types.js"; import { @@ -79,10 +79,16 @@ function extractFileOperations( */ function getMessageFromEntry(entry: SessionTreeEntry): AgentMessage | undefined { if (entry.type === "message") { - return entry.message; + return entry.message as AgentMessage; } if (entry.type === "custom_message") { - return createCustomMessage(entry.customType, entry.content, entry.display, entry.details, entry.timestamp); + return createCustomMessage( + entry.customType, + entry.content as string | (TextContent | ImageContent)[], + entry.display, + entry.details, + entry.timestamp, + ); } if (entry.type === "branch_summary") { return createBranchSummaryMessage(entry.summary, entry.fromId, entry.timestamp); @@ -158,7 +164,7 @@ export function getLastAssistantUsage(entries: SessionTreeEntry[]): Usage | unde for (let i = entries.length - 1; i >= 0; i--) { const entry = entries[i]; if (entry.type === "message") { - const usage = getAssistantUsage(entry.message); + const usage = getAssistantUsage(entry.message as AgentMessage); if (usage) return usage; } } @@ -405,7 +411,7 @@ export function findCutPoint( if (entry.type !== "message") continue; // Estimate this message's size - const messageTokens = estimateTokens(entry.message); + const messageTokens = estimateTokens(entry.message as AgentMessage); accumulatedTokens += messageTokens; // Check if we've exceeded the budget diff --git a/packages/agent/src/harness/session/jsonl-session-repo.ts b/packages/agent/src/harness/session/jsonl-session-repo.ts new file mode 100644 index 00000000..39a31ce0 --- /dev/null +++ b/packages/agent/src/harness/session/jsonl-session-repo.ts @@ -0,0 +1,142 @@ +import { constants } from "node:fs"; +import { access, mkdir, readdir, rm } from "node:fs/promises"; +import { join, resolve } from "node:path"; +import type { + JsonlSessionCreateOptions, + JsonlSessionListQuery, + 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"; + +async function exists(path: string): Promise { + try { + await access(path, constants.F_OK); + return true; + } catch { + return false; + } +} + +function encodeCwd(cwd: string): string { + return `--${cwd.replace(/^[/\\]/, "").replace(/[/\\:]/g, "-")}--`; +} + +export class JsonlSessionRepo implements JsonlSessionRepoApi { + private sessionsRoot: string; + + constructor(options: { sessionsRoot: string }) { + this.sessionsRoot = resolve(options.sessionsRoot); + } + + private getSessionDir(cwd: string): string { + return join(this.sessionsRoot, encodeCwd(cwd)); + } + + private createSessionFilePath(cwd: string, sessionId: string, timestamp: string): string { + 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(); + const createdAt = createTimestamp(); + const filePath = this.createSessionFilePath(options.cwd, id, createdAt); + const storage = await JsonlSessionStorage.create(filePath, { + cwd: options.cwd, + sessionId: id, + parentSessionPath: options.parentSessionPath, + }); + return toSession(storage); + } + + async open(ref: JsonlSessionRef): Promise> { + const filePath = this.refPath(ref); + if (!(await exists(filePath))) { + throw new Error(`Session not found: ${filePath}`); + } + const storage = await JsonlSessionStorage.open(filePath); + return toSession(storage); + } + + async list(query: JsonlSessionListQuery = {}): Promise { + const dirs = query.cwd ? [this.getSessionDir(query.cwd)] : await this.listSessionDirs(); + const sessions: JsonlSessionMetadata[] = []; + for (const dir of dirs) { + if (!(await exists(dir))) continue; + const files = (await readdir(dir)).filter((file) => file.endsWith(".jsonl")).map((file) => join(dir, file)); + for (const filePath of files) { + try { + sessions.push(await loadJsonlSessionMetadata(filePath)); + } catch { + // Ignore invalid session files when listing a directory. + } + } + } + sessions.sort((a, b) => new Date(b.createdAt).getTime() - new Date(a.createdAt).getTime()); + 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 fork( + ref: JsonlSessionRef, + 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 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, + }); + for (const entry of forkedEntries) { + await storage.appendEntry(entry); + } + return toSession(storage); + } + + private async listSessionDirs(): Promise { + if (!(await exists(this.sessionsRoot))) return []; + const entries = await readdir(this.sessionsRoot, { withFileTypes: true }); + return entries.filter((entry) => entry.isDirectory()).map((entry) => join(this.sessionsRoot, entry.name)); + } +} diff --git a/packages/agent/src/harness/session/jsonl-session-storage.ts b/packages/agent/src/harness/session/jsonl-session-storage.ts index ffc7a437..0ec40ed1 100644 --- a/packages/agent/src/harness/session/jsonl-session-storage.ts +++ b/packages/agent/src/harness/session/jsonl-session-storage.ts @@ -2,7 +2,7 @@ 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 { JsonlSessionInfo, SessionTreeEntry, SessionTreeStorage } from "../types.js"; +import type { JsonlSessionMetadata, SessionStorage, SessionTreeEntry } from "../types.js"; interface SessionHeader { type: "session"; @@ -13,7 +13,25 @@ interface SessionHeader { parentSession?: string; } -function headerToSessionInfo(header: SessionHeader, path: string): JsonlSessionInfo { +function updateLabelCache(labelsById: Map, entry: SessionTreeEntry): void { + if (entry.type !== "label") return; + const label = entry.label?.trim(); + if (label) { + labelsById.set(entry.targetId, label); + } else { + labelsById.delete(entry.targetId); + } +} + +function buildLabelsById(entries: SessionTreeEntry[]): Map { + const labelsById = new Map(); + for (const entry of entries) { + updateLabelCache(labelsById, entry); + } + return labelsById; +} + +function headerToSessionMetadata(header: SessionHeader, path: string): JsonlSessionMetadata { return { id: header.id, createdAt: header.timestamp, @@ -23,7 +41,7 @@ function headerToSessionInfo(header: SessionHeader, path: string): JsonlSessionI }; } -export async function loadJsonlSessionInfo(filePath: string): Promise { +export async function loadJsonlSessionMetadata(filePath: string): Promise { const stream = createReadStream(filePath, { encoding: "utf8" }); const lines = createInterface({ input: stream, crlfDelay: Infinity }); try { @@ -31,7 +49,7 @@ export async function loadJsonlSessionInfo(filePath: string): Promise { +export class JsonlSessionStorage implements SessionStorage { private readonly filePath: string; - private readonly header: SessionHeader; - private readonly sessionInfo: JsonlSessionInfo; + private readonly metadata: JsonlSessionMetadata; private entries: SessionTreeEntry[]; private byId: Map; + private labelsById: Map; private currentLeafId: string | null; - private headerWritten: boolean; - private constructor( - filePath: string, - header: SessionHeader, - entries: SessionTreeEntry[], - leafId: string | null, - headerWritten: boolean, - ) { + private constructor(filePath: string, header: SessionHeader, entries: SessionTreeEntry[], leafId: string | null) { this.filePath = resolve(filePath); - this.header = header; - this.sessionInfo = headerToSessionInfo(header, this.filePath); + this.metadata = headerToSessionMetadata(header, this.filePath); this.entries = entries; this.byId = new Map(entries.map((entry) => [entry.id, entry])); + this.labelsById = buildLabelsById(entries); this.currentLeafId = leafId; - this.headerWritten = headerWritten; } - static async open(filePath: string): Promise { + 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); + return new JsonlSessionStorage(resolvedPath, loaded.header, loaded.entries, loaded.leafId); } static async create( @@ -113,7 +123,7 @@ export class JsonlSessionTreeStorage implements SessionTreeStorage { + ): Promise { const resolvedPath = resolve(filePath); const header: SessionHeader = { type: "session", @@ -123,11 +133,13 @@ export class JsonlSessionTreeStorage implements SessionTreeStorage { - return this.sessionInfo; + async getMetadata(): Promise { + return this.metadata; } async getLeafId(): Promise { @@ -142,14 +154,10 @@ export class JsonlSessionTreeStorage implements SessionTreeStorage { - 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); + updateLabelCache(this.labelsById, entry); this.currentLeafId = entry.id; } @@ -157,6 +165,10 @@ export class JsonlSessionTreeStorage implements SessionTreeStorage { + return this.labelsById.get(id); + } + async getPathToRoot(leafId: string | null): Promise { if (leafId === null) return []; const path: SessionTreeEntry[] = []; diff --git a/packages/agent/src/harness/session/memory-session-repo.ts b/packages/agent/src/harness/session/memory-session-repo.ts new file mode 100644 index 00000000..318f6ed2 --- /dev/null +++ b/packages/agent/src/harness/session/memory-session-repo.ts @@ -0,0 +1,55 @@ +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/memory-session-storage.ts b/packages/agent/src/harness/session/memory-session-storage.ts index 4485d87b..cdf3e791 100644 --- a/packages/agent/src/harness/session/memory-session-storage.ts +++ b/packages/agent/src/harness/session/memory-session-storage.ts @@ -1,24 +1,44 @@ import { v7 as uuidv7 } from "uuid"; -import type { SessionInfo, SessionTreeEntry, SessionTreeStorage } from "../types.js"; +import type { SessionMetadata, SessionStorage, SessionTreeEntry } from "../types.js"; -export class InMemorySessionTreeStorage implements SessionTreeStorage { - private readonly sessionInfo: SessionInfo; +function updateLabelCache(labelsById: Map, entry: SessionTreeEntry): void { + if (entry.type !== "label") return; + const label = entry.label?.trim(); + if (label) { + labelsById.set(entry.targetId, label); + } else { + labelsById.delete(entry.targetId); + } +} + +function buildLabelsById(entries: SessionTreeEntry[]): Map { + const labelsById = new Map(); + for (const entry of entries) { + updateLabelCache(labelsById, entry); + } + return labelsById; +} + +export class InMemorySessionStorage implements SessionStorage { + private readonly metadata: SessionMetadata; private entries: SessionTreeEntry[]; private byId: Map; + private labelsById: Map; private leafId: string | null; - constructor(options?: { entries?: SessionTreeEntry[]; leafId?: string | null; sessionInfo?: SessionInfo }) { + constructor(options?: { entries?: SessionTreeEntry[]; leafId?: string | null; metadata?: SessionMetadata }) { this.entries = options?.entries ? [...options.entries] : []; this.byId = new Map(this.entries.map((entry) => [entry.id, entry])); + this.labelsById = buildLabelsById(this.entries); this.leafId = options?.leafId ?? this.entries[this.entries.length - 1]?.id ?? null; 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() }; + this.metadata = options?.metadata ?? { id: uuidv7(), createdAt: new Date().toISOString() }; } - async getSessionInfo(): Promise { - return this.sessionInfo; + async getMetadata(): Promise { + return this.metadata; } async getLeafId(): Promise { @@ -35,6 +55,7 @@ export class InMemorySessionTreeStorage implements SessionTreeStorage { async appendEntry(entry: SessionTreeEntry): Promise { this.entries.push(entry); this.byId.set(entry.id, entry); + updateLabelCache(this.labelsById, entry); this.leafId = entry.id; } @@ -42,6 +63,10 @@ export class InMemorySessionTreeStorage implements SessionTreeStorage { return this.byId.get(id); } + async getLabel(id: string): Promise { + return this.labelsById.get(id); + } + async getPathToRoot(leafId: string | null): Promise { if (leafId === null) return []; const path: SessionTreeEntry[] = []; diff --git a/packages/agent/src/harness/session/session-repo.ts b/packages/agent/src/harness/session/session-repo.ts index bde9c291..3362d4a0 100644 --- a/packages/agent/src/harness/session/session-repo.ts +++ b/packages/agent/src/harness/session/session-repo.ts @@ -1,41 +1,25 @@ -import { existsSync, mkdirSync, readdirSync, rmSync } from "node:fs"; -import { join, resolve } from "node:path"; import { v7 as uuidv7 } from "uuid"; -import type { - JsonlSessionInfo, - JsonlSessionRepo, - Session, - SessionInfo, - SessionRepo, - SessionTreeEntry, - SessionTreeStorage, -} from "../types.js"; -import { JsonlSessionTreeStorage } from "./jsonl-session-storage.js"; -import { InMemorySessionTreeStorage } from "./memory-session-storage.js"; -import { DefaultSessionTree } from "./session-tree.js"; +import type { Session, SessionMetadata, SessionStorage, SessionTreeEntry } from "../types.js"; +import { DefaultSession } from "./session-tree.js"; -function createSessionId(): string { +export function createSessionId(): string { return uuidv7(); } -function createTimestamp(): string { +export function createTimestamp(): string { return new Date().toISOString(); } -function toSession( - storage: SessionTreeStorage, - tree: DefaultSessionTree, -): Session { - return { storage, tree }; +export function toSession(storage: SessionStorage): Session { + return new DefaultSession(storage); } -function getPathEntriesToFork( - entries: SessionTreeEntry[], +export async function getPathEntriesToFork( + storage: SessionStorage, entryId: string, position: "before" | "at", -): SessionTreeEntry[] { - const byId = new Map(entries.map((entry) => [entry.id, entry])); - const target = byId.get(entryId); +): Promise { + const target = await storage.getEntry(entryId); if (!target) { throw new Error(`Entry ${entryId} not found`); } @@ -48,172 +32,5 @@ function getPathEntriesToFork( } effectiveLeafId = target.parentId; } - if (effectiveLeafId === null) { - return []; - } - const path: SessionTreeEntry[] = []; - let current = byId.get(effectiveLeafId); - while (current) { - path.unshift(current); - current = current.parentId ? byId.get(current.parentId) : undefined; - } - return path; -} - -export class InMemorySessionRepo implements SessionRepo { - private sessions = new Map>(); - - async create(options?: { id?: string }): Promise> { - const info: SessionInfo = { - id: options?.id ?? createSessionId(), - createdAt: createTimestamp(), - }; - const storage = new InMemorySessionTreeStorage({ sessionInfo: info }); - const session = toSession(storage, new DefaultSessionTree(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 [...this.sessions.values()]; - } - - 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 entries = await source.tree.getEntries(); - const forkedEntries = getPathEntriesToFork(entries, options.entryId, options.position ?? "before"); - const info: SessionInfo = { - id: options.id ?? createSessionId(), - createdAt: createTimestamp(), - }; - const leafId = forkedEntries[forkedEntries.length - 1]?.id ?? null; - const storage = new InMemorySessionTreeStorage({ sessionInfo: info, entries: forkedEntries, leafId }); - const session = toSession(storage, new DefaultSessionTree(storage)); - this.sessions.set(info.id, session); - return session; - } -} - -export class JsonlSessionFileRepo implements JsonlSessionRepo { - private sessionDir: string; - private cwd: string; - - constructor(options: { sessionDir: string; cwd: string }) { - this.sessionDir = resolve(options.sessionDir); - this.cwd = options.cwd; - mkdirSync(this.sessionDir, { recursive: true }); - } - - private createSessionFilePath(sessionId: string, timestamp: string): string { - return join(this.sessionDir, `${timestamp.replace(/[:.]/g, "-")}_${sessionId}.jsonl`); - } - - async create(options?: { id?: string; parentSessionPath?: string }): Promise> { - const id = options?.id ?? createSessionId(); - const createdAt = createTimestamp(); - const filePath = this.createSessionFilePath(id, createdAt); - const storage = await JsonlSessionTreeStorage.create(filePath, { - cwd: this.cwd, - sessionId: id, - parentSessionPath: options?.parentSessionPath, - }); - return toSession(storage, new DefaultSessionTree(storage)); - } - - 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 = await JsonlSessionTreeStorage.open(filePath); - return toSession(storage, new DefaultSessionTree(storage)); - } - - 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> = []; - for (const filePath of files) { - 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>> { - 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 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 { - const filePath = ref.includes("/") || ref.endsWith(".jsonl") ? resolve(ref) : join(this.sessionDir, ref); - if (existsSync(filePath)) { - rmSync(filePath, { force: true }); - } - } - - async fork( - ref: string, - options: { entryId: string; position?: "before" | "at"; id?: string }, - ): 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 sourceInfo = await source.storage.getSessionInfo(); - const storage = await JsonlSessionTreeStorage.create(filePath, { - cwd: sourceInfo.cwd, - sessionId: id, - parentSessionPath: sourceInfo.path, - }); - for (const entry of forkedEntries) { - await storage.appendEntry(entry); - } - if (forkedEntries.length === 0) { - await storage.getSessionInfo(); - } - return toSession(storage, new DefaultSessionTree(storage)); - } + return storage.getPathToRoot(effectiveLeafId); } diff --git a/packages/agent/src/harness/session/session-tree.ts b/packages/agent/src/harness/session/session-tree.ts index 0d2eb146..3ed24975 100644 --- a/packages/agent/src/harness/session/session-tree.ts +++ b/packages/agent/src/harness/session/session-tree.ts @@ -10,12 +10,12 @@ import type { LabelEntry, MessageEntry, ModelChangeEntry, + Session, SessionContext, - SessionInfo, SessionInfoEntry, - SessionTree, + SessionMetadata, + SessionStorage, SessionTreeEntry, - SessionTreeStorage, ThinkingLevelChangeEntry, } from "../types.js"; @@ -27,12 +27,12 @@ function generateId(byId: { has(id: string): boolean }): string { return randomUUID(); } -export function buildSessionContext(entries: SessionTreeEntry[]): SessionContext { +export function buildSessionContext(pathEntries: SessionTreeEntry[]): SessionContext { let thinkingLevel = "off"; let model: { provider: string; modelId: string } | null = null; let compaction: CompactionEntry | null = null; - for (const entry of entries) { + for (const entry of pathEntries) { if (entry.type === "thinking_level_change") { thinkingLevel = entry.thinkingLevel; } else if (entry.type === "model_change") { @@ -47,10 +47,16 @@ export function buildSessionContext(entries: SessionTreeEntry[]): SessionContext const messages: AgentMessage[] = []; const appendMessage = (entry: SessionTreeEntry) => { if (entry.type === "message") { - messages.push(entry.message); + messages.push(entry.message as AgentMessage); } else if (entry.type === "custom_message") { messages.push( - createCustomMessage(entry.customType, entry.content, entry.display, entry.details, entry.timestamp), + createCustomMessage( + entry.customType, + entry.content as string | (TextContent | ImageContent)[], + entry.display, + entry.details, + entry.timestamp, + ), ); } else if (entry.type === "branch_summary" && entry.summary) { messages.push(createBranchSummaryMessage(entry.summary, entry.fromId, entry.timestamp)); @@ -59,18 +65,18 @@ export function buildSessionContext(entries: SessionTreeEntry[]): SessionContext if (compaction) { messages.push(createCompactionSummaryMessage(compaction.summary, compaction.tokensBefore, compaction.timestamp)); - const compactionIdx = entries.findIndex((e) => e.type === "compaction" && e.id === compaction.id); + const compactionIdx = pathEntries.findIndex((e) => e.type === "compaction" && e.id === compaction.id); let foundFirstKept = false; for (let i = 0; i < compactionIdx; i++) { - const entry = entries[i]!; + const entry = pathEntries[i]!; if (entry.id === compaction.firstKeptEntryId) foundFirstKept = true; if (foundFirstKept) appendMessage(entry); } - for (let i = compactionIdx + 1; i < entries.length; i++) { - appendMessage(entries[i]!); + for (let i = compactionIdx + 1; i < pathEntries.length; i++) { + appendMessage(pathEntries[i]!); } } else { - for (const entry of entries) { + for (const entry of pathEntries) { appendMessage(entry); } } @@ -78,13 +84,21 @@ export function buildSessionContext(entries: SessionTreeEntry[]): SessionContext return { messages, thinkingLevel, model }; } -export class DefaultSessionTree implements SessionTree { - private storage: SessionTreeStorage; +export class DefaultSession implements Session { + private storage: SessionStorage; - constructor(storage: SessionTreeStorage) { + constructor(storage: SessionStorage) { this.storage = storage; } + getMetadata(): Promise { + return this.storage.getMetadata(); + } + + getStorage(): SessionStorage { + return this.storage; + } + getLeafId(): Promise { return this.storage.getLeafId(); } @@ -106,19 +120,8 @@ export class DefaultSessionTree impleme return buildSessionContext(await this.getBranch()); } - getSessionInfo(): Promise { - return this.storage.getSessionInfo(); - } - - async getLabel(id: string): Promise { - const entries = await this.storage.getEntries(); - for (let i = entries.length - 1; i >= 0; i--) { - const entry = entries[i]!; - if (entry.type === "label" && entry.targetId === id) { - return entry.label?.trim() || undefined; - } - } - return undefined; + getLabel(id: string): Promise { + return this.storage.getLabel(id); } async getSessionName(): Promise { @@ -193,24 +196,6 @@ export class DefaultSessionTree impleme } satisfies CompactionEntry); } - async appendBranchSummary( - fromId: string, - summary: string, - details?: T, - fromHook?: boolean, - ): Promise { - return this.appendTypedEntry({ - type: "branch_summary", - id: await this.makeEntryId(), - parentId: await this.storage.getLeafId(), - timestamp: new Date().toISOString(), - fromId, - summary, - details, - fromHook, - } satisfies BranchSummaryEntry); - } - async appendCustomEntry(customType: string, data?: unknown): Promise { return this.appendTypedEntry({ type: "custom", @@ -240,7 +225,10 @@ export class DefaultSessionTree impleme } satisfies CustomMessageEntry); } - async appendLabelChange(targetId: string, label: string | undefined): Promise { + async appendLabel(targetId: string, label: string | undefined): Promise { + if (!(await this.storage.getEntry(targetId))) { + throw new Error(`Entry ${targetId} not found`); + } return this.appendTypedEntry({ type: "label", id: await this.makeEntryId(), @@ -251,7 +239,7 @@ export class DefaultSessionTree impleme } satisfies LabelEntry); } - async appendSessionInfo(name: string): Promise { + async appendSessionName(name: string): Promise { return this.appendTypedEntry({ type: "session_info", id: await this.makeEntryId(), @@ -261,10 +249,24 @@ export class DefaultSessionTree impleme } satisfies SessionInfoEntry); } - async moveTo(entryId: string | null): Promise { + async moveTo( + entryId: string | null, + summary?: { summary: string; details?: unknown; fromHook?: boolean }, + ): Promise { if (entryId !== null && !(await this.storage.getEntry(entryId))) { throw new Error(`Entry ${entryId} not found`); } await this.storage.setLeafId(entryId); + if (!summary) return undefined; + return this.appendTypedEntry({ + type: "branch_summary", + id: await this.makeEntryId(), + parentId: entryId, + timestamp: new Date().toISOString(), + fromId: entryId ?? "root", + summary: summary.summary, + details: summary.details, + fromHook: summary.fromHook, + } satisfies BranchSummaryEntry); } } diff --git a/packages/agent/src/harness/types.ts b/packages/agent/src/harness/types.ts index 15f23853..7975690b 100644 --- a/packages/agent/src/harness/types.ts +++ b/packages/agent/src/harness/types.ts @@ -136,7 +136,7 @@ export interface LabelEntry extends SessionTreeEntryBase { } export interface SessionInfoEntry extends SessionTreeEntryBase { - type: "session_info"; + type: "session_info"; // legacy name, kept for backwards compatibility name?: string; } @@ -157,28 +157,32 @@ export interface SessionContext { model: { provider: string; modelId: string } | null; } -export interface SessionInfo { +export interface SessionMetadata { id: string; createdAt: string; } -export interface JsonlSessionInfo extends SessionInfo { +export interface JsonlSessionMetadata extends SessionMetadata { cwd: string; path: string; parentSessionPath?: string; } -export interface SessionTreeStorage { - getSessionInfo(): Promise; +export interface SessionStorage { + getMetadata(): Promise; getLeafId(): Promise; setLeafId(leafId: string | null): Promise; appendEntry(entry: SessionTreeEntry): Promise; getEntry(id: string): Promise; + getLabel(id: string): Promise; getPathToRoot(leafId: string | null): Promise; getEntries(): Promise; } -export interface SessionTree { +export interface Session { + getMetadata(): Promise; + getStorage(): SessionStorage; + getLeafId(): Promise; getEntry(id: string): Promise; getEntries(): Promise; @@ -197,7 +201,6 @@ export interface SessionTree { details?: T, fromHook?: boolean, ): Promise; - appendBranchSummary(fromId: string, summary: string, details?: T, fromHook?: boolean): Promise; appendCustomEntry(customType: string, data?: unknown): Promise; appendCustomMessageEntry( customType: string, @@ -205,15 +208,13 @@ export interface SessionTree { display: boolean, details?: T, ): Promise; - appendLabelChange(targetId: string, label: string | undefined): Promise; - appendSessionInfo(name: string): Promise; + appendLabel(targetId: string, label: string | undefined): Promise; + appendSessionName(name: string): Promise; - moveTo(entryId: string | null): Promise; -} - -export interface Session { - storage: SessionTreeStorage; - tree: SessionTree; + moveTo( + entryId: string | null, + summary?: { summary: string; details?: unknown; fromHook?: boolean }, + ): Promise; } export interface SessionCreateOptions { @@ -227,25 +228,38 @@ export interface SessionForkOptions { } export interface SessionRepo< - TRef = string, - TInfo extends SessionInfo = SessionInfo, + TMetadata extends SessionMetadata = SessionMetadata, TCreateOptions extends SessionCreateOptions = SessionCreateOptions, + TRef = string, + TListQuery = void, > { - create(options?: TCreateOptions): Promise>; - open(ref: TRef): Promise>; - list(): Promise>>; + create(options: TCreateOptions): Promise>; + open(ref: TRef): Promise>; + list(query?: TListQuery): Promise; delete(ref: TRef): Promise; - fork(ref: TRef, options: SessionForkOptions): Promise>; + fork(ref: TRef, options: SessionForkOptions & TCreateOptions): Promise>; } export interface JsonlSessionCreateOptions extends SessionCreateOptions { + cwd: string; parentSessionPath?: string; } -export interface JsonlSessionRepo - extends SessionRepo { - listByCwd(cwd: string): Promise>>; - getMostRecentByCwd(cwd: string): Promise | undefined>; +export type JsonlSessionRef = { path: string } | JsonlSessionMetadata; + +export interface JsonlSessionListQuery { + cwd?: string; +} + +export interface JsonlSessionResolveOptions { + cwd?: string; + searchAll?: boolean; +} + +export interface JsonlSessionRepoApi + extends SessionRepo { + resolve(ref: string, options?: JsonlSessionResolveOptions): Promise; + getMostRecent(query?: JsonlSessionListQuery): Promise; } export interface AgentHarnessPendingMutations { @@ -257,7 +271,7 @@ export interface AgentHarnessPendingMutations { } export interface AgentHarnessConversationState { - sessionTree: SessionTree; + session: Session; model: Model | undefined; thinkingLevel: ThinkingLevel; activeToolNames: string[]; @@ -284,8 +298,8 @@ export interface SavePointSnapshot { export interface AgentHarnessContext { env: ExecutionEnv; - conversation: Readonly; - operation: Readonly; + conversation: AgentHarnessConversationState; + operation: AgentHarnessOperationState; abortSignal?: AbortSignal; } @@ -549,7 +563,7 @@ export interface BranchSummaryResult { export interface AgentHarnessOptions { agent: Agent; env: ExecutionEnv; - sessionTree: SessionTree; + session: Session; promptTemplates?: PromptTemplate[]; skills?: Skill[]; requestAuth?: (model: Model) => Promise<{ apiKey: string; headers?: Record } | undefined>; @@ -562,8 +576,8 @@ export interface AgentHarnessOptions { export interface AgentHarness { readonly agent: Agent; readonly env: ExecutionEnv; - readonly conversation: Readonly; - readonly operation: Readonly; + readonly conversation: AgentHarnessConversationState; + readonly operation: AgentHarnessOperationState; prompt(text: string, options?: AgentHarnessPromptOptions): Promise; skill(name: string, args?: string): Promise; diff --git a/packages/agent/src/index.ts b/packages/agent/src/index.ts index 78222783..43557ef0 100644 --- a/packages/agent/src/index.ts +++ b/packages/agent/src/index.ts @@ -25,6 +25,8 @@ 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"; // Harness diff --git a/packages/agent/test/harness/session-tree.test.ts b/packages/agent/test/harness/session-tree.test.ts index ffcb01d4..17d578fb 100644 --- a/packages/agent/test/harness/session-tree.test.ts +++ b/packages/agent/test/harness/session-tree.test.ts @@ -3,10 +3,12 @@ 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, 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 { MessageEntry, SessionInfo, SessionTreeStorage } from "../../src/harness/types.js"; +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"; function createUserMessage(text: string): AgentMessage { return { @@ -56,12 +58,12 @@ afterEach(() => { async function runSessionTreeSuite( name: string, - createStorage: () => SessionTreeStorage | Promise, + createStorage: () => SessionStorage | Promise, inspect?: () => void, ) { describe(name, () => { it("appends messages and builds context in order", async () => { - const tree = new DefaultSessionTree(await createStorage()); + const tree = new DefaultSession(await createStorage()); await tree.appendMessage(createUserMessage("one")); await tree.appendMessage(createAssistantMessage("two")); const context = await tree.buildContext(); @@ -69,7 +71,7 @@ async function runSessionTreeSuite( }); it("tracks model and thinking level changes", async () => { - const tree = new DefaultSessionTree(await createStorage()); + const tree = new DefaultSession(await createStorage()); await tree.appendMessage(createUserMessage("one")); await tree.appendModelChange("openai", "gpt-4.1"); await tree.appendThinkingLevelChange("high"); @@ -79,7 +81,7 @@ async function runSessionTreeSuite( }); it("supports branching by moving the leaf and appending a new branch", async () => { - const tree = new DefaultSessionTree(await createStorage()); + 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")); @@ -93,7 +95,7 @@ async function runSessionTreeSuite( }); it("supports moving the leaf to root", async () => { - const tree = new DefaultSessionTree(await createStorage()); + const tree = new DefaultSession(await createStorage()); await tree.appendMessage(createUserMessage("one")); await tree.moveTo(null); expect(await tree.getLeafId()).toBeNull(); @@ -101,7 +103,7 @@ async function runSessionTreeSuite( }); it("reconstructs compaction summaries in context", async () => { - const tree = new DefaultSessionTree(await createStorage()); + const tree = new DefaultSession(await createStorage()); await tree.appendMessage(createUserMessage("one")); await tree.appendMessage(createAssistantMessage("two")); const user2 = await tree.appendMessage(createUserMessage("three")); @@ -113,16 +115,19 @@ async function runSessionTreeSuite( expect(context.messages).toHaveLength(4); }); - it("supports branch summary entries in context", async () => { - const tree = new DefaultSessionTree(await createStorage()); + it("supports moving with branch summary entries in context", async () => { + const tree = new DefaultSession(await createStorage()); const user1 = await tree.appendMessage(createUserMessage("one")); - await tree.appendBranchSummary(user1, "summary text"); + const summaryId = await tree.moveTo(user1, { summary: "summary text" }); + expect(summaryId).toBeTruthy(); + const summaryEntry = await tree.getEntry(summaryId!); + expect(summaryEntry).toMatchObject({ type: "branch_summary", parentId: user1, fromId: user1 }); const context = await tree.buildContext(); expect(context.messages[1]?.role).toBe("branchSummary"); }); it("supports custom message entries in context", async () => { - const tree = new DefaultSessionTree(await createStorage()); + const tree = new DefaultSession(await createStorage()); await tree.appendMessage(createUserMessage("one")); await tree.appendCustomMessageEntry("custom", "hello", true, { ok: true }); const context = await tree.buildContext(); @@ -130,10 +135,10 @@ async function runSessionTreeSuite( }); it("supports labels and session info entries without affecting context", async () => { - const tree = new DefaultSessionTree(await createStorage()); + const tree = new DefaultSession(await createStorage()); const user1 = await tree.appendMessage(createUserMessage("one")); - await tree.appendLabelChange(user1, "checkpoint"); - await tree.appendSessionInfo("name"); + await tree.appendLabel(user1, "checkpoint"); + await tree.appendSessionName("name"); const entries = await tree.getEntries(); expect(entries.some((entry) => entry.type === "label")).toBe(true); expect(entries.some((entry) => entry.type === "session_info")).toBe(true); @@ -142,16 +147,21 @@ async function runSessionTreeSuite( expect((await tree.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"); + }); + it("persists leaf changes and appended entries via storage", async () => { const storage = await createStorage(); - const tree = new DefaultSessionTree(storage); + const tree = new DefaultSession(storage); const user1 = await tree.appendMessage(createUserMessage("one")); await tree.appendMessage(createAssistantMessage("two")); - await tree.appendLabelChange(user1, "checkpoint"); - await tree.appendSessionInfo("name"); + await tree.appendLabel(user1, "checkpoint"); + await tree.appendSessionName("name"); await tree.moveTo(user1); await tree.appendMessage(createAssistantMessage("branched")); - const tree2 = new DefaultSessionTree(storage); + const tree2 = new DefaultSession(storage); const context = await tree2.buildContext(); expect(context.messages.map((message) => message.role)).toEqual(["user", "assistant"]); expect(await tree2.getLabel(user1)).toBe("checkpoint"); @@ -161,11 +171,11 @@ async function runSessionTreeSuite( }); } -describe("InMemorySessionTreeStorage", () => { +describe("InMemorySessionStorage", () => { 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); + 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 () => { @@ -177,7 +187,7 @@ describe("InMemorySessionTreeStorage", () => { message: createUserMessage("one"), }; const initialEntries = [entry]; - const storage = new InMemorySessionTreeStorage({ entries: initialEntries }); + 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"); @@ -186,9 +196,39 @@ describe("InMemorySessionTreeStorage", () => { }); it("rejects invalid leaf ids", async () => { - const storage = new InMemorySessionTreeStorage(); + const storage = new InMemorySessionStorage(); await expect(storage.setLeafId("missing")).rejects.toThrow("Entry missing not found"); - expect(() => new InMemorySessionTreeStorage({ leafId: "missing" })).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 () => { @@ -205,26 +245,70 @@ describe("InMemorySessionTreeStorage", () => { parentId: "root", message: createAssistantMessage("child"), }; - const storage = new InMemorySessionTreeStorage({ entries: [root, 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("SessionTree with in-memory storage", () => new InMemorySessionTreeStorage()); +runSessionTreeSuite("Session with in-memory storage", () => new InMemorySessionStorage()); -describe("JsonlSessionTreeStorage", () => { +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(JsonlSessionTreeStorage.open(filePath)).rejects.toMatchObject({ code: "ENOENT" }); + await expect(JsonlSessionStorage.open(filePath)).rejects.toMatchObject({ code: "ENOENT" }); }); - it("writes the header before the first appended entry", async () => { + it("writes the header on create", 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); + 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({ @@ -245,7 +329,7 @@ describe("JsonlSessionTreeStorage", () => { 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"); + await expect(JsonlSessionStorage.open(filePath)).rejects.toThrow("first line is not a valid session header"); }); it("ignores malformed entry lines", async () => { @@ -266,7 +350,7 @@ describe("JsonlSessionTreeStorage", () => { message: createUserMessage("one"), }; writeFileSync(filePath, `${JSON.stringify(header)}\nnot json\n${JSON.stringify(entry)}\n`); - const storage = await JsonlSessionTreeStorage.open(filePath); + const storage = await JsonlSessionStorage.open(filePath); expect((await storage.getEntries()).map((loadedEntry) => loadedEntry.id)).toEqual(["entry-1"]); expect(await storage.getLeafId()).toBe("entry-1"); }); @@ -274,19 +358,19 @@ describe("JsonlSessionTreeStorage", () => { 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, { + const storage = await JsonlSessionStorage.create(filePath, { cwd: dir, sessionId: "session-1", parentSessionPath: "/tmp/parent.jsonl", }); - const info = await storage.getSessionInfo(); + const info = await storage.getMetadata(); expect(info).toMatchObject({ id: "session-1", cwd: dir, path: filePath, parentSessionPath: "/tmp/parent.jsonl", }); - expect(existsSync(filePath)).toBe(false); + expect(existsSync(filePath)).toBe(true); await storage.appendEntry({ type: "message", id: "user-1", @@ -294,13 +378,13 @@ describe("JsonlSessionTreeStorage", () => { timestamp: "2026-01-01T00:00:00.000Z", message: createUserMessage("one"), }); - expect(await loadJsonlSessionInfo(filePath)).toEqual(info); + 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 JsonlSessionTreeStorage.create(filePath, { cwd: dir, sessionId: "session-1" }); + const storage = await JsonlSessionStorage.create(filePath, { cwd: dir, sessionId: "session-1" }); const root: MessageEntry = { type: "message", id: "root", @@ -316,12 +400,46 @@ describe("JsonlSessionTreeStorage", () => { }; await storage.appendEntry(root); await storage.appendEntry(child); - const loaded = await JsonlSessionTreeStorage.open(filePath); + 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"); @@ -334,7 +452,7 @@ describe("JsonlSessionTreeStorage", () => { }; const malformedSecondLine = "{".repeat(10000); writeFileSync(filePath, `${JSON.stringify(header)}\n${malformedSecondLine}\n`); - expect(await loadJsonlSessionInfo(filePath)).toEqual({ + expect(await loadJsonlSessionMetadata(filePath)).toEqual({ id: "session-1", createdAt: "2026-01-01T00:00:00.000Z", cwd: dir, @@ -345,10 +463,10 @@ describe("JsonlSessionTreeStorage", () => { }); runSessionTreeSuite( - "SessionTree with JSONL storage", + "Session with JSONL storage", async () => { const dir = createTempDir(); - return await JsonlSessionTreeStorage.create(join(dir, "session.jsonl"), { cwd: dir, sessionId: "session-1" }); + return await JsonlSessionStorage.create(join(dir, "session.jsonl"), { cwd: dir, sessionId: "session-1" }); }, () => { const dir = tempDirs[tempDirs.length - 1]!;