refactor: dedupe session storage indexing

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<string, string>, 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<string, string> {
-  const labelsById = new Map<string, string>();
-  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<string, unknown> {
   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<JsonlSessionMetadata> {
+export class JsonlSessionStorage extends BaseSessionStorage<JsonlSessionMetadata> {
   private readonly fs: JsonlSessionStorageFileSystem;
   private readonly filePath: string;
-  private readonly metadata: JsonlSessionMetadata;
-  private entries: SessionTreeEntry[];
-  private byId: Map<string, SessionTreeEntry>;
-  private labelsById: Map<string, string>;
-  private currentLeafId: string | null;
 
   private constructor(
     fs: JsonlSessionStorageFileSystem,
@@ -217,13 +172,9 @@ export class JsonlSessionStorage implements SessionStorage<JsonlSessionMetadata>
     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<JsonlSessionMetadata>
     return new JsonlSessionStorage(fs, filePath, header, [], null);
   }
 
-  async getMetadata(): Promise<JsonlSessionMetadata> {
-    return this.metadata;
-  }
-
-  async getLeafId(): Promise<string | null> {
-    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<void> {
-    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<void> {
+    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<string> {
-    return generateEntryId(this.byId);
-  }
-
-  async appendEntry(entry: SessionTreeEntry): Promise<void> {
+  override async appendEntry(entry: SessionTreeEntry): Promise<void> {
     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<SessionTreeEntry | undefined> {
-    return this.byId.get(id);
-  }
-
-  async findEntries<TType extends SessionTreeEntry["type"]>(
-    type: TType,
-  ): Promise<Array<Extract<SessionTreeEntry, { type: TType }>>> {
-    return this.entries.filter(
-      (entry): entry is Extract<SessionTreeEntry, { type: TType }> => entry.type === type,
-    );
-  }
-
-  async getLabel(id: string): Promise<string | undefined> {
-    return this.labelsById.get(id);
-  }
-
-  async getPathToRoot(leafId: string | null): Promise<SessionTreeEntry[]> {
-    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<SessionTreeEntry[]> {
-    return [...this.entries];
+    this.recordEntry(entry);
   }
 }

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",
+    });
+  });
+});

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<string, string>, 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<string, string> {
-  const labelsById = new Map<string, string>();
-  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<TMetadata> {
-  private readonly metadata: TMetadata;
-  private entries: SessionTreeEntry[];
-  private byId: Map<string, SessionTreeEntry>;
-  private labelsById: Map<string, string>;
-  private leafId: string | null;
-
+> extends BaseSessionStorage<TMetadata> {
   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<TMetadata> {
-    return this.metadata;
-  }
-
-  async getLeafId(): Promise<string | null> {
-    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<void> {
-    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<string> {
-    return generateEntryId(this.byId);
-  }
-
-  async appendEntry(entry: SessionTreeEntry): Promise<void> {
-    this.entries.push(entry);
-    this.byId.set(entry.id, entry);
-    updateLabelCache(this.labelsById, entry);
-    this.leafId = leafIdAfterEntry(entry);
-  }
-
-  async getEntry(id: string): Promise<SessionTreeEntry | undefined> {
-    return this.byId.get(id);
-  }
-
-  async findEntries<TType extends SessionTreeEntry["type"]>(
-    type: TType,
-  ): Promise<Array<Extract<SessionTreeEntry, { type: TType }>>> {
-    return this.entries.filter(
-      (entry): entry is Extract<SessionTreeEntry, { type: TType }> => entry.type === type,
+    super(
+      options?.metadata ?? ({ id: uuidv7(), createdAt: new Date().toISOString() } as TMetadata),
+      options?.entries ? [...options.entries] : [],
     );
   }
 
-  async getLabel(id: string): Promise<string | undefined> {
-    return this.labelsById.get(id);
+  override async setLeafId(leafId: string | null): Promise<void> {
+    this.recordEntry(this.createLeafEntry(leafId));
   }
 
-  async getPathToRoot(leafId: string | null): Promise<SessionTreeEntry[]> {
-    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<SessionTreeEntry[]> {
-    return [...this.entries];
+  override async appendEntry(entry: SessionTreeEntry): Promise<void> {
+    this.recordEntry(entry);
   }
 }

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<string, string>, 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<string, string> {
+  const labelsById = new Map<string, string>();
+  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<TMetadata> {
+  private readonly metadata: TMetadata;
+  private readonly entries: SessionTreeEntry[];
+  private readonly byId: Map<string, SessionTreeEntry>;
+  private readonly labelsById: Map<string, string>;
+  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<TMetadata> {
+    return this.metadata;
+  }
+
+  async getLeafId(): Promise<string | null> {
+    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<string> {
+    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<SessionTreeEntry | undefined> {
+    return this.byId.get(id);
+  }
+
+  async findEntries<TType extends SessionTreeEntry["type"]>(
+    type: TType,
+  ): Promise<Array<Extract<SessionTreeEntry, { type: TType }>>> {
+    return this.entries.filter(
+      (entry): entry is Extract<SessionTreeEntry, { type: TType }> => entry.type === type,
+    );
+  }
+
+  async getLabel(id: string): Promise<string | undefined> {
+    return this.labelsById.get(id);
+  }
+
+  async getPathToRoot(leafId: string | null): Promise<SessionTreeEntry[]> {
+    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<SessionTreeEntry[]> {
+    return [...this.entries];
+  }
+
+  abstract setLeafId(leafId: string | null): Promise<void>;
+  abstract appendEntry(entry: SessionTreeEntry): Promise<void>;
+}