diff --git a/packages/agent/src/harness/session/array-session-reader.ts b/packages/agent/src/harness/session/array-session-reader.ts deleted file mode 100644 index 696a7b14f..000000000 --- a/packages/agent/src/harness/session/array-session-reader.ts +++ /dev/null @@ -1,107 +0,0 @@ -import type { SessionEntryCursorOptions, SessionMetadata, SessionReader, SessionTreeEntry } from "../types.ts"; -import { SessionError } from "../types.ts"; - -export function createArraySessionReader( - metadata: TMetadata, - readCurrentEntries: () => readonly SessionTreeEntry[], -): SessionReader { - let indexedCount = 0; - let indexedEntries: readonly SessionTreeEntry[] | undefined; - let leafId: string | null = null; - const byId = new Map(); - - const updateIndex = (): readonly SessionTreeEntry[] => { - const entries = readCurrentEntries(); - if (entries !== indexedEntries || entries.length < indexedCount) { - indexedCount = 0; - leafId = null; - byId.clear(); - indexedEntries = entries; - } - for (; indexedCount < entries.length; indexedCount++) { - const entry = entries[indexedCount]!; - byId.set(entry.id, entry); - leafId = entry.type === "leaf" ? entry.targetId : entry.id; - } - return entries; - }; - - return { - metadata, - async readHead() { - updateIndex(); - if (leafId !== null && !byId.has(leafId)) { - throw new SessionError("invalid_session", `Entry ${leafId} not found`); - } - return { leafId }; - }, - async readEntry(id) { - updateIndex(); - return byId.get(id); - }, - async readEntries(options?: SessionEntryCursorOptions) { - const entries = updateIndex(); - const start = options?.afterEntrySeq ?? 0; - const end = options?.limit === undefined ? undefined : start + options.limit; - return entries.slice(start, end); - }, - async findEntriesOnBranch(query) { - updateIndex(); - if (query.limit !== undefined && (!Number.isInteger(query.limit) || query.limit <= 0)) { - throw new RangeError("Session branch query limit must be a positive integer"); - } - if (query.start === null) return []; - const pathFromStart: SessionTreeEntry[] = []; - const visited = new Set(); - let current = byId.get(query.start); - if (!current) throw new SessionError("not_found", `Entry ${query.start} not found`); - while (current) { - if (visited.has(current.id)) { - throw new SessionError("invalid_session", `Session branch contains a cycle at ${current.id}`); - } - visited.add(current.id); - pathFromStart.push(current); - if (query.order !== "oldestFirst" && (current.id === query.stopAtId || current.type === query.stopAtType)) { - break; - } - if (!current.parentId) break; - const parent = byId.get(current.parentId); - if (!parent) throw new SessionError("invalid_session", `Entry ${current.parentId} not found`); - current = parent; - } - const traversal = query.order === "oldestFirst" ? pathFromStart.reverse() : pathFromStart; - const stopIndex = - query.order === "oldestFirst" - ? traversal.findIndex((entry) => entry.id === query.stopAtId || entry.type === query.stopAtType) - : -1; - const bounded = stopIndex === -1 ? traversal : traversal.slice(0, stopIndex + 1); - const entries = bounded.filter( - (entry) => - (query.type === undefined || entry.type === query.type) && - (query.customType === undefined || (entry.type === "custom" && entry.customType === query.customType)), - ); - return query.limit === undefined ? entries : entries.slice(0, query.limit); - }, - async readPathToRootOrCompaction(requestedLeafId) { - updateIndex(); - if (requestedLeafId === null) return []; - const path: SessionTreeEntry[] = []; - let stopAtEntryId: string | null = null; - let current = byId.get(requestedLeafId); - if (!current) throw new SessionError("not_found", `Entry ${requestedLeafId} not found`); - while (current) { - path.push(current); - if (stopAtEntryId !== null && current.id === stopAtEntryId) break; - if (current.type === "compaction") { - if (current.retainedTail) break; - stopAtEntryId = current.firstKeptEntryId ?? null; - } - if (!current.parentId) break; - const parent = byId.get(current.parentId); - if (!parent) throw new SessionError("invalid_session", `Entry ${current.parentId} not found`); - current = parent; - } - return path.reverse(); - }, - }; -} diff --git a/packages/agent/src/harness/session/jsonl-store.ts b/packages/agent/src/harness/session/jsonl-store.ts index 455548cdc..664d21be7 100644 --- a/packages/agent/src/harness/session/jsonl-store.ts +++ b/packages/agent/src/harness/session/jsonl-store.ts @@ -9,10 +9,10 @@ import type { SessionTreeEntry, } from "../types.ts"; import { SessionError, toError } from "../types.ts"; -import { createArraySessionReader } from "./array-session-reader.ts"; import { readSessionEntriesForFork } from "./fork.ts"; import { KeyedOperationQueue } from "./keyed-operation-queue.ts"; import { createSessionId, createTimestamp, getFileSystemResultOrThrow } from "./repository.ts"; +import { createSessionEntryCollectionReader, SessionEntryCollection } from "./session-entry-collection.ts"; export interface JsonlSessionStoreOptions { fs: JsonlSessionStoreFileSystem; @@ -117,7 +117,13 @@ function parseEntry(line: string, path: string, lineNumber: number): SessionTree } if (typeof value !== "object" || value === null) throw invalidEntry(path, lineNumber, "is not a valid session entry"); - const entry = value as { type?: unknown; id?: unknown; parentId?: unknown; timestamp?: unknown; targetId?: unknown }; + const entry = value as { + type?: unknown; + id?: unknown; + parentId?: unknown; + timestamp?: unknown; + targetId?: unknown; + }; if (typeof entry.type !== "string") throw invalidEntry(path, lineNumber, "is missing entry type"); if (typeof entry.id !== "string" || !entry.id) throw invalidEntry(path, lineNumber, "is missing entry id"); if (entry.parentId !== null && typeof entry.parentId !== "string") @@ -199,8 +205,7 @@ class JsonlSessionStore private readonly fs: JsonlSessionStoreFileSystem; private readonly sessionsRootInput: string; private sessionsRoot: string | undefined; - private readonly entryIdsByPath = new Map>(); - private readonly entriesByPath = new Map(); + private readonly entryCollectionsByPath = new Map(); private readonly operationKeysByPath = new Map(); private readonly operations: KeyedOperationQueue; private disposed = false; @@ -234,8 +239,9 @@ class JsonlSessionStore throw new SessionError("not_found", `Session not found: ${metadata.path}`); } const document = await loadJsonlSession(this.fs, metadata.path); - this.entriesByPath.set(metadata.path, document.entries); - this.entryIdsByPath.set(metadata.path, new Set(document.entries.map((entry) => entry.id))); + const entries = this.entryCollectionsByPath.get(metadata.path); + if (entries) entries.replace(document.entries); + else this.entryCollectionsByPath.set(metadata.path, new SessionEntryCollection(document.entries)); return this.reader(document.metadata); } @@ -267,19 +273,17 @@ class JsonlSessionStore ) { throw new SessionError("not_found", `Session not found: ${metadata.path}`); } - let entryIds = this.entryIdsByPath.get(metadata.path); - if (!entryIds) { + let entries = this.entryCollectionsByPath.get(metadata.path); + if (!entries) { await this.loadDocument(metadata); - entryIds = this.entryIdsByPath.get(metadata.path)!; + entries = this.entryCollectionsByPath.get(metadata.path)!; } - if (entryIds.has(entry.id)) throw new SessionError("invalid_entry", `Entry ${entry.id} already exists`); + if (entries.has(entry.id)) throw new SessionError("invalid_entry", `Entry ${entry.id} already exists`); getFileSystemResultOrThrow( await this.fs.appendFile(metadata.path, `${JSON.stringify(entry)}\n`), `Failed to append session entry ${entry.id}`, ); - entryIds.add(entry.id); - this.entryIdsByPath.set(metadata.path, entryIds); - this.entriesByPath.get(metadata.path)!.push(entry); + entries.append(entry); }); } @@ -290,8 +294,7 @@ class JsonlSessionStore await this.fs.remove(metadata.path, { force: true }), `Failed to delete session ${metadata.path}`, ); - this.entryIdsByPath.delete(metadata.path); - this.entriesByPath.delete(metadata.path); + this.entryCollectionsByPath.delete(metadata.path); this.operationKeysByPath.delete(metadata.path); }); } @@ -308,10 +311,11 @@ class JsonlSessionStore throw new SessionError("not_found", `Session not found: ${source.path}`); } const document = await loadJsonlSession(this.fs, source.path); - this.entriesByPath.set(source.path, document.entries); - this.entryIdsByPath.set(source.path, new Set(document.entries.map((entry) => entry.id))); + const entries = this.entryCollectionsByPath.get(source.path); + if (entries) entries.replace(document.entries); + else this.entryCollectionsByPath.set(source.path, new SessionEntryCollection(document.entries)); return readSessionEntriesForFork( - createArraySessionReader(document.metadata, () => document.entries), + createSessionEntryCollectionReader(document.metadata, this.entryCollectionsByPath.get(source.path)!), selection, ); }); @@ -372,45 +376,50 @@ class JsonlSessionStore }; const content = [JSON.stringify(header), ...entries.map((entry) => JSON.stringify(entry)), ""].join("\n"); getFileSystemResultOrThrow(await this.fs.writeFile(path, content), `Failed to create session ${path}`); - const storedEntries = [...entries]; - this.entriesByPath.set(path, storedEntries); - this.entryIdsByPath.set(path, new Set(storedEntries.map((entry) => entry.id))); + this.entryCollectionsByPath.set(path, new SessionEntryCollection(entries)); this.operationKeysByPath.set(path, descriptor.operationKey); return this.reader(metadataFromHeader(header, path)); } private reader(metadata: JsonlSessionMetadata): SessionReader { - const reader = createArraySessionReader(metadata, () => { - const entries = this.entriesByPath.get(metadata.path); - if (!entries) throw new SessionError("not_found", `Session not found: ${metadata.path}`); - return entries; - }); const operationKey = this.operationKey(metadata); return { - metadata: reader.metadata, + metadata, readHead: () => { this.assertOpen(); - return this.operations.enqueue(operationKey, () => reader.readHead()); + return this.operations.enqueue(operationKey, () => this.entryCollection(metadata.path).readHead()); }, readEntry: (id) => { this.assertOpen(); - return this.operations.enqueue(operationKey, () => reader.readEntry(id)); + return this.operations.enqueue(operationKey, () => this.entryCollection(metadata.path).readEntry(id)); }, readEntries: (options) => { this.assertOpen(); - return this.operations.enqueue(operationKey, () => reader.readEntries(options)); + return this.operations.enqueue(operationKey, () => + this.entryCollection(metadata.path).readEntries(options), + ); }, findEntriesOnBranch: (query) => { this.assertOpen(); - return this.operations.enqueue(operationKey, () => reader.findEntriesOnBranch(query)); + return this.operations.enqueue(operationKey, () => + this.entryCollection(metadata.path).findEntriesOnBranch(query), + ); }, readPathToRootOrCompaction: (leafId) => { this.assertOpen(); - return this.operations.enqueue(operationKey, () => reader.readPathToRootOrCompaction(leafId)); + return this.operations.enqueue(operationKey, () => + this.entryCollection(metadata.path).readPathToRootOrCompaction(leafId), + ); }, }; } + private entryCollection(path: string): SessionEntryCollection { + const entries = this.entryCollectionsByPath.get(path); + if (!entries) throw new SessionError("not_found", `Session not found: ${path}`); + return entries; + } + private async getSessionsRoot(): Promise { this.sessionsRoot ??= getFileSystemResultOrThrow( await this.fs.absolutePath(this.sessionsRootInput), diff --git a/packages/agent/src/harness/session/memory-store.ts b/packages/agent/src/harness/session/memory-store.ts index f69a9bb6f..1d44b1830 100644 --- a/packages/agent/src/harness/session/memory-store.ts +++ b/packages/agent/src/harness/session/memory-store.ts @@ -1,15 +1,15 @@ import type { SessionForkSelection, SessionMetadata, SessionReader, SessionStore, SessionTreeEntry } from "../types.ts"; import { SessionError } from "../types.ts"; -import { createArraySessionReader } from "./array-session-reader.ts"; import { readSessionEntriesForFork } from "./fork.ts"; import { KeyedOperationQueue } from "./keyed-operation-queue.ts"; import { createSessionId, createTimestamp } from "./repository.ts"; +import { createSessionEntryCollectionReader, SessionEntryCollection } from "./session-entry-collection.ts"; export type InMemorySessionCreateOptions = { id?: string }; interface InMemorySessionState { metadata: SessionMetadata; - entries: SessionTreeEntry[]; + entries: SessionEntryCollection; } class InMemorySessionStore implements SessionStore { @@ -24,7 +24,7 @@ class InMemorySessionStore implements SessionStore { const state: InMemorySessionState = { metadata: { id, createdAt: createTimestamp() }, - entries: [], + entries: new SessionEntryCollection(), }; this.sessions.set(state.metadata.id, state); return this.reader(state); @@ -45,10 +45,7 @@ class InMemorySessionStore implements SessionStore { const state = this.getState(metadata); - if (state.entries.some((existing) => existing.id === entry.id)) { - throw new SessionError("invalid_entry", `Entry ${entry.id} already exists`); - } - state.entries.push(entry); + state.entries.append(entry); }); } @@ -69,14 +66,14 @@ class InMemorySessionStore implements SessionStore { const sourceState = this.getState(source); return readSessionEntriesForFork( - createArraySessionReader(sourceState.metadata, () => sourceState.entries), + createSessionEntryCollectionReader(sourceState.metadata, sourceState.entries), selection, ); }); return this.operations.enqueue(id, async () => { const state: InMemorySessionState = { metadata: { id, createdAt: createTimestamp() }, - entries: [...(await sourceEntries)], + entries: new SessionEntryCollection(await sourceEntries), }; this.sessions.set(state.metadata.id, state); return this.reader(state); @@ -102,7 +99,7 @@ class InMemorySessionStore implements SessionStore { - const reader = createArraySessionReader(state.metadata, () => state.entries); + const reader = createSessionEntryCollectionReader(state.metadata, state.entries); return { metadata: reader.metadata, readHead: () => { diff --git a/packages/agent/src/harness/session/session-entry-collection.ts b/packages/agent/src/harness/session/session-entry-collection.ts new file mode 100644 index 000000000..47157f299 --- /dev/null +++ b/packages/agent/src/harness/session/session-entry-collection.ts @@ -0,0 +1,148 @@ +import type { + SessionBranchQuery, + SessionEntryCursorOptions, + SessionHead, + SessionMetadata, + SessionReader, + SessionTreeEntry, +} from "../types.ts"; +import { SessionError } from "../types.ts"; + +/** Ordered entries and derived indexes for an array-backed session store. */ +export class SessionEntryCollection { + private entries: SessionTreeEntry[] = []; + private byId = new Map(); + private leafId: string | null = null; + + constructor(entries: readonly SessionTreeEntry[] = []) { + this.replace(entries); + } + + has(id: string): boolean { + return this.byId.has(id); + } + + append(entry: SessionTreeEntry): void { + if (this.byId.has(entry.id)) { + throw new SessionError("invalid_entry", `Entry ${entry.id} already exists`); + } + this.entries.push(entry); + this.byId.set(entry.id, entry); + this.leafId = entry.type === "leaf" ? entry.targetId : entry.id; + } + + replace(entries: readonly SessionTreeEntry[]): void { + const nextEntries = [...entries]; + const nextById = new Map(); + let nextLeafId: string | null = null; + for (const entry of nextEntries) { + if (nextById.has(entry.id)) { + throw new SessionError("invalid_entry", `Entry ${entry.id} already exists`); + } + nextById.set(entry.id, entry); + nextLeafId = entry.type === "leaf" ? entry.targetId : entry.id; + } + this.entries = nextEntries; + this.byId = nextById; + this.leafId = nextLeafId; + } + + readHead(): SessionHead { + if (this.leafId !== null && !this.byId.has(this.leafId)) { + throw new SessionError("invalid_session", `Entry ${this.leafId} not found`); + } + return { leafId: this.leafId }; + } + + readEntry(id: string): SessionTreeEntry | undefined { + return this.byId.get(id); + } + + readEntries(options?: SessionEntryCursorOptions): readonly SessionTreeEntry[] { + const start = options?.afterEntrySeq ?? 0; + const end = options?.limit === undefined ? undefined : start + options.limit; + return this.entries.slice(start, end); + } + + findEntriesOnBranch(query: SessionBranchQuery & { start: string | null }): readonly SessionTreeEntry[] { + if (query.limit !== undefined && (!Number.isInteger(query.limit) || query.limit <= 0)) { + throw new RangeError("Session branch query limit must be a positive integer"); + } + if (query.start === null) return []; + const pathFromStart: SessionTreeEntry[] = []; + const visited = new Set(); + let current = this.byId.get(query.start); + if (!current) throw new SessionError("not_found", `Entry ${query.start} not found`); + while (current) { + if (visited.has(current.id)) { + throw new SessionError("invalid_session", `Session branch contains a cycle at ${current.id}`); + } + visited.add(current.id); + pathFromStart.push(current); + if (query.order !== "oldestFirst" && (current.id === query.stopAtId || current.type === query.stopAtType)) { + break; + } + 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; + } + const traversal = query.order === "oldestFirst" ? pathFromStart.reverse() : pathFromStart; + const stopIndex = + query.order === "oldestFirst" + ? traversal.findIndex((entry) => entry.id === query.stopAtId || entry.type === query.stopAtType) + : -1; + const bounded = stopIndex === -1 ? traversal : traversal.slice(0, stopIndex + 1); + const entries = bounded.filter( + (entry) => + (query.type === undefined || entry.type === query.type) && + (query.customType === undefined || (entry.type === "custom" && entry.customType === query.customType)), + ); + return query.limit === undefined ? entries : entries.slice(0, query.limit); + } + + readPathToRootOrCompaction(requestedLeafId: string | null): readonly SessionTreeEntry[] { + if (requestedLeafId === null) return []; + const path: SessionTreeEntry[] = []; + let stopAtEntryId: string | null = null; + let current = this.byId.get(requestedLeafId); + if (!current) throw new SessionError("not_found", `Entry ${requestedLeafId} not found`); + while (current) { + path.push(current); + if (stopAtEntryId !== null && current.id === stopAtEntryId) break; + if (current.type === "compaction") { + if (current.retainedTail) break; + stopAtEntryId = current.firstKeptEntryId ?? null; + } + 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.reverse(); + } +} + +export function createSessionEntryCollectionReader( + metadata: TMetadata, + entries: SessionEntryCollection, +): SessionReader { + return { + metadata, + async readHead() { + return entries.readHead(); + }, + async readEntry(id) { + return entries.readEntry(id); + }, + async readEntries(options) { + return entries.readEntries(options); + }, + async findEntriesOnBranch(query) { + return entries.findEntriesOnBranch(query); + }, + async readPathToRootOrCompaction(leafId) { + return entries.readPathToRootOrCompaction(leafId); + }, + }; +} diff --git a/packages/agent/test/harness/branch-query.test.ts b/packages/agent/test/harness/branch-query.test.ts index 8810a27fa..7483382f2 100644 --- a/packages/agent/test/harness/branch-query.test.ts +++ b/packages/agent/test/harness/branch-query.test.ts @@ -101,7 +101,7 @@ describe("bounded session branch queries", () => { ).toEqual(expected.fullPath); }); - it("rejects corrupt parent chains in array-backed readers", async () => { + it("rejects corrupt parent chains in array-backed stores", async () => { const store = createInMemorySessionStore(); ownedStores.push(store); const reader = await store.create({ id: "corrupt-memory" }); @@ -156,6 +156,20 @@ describe("bounded session branch queries", () => { ).toEqual(expected.fullPath); }); + it("keeps existing JSONL readers synchronized when a session is reloaded", async () => { + const root = createTempDir(); + const store = createJsonlSessionStore({ fs: new NodeExecutionEnv({ cwd: root }), sessionsRoot: root }); + ownedStores.push(store); + const repo = createSessionRepository({ store }); + const original = await repo.create({ id: "jsonl-reload", cwd: root }); + await original.appendMessage(createUserMessage("root")); + + const reopened = await repo.open(await original.getMetadata()); + const appendedId = await reopened.appendMessage(createAssistantMessage("after reload")); + + expect(await original.getEntry(appendedId)).toMatchObject({ id: appendedId }); + }); + it("does not decode SQLite branch entries outside query bounds", async () => { const root = createTempDir(); const databasePath = join(root, "sessions.sqlite");