perf(ai-vault): bound per-row bookkeeping in unlimited session scans (#20291)

* perf: deduplicate unlimited vault scans once * perf: release discarded vault aliases during unlimited scans * fix: bound per-session bookkeeping in unlimited vault scans - Drop the per-session alias-key string, wrapper object and positions array the accumulator retained for every parsed row; index winning positions by the row's own sessionId instead (~430 B -> ~45 B per session at 50k rows). - Add a --expose-gc retention test asserting a 50k mostly-unique load-all corpus stays under 128 B of bookkeeping per session while matching dedupeCodexSessionsBySessionId exactly. * perf(ai-vault): bound per-row bookkeeping in CodexSessionCollection Key winners by the row's own sessionId string so an unlimited scan retains no alias-key string per live row (301 -> ~115 B/row measured over 50k rows), and split into a per-alias-key map only for the rare id that spans several hosts, namespaces, or rollout names, so admission stays O(1). Fold the PR's CodexSessionAccumulator into the collection main already routes every scan through, and rerun its scanner-level tests against that single class.
2026-09-22 16:02:32 +00:00 · 2026-09-12 22:10:30 -07:00
parent e86cba888b
commit f2b6434fe6
2 changed files with 313 additions and 26 deletions
@@ -252,12 +252,17 @@ export function dedupeCodexSessionsBySessionId(
  })
 }

