added cached id->file index and made resume path async => freezing on resume fixed

pi-extension/subagents/index.ts

@@ -28,11 +28,12 @@ import {
 } from "./cmux.ts";
 
 import {
+  countSessionEntryLines,
   findLastAssistantMessage,
   getNewEntries,
   getSessionId,
   readSubagentLoadout,
-  resolveSessionFileById,
+  resolveSessionFileByIdAsync,
   seedSubagentSessionFile,
   summarizeSessionStats,
   writeSubagentLoadout,
@@ -2083,12 +2084,14 @@ export default function subagentsExtension(pi: ExtensionAPI) {
         const startTime = Date.now();
         const id = Math.random().toString(16).slice(2, 10);
 
-        // Resolve the id (or path) to a concrete session file.
+        // Resolve the id (or path) to a concrete session file. Use the async
+        // resolver so a cold index build (thousands of sessions) doesn't block
+        // the event loop and freeze the terminal UI.
         const sessionsRoot = dirname(ctx.sessionManager.getSessionDir());
         const sessionPath =
           requestedId.includes("/") || requestedId.includes("\\") || requestedId.endsWith(".jsonl")
             ? requestedId
-            : resolveSessionFileById(requestedId, sessionsRoot);
+            : await resolveSessionFileByIdAsync(requestedId, sessionsRoot);
 
         if (!sessionPath || !existsSync(sessionPath)) {
           const err = `No session found for id "${requestedId}". Use the session id returned in a completed subagent's result.`;
@@ -2119,8 +2122,10 @@ export default function subagentsExtension(pi: ExtensionAPI) {
 
         const resumedSessionId = getSessionId(sessionPath) ?? requestedId;
 
-        // Record entry count before resuming so we can extract new messages
+        // Record entry count before resuming so we can extract new messages.
-        const entryCountBefore = getNewEntries(sessionPath, 0).length;
+        // Count lines cheaply (no per-line JSON.parse) so resuming a large
+        // transcript doesn't block the UI.
+        const entryCountBefore = countSessionEntryLines(sessionPath);
 
         const surface = createSurface(name);
         await new Promise<void>((resolve) => setTimeout(resolve, getShellReadyDelayMs()));

pi-extension/subagents/session.ts

@@ -217,46 +217,176 @@ function readHeaderId(sessionFile: string): string | null {
  * resolution order: exact match first, then prefix match. Most recently
  * modified file wins on ties. Returns null when nothing matches.
  */
-export function resolveSessionFileById(sessionId: string, sessionsRoot: string): string | null {
+/**
-  if (!sessionId || !existsSync(sessionsRoot)) return null;
+ * In-process index of session id → session file, per sessions root.
+ *
+ * Resolving a session id naively walks every `.jsonl` under the sessions tree
+ * and reads each header. With a few thousand sessions that is thousands of
+ * synchronous open/read/stat syscalls — on the extension host's single thread
+ * that blocks the entire terminal UI for many seconds (measured ~67s on a
+ * 2010-file tree). To avoid that, we build the index once per root and cache
+ * it; subsequent lookups are O(1). The cache is validated cheaply (a directory
+ * listing plus statSync-only mtime checks) on every call, so new sessions are
+ * picked up without re-reading unchanged headers and without ever freezing the
+ * UI again.
+ */
+interface SessionIndex {
+  idToFile: Map<string, { path: string; mtime: number }>;
+  /** file path → mtime when indexed (staleness detection). */
+  files: Map<string, number>;
+  /** top-level dir signature used to detect newly added cwd dirs. */
+  topSig: string;
+}
+const sessionIndexCache = new Map<string, SessionIndex>();
 
-  const candidates: Array<{ path: string; id: string; mtime: number }> = [];
+function topLevelSignature(root: string): string {
-  const walk = (dir: string) => {
+  const parts: string[] = [];
-    let entries: import("node:fs").Dirent[];
+  let entries: import("node:fs").Dirent[];
-    try {
+  try {
-      entries = readdirSync(dir, { withFileTypes: true });
+    entries = readdirSync(root, { withFileTypes: true });
-    } catch {
+  } catch {
-      return;
+    return "";
+  }
+  for (const e of entries) {
+    const full = join(root, e.name);
+    if (e.isDirectory()) {
+      let m = 0;
+      try {
+        m = statSync(full).mtimeMs;
+      } catch {
+        /* ignore */
+      }
+      parts.push(`d:${e.name}:${m}`);
+    } else if (e.isFile() && e.name.endsWith(".jsonl")) {
+      parts.push(`f:${e.name}`);
     }
-    for (const entry of entries) {
+  }
-      const full = join(dir, entry.name);
+  parts.sort();
-      if (entry.isDirectory()) {
+  return parts.join("|");
-        walk(full);
+}
-      } else if (entry.isFile() && entry.name.endsWith(".jsonl")) {
+
-        const id = readHeaderId(full);
+/** Recursively index new/changed .jsonl files under dir into idx. */
-        if (!id) continue;
+function indexDir(dir: string, idx: SessionIndex): void {
-        let mtime = 0;
+  let entries: import("node:fs").Dirent[];
-        try {
+  try {
-          mtime = statSync(full).mtimeMs;
+    entries = readdirSync(dir, { withFileTypes: true });
-        } catch {
+  } catch {
-          /* ignore */
+    return;
-        }
+  }
-        candidates.push({ path: full, id, mtime });
+  for (const entry of entries) {
+    const full = join(dir, entry.name);
+    if (entry.isDirectory()) {
+      indexDir(full, idx);
+    } else if (entry.isFile() && entry.name.endsWith(".jsonl")) {
+      let mtime = 0;
+      try {
+        mtime = statSync(full).mtimeMs;
+      } catch {
+        continue;
+      }
+      const known = idx.files.get(full);
+      if (known !== undefined && known === mtime) continue; // unchanged
+      const id = readHeaderId(full); // only read headers for new/changed files
+      idx.files.set(full, mtime);
+      if (!id) continue;
+      const prev = idx.idToFile.get(id);
+      if (!prev || mtime >= prev.mtime) {
+        idx.idToFile.set(id, { path: full, mtime });
       }
     }
-  };
+  }
-  walk(sessionsRoot);
+}
 
-  candidates.sort((a, b) => b.mtime - a.mtime);
+function getSessionIndex(sessionsRoot: string): SessionIndex {
-  const exact = candidates.find((c) => c.id === sessionId);
+  let idx = sessionIndexCache.get(sessionsRoot);
-  if (exact) return exact.path;
+  const sig = topLevelSignature(sessionsRoot);
-  const prefix = candidates.find((c) => c.id.startsWith(sessionId));
+  if (!idx) {
-  return prefix ? prefix.path : null;
+    idx = { idToFile: new Map(), files: new Map(), topSig: sig };
+    sessionIndexCache.set(sessionsRoot, idx);
+    indexDir(sessionsRoot, idx); // first build: full scan, once per process
+  } else if (idx.topSig !== sig) {
+    idx.topSig = sig;
+    indexDir(sessionsRoot, idx); // a cwd dir was added/changed: incremental rescan
+  } else {
+    indexDir(sessionsRoot, idx); // cheap: stats files, reads only new/changed headers
+  }
+  return idx;
+}
+
+export function resolveSessionFileById(sessionId: string, sessionsRoot: string): string | null {
+  if (!sessionId || !existsSync(sessionsRoot)) return null;
+  const idx = getSessionIndex(sessionsRoot);
+  return lookupSessionIndex(idx, sessionId);
+}
+
+function lookupSessionIndex(
+  idx: { idToFile: Map<string, { path: string; mtime: number }> },
+  sessionId: string,
+): string | null {
+  // Exact match first.
+  const exact = idx.idToFile.get(sessionId);
+  if (exact && existsSync(exact.path)) return exact.path;
+
+  // Prefix match: most recently modified wins (ids are unique in practice, so
+  // this is only a convenience for hand-typed short prefixes).
+  let best: { path: string; mtime: number } | null = null;
+  for (const [id, rec] of idx.idToFile) {
+    if (!id.startsWith(sessionId)) continue;
+    if (!existsSync(rec.path)) continue;
+    if (!best || rec.mtime > best.mtime) best = rec;
+  }
+  return best ? best.path : null;
+}
+
+/**
+ * Async variant used by the interactive resume path. Index building/refresh is
+ * synchronous I/O, which can take many seconds on a cold OS page cache with a
+ * few thousand sessions; running it synchronously would block the extension
+ * host's single thread and freeze the terminal UI. Deferring to a macrotask
+ * keeps the event loop responsive. The heavy work only happens on the first
+ * resolution per process (and incrementally thereafter); warm lookups are ~50ms.
+ */
+export async function resolveSessionFileByIdAsync(
+  sessionId: string,
+  sessionsRoot: string,
+): Promise<string | null> {
+  if (!sessionId || !existsSync(sessionsRoot)) return null;
+  // Let the event loop breathe (and the UI repaint) before the sync scan.
+  await new Promise<void>((r) => setImmediate(r));
+  const idx = getSessionIndex(sessionsRoot);
+  return lookupSessionIndex(idx, sessionId);
+}
+
+/** Test hook: drop the cached session index so tests start clean. */
+export function resetSessionIndexCache(): void {
+  sessionIndexCache.clear();
 }
 
 /**
  * Return entries added after `afterLine` (1-indexed count of existing entries).
  */
+/**
+ * Count the number of entry lines in a session file without parsing each line
+ * into an object. Used by the resume path, which only needs the *count* of
+ * pre-existing entries (so it can later slice out the new ones). Parsing every
+ * line of a large resumed transcript synchronously at resume time would block
+ * the UI; counting newlines is dramatically cheaper.
+ */
+export function countSessionEntryLines(sessionFile: string): number {
+  try {
+    const raw = readFileSync(sessionFile, "utf8");
+    // Count non-blank lines, mirroring getNewEntries' `.filter(line => line.trim())`
+    // but skipping the per-line JSON.parse that makes resume slow on big files.
+    let count = 0;
+    for (const line of raw.split("\n")) {
+      if (line.trim()) count++;
+    }
+    return count;
+  } catch {
+    return 0;
+  }
+}
+
 export function getNewEntries(sessionFile: string, afterLine: number): SessionEntry[] {
   const raw = readFileSync(sessionFile, "utf8");
   const lines = raw.split("\n").filter((line) => line.trim());

test/test.ts

@@ -1,4 +1,4 @@
-import { describe, it, before, after } from "node:test";
+import { describe, it, before, after, beforeEach } from "node:test";
 import assert from "node:assert/strict";
 import { mkdtempSync, writeFileSync, readFileSync, mkdirSync, rmSync, existsSync } from "node:fs";
 import { join } from "node:path";
@@ -10,11 +10,13 @@ import * as subagentsModule from "../pi-extension/subagents/index.ts";
 import {
   getLeafId,
   getNewEntries,
+  countSessionEntryLines,
   getSessionId,
   readSubagentLoadout,
   writeSubagentLoadout,
   loadoutSidecarPath,
   type SubagentLoadout,
+  resetSessionIndexCache,
   resolveSessionFileById,
   findLastAssistantMessage,
   appendBranchSummary,
@@ -265,6 +267,19 @@ describe("session.ts", () => {
       const entries = getNewEntries(file, 2);
       assert.equal(entries.length, 0);
     });
+
+    it("countSessionEntryLines matches getNewEntries(0).length without parsing", () => {
+      const file = createSessionFile(dir, [SESSION_HEADER, MODEL_CHANGE, USER_MSG, ASSISTANT_MSG]);
+      assert.equal(countSessionEntryLines(file), getNewEntries(file, 0).length);
+      assert.equal(countSessionEntryLines(file), 4);
+    });
+
+    it("countSessionEntryLines ignores blank lines and returns 0 for missing files", () => {
+      const file = join(dir, "blanks.jsonl");
+      writeFileSync(file, JSON.stringify({ type: "session", id: "x" }) + "\n\n\n");
+      assert.equal(countSessionEntryLines(file), 1);
+      assert.equal(countSessionEntryLines(join(dir, "does-not-exist.jsonl")), 0);
+    });
   });
 
   describe("getSessionId / resolveSessionFileById", () => {
@@ -274,6 +289,12 @@ describe("session.ts", () => {
       return p;
     }
 
+    // The resolver caches an id→file index per root; reset it so each test
+    // builds a fresh index from the current on-disk state.
+    beforeEach(() => {
+      resetSessionIndexCache();
+    });
+
     it("reads the header id from a session file", () => {
       const file = createSessionFile(dir, [SESSION_HEADER, MODEL_CHANGE, USER_MSG]);
       assert.equal(getSessionId(file), "sess-001");
@@ -299,6 +320,16 @@ describe("session.ts", () => {
       writeSession(dir, "c.jsonl", "abc");
       assert.equal(resolveSessionFileById("zzz", dir), null);
     });
+
+    it("picks up newly added sessions on repeat calls without a reset", () => {
+      // Prime the index (first call builds it).
+      writeSession(dir, "first.jsonl", "id-first");
+      assert.equal(resolveSessionFileById("id-first", dir) !== null, true);
+      // Add a new session AFTER the index was built — no reset. The resolver's
+      // cheap refresh should index it.
+      const b = writeSession(dir, "second.jsonl", "id-second");
+      assert.equal(resolveSessionFileById("id-second", dir), b);
+    });
   });
 
   describe("subagent loadout snapshot", () => {