From 36e1291aef585d4b774791547dcba56c4d27fbb4 Mon Sep 17 00:00:00 2001 From: Brennan Benson <79079362+brennanb2025@users.noreply.github.com> Date: Tue, 13 Oct 2026 17:17:56 -0700 Subject: [PATCH] perf(session-history): end a scan only when it stops making progress (#28434) * perf(session-history): end a scan only when it stops making progress A cold Session History scan of a large history was killed 130 seconds in every time, however healthy it was, and every Refresh died at the same place. The scanner child now reports progress for a running scan, at most once a second and only when real work advanced: a directory listed, a batch of files statted, a transcript chunk read, a candidate file parsed. Progress is keyed to the scan's own cancel signal, so the search index or an interactive read in the same child never extends it. The parent re-arms the 130 second window only for an active, uncancelled scan whose counter grew. When a request deadline does fire, the child gets the shutdown frame and the existing two-second grace to abort and save its parse cache before the kill. Exit, disconnect, malformed-message and cancel-timeout faults still kill at once. Startup, cancel and interactive bounds are unchanged. * fix(session-history): keep the deadline module off child_process and let the orcad smoke skip progress frames --- .../scripts/session-scanner-service-smoke.mjs | 3 +- .../ai-vault/ai-vault-scan-progress.test.ts | 134 +++++++++++ src/main/ai-vault/ai-vault-scan-progress.ts | 41 ++++ .../ai-vault/session-scanner-discovery.ts | 3 + .../ai-vault/session-scanner-jsonl-reader.ts | 2 + .../ai-vault/session-scanner-parse-cache.ts | 5 +- .../session-scanner-service-client-state.ts | 69 ++++-- .../session-scanner-service-client.ts | 39 +-- .../session-scanner-service-entry.test.ts | 39 ++- .../ai-vault/session-scanner-service-entry.ts | 18 +- .../session-scanner-service-protocol.ts | 10 + ...n-scanner-service-request-deadline.test.ts | 227 ++++++++++++++++++ ...ession-scanner-service-request-deadline.ts | 63 +++++ 13 files changed, 606 insertions(+), 47 deletions(-) create mode 100644 src/main/ai-vault/ai-vault-scan-progress.test.ts create mode 100644 src/main/ai-vault/ai-vault-scan-progress.ts create mode 100644 src/main/ai-vault/session-scanner-service-request-deadline.test.ts create mode 100644 src/main/ai-vault/session-scanner-service-request-deadline.ts diff --git a/config/scripts/session-scanner-service-smoke.mjs b/config/scripts/session-scanner-service-smoke.mjs index 146a88ad17b..f67766bd3d5 100644 --- a/config/scripts/session-scanner-service-smoke.mjs +++ b/config/scripts/session-scanner-service-smoke.mjs @@ -30,7 +30,8 @@ child.on('message', (message) => { child.send({ type: 'request', id: 1, operation: 'scan', options: { limit: 5 } }) return } - if (!message || message.id !== 1) { + // Progress frames share the request id; only the result or error answers it. + if (!message || message.id !== 1 || (message.type !== 'result' && message.type !== 'error')) { return } const sessions = message.type === 'result' ? message.value.result.sessions : [] diff --git a/src/main/ai-vault/ai-vault-scan-progress.test.ts b/src/main/ai-vault/ai-vault-scan-progress.test.ts new file mode 100644 index 00000000000..01249bd737c --- /dev/null +++ b/src/main/ai-vault/ai-vault-scan-progress.test.ts @@ -0,0 +1,134 @@ +import { mkdtemp, rm, writeFile } from 'node:fs/promises' +import { tmpdir } from 'node:os' +import { join } from 'node:path' +import { afterEach, describe, expect, it } from 'vitest' +import { + AI_VAULT_SCAN_PROGRESS_INTERVAL_MS, + noteAiVaultScanProgress, + trackAiVaultScanProgress +} from './ai-vault-scan-progress' +import { scanAiVaultSessions } from './session-scanner' +import { consumeCompleteJsonlLines } from './session-scanner-jsonl-reader' +import { resetSessionParseCacheForTests } from './session-scanner-parse-cache' +import { isolatedScanRoots, writeJsonlFile } from './session-scanner-test-fixtures' + +let tempRoots: string[] = [] + +afterEach(async () => { + resetSessionParseCacheForTests() + await Promise.all(tempRoots.map((root) => rm(root, { recursive: true, force: true }))) + tempRoots = [] +}) + +async function tempRoot(): Promise { + const root = await mkdtemp(join(tmpdir(), 'orca-scan-progress-')) + tempRoots.push(root) + return root +} + +/** A clock that has always moved a full interval, so every advance is reported. */ +function steppingClock(): () => number { + let now = 0 + return () => (now += AI_VAULT_SCAN_PROGRESS_INTERVAL_MS) +} + +describe('scan progress reporting', () => { + it('reports at most once per interval, and only after the counter moved', () => { + let now = 0 + const reports: number[] = [] + const controller = new AbortController() + const stop = trackAiVaultScanProgress( + controller.signal, + (units) => reports.push(units), + () => now + ) + + noteAiVaultScanProgress(controller.signal) + noteAiVaultScanProgress(controller.signal) + now += AI_VAULT_SCAN_PROGRESS_INTERVAL_MS - 1 + noteAiVaultScanProgress(controller.signal) + expect(reports).toEqual([1]) + + // Time alone is not progress: nothing is sent until work advances again. + now += 60_000 + expect(reports).toEqual([1]) + noteAiVaultScanProgress(controller.signal) + expect(reports).toEqual([1, 4]) + + stop() + now += 60_000 + noteAiVaultScanProgress(controller.signal) + expect(reports).toEqual([1, 4]) + }) + + it('never credits work done under another signal', () => { + const reports: number[] = [] + const scan = new AbortController() + trackAiVaultScanProgress(scan.signal, (units) => reports.push(units), steppingClock()) + + noteAiVaultScanProgress(new AbortController().signal) + noteAiVaultScanProgress(undefined) + expect(reports).toEqual([]) + }) + + it('counts every chunk a long transcript read consumes', async () => { + const root = await tempRoot() + const path = join(root, 'long.jsonl') + const line = JSON.stringify({ text: 'x'.repeat(1_000) }) + await writeFile(path, `${Array.from({ length: 400 }, () => line).join('\n')}\n`) + const reports: number[] = [] + const controller = new AbortController() + trackAiVaultScanProgress(controller.signal, (units) => reports.push(units), steppingClock()) + + await consumeCompleteJsonlLines({ + path, + start: 0, + onLine: () => undefined, + signal: controller.signal + }) + + // ~400 KB arrives in several stream chunks; one unit per chunk. + expect(reports.length).toBeGreaterThan(3) + expect(reports).toEqual(reports.map((_, index) => index + 1)) + }) + + it('reports discovery and parsing for its own scan, never for an unrelated one', async () => { + const root = await tempRoot() + const roots = isolatedScanRoots(root) + for (const name of ['a', 'b', 'c']) { + await writeJsonlFile(join(roots.claudeProjectsDir, 'project', `${name}.jsonl`), [ + { + type: 'user', + sessionId: name, + timestamp: '2026-05-01T10:00:00.000Z', + cwd: '/tmp/claude', + message: { role: 'user', content: `ask ${name}` } + } + ]) + } + const reports: number[] = [] + const tracked = new AbortController() + trackAiVaultScanProgress(tracked.signal, (units) => reports.push(units), steppingClock()) + + // Same readers, another owner's signal (the index loop or an interactive read). + await scanAiVaultSessions({ + ...roots, + platform: 'darwin', + limit: 20, + signal: new AbortController().signal + }) + expect(reports).toEqual([]) + + resetSessionParseCacheForTests() + const result = await scanAiVaultSessions({ + ...roots, + platform: 'darwin', + limit: 20, + signal: tracked.signal + }) + expect(result.sessions).toHaveLength(3) + // Two directory listings, a stat batch, then a chunk and a parse per file. + expect(reports.length).toBeGreaterThanOrEqual(2 + 1 + 3 * 2) + expect(reports).toEqual(reports.map((_, index) => index + 1)) + }) +}) diff --git a/src/main/ai-vault/ai-vault-scan-progress.ts b/src/main/ai-vault/ai-vault-scan-progress.ts new file mode 100644 index 00000000000..16551ce0a3d --- /dev/null +++ b/src/main/ai-vault/ai-vault-scan-progress.ts @@ -0,0 +1,41 @@ +// Why: the parent's scan deadline is a no-progress window, so the child has to +// say when a scan's work advanced. Keyed by the request's cancel signal, which +// every work site already receives, so the index loop or an interactive read +// running through the same readers can never extend a scan's deadline. +const reporters = new WeakMap void>() + +// Rate-limited because a cold scan consumes thousands of chunks and stats, and +// the parent only reads the counter as "still moving" inside a 130 s window. +export const AI_VAULT_SCAN_PROGRESS_INTERVAL_MS = 1_000 + +/** + * Reports advancing work for the scan that owns `signal`, at most once per + * interval and only when the counter moved, until the returned stop is called. + */ +export function trackAiVaultScanProgress( + signal: AbortSignal, + report: (units: number) => void, + now: () => number = () => performance.now() +): () => void { + let units = 0 + let lastReportedAt = Number.NEGATIVE_INFINITY + reporters.set(signal, () => { + units++ + const at = now() + if (at - lastReportedAt < AI_VAULT_SCAN_PROGRESS_INTERVAL_MS) { + return + } + lastReportedAt = at + report(units) + }) + return () => { + reporters.delete(signal) + } +} + +/** Counts one unit of real work for the scan that owns `signal`; a no-op for any other work. */ +export function noteAiVaultScanProgress(signal: AbortSignal | undefined): void { + if (signal) { + reporters.get(signal)?.() + } +} diff --git a/src/main/ai-vault/session-scanner-discovery.ts b/src/main/ai-vault/session-scanner-discovery.ts index aa766c85647..24ba97ebf4c 100644 --- a/src/main/ai-vault/session-scanner-discovery.ts +++ b/src/main/ai-vault/session-scanner-discovery.ts @@ -12,6 +12,7 @@ import { recordSessionScanIssue } from './session-scan-issues' import type { FileWithMtime, SessionFileDiscovery } from './session-scanner-types' import { errorMessage } from './session-scanner-values' import { throwIfAiVaultScanCancelled } from './ai-vault-scan-cancellation' +import { noteAiVaultScanProgress } from './ai-vault-scan-progress' const NATIVE_DISCOVERY_CONCURRENCY = 4 const NATIVE_DISCOVERY_BATCH_SIZE = 16 @@ -50,6 +51,7 @@ export async function discoverFiles(args: { function consume(observations: readonly SessionFileObservation[]): void { // A stat the cancel cut short is not an unreadable file; never record it. throwIfAiVaultScanCancelled(args.signal) + noteAiVaultScanProgress(args.signal) for (const observation of observations) { if ('issue' in observation) { recordSessionScanIssue(args.issues, { agent: args.agent, ...observation.issue }) @@ -274,6 +276,7 @@ export async function forEachSessionFile( } return } + noteAiVaultScanProgress(options.signal) for (const entry of entries) { options.signal?.throwIfAborted() diff --git a/src/main/ai-vault/session-scanner-jsonl-reader.ts b/src/main/ai-vault/session-scanner-jsonl-reader.ts index ff0a9bf4dbc..68db27ab5cc 100644 --- a/src/main/ai-vault/session-scanner-jsonl-reader.ts +++ b/src/main/ai-vault/session-scanner-jsonl-reader.ts @@ -1,5 +1,6 @@ import { openTranscriptReadStream } from '../native-chat/wsl-transcript-fs-access' import { throwIfAiVaultScanCancelled } from './ai-vault-scan-cancellation' +import { noteAiVaultScanProgress } from './ai-vault-scan-progress' import { MAX_SESSION_TRANSCRIPT_RECORD_BYTES, type SkippedTranscriptRecord @@ -57,6 +58,7 @@ export async function consumeCompleteJsonlLines(args: { // Per chunk, so a cancel stops a giant (or over-budget) record within one // chunk; leaving the loop destroys the stream and releases its handle. throwIfAiVaultScanCancelled(args.signal) + noteAiVaultScanProgress(args.signal) bytesRead += chunk.length const chunkEnd = chunkStart + chunk.length let lineStart = 0 diff --git a/src/main/ai-vault/session-scanner-parse-cache.ts b/src/main/ai-vault/session-scanner-parse-cache.ts index 1f6b029ca36..0598204c4be 100644 --- a/src/main/ai-vault/session-scanner-parse-cache.ts +++ b/src/main/ai-vault/session-scanner-parse-cache.ts @@ -1,6 +1,7 @@ import { createQoderSessionResumeState } from './session-scanner-qoder-parser' import type { AiVaultSession } from '../../shared/ai-vault-types' import { inSessionParseFileLane } from './session-parse-file-lane' +import { noteAiVaultScanProgress } from './ai-vault-scan-progress' import { createAntigravitySessionResumeState } from './session-scanner-antigravity-parser' import { createCodexSessionResumeState } from './session-scanner-codex-parser' import { createDroidSessionResumeState } from './session-scanner-droid-parser' @@ -128,9 +129,11 @@ export async function parseAgentSessionFileCached( ): Promise { // The whole lookup-read-store sequence runs in the lane: a concurrent parse of // the same path shares this entry's resume point and its message channel. - return inSessionParseFileLane(candidate.file.path, () => + const session = await inSessionParseFileLane(candidate.file.path, () => parseCachedInLane(candidate, platform, stats, requireRead, signal) ) + noteAiVaultScanProgress(signal) + return session } /** diff --git a/src/main/ai-vault/session-scanner-service-client-state.ts b/src/main/ai-vault/session-scanner-service-client-state.ts index cc5b514988b..f758595d560 100644 --- a/src/main/ai-vault/session-scanner-service-client-state.ts +++ b/src/main/ai-vault/session-scanner-service-client-state.ts @@ -3,6 +3,7 @@ import type { ChildProcess } from 'node:child_process' import { createAiVaultScanCancelledError } from './ai-vault-scan-cancellation' import { AI_VAULT_SERVICE_PROTOCOL_VERSION, + aiVaultServiceLane, type AiVaultServiceInit, type AiVaultServiceLane, type AiVaultServiceRequest, @@ -10,8 +11,6 @@ import { } from './session-scanner-service-protocol' export const AI_VAULT_SERVICE_READY_TIMEOUT_MS = 5_000 -export const AI_VAULT_SERVICE_SCAN_TIMEOUT_MS = 130_000 -export const AI_VAULT_SERVICE_INTERACTIVE_TIMEOUT_MS = 15_000 export const AI_VAULT_SERVICE_MAX_CALLS = 16 export const AI_VAULT_SERVICE_IDLE_TIMEOUT_MS = 10 * 60_000 export const AI_VAULT_SERVICE_SHUTDOWN_TIMEOUT_MS = 2_000 @@ -124,6 +123,48 @@ export function retireAiVaultServiceChild(child: ChildProcess): void { child.unref() } +/** + * `retire`: a request deadline sends the shutdown frame and allows the usual + * grace to abort and save the parse cache. A broken channel or protocol is + * killed at once. + */ +export function stopAiVaultServiceChild(child: ChildProcess, retire: boolean): void { + if (retire) { + retireAiVaultServiceChild(child) + return + } + child.removeAllListeners() + child.kill() +} + +/** A queued call, already listening for its caller's abort. */ +export function createAiVaultServicePendingCall( + request: AiVaultServiceRequest, + signal: AbortSignal | undefined, + resolve: (value: unknown) => void, + reject: (error: Error) => void, + cancel: (call: AiVaultServicePendingCall) => void +): AiVaultServicePendingCall { + const call: AiVaultServicePendingCall = { + request, + lane: aiVaultServiceLane(request.operation), + signal, + resolve, + reject, + timer: null, + onAbort: null, + cancelled: false, + sent: false, + startRetried: false, + progressUnits: 0 + } + if (signal) { + call.onAbort = () => cancel(call) + signal.addEventListener('abort', call.onAbort, { once: true }) + } + return call +} + function armAiVaultServiceCancellationTimeout( call: AiVaultServicePendingCall, onExpired: () => void @@ -249,6 +290,8 @@ export type AiVaultServicePendingCall = { /** Whether the child received the request; an unsent call gets no reply. */ sent: boolean startRetried: boolean + /** Highest progress counter the child reported for this scan. */ + progressUnits: number } export type AiVaultServiceReadyWaiter = { @@ -306,25 +349,3 @@ export class AiVaultServiceSessionSearchHold { return this.enabled } } - -/** Starts the request deadline only once the child is ready to receive it. */ -export function sendAiVaultServiceCall( - child: ChildProcess, - call: AiVaultServicePendingCall, - isActive: () => boolean, - onFault: (error: Error) => void -): void { - if (call.cancelled || !isActive()) { - return - } - const timeoutMs = - call.request.operation === 'scan' - ? AI_VAULT_SERVICE_SCAN_TIMEOUT_MS - : AI_VAULT_SERVICE_INTERACTIVE_TIMEOUT_MS - call.timer = setTimeout(() => { - onFault(new Error(`AI Vault service timed out after ${timeoutMs}ms.`)) - }, timeoutMs) - call.timer.unref?.() - call.sent = true - child.send(call.request) -} diff --git a/src/main/ai-vault/session-scanner-service-client.ts b/src/main/ai-vault/session-scanner-service-client.ts index e547ee1dee7..5367c6fbc0d 100644 --- a/src/main/ai-vault/session-scanner-service-client.ts +++ b/src/main/ai-vault/session-scanner-service-client.ts @@ -11,18 +11,22 @@ import { attachAiVaultServiceChild, cancelAiVaultServiceCall, clearAiVaultServiceCall, + createAiVaultServicePendingCall, createAiVaultServiceReadyWaiter, rejectAiVaultServiceCall, requeueOrRejectAiVaultServiceStart, retireAiVaultServiceChild, - sendAiVaultServiceCall, + stopAiVaultServiceChild, type AiVaultServiceClientOptions, type AiVaultServicePendingCall, type AiVaultServiceReadyWaiter } from './session-scanner-service-client-state' +import { + noteAiVaultServiceScanProgress, + sendAiVaultServiceCall +} from './session-scanner-service-request-deadline' import { AiVaultServiceRestartPolicy } from './session-scanner-service-restart-policy' import { - aiVaultServiceLane, isAiVaultServiceChildMessage, type AiVaultSessionSearchInit, type AiVaultServiceRequest, @@ -56,22 +60,14 @@ export class AiVaultScannerServiceClient { } const request = { ...body, id: this.nextId++ } as AiVaultServiceRequest return new Promise((resolve, reject) => { - const call: AiVaultServicePendingCall = { + const call = createAiVaultServicePendingCall( request, - lane: aiVaultServiceLane(request.operation), signal, - resolve: resolve as (value: unknown) => void, + // oxlint-disable-next-line typescript/consistent-type-assertions -- SAFETY: callers choose T as the result value of the operation they send, which is what the child answers with. + resolve as (value: unknown) => void, reject, - timer: null, - onAbort: null, - cancelled: false, - sent: false, - startRetried: false - } - if (signal) { - call.onAbort = () => this.cancel(call) - signal.addEventListener('abort', call.onAbort, { once: true }) - } + (aborted) => this.cancel(aborted) + ) this.queue.push(call) this.idleRetirement.clear() this.pump() @@ -145,7 +141,7 @@ export class AiVaultScannerServiceClient { child, call, () => this.active.get(call.lane) === call, - (error) => this.onFault(error) + this.onDeadline ), (error: Error) => { if (this.active.get(lane) !== call) { @@ -254,6 +250,10 @@ export class AiVaultScannerServiceClient { if (!call) { return } + if (message.type === 'progress') { + noteAiVaultServiceScanProgress(call, message.units, this.onDeadline) + return + } this.active.delete(call.lane) clearAiVaultServiceCall(call) if (!call.cancelled) { @@ -276,14 +276,15 @@ export class AiVaultScannerServiceClient { }) } - private onFault(error: Error): void { + private readonly onDeadline = (error: Error): void => this.onFault(error, true) + + private onFault(error: Error, retire = false): void { const child = this.child if (!child) { return } this.child = null - child.removeAllListeners() - child.kill() + stopAiVaultServiceChild(child, retire) if (this.readyWaiter) { clearTimeout(this.readyWaiter.timer) this.readyWaiter.reject(error) diff --git a/src/main/ai-vault/session-scanner-service-entry.test.ts b/src/main/ai-vault/session-scanner-service-entry.test.ts index 70dae0d764b..458ba6f5e0c 100644 --- a/src/main/ai-vault/session-scanner-service-entry.test.ts +++ b/src/main/ai-vault/session-scanner-service-entry.test.ts @@ -1,4 +1,5 @@ import { beforeAll, beforeEach, describe, expect, it, vi } from 'vitest' +import { noteAiVaultScanProgress } from './ai-vault-scan-progress' import { AI_VAULT_SERVICE_PROTOCOL_VERSION } from './session-scanner-service-protocol' const invalidateSessionParseCacheEntry = vi.hoisted(() => vi.fn()) @@ -18,7 +19,7 @@ vi.mock('./session-subagent-reader', () => ({ listLocalAiVaultSubagentSessions: vi.fn(() => Promise.resolve({ sessions: [], issues: [] })) })) -const sent: { type: string; id?: number }[] = [] +const sent: { type: string; id?: number; units?: number }[] = [] function emit(message: unknown): void { process.emit('message', message as never, undefined as never) @@ -91,4 +92,40 @@ describe('AI Vault service entry cache invalidation', () => { await runScan(12) expect(invalidationsFor('/transcripts/b.jsonl')).toBe(3) }) + + it('sends rate-limited progress under the running scan id, and none once it settles', async () => { + let now = 0 + const clock = vi.spyOn(performance, 'now').mockImplementation(() => now) + let signal: AbortSignal | undefined + let releaseScan: (() => void) | undefined + scanAiVaultSessions.mockImplementationOnce( + (options: { signal: AbortSignal }) => + new Promise((resolve) => { + signal = options.signal + releaseScan = () => resolve({ sessions: [], issues: [], scannedAt: '2026-08-10' }) + }) + ) + try { + emit({ type: 'request', id: 20, operation: 'scan', options: {} }) + await vi.waitFor(() => expect(signal).toBeDefined()) + + noteAiVaultScanProgress(signal) + noteAiVaultScanProgress(signal) + now += 1_000 + noteAiVaultScanProgress(signal) + const progress = (): typeof sent => sent.filter((message) => message.type === 'progress') + expect(progress()).toEqual([ + { type: 'progress', id: 20, units: 1 }, + { type: 'progress', id: 20, units: 3 } + ]) + + releaseScan?.() + await vi.waitFor(() => expect(sent.some((message) => message.type === 'result')).toBe(true)) + now += 10_000 + noteAiVaultScanProgress(signal) + expect(progress()).toHaveLength(2) + } finally { + clock.mockRestore() + } + }) }) diff --git a/src/main/ai-vault/session-scanner-service-entry.ts b/src/main/ai-vault/session-scanner-service-entry.ts index d1a9da14a61..b1069757ae9 100644 --- a/src/main/ai-vault/session-scanner-service-entry.ts +++ b/src/main/ai-vault/session-scanner-service-entry.ts @@ -1,5 +1,7 @@ import { requestSessionSearchRoots } from './session-scanner-service-root-request' +import { trackAiVaultScanProgress } from './ai-vault-scan-progress' import type { AiVaultSessionTitle } from '../../shared/ai-vault-session-title' +import type { AiVaultListResult } from '../../shared/ai-vault-types' import { readAiVaultFirstUserPrompt } from './session-first-user-prompt-read' import { flushSessionParseCachePersist, @@ -45,6 +47,20 @@ function titleKey(request: { agent: string; sessionId: string }): string { return `${request.agent}\0${request.sessionId}` } +async function scanReportingProgress( + request: Extract, + signal: AbortSignal +): Promise { + const stopProgress = trackAiVaultScanProgress(signal, (units) => + send({ type: 'progress', id: request.id, units }) + ) + try { + return await scanAiVaultSessions({ ...request.options, signal }) + } finally { + stopProgress() + } +} + async function executeRequest(request: AiVaultServiceRequest): Promise { if (sessionSearch.handles(request)) { try { @@ -90,7 +106,7 @@ async function executeRequest(request: AiVaultServiceRequest): Promise 0 + ) + } return (message.type === 'result' || message.type === 'error') && Number.isSafeInteger(message.id) } diff --git a/src/main/ai-vault/session-scanner-service-request-deadline.test.ts b/src/main/ai-vault/session-scanner-service-request-deadline.test.ts new file mode 100644 index 00000000000..164ed0a67c8 --- /dev/null +++ b/src/main/ai-vault/session-scanner-service-request-deadline.test.ts @@ -0,0 +1,227 @@ +import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest' +import { AiVaultScannerServiceClient } from './session-scanner-service-client' +import { + AI_VAULT_SERVICE_READY_TIMEOUT_MS, + AI_VAULT_SERVICE_SHUTDOWN_TIMEOUT_MS +} from './session-scanner-service-client-state' +import { + AI_VAULT_SERVICE_INTERACTIVE_TIMEOUT_MS, + AI_VAULT_SERVICE_SCAN_TIMEOUT_MS +} from './session-scanner-service-request-deadline' +import { + AiVaultServiceTestChild, + aiVaultServiceRequestId, + readyAiVaultServiceChild +} from './session-scanner-service-test-child' + +const SCAN = { type: 'request', operation: 'scan', options: {} } as const +const SCAN_VALUE = { result: { sessions: [], issues: [], scannedAt: '2026-10-12' }, durationMs: 1 } +const TIMED_OUT = `AI Vault service timed out after ${AI_VAULT_SERVICE_SCAN_TIMEOUT_MS}ms.` + +function setup(): { children: AiVaultServiceTestChild[]; client: AiVaultScannerServiceClient } { + const children: AiVaultServiceTestChild[] = [] + const client = new AiVaultScannerServiceClient({ + processFactory: () => { + const child = new AiVaultServiceTestChild(12_345 + children.length) + children.push(child) + return child.asChildProcess() + }, + init: () => ({ sessionParseCache: null, sessionSearch: null }) + }) + return { children, client } +} + +/** Starts one scan on a ready child and returns its request id. */ +async function startScan( + client: AiVaultScannerServiceClient, + children: AiVaultServiceTestChild[], + signal?: AbortSignal +): Promise<{ scan: Promise; child: AiVaultServiceTestChild; id: number }> { + const scan = client.request(SCAN, signal) + const child = children.at(-1)! + readyAiVaultServiceChild(child) + await Promise.resolve() + return { scan, child, id: aiVaultServiceRequestId(child, 'scan') } +} + +function progress(child: AiVaultServiceTestChild, id: number, units: number): void { + child.emit('message', { type: 'progress', id, units }) +} + +beforeEach(() => { + vi.useFakeTimers() +}) + +afterEach(() => { + vi.useRealTimers() +}) + +describe('AI Vault service scan watchdog', () => { + it('keeps a scan that keeps advancing alive past several windows', async () => { + const { children, client } = setup() + const { scan, child, id } = await startScan(client, children) + + for (let units = 1; units <= 4; units++) { + vi.advanceTimersByTime(AI_VAULT_SERVICE_SCAN_TIMEOUT_MS - 1) + progress(child, id, units) + } + expect(child.sent).not.toContainEqual({ type: 'shutdown' }) + child.emit('message', { type: 'result', id, operation: 'scan', value: SCAN_VALUE }) + + await expect(scan).resolves.toEqual(SCAN_VALUE) + expect(child.killed).toBe(false) + client.dispose() + }) + + it('faults a silent scan after the window', async () => { + const { children, client } = setup() + const { scan, child } = await startScan(client, children) + const rejected = expect(scan).rejects.toThrow(TIMED_OUT) + + vi.advanceTimersByTime(AI_VAULT_SERVICE_SCAN_TIMEOUT_MS - 1) + expect(child.sent).not.toContainEqual({ type: 'shutdown' }) + vi.advanceTimersByTime(1) + + await rejected + expect(child.sent).toContainEqual({ type: 'shutdown' }) + client.dispose() + }) + + it('faults a scan one window after its counter last advanced', async () => { + const { children, client } = setup() + const { scan, child, id } = await startScan(client, children) + const rejected = expect(scan).rejects.toThrow(TIMED_OUT) + + vi.advanceTimersByTime(10_000) + progress(child, id, 5) + for (let repeat = 0; repeat < 3; repeat++) { + vi.advanceTimersByTime(40_000) + progress(child, id, 5) + progress(child, id, 2) + } + vi.advanceTimersByTime(AI_VAULT_SERVICE_SCAN_TIMEOUT_MS - 120_001) + expect(child.sent).not.toContainEqual({ type: 'shutdown' }) + vi.advanceTimersByTime(1) + + await rejected + expect(child.sent).toContainEqual({ type: 'shutdown' }) + client.dispose() + }) + + it('ignores progress for another request id', async () => { + const { children, client } = setup() + const { scan, child, id } = await startScan(client, children) + const rejected = expect(scan).rejects.toThrow(TIMED_OUT) + + vi.advanceTimersByTime(AI_VAULT_SERVICE_SCAN_TIMEOUT_MS - 1) + progress(child, id + 100, 1) + vi.advanceTimersByTime(1) + + await rejected + client.dispose() + }) + + it('keeps the fixed bound on a non-scan call that reports progress', async () => { + const { children, client } = setup() + const subagents = client.request({ + type: 'request', + operation: 'subagents', + request: { agent: 'claude', parentFilePath: '/tmp/parent.jsonl' } + }) + const child = children[0]! + readyAiVaultServiceChild(child) + await Promise.resolve() + const id = aiVaultServiceRequestId(child, 'subagents') + const rejected = expect(subagents).rejects.toThrow( + `timed out after ${AI_VAULT_SERVICE_INTERACTIVE_TIMEOUT_MS}ms` + ) + + vi.advanceTimersByTime(AI_VAULT_SERVICE_INTERACTIVE_TIMEOUT_MS - 1) + progress(child, id, 1) + vi.advanceTimersByTime(1) + + await rejected + client.dispose() + }) + + it('does not extend the cancel acknowledgement bound', async () => { + const { children, client } = setup() + const controller = new AbortController() + const { scan, child, id } = await startScan(client, children, controller.signal) + + controller.abort() + await expect(scan).rejects.toMatchObject({ name: 'AbortError' }) + expect(child.sent).toContainEqual({ type: 'cancel', id }) + vi.advanceTimersByTime(AI_VAULT_SERVICE_SHUTDOWN_TIMEOUT_MS - 1) + progress(child, id, 1) + expect(child.killed).toBe(false) + vi.advanceTimersByTime(1) + + expect(child.killed).toBe(true) + client.dispose() + }) + + it('keeps the startup bound when progress arrives before ready', async () => { + const { children, client } = setup() + const scan = client.request(SCAN) + const rejected = expect(scan).rejects.toThrow('did not become ready') + const child = children[0]! + + vi.advanceTimersByTime(AI_VAULT_SERVICE_READY_TIMEOUT_MS - 1) + progress(child, 1, 1) + vi.advanceTimersByTime(1) + expect(child.killed).toBe(true) + // The cold-start retry gets the same bound on the respawned child. + vi.advanceTimersByTime(250 + AI_VAULT_SERVICE_READY_TIMEOUT_MS) + + await rejected + client.dispose() + }) + + it('retires a stalled child with shutdown first, records the fault, and never retries the scan', async () => { + const { children, client } = setup() + const { scan, child, id } = await startScan(client, children) + const rejected = expect(scan).rejects.toThrow(TIMED_OUT) + + vi.advanceTimersByTime(AI_VAULT_SERVICE_SCAN_TIMEOUT_MS) + await rejected + expect(child.sent).toContainEqual({ type: 'shutdown' }) + expect(child.killed).toBe(false) + + // The fault went through the restart policy: the next start waits out its delay. + const next = client.request(SCAN) + vi.advanceTimersByTime(249) + expect(children).toHaveLength(1) + vi.advanceTimersByTime(1) + expect(children).toHaveLength(2) + readyAiVaultServiceChild(children[1]!) + await Promise.resolve() + expect(children[1]!.sent).toContainEqual( + expect.objectContaining({ operation: 'scan', id: id + 1 }) + ) + expect(children[1]!.sent).not.toContainEqual(expect.objectContaining({ id })) + + // A late frame from the retiring child reaches nothing; the kill waits out the grace. + progress(child, id, 9) + vi.advanceTimersByTime(AI_VAULT_SERVICE_SHUTDOWN_TIMEOUT_MS - 251) + expect(child.killed).toBe(false) + vi.advanceTimersByTime(1) + expect(child.killed).toBe(true) + + client.dispose() + await expect(next).rejects.toThrow('disposed') + }) + + it('kills a child at once on a malformed progress frame', async () => { + const { children, client } = setup() + const { scan, child, id } = await startScan(client, children) + const rejected = expect(scan).rejects.toThrow('malformed') + + progress(child, id, 0) + + await rejected + expect(child.killed).toBe(true) + expect(child.sent).not.toContainEqual({ type: 'shutdown' }) + client.dispose() + }) +}) diff --git a/src/main/ai-vault/session-scanner-service-request-deadline.ts b/src/main/ai-vault/session-scanner-service-request-deadline.ts new file mode 100644 index 00000000000..d8816b81249 --- /dev/null +++ b/src/main/ai-vault/session-scanner-service-request-deadline.ts @@ -0,0 +1,63 @@ +import type { AiVaultServicePendingCall } from './session-scanner-service-client-state' +import type { AiVaultServiceRequest } from './session-scanner-service-protocol' + +/** No-progress window: a scan whose work keeps advancing re-arms it, however long it runs. */ +export const AI_VAULT_SERVICE_SCAN_TIMEOUT_MS = 130_000 +export const AI_VAULT_SERVICE_INTERACTIVE_TIMEOUT_MS = 15_000 + +function armAiVaultServiceRequestDeadline( + call: AiVaultServicePendingCall, + onTimeout: (error: Error) => void +): void { + const timeoutMs = + call.request.operation === 'scan' + ? AI_VAULT_SERVICE_SCAN_TIMEOUT_MS + : AI_VAULT_SERVICE_INTERACTIVE_TIMEOUT_MS + if (call.timer) { + clearTimeout(call.timer) + } + // Same text for both meanings: remote clients humanize it, so a new wording + // would reach older ones as "stopped unexpectedly". + call.timer = setTimeout(() => { + onTimeout(new Error(`AI Vault service timed out after ${timeoutMs}ms.`)) + }, timeoutMs) + call.timer.unref?.() +} + +/** Starts the request deadline only once the child is ready to receive it. */ +export function sendAiVaultServiceCall( + // Only the send this needs, so the module stays off the direct child_process import list. + child: { send(request: AiVaultServiceRequest): unknown }, + call: AiVaultServicePendingCall, + isActive: () => boolean, + onTimeout: (error: Error) => void +): void { + if (call.cancelled || !isActive()) { + return + } + armAiVaultServiceRequestDeadline(call, onTimeout) + call.sent = true + child.send(call.request) +} + +/** + * Restarts a running scan's no-progress window when its counter advanced. + * A cancelled call keeps its two-second acknowledgement bound, and interactive + * calls keep their fixed one, so only a live scan's own work can extend it. + */ +export function noteAiVaultServiceScanProgress( + call: AiVaultServicePendingCall, + units: number, + onTimeout: (error: Error) => void +): void { + if ( + call.cancelled || + !call.sent || + call.request.operation !== 'scan' || + units <= call.progressUnits + ) { + return + } + call.progressUnits = units + armAiVaultServiceRequestDeadline(call, onTimeout) +}