fix(agent): synchronize array-backed entry indexes eagerly

packages/agent/src/harness/session/array-session-reader.ts

@@ -1,107 +0,0 @@
-import type { SessionEntryCursorOptions, SessionMetadata, SessionReader, SessionTreeEntry } from "../types.ts";
-import { SessionError } from "../types.ts";
-
-export function createArraySessionReader<TMetadata extends SessionMetadata>(
-	metadata: TMetadata,
-	readCurrentEntries: () => readonly SessionTreeEntry[],
-): SessionReader<TMetadata> {
-	let indexedCount = 0;
-	let indexedEntries: readonly SessionTreeEntry[] | undefined;
-	let leafId: string | null = null;
-	const byId = new Map<string, SessionTreeEntry>();
-
-	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<string>();
-			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();
-		},
-	};
-}

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<string, Set<string>>();
-	private readonly entriesByPath = new Map<string, SessionTreeEntry[]>();
+	private readonly entryCollectionsByPath = new Map<string, SessionEntryCollection>();
 	private readonly operationKeysByPath = new Map<string, string>();
 	private readonly operations: KeyedOperationQueue<string>;
 	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<JsonlSessionMetadata> {
-		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<string> {
 		this.sessionsRoot ??= getFileSystemResultOrThrow(
 			await this.fs.absolutePath(this.sessionsRootInput),

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<SessionMetadata, InMemorySessionCreateOptions, void> {
@@ -24,7 +24,7 @@ class InMemorySessionStore implements SessionStore<SessionMetadata, InMemorySess
 		return this.operations.enqueue(id, () => {
 			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<SessionMetadata, InMemorySess
 		this.assertOpen();
 		return this.operations.enqueue(metadata.id, () => {
 			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<SessionMetadata, InMemorySess
 		const sourceEntries = this.operations.enqueue(source.id, () => {
 			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<SessionMetadata, InMemorySess
 	}
 
 	private reader(state: InMemorySessionState): SessionReader<SessionMetadata> {
-		const reader = createArraySessionReader(state.metadata, () => state.entries);
+		const reader = createSessionEntryCollectionReader(state.metadata, state.entries);
 		return {
 			metadata: reader.metadata,
 			readHead: () => {

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<string, SessionTreeEntry>();
+	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<string, SessionTreeEntry>();
+		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<string>();
+		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<TMetadata extends SessionMetadata>(
+	metadata: TMetadata,
+	entries: SessionEntryCollection,
+): SessionReader<TMetadata> {
+	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);
+		},
+	};
+}

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");