From 4d3902c52cc40193a83faea93d7accd27e5e4ca9 Mon Sep 17 00:00:00 2001 From: Brennan Benson <79079362+brennanb2025@users.noreply.github.com> Date: Fri, 14 Aug 2026 13:08:36 -0700 Subject: [PATCH] perf(ai-vault): bound repeated runtime SSH fanout --- .../runtime-owned-ssh-session-list.test.ts | 33 +++++++++++++++ .../runtime-owned-ssh-session-list.ts | 39 ++++++++++++++++++ .../ai-vault-runtime-ssh-scope.repro.test.ts | 32 +++++++++++++++ src/main/runtime/rpc/methods/ai-vault.ts | 40 ++++++++++++++----- 4 files changed, 135 insertions(+), 9 deletions(-) diff --git a/src/main/ai-vault/runtime-owned-ssh-session-list.test.ts b/src/main/ai-vault/runtime-owned-ssh-session-list.test.ts index 2effc916147..d4997c1ecd1 100644 --- a/src/main/ai-vault/runtime-owned-ssh-session-list.test.ts +++ b/src/main/ai-vault/runtime-owned-ssh-session-list.test.ts @@ -17,11 +17,13 @@ vi.mock('../ipc/runtime-environment-transport-routing', () => ({ const { findRuntimeOwningSshAiVaultHost, listRuntimeOwnedSshAiVaultTargets, + resetRuntimeOwnedSshOwnerCacheForTests, scanRuntimeOwnedSshAiVaultSessions } = await import('./runtime-owned-ssh-session-list') beforeEach(() => { vi.clearAllMocks() + resetRuntimeOwnedSshOwnerCacheForTests() mocks.listEnvironments.mockReturnValue([{ id: 'hub-runtime' }]) mocks.callRuntimeEnvironment.mockResolvedValue({ ok: true, @@ -73,6 +75,37 @@ describe('runtime-owned SSH AI Vault inventory', () => { ) }) + it('reuses a bounded owner lookup for repeated title batches', async () => { + await findRuntimeOwningSshAiVaultHost('/user-data', 'hub-owned-host') + await findRuntimeOwningSshAiVaultHost('/user-data', 'hub-owned-host') + + expect(mocks.callRuntimeEnvironment).toHaveBeenCalledTimes(1) + }) + + it('expires cached owner lookups after one minute', async () => { + vi.useFakeTimers({ now: new Date('2026-08-14T00:00:00.000Z') }) + try { + await findRuntimeOwningSshAiVaultHost('/user-data', 'hub-owned-host') + await vi.advanceTimersByTimeAsync(60_000) + await findRuntimeOwningSshAiVaultHost('/user-data', 'hub-owned-host') + + expect(mocks.callRuntimeEnvironment).toHaveBeenCalledTimes(2) + } finally { + vi.useRealTimers() + } + }) + + it('does not share owner lookups across runtime inventories', async () => { + await findRuntimeOwningSshAiVaultHost('/user-data', 'hub-owned-host') + mocks.listEnvironments.mockReturnValue([{ id: 'other-runtime' }]) + mocks.callRuntimeEnvironment.mockResolvedValueOnce({ ok: true, result: { targets: [] } }) + + await expect( + findRuntimeOwningSshAiVaultHost('/user-data', 'hub-owned-host') + ).resolves.toBeNull() + expect(mocks.callRuntimeEnvironment).toHaveBeenCalledTimes(2) + }) + it('does not treat recipe-VM targets as runtime-owned SSH history hosts', async () => { await expect(findRuntimeOwningSshAiVaultHost('/user-data', 'runtime-ssh-recipe')).resolves.toBe( null diff --git a/src/main/ai-vault/runtime-owned-ssh-session-list.ts b/src/main/ai-vault/runtime-owned-ssh-session-list.ts index 1b49cb37ccb..268d5c97f35 100644 --- a/src/main/ai-vault/runtime-owned-ssh-session-list.ts +++ b/src/main/ai-vault/runtime-owned-ssh-session-list.ts @@ -30,6 +30,15 @@ export type RuntimeOwnedSshAiVaultScanOptions = { } const UNSUPPORTED_SSH_HOST_PARAM = /invalid (runtime )?execution host id/i +const OWNER_CACHE_TTL_MS = 60_000 +const OWNER_CACHE_MAX_ENTRIES = 256 + +type CachedRuntimeOwnedSshHost = { + owner: RuntimeOwnedSshAiVaultHost + expiresAt: number +} + +const cachedOwners = new Map() export async function listRuntimeOwnedSshAiVaultTargets( userDataPath: string, @@ -78,6 +87,15 @@ export async function findRuntimeOwningSshAiVaultHost( return null } const environments = listEnvironments(userDataPath) + const cacheKey = JSON.stringify({ + userDataPath, + targetId, + environmentIds: environments.map((environment) => environment.id).sort() + }) + const cached = cachedOwners.get(cacheKey) + if (cached && cached.expiresAt > Date.now()) { + return cached.owner + } const inventories = await Promise.all( environments.map((environment) => listRuntimeOwnedSshAiVaultTargets(userDataPath, environment.id, options) @@ -86,12 +104,33 @@ export async function findRuntimeOwningSshAiVaultHost( for (const hosts of inventories) { const match = hosts.find((host) => host.targetId === targetId) if (match) { + cacheRuntimeOwnedSshHost(cacheKey, match) return match } } return null } +function cacheRuntimeOwnedSshHost(cacheKey: string, owner: RuntimeOwnedSshAiVaultHost): void { + const now = Date.now() + for (const [key, entry] of cachedOwners) { + if (entry.expiresAt <= now) { + cachedOwners.delete(key) + } + } + if (!cachedOwners.has(cacheKey) && cachedOwners.size >= OWNER_CACHE_MAX_ENTRIES) { + const oldestKey = cachedOwners.keys().next().value + if (oldestKey) { + cachedOwners.delete(oldestKey) + } + } + cachedOwners.set(cacheKey, { owner, expiresAt: now + OWNER_CACHE_TTL_MS }) +} + +export function resetRuntimeOwnedSshOwnerCacheForTests(): void { + cachedOwners.clear() +} + export async function scanRuntimeOwnedSshAiVaultSessions( userDataPath: string, environmentId: string, diff --git a/src/main/runtime/rpc/methods/ai-vault-runtime-ssh-scope.repro.test.ts b/src/main/runtime/rpc/methods/ai-vault-runtime-ssh-scope.repro.test.ts index b793bf9119a..6834a8398f5 100644 --- a/src/main/runtime/rpc/methods/ai-vault-runtime-ssh-scope.repro.test.ts +++ b/src/main/runtime/rpc/methods/ai-vault-runtime-ssh-scope.repro.test.ts @@ -38,6 +38,7 @@ import { listAiVaultSessions, resetAiVaultSessionListCacheForTests } from '../../../ai-vault/cached-session-list' +import { resetAiVaultHostLegCacheForTests } from '../../../ipc/ai-vault-host-leg-cache' const SCANNED_AT = '2026-08-12T00:00:00.000Z' const SSH_HOST_OWNED_BY_RUNTIME = 'ssh:hub-owned-host' @@ -98,6 +99,7 @@ function makeDispatcher(): RpcDispatcher { describe('a runtime can scan an SSH host that it owns', () => { beforeEach(() => { resetAiVaultSessionListCacheForTests() + resetAiVaultHostLegCacheForTests() scanAiVaultSessionsInWorker.mockReset() scanAiVaultSessionsInWorker.mockResolvedValue({ sessions: [makeRuntimeLocalSession()], @@ -122,6 +124,7 @@ describe('a runtime can scan an SSH host that it owns', () => { afterEach(() => { resetAiVaultSessionListCacheForTests() + resetAiVaultHostLegCacheForTests() }) it('routes an SSH execution host to that host instead of rejecting it', async () => { @@ -205,6 +208,35 @@ describe('a runtime can scan an SSH host that it owns', () => { ]) }) + it('reuses each SSH host leg across repeated non-force all-host scans', async () => { + const dispatcher = makeDispatcher() + const params = { + limit: 500, + executionHostId: 'runtime:hub-runtime', + includeOwnedSshHosts: true + } + + await dispatcher.dispatch(makeRequest('aiVault.listSessions', params)) + await dispatcher.dispatch(makeRequest('aiVault.listSessions', params)) + + expect(scanSshAiVaultSessions).toHaveBeenCalledTimes(1) + }) + + it('bypasses the SSH host-leg cache for a forced refresh', async () => { + const dispatcher = makeDispatcher() + const params = { + limit: 500, + force: true, + executionHostId: 'runtime:hub-runtime', + includeOwnedSshHosts: true + } + + await dispatcher.dispatch(makeRequest('aiVault.listSessions', params)) + await dispatcher.dispatch(makeRequest('aiVault.listSessions', params)) + + expect(scanSshAiVaultSessions).toHaveBeenCalledTimes(2) + }) + it('resolves SSH session titles on the transcript-owning host', async () => { const dispatcher = makeDispatcher() const requests = [{ agent: 'claude', sessionId: 'sess-ssh' }] diff --git a/src/main/runtime/rpc/methods/ai-vault.ts b/src/main/runtime/rpc/methods/ai-vault.ts index 74b214720f8..22556797386 100644 --- a/src/main/runtime/rpc/methods/ai-vault.ts +++ b/src/main/runtime/rpc/methods/ai-vault.ts @@ -11,7 +11,14 @@ import { getActiveSshAiVaultHostInfos, requestActiveSshAiVaultSessionTitles } from '../../../ipc/ssh' -import { AI_VAULT_AGENTS, AI_VAULT_SCOPE_PATHS_MAX_COUNT } from '../../../../shared/ai-vault-types' +import { scanHostLegWithCache } from '../../../ipc/ai-vault-host-leg-cache' +import { requestedAiVaultSessionDepth } from '../../../../shared/ai-vault-session-depth' +import { + AI_VAULT_AGENTS, + AI_VAULT_SCOPE_PATHS_MAX_COUNT, + type AiVaultListArgs, + type AiVaultListResult +} from '../../../../shared/ai-vault-types' import { AI_VAULT_SESSION_TITLE_REQUEST_MAX_COUNT } from '../../../../shared/ai-vault-session-title' import { LOCAL_EXECUTION_HOST_ID, parseExecutionHostId } from '../../../../shared/execution-host' @@ -123,19 +130,13 @@ export const AI_VAULT_METHODS: RpcMethod[] = [ scopePaths: params.scopePaths } if (parsed?.kind === 'ssh') { - return scanSshAiVaultSessions(parsed.targetId, listArgs, { - timeoutMs: OWNED_SSH_SCAN_TIMEOUT_MS - }) + return scanRuntimeSshAiVaultSessions(parsed.targetId, listArgs) } const sshHosts = params.includeOwnedSshHosts === true ? getActiveSshAiVaultHostInfos() : [] const sshPromise = sshHosts.length > 0 ? Promise.all( - sshHosts.map((host) => - scanSshAiVaultSessions(host.targetId, listArgs, { - timeoutMs: OWNED_SSH_SCAN_TIMEOUT_MS - }) - ) + sshHosts.map((host) => scanRuntimeSshAiVaultSessions(host.targetId, listArgs)) ) : Promise.resolve([]) const result = await runtime.listAiVaultSessions(listArgs) @@ -164,3 +165,24 @@ export const AI_VAULT_METHODS: RpcMethod[] = [ }) }) ] + +function scanRuntimeSshAiVaultSessions( + targetId: string, + args: AiVaultListArgs +): Promise { + const scopePaths = args.scopePaths ?? [] + return scanHostLegWithCache({ + cacheKey: JSON.stringify({ + route: 'runtime-rpc-ssh', + targetId, + scopePaths: [...new Set(scopePaths)].sort() + }), + depth: requestedAiVaultSessionDepth(args), + scopePaths, + force: args.force === true, + scan: () => + scanSshAiVaultSessions(targetId, args, { + timeoutMs: OWNED_SSH_SCAN_TIMEOUT_MS + }) + }) +}