pi-sessions 0.3.2 → 0.4.0

package/extensions/session-search/hooks.ts

@@ -1,4 +1,4 @@
-import { existsSync } from "node:fs";
+import { existsSync, type Stats, statSync } from "node:fs";
 import type { ImageContent, TextContent } from "@earendil-works/pi-ai";
 import {
   isToolCallEventType,
@@ -6,25 +6,34 @@ import {
   type ToolResultEvent,
 } from "@earendil-works/pi-coding-agent";
 import {
+  clearSessionChunksBySourceKind,
   clearSessionIndexedData,
-  getIndexStatus,
+  getMetadata,
   getSessionById,
+  getSessionRowByPath,
   INDEX_SCHEMA_VERSION,
   insertSessionFileTouch,
   insertTextChunk,
   openIndexDatabase,
-  rebuildSessionLineageRelations,
+  refreshSessionLineageRelationsFor,
+  type SessionIndexDatabase,
   type SessionLineageRow,
   type SessionOrigin,
+  type SessionRow,
   setMetadata,
   upsertSession,
 } from "../shared/session-index/index.js";
-import { type ExtractedSessionRecord, extractSessionRecord } from "./extract.js";
+import {
+  createSessionNameChunk,
+  type ExtractedSessionRecord,
+  type ExtractedSessionTail,
+  extractSessionRecord,
+  extractSessionTail,
+  type SessionFileTouch,
+} from "./extract.js";
+import { deriveSessionRepoRoots } from "./normalize.js";
 
 const TOOL_RESULT_TEXT_LIMIT = 500;
-const HOOK_INDEX_WRITE_TIMEOUT_MS = 500;
-const HOOK_INDEX_WRITE_MAX_ATTEMPTS = 4;
-const HOOK_INDEX_WRITE_RETRY_DELAYS_MS = [150, 400, 1000];
 
 type TrackedToolName = "read" | "edit" | "write";
 
@@ -99,24 +108,15 @@ export function createSessionHookController(options: { indexPath: string }): Ses
       return syncAttachedSession(indexPath, state, "session_start");
     },
     async handleSessionSwitch(previousSessionFile, sessionFile, cwd, sessionOrigin) {
-      const previousSynced = await syncSessionFile(
-        indexPath,
-        previousSessionFile,
-        "session_switch",
-      );
+      const previousSynced = syncSessionFile(indexPath, previousSessionFile, "session_switch");
       attachSession(state, sessionFile, cwd);
-      const currentSynced = await syncAttachedSession(
-        indexPath,
-        state,
-        "session_switch",
-        sessionOrigin,
-      );
+      const currentSynced = syncAttachedSession(indexPath, state, "session_switch", sessionOrigin);
       return previousSynced || currentSynced;
     },
     async handleSessionFork(previousSessionFile, sessionFile, cwd) {
-      const previousSynced = await syncSessionFile(indexPath, previousSessionFile, "session_fork");
+      const previousSynced = syncSessionFile(indexPath, previousSessionFile, "session_fork");
       attachSession(state, sessionFile, cwd);
-      const currentSynced = await syncAttachedSession(indexPath, state, "session_fork", "fork");
+      const currentSynced = syncAttachedSession(indexPath, state, "session_fork", "fork");
       return previousSynced || currentSynced;
     },
     handleToolCall(event, sessionFile, cwd) {
@@ -144,7 +144,7 @@ export function createSessionHookController(options: { indexPath: string }): Ses
     async handleTurnEnd(sessionFile, cwd) {
       attachSession(state, sessionFile, cwd);
       try {
-        return await syncAttachedSession(indexPath, state, "turn_end");
+        return syncAttachedSession(indexPath, state, "turn_end");
       } finally {
         clearTurnState(state);
       }
@@ -160,7 +160,7 @@ export function createSessionHookController(options: { indexPath: string }): Ses
     async handleSessionShutdown(sessionFile, cwd) {
       attachSession(state, sessionFile, cwd);
       try {
-        return await syncAttachedSession(indexPath, state, "session_shutdown");
+        return syncAttachedSession(indexPath, state, "session_shutdown");
       } finally {
         clearTurnState(state);
         state.currentSessionFile = undefined;
@@ -182,73 +182,67 @@ function attachSession(
   state.currentCwd = cwd;
 }
 
-async function retryIndexWrite<T>(action: () => T): Promise<T> {
-  for (let attempt = 1; attempt <= HOOK_INDEX_WRITE_MAX_ATTEMPTS; attempt += 1) {
-    try {
-      return action();
-    } catch (error) {
-      if (!isRetryableSqliteLock(error) || attempt >= HOOK_INDEX_WRITE_MAX_ATTEMPTS) {
-        throw error;
-      }
-
-      await sleep(getRetryDelayMs(attempt));
-    }
-  }
-
-  return action();
-}
-
-function getRetryDelayMs(attempt: number): number {
-  return (
-    HOOK_INDEX_WRITE_RETRY_DELAYS_MS[attempt - 1] ??
-    HOOK_INDEX_WRITE_RETRY_DELAYS_MS[HOOK_INDEX_WRITE_RETRY_DELAYS_MS.length - 1] ??
-    0
-  );
-}
-
-function isRetryableSqliteLock(error: unknown): boolean {
-  if (!(error instanceof Error)) {
-    return false;
-  }
-
-  const code = "code" in error && typeof error.code === "string" ? error.code : undefined;
-  const message = error.message.toLowerCase();
-  return (
-    code === "SQLITE_BUSY" ||
-    code === "SQLITE_LOCKED" ||
-    code?.startsWith("SQLITE_BUSY_") === true ||
-    code?.startsWith("SQLITE_LOCKED_") === true ||
-    message.includes("database is locked")
-  );
-}
-
-function sleep(ms: number): Promise<void> {
-  return new Promise((resolve) => setTimeout(resolve, ms));
-}
-
 function syncAttachedSession(
   indexPath: string,
   state: SessionHookState,
   eventType: string,
   sessionOrigin?: SessionOrigin,
-): Promise<boolean> {
+): boolean {
   return syncSessionFile(indexPath, state.currentSessionFile, eventType, state, sessionOrigin);
 }
 
-async function syncSessionFile(
+function syncSessionFile(
   indexPath: string,
   sessionFile: string | undefined,
   eventType: string,
   state?: SessionHookState,
   sessionOrigin?: SessionOrigin,
-): Promise<boolean> {
-  if (!sessionFile || !existsSync(sessionFile)) {
+): boolean {
+  if (!sessionFile || !existsSync(sessionFile) || !existsSync(indexPath)) {
     return false;
   }
 
-  const status = getIndexStatus(indexPath);
-  if (!status.exists || status.schemaVersion !== INDEX_SCHEMA_VERSION) {
-    return false;
+  const db = openIndexDatabase(indexPath, { create: false });
+  try {
+    if (readIndexSchemaVersion(db) !== INDEX_SCHEMA_VERSION) {
+      return false;
+    }
+
+    const synced = syncSessionFileWithDb(db, sessionFile, eventType, sessionOrigin);
+    if (synced && state) {
+      state.lastFlushedSessionFile = sessionFile;
+    }
+    return synced;
+  } finally {
+    db.close();
+  }
+}
+
+interface TailSyncBaseline extends SessionRow {
+  indexedFileSize: number;
+  indexedFileMtimeMs: number;
+  indexedFileAnchor: string;
+}
+
+function syncSessionFileWithDb(
+  db: SessionIndexDatabase,
+  sessionFile: string,
+  eventType: string,
+  sessionOrigin?: SessionOrigin,
+): boolean {
+  const baseline = asTailSyncBaseline(getSessionRowByPath(db, sessionFile));
+  const stat = statSync(sessionFile);
+
+  if (baseline && isIndexCurrent(baseline, stat, sessionOrigin)) {
+    db.transaction(() => writeHookSyncMetadata(db, eventType)).immediate();
+    return true;
+  }
+
+  if (baseline && stat.size > baseline.indexedFileSize) {
+    const tail = extractSessionTail(sessionFile, baseline);
+    if (tail && applyTailSync(db, baseline, tail, eventType, sessionOrigin)) {
+      return true;
+    }
   }
 
   const extracted = extractSessionRecord(sessionFile);
@@ -256,42 +250,198 @@ async function syncSessionFile(
     return false;
   }
 
-  await retryIndexWrite(() => {
-    const db = openIndexDatabase(indexPath, {
-      create: false,
-      timeoutMs: HOOK_INDEX_WRITE_TIMEOUT_MS,
-    });
-    try {
-      db.transaction(() => {
-        const existingSession = getSessionById(db, extracted.sessionId);
-        const sessionRow = mergeSessionLineage(extracted, existingSession, sessionOrigin);
-        upsertSession(db, sessionRow, "hook");
-        if (shouldRefreshLineageRelations(existingSession, sessionRow)) {
-          rebuildSessionLineageRelations(db);
-        }
-        clearSessionIndexedData(db, extracted.sessionId);
+  applyFullSync(db, extracted, eventType, sessionOrigin);
+  return true;
+}
 
-        for (const chunk of extracted.chunks) {
-          insertTextChunk(db, { sessionId: extracted.sessionId, ...chunk });
-        }
+function asTailSyncBaseline(row: SessionRow | undefined): TailSyncBaseline | undefined {
+  if (
+    !row ||
+    row.indexedFileSize === undefined ||
+    row.indexedFileMtimeMs === undefined ||
+    row.indexedFileAnchor === undefined ||
+    row.indexedFileAnchor.length === 0
+  ) {
+    return undefined;
+  }
 
-        for (const fileTouch of extracted.fileTouches) {
-          insertSessionFileTouch(db, { sessionId: extracted.sessionId, ...fileTouch });
-        }
+  return {
+    ...row,
+    indexedFileSize: row.indexedFileSize,
+    indexedFileMtimeMs: row.indexedFileMtimeMs,
+    indexedFileAnchor: row.indexedFileAnchor,
+  };
+}
+
+function isIndexCurrent(
+  baseline: TailSyncBaseline,
+  stat: Stats,
+  sessionOrigin?: SessionOrigin,
+): boolean {
+  return (
+    stat.size === baseline.indexedFileSize &&
+    Math.trunc(stat.mtimeMs) === baseline.indexedFileMtimeMs &&
+    (sessionOrigin === undefined || sessionOrigin === baseline.sessionOrigin)
+  );
+}
+
+// All write transactions are immediate: they read before they write, and a
+// deferred transaction whose snapshot goes stale fails with SQLITE_BUSY on the
+// read-to-write upgrade without ever invoking the busy handler. Immediate mode
+// takes the write lock at BEGIN, where busy_timeout queues us behind concurrent
+// writers from other pi processes.
+function applyFullSync(
+  db: SessionIndexDatabase,
+  extracted: ExtractedSessionRecord,
+  eventType: string,
+  sessionOrigin?: SessionOrigin,
+): void {
+  db.transaction(() => {
+    const existingSession = getSessionById(db, extracted.sessionId);
+    const sessionRow = mergeSessionLineage(extracted, existingSession, sessionOrigin);
+    clearSessionIndexedData(db, extracted.sessionId);
+    upsertSession(db, sessionRow, "hook");
+    if (shouldRefreshLineageRelations(existingSession, sessionRow)) {
+      refreshSessionLineageRelationsFor(db, [
+        extracted.sessionId,
+        existingSession?.parentSessionId,
+        sessionRow.parentSessionId,
+      ]);
+    }
 
-        setMetadata(db, "hook_updated_at", new Date().toISOString());
-        setMetadata(db, "hook_last_event", eventType);
-      })();
-    } finally {
-      db.close();
+    for (const chunk of extracted.chunks) {
+      insertTextChunk(db, { sessionId: extracted.sessionId, ...chunk });
     }
-  });
 
-  if (state) {
-    state.lastFlushedSessionFile = sessionFile;
-  }
+    for (const fileTouch of extracted.fileTouches) {
+      insertSessionFileTouch(db, { sessionId: extracted.sessionId, ...fileTouch });
+    }
 
-  return true;
+    writeHookSyncMetadata(db, eventType);
+  }).immediate();
+}
+
+function applyTailSync(
+  db: SessionIndexDatabase,
+  baseline: TailSyncBaseline,
+  tail: ExtractedSessionTail,
+  eventType: string,
+  sessionOrigin?: SessionOrigin,
+): boolean {
+  return db
+    .transaction((): boolean => {
+      // Another process may have advanced the index between our baseline read
+      // and this transaction; the tail deltas would then double-count.
+      const current = getSessionRowByPath(db, baseline.sessionPath);
+      if (
+        !current ||
+        current.sessionId !== baseline.sessionId ||
+        current.indexedFileSize !== baseline.indexedFileSize
+      ) {
+        return false;
+      }
+
+      const scan = tail.scan;
+      upsertSession(db, buildTailSessionRow(baseline, tail, sessionOrigin), "hook");
+
+      if (scan.sessionName !== undefined && scan.sessionName !== baseline.sessionName) {
+        clearSessionChunksBySourceKind(db, baseline.sessionId, "session_name");
+        if (scan.sessionName) {
+          insertTextChunk(db, {
+            sessionId: baseline.sessionId,
+            ...createSessionNameChunk(scan.sessionName, baseline.startedAt),
+          });
+        }
+      }
+
+      if (!baseline.handoffGoal && scan.handoffMetadata) {
+        const { entryId, ts, metadata } = scan.handoffMetadata;
+        insertTextChunk(db, {
+          sessionId: baseline.sessionId,
+          entryId,
+          entryType: "custom",
+          ts,
+          sourceKind: "handoff_goal",
+          text: metadata.goal,
+        });
+        insertTextChunk(db, {
+          sessionId: baseline.sessionId,
+          entryId,
+          entryType: "custom",
+          ts,
+          sourceKind: "handoff_next_task",
+          text: metadata.nextTask,
+        });
+      }
+
+      for (const chunk of scan.chunks) {
+        insertTextChunk(db, { sessionId: baseline.sessionId, ...chunk });
+      }
+
+      for (const fileTouch of scan.fileTouches) {
+        insertSessionFileTouch(db, { sessionId: baseline.sessionId, ...fileTouch });
+      }
+
+      writeHookSyncMetadata(db, eventType);
+      return true;
+    })
+    .immediate();
+}
+
+function buildTailSessionRow(
+  baseline: TailSyncBaseline,
+  tail: ExtractedSessionTail,
+  sessionOrigin?: SessionOrigin,
+): SessionRow {
+  const scan = tail.scan;
+  const tailHandoffMetadata = baseline.handoffGoal ? undefined : scan.handoffMetadata?.metadata;
+  const tailOrigin =
+    baseline.parentSessionPath && tailHandoffMetadata?.origin === "handoff"
+      ? ("handoff" as const)
+      : undefined;
+  const nextOrigin = resolveSessionOrigin(sessionOrigin, tailOrigin, baseline.sessionOrigin);
+
+  return {
+    sessionId: baseline.sessionId,
+    sessionPath: baseline.sessionPath,
+    sessionName: scan.sessionName ?? baseline.sessionName,
+    firstUserPrompt: baseline.firstUserPrompt || (scan.firstUserPrompt ?? ""),
+    cwd: baseline.cwd,
+    repoRoots: mergeRepoRoots(baseline, scan.fileTouches),
+    startedAt: baseline.startedAt,
+    modifiedAt:
+      scan.maxEntryTs !== undefined && scan.maxEntryTs > baseline.modifiedAt
+        ? scan.maxEntryTs
+        : baseline.modifiedAt,
+    messageCount: baseline.messageCount + scan.messageCount,
+    entryCount: baseline.entryCount + scan.entryCount,
+    parentSessionPath: baseline.parentSessionPath,
+    parentSessionId: baseline.parentSessionId,
+    sessionOrigin: baseline.parentSessionPath ? (nextOrigin ?? "unknown_child") : undefined,
+    handoffGoal: baseline.handoffGoal ?? tailHandoffMetadata?.goal,
+    handoffNextTask: baseline.handoffNextTask ?? tailHandoffMetadata?.nextTask,
+    indexedFileSize: tail.indexedFileSize,
+    indexedFileMtimeMs: tail.indexedFileMtimeMs,
+    indexedFileAnchor: tail.indexedFileAnchor,
+  };
+}
+
+function mergeRepoRoots(baseline: TailSyncBaseline, fileTouches: SessionFileTouch[]): string[] {
+  const merged = new Set([
+    ...baseline.repoRoots,
+    ...deriveSessionRepoRoots(baseline.cwd, fileTouches),
+  ]);
+  return [...merged].sort();
+}
+
+function writeHookSyncMetadata(db: SessionIndexDatabase, eventType: string): void {
+  setMetadata(db, "hook_updated_at", new Date().toISOString());
+  setMetadata(db, "hook_last_event", eventType);
+}
+
+function readIndexSchemaVersion(db: SessionIndexDatabase): number | undefined {
+  const raw = getMetadata(db, "schema_version");
+  return raw === undefined ? undefined : Number(raw);
 }
 
 function mergeSessionLineage(

package/extensions/session-search/reindex.ts

@@ -1,17 +1,18 @@
-import { renameSync, rmSync } from "node:fs";
+import path from "node:path";
 import { SessionManager } from "@earendil-works/pi-coding-agent";
 import {
-  createTempIndexPath,
+  clearSessionIndexedData,
+  ensureIndexDir,
   initializeSchema,
-  insertSession,
   insertSessionFileTouch,
   insertTextChunk,
   openIndexDatabase,
   rebuildSessionLineageRelations,
   type SessionIndexDatabase,
   setMetadata,
+  upsertSession,
 } from "../shared/session-index/index.js";
-import { extractSessionRecord, type SearchTextChunk, type SessionFileTouch } from "./extract.js";
+import { type ExtractedSessionRecord, extractSessionRecord } from "./extract.js";
 
 export interface ReindexOptions {
   indexPath: string;
@@ -24,75 +25,67 @@ export interface ReindexResult {
 }
 
 export async function rebuildSessionIndex(options: ReindexOptions): Promise<ReindexResult> {
-  const finalIndexPath = options.indexPath;
-  const tempIndexPath = createTempIndexPath(finalIndexPath);
+  const indexPath = options.indexPath;
+  ensureIndexDir(path.dirname(indexPath));
   const sessionFiles = (await SessionManager.listAll()).map((session) => session.path);
 
-  const db = openIndexDatabase(tempIndexPath, { create: true });
-  let sessionCount: number;
-  let chunkCount: number;
+  const db = openIndexDatabase(indexPath, { create: true });
   try {
-    initializeSchema(db);
-    ({ sessionCount, chunkCount } = indexSessionFiles(db, sessionFiles));
-    // Fold the WAL into the main file so the single-file rename below is
-    // self-contained. bun:sqlite does not checkpoint on close.
-    db.exec("PRAGMA wal_checkpoint(TRUNCATE)");
-  } finally {
-    db.close();
-  }
+    db.transaction(() => {
+      dropIndexTables(db);
+      initializeSchema(db);
+    }).immediate();
 
-  renameSync(tempIndexPath, finalIndexPath);
-  rmSync(`${tempIndexPath}-wal`, { force: true });
-  rmSync(`${tempIndexPath}-shm`, { force: true });
-  return { sessionCount, chunkCount, indexPath: finalIndexPath };
-}
-
-function indexSessionFiles(
-  db: SessionIndexDatabase,
-  sessionFiles: string[],
-): { sessionCount: number; chunkCount: number } {
-  return db.transaction((files: string[]) => {
     let sessionCount = 0;
     let chunkCount = 0;
-
-    for (const sessionFile of files) {
+    for (const sessionFile of sessionFiles) {
       const extracted = extractSessionRecord(sessionFile);
       if (!extracted) {
         continue;
       }
 
-      insertSession(db, extracted, "full_reindex");
+      db.transaction(() => indexSession(db, extracted)).immediate();
       sessionCount += 1;
-      chunkCount += insertSessionChunks(db, extracted.sessionId, extracted.chunks);
-      insertSessionFileTouches(db, extracted.sessionId, extracted.fileTouches);
+      chunkCount += extracted.chunks.length;
     }
 
-    rebuildSessionLineageRelations(db);
-    setMetadata(db, "indexed_at", new Date().toISOString());
-    setMetadata(db, "session_source", "SessionManager.listAll()");
-
-    return { sessionCount, chunkCount };
-  })(sessionFiles);
-}
+    db.transaction(() => {
+      rebuildSessionLineageRelations(db);
+      setMetadata(db, "indexed_at", new Date().toISOString());
+      setMetadata(db, "session_source", "SessionManager.listAll()");
+    }).immediate();
 
-function insertSessionChunks(
-  db: SessionIndexDatabase,
-  sessionId: string,
-  chunks: SearchTextChunk[],
-): number {
-  for (const chunk of chunks) {
-    insertTextChunk(db, { sessionId, ...chunk });
+    return { sessionCount, chunkCount, indexPath };
+  } finally {
+    db.close();
   }
+}
 
-  return chunks.length;
+function dropIndexTables(db: SessionIndexDatabase): void {
+  // Children before parents: foreign_keys is ON and DROP TABLE runs an
+  // implicit DELETE that parent-side constraints would reject.
+  db.exec(`
+    DROP TABLE IF EXISTS session_lineage_relations;
+    DROP TABLE IF EXISTS session_text_chunks;
+    DROP TABLE IF EXISTS session_file_touches;
+    DROP TABLE IF EXISTS session_text_chunks_fts;
+    DROP TABLE IF EXISTS sessions;
+    DROP TABLE IF EXISTS metadata;
+  `);
 }
 
-function insertSessionFileTouches(
-  db: SessionIndexDatabase,
-  sessionId: string,
-  fileTouches: SessionFileTouch[],
-): void {
-  for (const fileTouch of fileTouches) {
-    insertSessionFileTouch(db, { sessionId, ...fileTouch });
+function indexSession(db: SessionIndexDatabase, extracted: ExtractedSessionRecord): void {
+  clearSessionIndexedData(db, extracted.sessionId);
+  upsertSession(db, extracted, "full_reindex");
+
+  for (const chunk of extracted.chunks) {
+    insertTextChunk(db, { sessionId: extracted.sessionId, ...chunk });
+  }
+
+  for (const fileTouch of extracted.fileTouches) {
+    insertSessionFileTouch(db, {
+      sessionId: extracted.sessionId,
+      ...fileTouch,
+    });
   }
 }

package/extensions/shared/session-index/common.ts

@@ -3,7 +3,7 @@ import type { FileTouchOp, FileTouchSource, PathScope } from "../../session-sear
 import { safeParseTypeBoxJson } from "../typebox.js";
 import type { SqliteDatabase } from "./sqlite.js";
 
-export const INDEX_SCHEMA_VERSION = 7;
+export const INDEX_SCHEMA_VERSION = 8;
 
 export type SessionOrigin = "handoff" | "fork" | "unknown_child";
 export type SessionLineageRelation =
@@ -30,6 +30,9 @@ export interface SessionRow {
   sessionOrigin?: SessionOrigin | undefined;
   handoffGoal?: string | undefined;
   handoffNextTask?: string | undefined;
+  indexedFileSize?: number | undefined;
+  indexedFileMtimeMs?: number | undefined;
+  indexedFileAnchor?: string | undefined;
 }
 
 export interface SessionLineageRow {

package/extensions/shared/session-index/lineage.ts

@@ -185,9 +185,39 @@ export function getSiblingSessions(
   return queryRelatedSessions(db, sessionId, ["sibling"]);
 }
 
+interface SessionGraph {
+  nodes: Map<string, SessionGraphNode>;
+  childrenByParent: Map<string, string[]>;
+}
+
 export function rebuildSessionLineageRelations(db: SessionIndexDatabase): void {
   db.prepare(`DELETE FROM session_lineage_relations`).run();
 
+  const graph = buildSessionGraph(db);
+  insertLineageRelationRows(db, graph.nodes.keys(), graph);
+}
+
+// Recomputes lineage for just the connected component(s) containing the seed
+// sessions. Lineage relations only ever link sessions connected through parent
+// edges, so the rest of the table is untouched.
+export function refreshSessionLineageRelationsFor(
+  db: SessionIndexDatabase,
+  seedSessionIds: Array<string | undefined>,
+): void {
+  const graph = buildSessionGraph(db);
+  const component = collectLineageComponent(seedSessionIds, graph);
+  if (component.size === 0) {
+    return;
+  }
+
+  const placeholders = [...component].map(() => "?").join(", ");
+  db.prepare(`DELETE FROM session_lineage_relations WHERE session_id IN (${placeholders})`).run(
+    ...component,
+  );
+  insertLineageRelationRows(db, component, graph);
+}
+
+function buildSessionGraph(db: SessionIndexDatabase): SessionGraph {
   const rows = parseTypeBoxRows(
     SESSION_GRAPH_ROW_SCHEMA,
     db
@@ -229,6 +259,40 @@ export function rebuildSessionLineageRelations(db: SessionIndexDatabase): void {
     childrenByParent.set(node.resolvedParentSessionId, children);
   }
 
+  return { nodes, childrenByParent };
+}
+
+function collectLineageComponent(
+  seedSessionIds: Array<string | undefined>,
+  graph: SessionGraph,
+): Set<string> {
+  const component = new Set<string>();
+  const queue = seedSessionIds.filter(
+    (sessionId): sessionId is string => sessionId !== undefined && graph.nodes.has(sessionId),
+  );
+
+  while (queue.length > 0) {
+    const sessionId = queue.pop();
+    if (!sessionId || component.has(sessionId)) {
+      continue;
+    }
+
+    component.add(sessionId);
+    const parentId = graph.nodes.get(sessionId)?.resolvedParentSessionId;
+    if (parentId) {
+      queue.push(parentId);
+    }
+    queue.push(...(graph.childrenByParent.get(sessionId) ?? []));
+  }
+
+  return component;
+}
+
+function insertLineageRelationRows(
+  db: SessionIndexDatabase,
+  sessionIds: Iterable<string>,
+  graph: SessionGraph,
+): void {
   const insertRelation = db.prepare(
     `
       INSERT INTO session_lineage_relations(session_id, related_session_id, relation, distance)
@@ -236,8 +300,12 @@ export function rebuildSessionLineageRelations(db: SessionIndexDatabase): void {
     `,
   );
 
-  for (const node of nodes.values()) {
-    const relations = collectMaterializedLineageRows(node.sessionId, nodes, childrenByParent);
+  for (const sessionId of sessionIds) {
+    const relations = collectMaterializedLineageRows(
+      sessionId,
+      graph.nodes,
+      graph.childrenByParent,
+    );
     for (const relation of relations.values()) {
       insertRelation.run(
         relation.sessionId,