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

This commit is contained in:
Amos
2026-07-21 22:28:46 +03:00
parent a541b3d98b
commit 67495c38fb
3 changed files with 202 additions and 36 deletions
+10 -5
View File
@@ -28,11 +28,12 @@ import {
} from "./cmux.ts"; } from "./cmux.ts";
import { import {
countSessionEntryLines,
findLastAssistantMessage, findLastAssistantMessage,
getNewEntries, getNewEntries,
getSessionId, getSessionId,
readSubagentLoadout, readSubagentLoadout,
resolveSessionFileById, resolveSessionFileByIdAsync,
seedSubagentSessionFile, seedSubagentSessionFile,
summarizeSessionStats, summarizeSessionStats,
writeSubagentLoadout, writeSubagentLoadout,
@@ -2083,12 +2084,14 @@ export default function subagentsExtension(pi: ExtensionAPI) {
const startTime = Date.now(); const startTime = Date.now();
const id = Math.random().toString(16).slice(2, 10); 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 sessionsRoot = dirname(ctx.sessionManager.getSessionDir());
const sessionPath = const sessionPath =
requestedId.includes("/") || requestedId.includes("\\") || requestedId.endsWith(".jsonl") requestedId.includes("/") || requestedId.includes("\\") || requestedId.endsWith(".jsonl")
? requestedId ? requestedId
: resolveSessionFileById(requestedId, sessionsRoot); : await resolveSessionFileByIdAsync(requestedId, sessionsRoot);
if (!sessionPath || !existsSync(sessionPath)) { if (!sessionPath || !existsSync(sessionPath)) {
const err = `No session found for id "${requestedId}". Use the session id returned in a completed subagent's result.`; 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; 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); const surface = createSurface(name);
await new Promise<void>((resolve) => setTimeout(resolve, getShellReadyDelayMs())); await new Promise<void>((resolve) => setTimeout(resolve, getShellReadyDelayMs()));
+146 -16
View File
@@ -217,11 +217,56 @@ function readHeaderId(sessionFile: string): string | null {
* resolution order: exact match first, then prefix match. Most recently * resolution order: exact match first, then prefix match. Most recently
* modified file wins on ties. Returns null when nothing matches. * 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[];
try {
entries = readdirSync(root, { withFileTypes: true });
} catch {
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}`);
}
}
parts.sort();
return parts.join("|");
}
/** Recursively index new/changed .jsonl files under dir into idx. */
function indexDir(dir: string, idx: SessionIndex): void {
let entries: import("node:fs").Dirent[]; let entries: import("node:fs").Dirent[];
try { try {
entries = readdirSync(dir, { withFileTypes: true }); entries = readdirSync(dir, { withFileTypes: true });
@@ -231,32 +276,117 @@ export function resolveSessionFileById(sessionId: string, sessionsRoot: string):
for (const entry of entries) { for (const entry of entries) {
const full = join(dir, entry.name); const full = join(dir, entry.name);
if (entry.isDirectory()) { if (entry.isDirectory()) {
walk(full); indexDir(full, idx);
} else if (entry.isFile() && entry.name.endsWith(".jsonl")) { } else if (entry.isFile() && entry.name.endsWith(".jsonl")) {
const id = readHeaderId(full);
if (!id) continue;
let mtime = 0; let mtime = 0;
try { try {
mtime = statSync(full).mtimeMs; mtime = statSync(full).mtimeMs;
} catch { } catch {
/* ignore */ continue;
} }
candidates.push({ path: full, id, mtime }); 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). * 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[] { export function getNewEntries(sessionFile: string, afterLine: number): SessionEntry[] {
const raw = readFileSync(sessionFile, "utf8"); const raw = readFileSync(sessionFile, "utf8");
const lines = raw.split("\n").filter((line) => line.trim()); const lines = raw.split("\n").filter((line) => line.trim());
+32 -1
View File
@@ -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 assert from "node:assert/strict";
import { mkdtempSync, writeFileSync, readFileSync, mkdirSync, rmSync, existsSync } from "node:fs"; import { mkdtempSync, writeFileSync, readFileSync, mkdirSync, rmSync, existsSync } from "node:fs";
import { join } from "node:path"; import { join } from "node:path";
@@ -10,11 +10,13 @@ import * as subagentsModule from "../pi-extension/subagents/index.ts";
import { import {
getLeafId, getLeafId,
getNewEntries, getNewEntries,
countSessionEntryLines,
getSessionId, getSessionId,
readSubagentLoadout, readSubagentLoadout,
writeSubagentLoadout, writeSubagentLoadout,
loadoutSidecarPath, loadoutSidecarPath,
type SubagentLoadout, type SubagentLoadout,
resetSessionIndexCache,
resolveSessionFileById, resolveSessionFileById,
findLastAssistantMessage, findLastAssistantMessage,
appendBranchSummary, appendBranchSummary,
@@ -265,6 +267,19 @@ describe("session.ts", () => {
const entries = getNewEntries(file, 2); const entries = getNewEntries(file, 2);
assert.equal(entries.length, 0); 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", () => { describe("getSessionId / resolveSessionFileById", () => {
@@ -274,6 +289,12 @@ describe("session.ts", () => {
return p; 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", () => { it("reads the header id from a session file", () => {
const file = createSessionFile(dir, [SESSION_HEADER, MODEL_CHANGE, USER_MSG]); const file = createSessionFile(dir, [SESSION_HEADER, MODEL_CHANGE, USER_MSG]);
assert.equal(getSessionId(file), "sess-001"); assert.equal(getSessionId(file), "sess-001");
@@ -299,6 +320,16 @@ describe("session.ts", () => {
writeSession(dir, "c.jsonl", "abc"); writeSession(dir, "c.jsonl", "abc");
assert.equal(resolveSessionFileById("zzz", dir), null); 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", () => { describe("subagent loadout snapshot", () => {