diff --git a/packages/agent-core/src/harness/session/jsonl-storage.ts b/packages/agent-core/src/harness/session/jsonl-storage.ts index 805a84c83a10..0cadc35a53ee 100644 --- a/packages/agent-core/src/harness/session/jsonl-storage.ts +++ b/packages/agent-core/src/harness/session/jsonl-storage.ts @@ -1,13 +1,7 @@ -import type { - FileSystem, - JsonlSessionMetadata, - LeafEntry, - SessionStorage, - SessionTreeEntry, -} from "../types.js"; +import type { FileSystem, JsonlSessionMetadata, SessionTreeEntry } from "../types.js"; import { SessionError, toError } from "../types.js"; import { getFileSystemResultOrThrow } from "./repo-utils.js"; -import { uuidv7 } from "./uuid.js"; +import { BaseSessionStorage, leafIdAfterEntry } from "./storage-base.js"; type JsonlSessionStorageFileSystem = Pick< FileSystem, @@ -23,36 +17,6 @@ interface SessionHeader { parentSession?: string; } -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 generateEntryId(byId: { has(id: string): boolean }): string { - for (let i = 0; i < 100; i++) { - const id = uuidv7().slice(0, 8); - if (!byId.has(id)) { - return id; - } - } - return uuidv7(); -} - function isRecord(value: unknown): value is Record { return typeof value === "object" && value !== null; } @@ -144,10 +108,6 @@ function parseEntryLine(line: string, filePath: string, lineNumber: number): Ses return parsed as unknown as SessionTreeEntry; } -function leafIdAfterEntry(entry: SessionTreeEntry): string | null { - return entry.type === "leaf" ? entry.targetId : entry.id; -} - function headerToSessionMetadata(header: SessionHeader, path: string): JsonlSessionMetadata { return { id: header.id, @@ -201,14 +161,9 @@ async function loadJsonlStorage( return { header, entries, leafId }; } -export class JsonlSessionStorage implements SessionStorage { +export class JsonlSessionStorage extends BaseSessionStorage { private readonly fs: JsonlSessionStorageFileSystem; private readonly filePath: string; - private readonly metadata: JsonlSessionMetadata; - private entries: SessionTreeEntry[]; - private byId: Map; - private labelsById: Map; - private currentLeafId: string | null; private constructor( fs: JsonlSessionStorageFileSystem, @@ -217,13 +172,9 @@ export class JsonlSessionStorage implements SessionStorage entries: SessionTreeEntry[], leafId: string | null, ) { + super(headerToSessionMetadata(header, filePath), entries, leafId); this.fs = fs; this.filePath = 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; } static async open( @@ -258,92 +209,20 @@ export class JsonlSessionStorage implements SessionStorage return new JsonlSessionStorage(fs, filePath, header, [], null); } - async getMetadata(): Promise { - return this.metadata; - } - - async getLeafId(): Promise { - if (this.currentLeafId !== null && !this.byId.has(this.currentLeafId)) { - throw new SessionError("invalid_session", `Entry ${this.currentLeafId} not found`); - } - return this.currentLeafId; - } - - async setLeafId(leafId: string | null): Promise { - if (leafId !== null && !this.byId.has(leafId)) { - throw new SessionError("not_found", `Entry ${leafId} not found`); - } - const entry: LeafEntry = { - type: "leaf", - id: generateEntryId(this.byId), - parentId: this.currentLeafId, - timestamp: new Date().toISOString(), - targetId: leafId, - }; + override async setLeafId(leafId: string | null): Promise { + const entry = this.createLeafEntry(leafId); getFileSystemResultOrThrow( await this.fs.appendFile(this.filePath, `${JSON.stringify(entry)}\n`), `Failed to append session leaf ${entry.id}`, ); - this.entries.push(entry); - this.byId.set(entry.id, entry); - this.currentLeafId = leafId; + this.recordEntry(entry); } - async createEntryId(): Promise { - return generateEntryId(this.byId); - } - - async appendEntry(entry: SessionTreeEntry): Promise { + override async appendEntry(entry: SessionTreeEntry): Promise { getFileSystemResultOrThrow( await this.fs.appendFile(this.filePath, `${JSON.stringify(entry)}\n`), `Failed to append session entry ${entry.id}`, ); - this.entries.push(entry); - this.byId.set(entry.id, entry); - updateLabelCache(this.labelsById, entry); - this.currentLeafId = leafIdAfterEntry(entry); - } - - async getEntry(id: string): Promise { - 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); - } - - async getPathToRoot(leafId: string | null): Promise { - if (leafId === null) { - return []; - } - const path: SessionTreeEntry[] = []; - let current = this.byId.get(leafId); - if (!current) { - throw new SessionError("not_found", `Entry ${leafId} not found`); - } - while (current) { - path.unshift(current); - if (!current.parentId) { - break; - } - const parent = this.byId.get(current.parentId); - if (!parent) { - throw new SessionError("invalid_session", `Entry ${current.parentId} not found`); - } - current = parent; - } - return path; - } - - async getEntries(): Promise { - return [...this.entries]; + this.recordEntry(entry); } } diff --git a/packages/agent-core/src/harness/session/memory-storage.test.ts b/packages/agent-core/src/harness/session/memory-storage.test.ts new file mode 100644 index 000000000000..764956ec1c36 --- /dev/null +++ b/packages/agent-core/src/harness/session/memory-storage.test.ts @@ -0,0 +1,62 @@ +import { describe, expect, it } from "vitest"; +import type { SessionTreeEntry } from "../types.js"; +import { InMemorySessionStorage } from "./memory-storage.js"; + +const rootEntry: SessionTreeEntry = { + type: "custom", + id: "root", + parentId: null, + timestamp: "2026-01-01T00:00:00.000Z", + customType: "root", +}; + +const childEntry: SessionTreeEntry = { + type: "custom", + id: "child", + parentId: "root", + timestamp: "2026-01-01T00:00:01.000Z", + customType: "child", +}; + +describe("InMemorySessionStorage", () => { + it("uses shared entry indexes for labels, leaves, and paths", async () => { + const storage = new InMemorySessionStorage({ + entries: [ + rootEntry, + childEntry, + { + type: "label", + id: "label-1", + parentId: "child", + timestamp: "2026-01-01T00:00:02.000Z", + targetId: "child", + label: " latest ", + }, + ], + }); + + expect(await storage.getLeafId()).toBe("label-1"); + expect(await storage.getLabel("child")).toBe("latest"); + expect((await storage.getPathToRoot("child")).map((entry) => entry.id)).toEqual([ + "root", + "child", + ]); + }); + + it("records explicit leaf updates through the shared storage path", async () => { + const storage = new InMemorySessionStorage({ + entries: [rootEntry, childEntry], + }); + + await storage.setLeafId("root"); + + const entries = await storage.getEntries(); + const leaf = entries.at(-1); + expect(await storage.getLeafId()).toBe("root"); + expect(leaf).toMatchObject({ + type: "leaf", + parentId: "child", + targetId: "root", + }); + }); +}); diff --git a/packages/agent-core/src/harness/session/memory-storage.ts b/packages/agent-core/src/harness/session/memory-storage.ts index ab197c0a350c..d385e26ea5d9 100644 --- a/packages/agent-core/src/harness/session/memory-storage.ts +++ b/packages/agent-core/src/harness/session/memory-storage.ts @@ -1,148 +1,22 @@ -import { - type LeafEntry, - SessionError, - type SessionMetadata, - type SessionStorage, - type SessionTreeEntry, -} from "../types.js"; +import { type SessionMetadata, type SessionTreeEntry } from "../types.js"; +import { BaseSessionStorage } from "./storage-base.js"; import { uuidv7 } from "./uuid.js"; -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 generateEntryId(byId: { has(id: string): boolean }): string { - for (let i = 0; i < 100; i++) { - const id = uuidv7().slice(0, 8); - if (!byId.has(id)) { - return id; - } - } - return uuidv7(); -} - -function leafIdAfterEntry(entry: SessionTreeEntry): string | null { - return entry.type === "leaf" ? entry.targetId : entry.id; -} - export class InMemorySessionStorage< TMetadata extends SessionMetadata = SessionMetadata, -> implements SessionStorage { - private readonly metadata: TMetadata; - private entries: SessionTreeEntry[]; - private byId: Map; - private labelsById: Map; - private leafId: string | null; - +> extends BaseSessionStorage { constructor(options?: { entries?: SessionTreeEntry[]; metadata?: TMetadata }) { - this.entries = options?.entries ? [...options.entries] : []; - this.byId = new Map(this.entries.map((entry) => [entry.id, entry])); - this.labelsById = buildLabelsById(this.entries); - this.leafId = null; - for (const entry of this.entries) { - this.leafId = leafIdAfterEntry(entry); - } - if (this.leafId !== null && !this.byId.has(this.leafId)) { - throw new SessionError("invalid_session", `Entry ${this.leafId} not found`); - } - this.metadata = - options?.metadata ?? ({ id: uuidv7(), createdAt: new Date().toISOString() } as TMetadata); - } - - async getMetadata(): Promise { - return this.metadata; - } - - async getLeafId(): Promise { - if (this.leafId !== null && !this.byId.has(this.leafId)) { - throw new SessionError("invalid_session", `Entry ${this.leafId} not found`); - } - return this.leafId; - } - - async setLeafId(leafId: string | null): Promise { - if (leafId !== null && !this.byId.has(leafId)) { - throw new SessionError("not_found", `Entry ${leafId} not found`); - } - const entry: LeafEntry = { - type: "leaf", - id: generateEntryId(this.byId), - parentId: this.leafId, - timestamp: new Date().toISOString(), - targetId: leafId, - }; - this.entries.push(entry); - this.byId.set(entry.id, entry); - 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); - updateLabelCache(this.labelsById, entry); - this.leafId = leafIdAfterEntry(entry); - } - - async getEntry(id: string): Promise { - return this.byId.get(id); - } - - async findEntries( - type: TType, - ): Promise>> { - return this.entries.filter( - (entry): entry is Extract => entry.type === type, + super( + options?.metadata ?? ({ id: uuidv7(), createdAt: new Date().toISOString() } as TMetadata), + options?.entries ? [...options.entries] : [], ); } - async getLabel(id: string): Promise { - return this.labelsById.get(id); + override async setLeafId(leafId: string | null): Promise { + this.recordEntry(this.createLeafEntry(leafId)); } - async getPathToRoot(leafId: string | null): Promise { - if (leafId === null) { - return []; - } - const path: SessionTreeEntry[] = []; - let current = this.byId.get(leafId); - if (!current) { - throw new SessionError("not_found", `Entry ${leafId} not found`); - } - while (current) { - path.unshift(current); - if (!current.parentId) { - break; - } - const parent = this.byId.get(current.parentId); - if (!parent) { - throw new SessionError("invalid_session", `Entry ${current.parentId} not found`); - } - current = parent; - } - return path; - } - - async getEntries(): Promise { - return [...this.entries]; + override async appendEntry(entry: SessionTreeEntry): Promise { + this.recordEntry(entry); } } diff --git a/packages/agent-core/src/harness/session/storage-base.ts b/packages/agent-core/src/harness/session/storage-base.ts new file mode 100644 index 000000000000..e8a311c81e0b --- /dev/null +++ b/packages/agent-core/src/harness/session/storage-base.ts @@ -0,0 +1,156 @@ +import { + type LeafEntry, + SessionError, + type SessionMetadata, + type SessionStorage, + type SessionTreeEntry, +} from "../types.js"; +import { uuidv7 } from "./uuid.js"; + +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 generateEntryId(byId: { has(id: string): boolean }): string { + for (let i = 0; i < 100; i++) { + const id = uuidv7().slice(0, 8); + if (!byId.has(id)) { + return id; + } + } + return uuidv7(); +} + +export function leafIdAfterEntry(entry: SessionTreeEntry): string | null { + return entry.type === "leaf" ? entry.targetId : entry.id; +} + +function resolveLeafId(entries: readonly SessionTreeEntry[]): string | null { + let leafId: string | null = null; + for (const entry of entries) { + leafId = leafIdAfterEntry(entry); + } + return leafId; +} + +export abstract class BaseSessionStorage< + TMetadata extends SessionMetadata = SessionMetadata, +> implements SessionStorage { + private readonly metadata: TMetadata; + private readonly entries: SessionTreeEntry[]; + private readonly byId: Map; + private readonly labelsById: Map; + private leafId: string | null; + + protected constructor( + metadata: TMetadata, + entries: SessionTreeEntry[], + leafId: string | null = resolveLeafId(entries), + ) { + this.metadata = metadata; + this.entries = entries; + this.byId = new Map(entries.map((entry) => [entry.id, entry])); + this.labelsById = buildLabelsById(entries); + this.leafId = leafId; + if (this.leafId !== null && !this.byId.has(this.leafId)) { + throw new SessionError("invalid_session", `Entry ${this.leafId} not found`); + } + } + + async getMetadata(): Promise { + return this.metadata; + } + + async getLeafId(): Promise { + if (this.leafId !== null && !this.byId.has(this.leafId)) { + throw new SessionError("invalid_session", `Entry ${this.leafId} not found`); + } + return this.leafId; + } + + protected createLeafEntry(leafId: string | null): LeafEntry { + if (leafId !== null && !this.byId.has(leafId)) { + throw new SessionError("not_found", `Entry ${leafId} not found`); + } + return { + type: "leaf", + id: generateEntryId(this.byId), + parentId: this.leafId, + timestamp: new Date().toISOString(), + targetId: leafId, + }; + } + + async createEntryId(): Promise { + return generateEntryId(this.byId); + } + + protected recordEntry(entry: SessionTreeEntry): void { + this.entries.push(entry); + this.byId.set(entry.id, entry); + updateLabelCache(this.labelsById, entry); + this.leafId = leafIdAfterEntry(entry); + } + + async getEntry(id: string): Promise { + 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); + } + + async getPathToRoot(leafId: string | null): Promise { + if (leafId === null) { + return []; + } + const path: SessionTreeEntry[] = []; + let current = this.byId.get(leafId); + if (!current) { + throw new SessionError("not_found", `Entry ${leafId} not found`); + } + while (current) { + path.unshift(current); + if (!current.parentId) { + break; + } + const parent = this.byId.get(current.parentId); + if (!parent) { + throw new SessionError("invalid_session", `Entry ${current.parentId} not found`); + } + current = parent; + } + return path; + } + + async getEntries(): Promise { + return [...this.entries]; + } + + abstract setLeafId(leafId: string | null): Promise; + abstract appendEntry(entry: SessionTreeEntry): Promise; +}