+type CodexSessionWinner = { session: AiVaultSession; indices: number | number[] }
+
 /** Scan-local accumulation; parsed rows must not be mutated after admission. */
 export class CodexSessionCollection {
  private readonly sessions = new Map<number, AiVaultSession>()
-  private readonly bestByKey = new Map<
+  // Keyed by the row's own sessionId string, so an unlimited scan retains no
+  // alias key per live row; a per-alias-key map appears only for the rare id
+  // that spans several hosts, namespaces, or rollout names.
+  private readonly winnersBySessionId = new Map<
    string,
-    { session: AiVaultSession; indices: number | number[] }
+    CodexSessionWinner | Map<string, CodexSessionWinner>
  >()
  private nextIndex = 0

@@ -272,33 +277,70 @@ export class CodexSessionCollection {
  add(session: AiVaultSession): void {
    const key = codexSessionAliasKey(session)
    const index = this.nextIndex++
-    if (key) {
-      const best = this.bestByKey.get(key)
-      if (best?.session === session) {
-        // The batch filter retains every occurrence of the winning object.
-        if (typeof best.indices === 'number') {
-          best.indices = [best.indices, index]
-        } else {
-          best.indices.push(index)
-        }
-      } else {
-        if (best) {
-          if (!codexSessionAliasBeats(session, best.session)) {
-            return
-          }
-          if (typeof best.indices === 'number') {
-            this.sessions.delete(best.indices)
-          } else {
-            for (const previousIndex of best.indices) {
-              this.sessions.delete(previousIndex)
-            }
-          }
-        }
-        this.bestByKey.set(key, { session, indices: index })
-      }
+    if (key && !this.admit(session, key, index)) {
+      return
    }
    this.sessions.set(index, session)
  }
+
+  /** Whether the row is retained; a losing alias is dropped. */
+  private admit(session: AiVaultSession, key: string, index: number): boolean {
+    const bucket = this.winnersBySessionId.get(session.sessionId)
+    if (bucket instanceof Map) {
+      const winner = this.contest(bucket.get(key), session, index)
+      if (winner) {
+        bucket.set(key, winner)
+      }
+      return winner !== null
+    }
+    const bucketKey = bucket && codexSessionAliasKey(bucket.session)
+    if (bucket && bucketKey && bucketKey !== key) {
+      this.winnersBySessionId.set(
+        session.sessionId,
+        new Map([
+          [bucketKey, bucket],
+          [key, { session, indices: index }]
+        ])
+      )
+      return true
+    }
+    const winner = this.contest(bucket, session, index)
+    if (winner) {
+      this.winnersBySessionId.set(session.sessionId, winner)
+    }
+    return winner !== null
+  }
+
+  /** The alias key's winner after this row, or null when the row loses. */
+  private contest(
+    best: CodexSessionWinner | undefined,
+    session: AiVaultSession,
+    index: number
+  ): CodexSessionWinner | null {
+    if (!best) {
+      return { session, indices: index }
+    }
+    if (best.session === session) {
+      // The batch filter retains every occurrence of the winning object.
+      if (typeof best.indices === 'number') {
+        best.indices = [best.indices, index]
+      } else {
+        best.indices.push(index)
+      }
+      return best
+    }
+    if (!codexSessionAliasBeats(session, best.session)) {
+      return null
+    }
+    if (typeof best.indices === 'number') {
+      this.sessions.delete(best.indices)
+    } else {
+      for (const previousIndex of best.indices) {
+        this.sessions.delete(previousIndex)
+      }
+    }
+    return { session, indices: index }
+  }
 }

 function codexSessionAliasKey(session: AiVaultSession): string | null {
@@ -0,0 +1,245 @@
+import { beforeEach, expect, it, vi } from 'vitest'
+import type * as CodexDedup from './codex-session-root-dedup'
+import type { AiVaultSession } from '../../shared/ai-vault-types'
+
+const fixture = vi.hoisted((): { sessions: AiVaultSession[]; visits: number } => ({
+  sessions: [],
+  visits: 0
+}))
+vi.mock('./session-scanner-source-discovery', () => ({
+  discoverAiVaultSessionSources: async () => [],
+  DEFAULT_CODEX_HOME_DIR: '/fixture'
+}))
+vi.mock('./session-scanner-candidates', () => ({
+  sessionCandidatesFromDiscoveries: async () => candidates()
+}))
+vi.mock('./session-parse-cache-persistence', () => ({
+  ensureSessionParseCacheLoaded: async () => {},
+  scheduleSessionParseCachePersist: () => {}
+}))
+vi.mock('./session-scanner-parse-cache', () => ({
+  createSessionParseStats: () => ({
+    reused: 0,
+    incremental: 0,
+    fullParses: 0,
+    earlyStopped: 0,
+    bytesRead: 0
+  }),
+  parseAgentSessionFileCached: async (candidate: { session: AiVaultSession }) => candidate.session
+}))
+vi.mock('./remote-session-scanner-sources', () => ({ remoteSessionSources: () => [{}] }))
+vi.mock('./remote-session-scanner-discovery', () => ({
+  discoverRemoteSourceCandidates: async () => candidates()
+}))
+vi.mock('./remote-session-parse-cache', () => ({
+  remoteSessionParseHostKey: () => 'fixture',
+  parseRemoteSessionFileCached: async ({ candidate }: { candidate: { session: AiVaultSession } }) =>
+    candidate.session
+}))
+vi.mock('./codex-session-root-dedup', async (original) => {
+  const actual = await original<typeof CodexDedup>()
+  return {
+    ...actual,
+    dedupeCodexSessionsBySessionId: (sessions: AiVaultSession[]) => {
+      fixture.visits += sessions.length
+      return actual.dedupeCodexSessionsBySessionId(sessions)
+    }
+  }
+})
+
+import { scanAiVaultSessions } from './session-scanner'
+import { scanRemoteAiVaultSessions } from './remote-session-scanner'
+import { CodexSessionCollection, dedupeCodexSessionsBySessionId } from './codex-session-root-dedup'
+
+function candidates() {
+  return fixture.sessions.map((session) => ({
+    agent: session.agent,
+    file: { path: session.filePath, mtimeMs: Date.parse(session.modifiedAt) },
+    codexHome: session.codexHome,
+    session,
+    source: { agent: session.agent }
+  }))
+}
+
+function session(index: number): AiVaultSession {
+  return {
+    id: String(index),
+    executionHostId: 'local',
+    agent: 'codex',
+    sessionId: String(index),
+    title: 'fixture',
+    cwd: '/fixture',
+    branch: null,
+    model: null,
+    filePath: `/fixture/rollout-${index}.jsonl`,
+    codexHome: null,
+    createdAt: null,
+    updatedAt: null,
+    modifiedAt: '2026-01-01T00:00:00.000Z',
+    messageCount: 1,
+    totalTokens: 0,
+    previewMessages: [],
+    queuedMessageCount: 0,
+    subagentTranscriptCount: 0,
+    resumeCommand: '',
+    subagent: null
+  }
+}
+
+beforeEach(() => {
+  fixture.sessions = []
+  fixture.visits = 0
+})
+
+for (const host of ['local', 'remote'] as const) {
+  const scan = (unlimited: boolean, limit?: number) =>
+    host === 'local'
+      ? scanAiVaultSessions({ unlimited, limit })
+      : scanRemoteAiVaultSessions({
+          unlimited,
+          limit,
+          provider: { readDir: vi.fn(), readFile: vi.fn(), stat: vi.fn() },
+          executionHostId: 'local',
+          remoteHome: '/fixture',
+          hostPlatform: {
+            relayPlatform: 'linux-x64',
+            os: 'linux',
+            arch: 'x64',
+            pathFlavor: 'posix',
+            commandDialect: 'posix',
+            pathSeparator: '/',
+            pathDelimiter: ':'
+          }
+        })
+
+  it(`${host}: load-all processes deduplication linearly and retains late canonical aliases`, async () => {
+    fixture.sessions = Array.from({ length: 10000 }, (_, i) => session(i))
+    fixture.sessions[0] = {
+      ...fixture.sessions[0]!,
+      codexHome: '/custom',
+      filePath: '/custom/rollout-0.jsonl'
+    }
+    fixture.sessions.push(session(0))
+    const expected = dedupeCodexSessionsBySessionId(fixture.sessions)
+    fixture.visits = 0
+    const started = performance.now()
+    const result = await scan(true)
+    process.stdout.write(
+      `${JSON.stringify({ host, candidates: fixture.sessions.length, scanMs: performance.now() - started, dedupVisits: fixture.visits })}\n`
+    )
+    expect(result.issues).toEqual([])
+    expect(result.sessions).toEqual(expected)
+    expect(fixture.visits).toBeLessThanOrEqual(fixture.sessions.length * 2)
+  }, 30000)
+
+  it(`${host}: capped scans still fill the unique-session budget`, async () => {
+    fixture.sessions = [
+      session(0),
+      ...Array.from({ length: 8 }, () => ({
+        ...session(0),
+        filePath: '/custom/rollout-0.jsonl',
+        codexHome: '/custom'
+      })),
+      ...Array.from({ length: 10 }, (_, i) => session(i + 1))
+    ]
+    const result = await scan(false, 10)
+    expect(result.sessions).toHaveLength(10)
+    expect(new Set(result.sessions.map((row) => row.sessionId)).size).toBe(10)
+  })
+}
+
+it('incremental canonical selection preserves winner occurrence order, ties and repeated references', () => {
+  const collection = new CodexSessionCollection()
+  const same = session(0)
+  const rows: AiVaultSession[] = []
+  const variants: AiVaultSession[] = [
+    same,
+    same,
+    { ...same, codexHome: '/custom', filePath: '/custom/rollout-0.jsonl' },
+    { ...same, agent: 'claude' as const },
+    { ...same, executionHostId: 'ssh:fixture' },
+    { ...same, modifiedAt: '2026-02-01T00:00:00.000Z' },
+    { ...same, filePath: '/aaa/rollout-0.jsonl' },
+    { ...same, filePath: '/fixture/rollout-0-fork.jsonl', modifiedAt: 'invalid' },
+    session(1)
+  ]
+  let seed = 42
+  for (let index = 0; index < 2000; index++) {
+    seed = (Math.imul(seed, 1664525) + 1013904223) >>> 0
+    const row = variants[seed % variants.length]!
+    rows.push(row)
+    collection.add(row)
+    expect([...collection.values()]).toEqual(dedupeCodexSessionsBySessionId(rows))
+  }
+})
+
+it('retains only canonical rows during duplicate-heavy load-all scans', () => {
+  const collection = new CodexSessionCollection()
+  for (let index = 0; index < 10000; index++) {
+    const row = session(index % 100)
+    collection.add({
+      ...row,
+      codexHome: '/custom',
+      filePath: `/custom/rollout-${index % 100}.jsonl`
+    })
+    expect(collection.size).toBeLessThanOrEqual(100)
+  }
+  for (let index = 0; index < 100; index++) {
+    collection.add(session(index))
+  }
+  expect(collection.size).toBe(100)
+  expect([...collection.values()].every((row) => row.codexHome === null)).toBe(true)
+})
+
+it('admits rows sharing one session id across rollout names without rescanning', () => {
+  const count = 4000
+  let pathReads = 0
+  const collection = new CodexSessionCollection()
+  for (let index = 0; index < count; index++) {
+    const row = { ...session(index), sessionId: 'shared' }
+    collection.add({
+      ...row,
+      get filePath() {
+        pathReads++
+        return row.filePath
+      }
+    })
+  }
+  expect(collection.size).toBe(count)
+  expect(pathReads).toBeLessThanOrEqual(count * 4)
+})
+
+it('bounds per-session bookkeeping for a large mostly-unique load-all corpus', () => {
+  const gc = globalThis.gc
+  if (!gc) {
+    throw new Error('Retention test requires --expose-gc (config/vitest.config.ts)')
+  }
+  const heapUsed = () => {
+    gc()
+    gc()
+    return process.memoryUsage().heapUsed
+  }
+  const count = 50000
+  // Why pre-build: the corpus itself must not count against the collection.
+  const corpus = Array.from({ length: count }, (_, index) =>
+    index % 100 === 99
+      ? {
+          ...session(index - 1),
+          codexHome: '/custom',
+          filePath: `/custom/rollout-${index - 1}.jsonl`
+        }
+      : session(index)
+  )
+  const expected = dedupeCodexSessionsBySessionId(corpus)
+  const before = heapUsed()
+  const collection = new CodexSessionCollection()
+  for (const row of corpus) {
+    collection.add(row)
+  }
+  const retained = heapUsed() - before
+
+  expect([...collection.values()]).toEqual(expected)
+  // Two map entries plus one winner record per live row measure ~115 B; an
+  // alias-key string per live row measured ~300 B.
+  expect(retained).toBeLessThan(count * 160)
+})