From a7d3f2015385db7e3853b5df091458723d5cbb45 Mon Sep 17 00:00:00 2001 From: Brennan Benson <79079362+brennanb2025@users.noreply.github.com> Date: Fri, 14 Aug 2026 13:18:36 -0700 Subject: [PATCH] perf(ai-vault): deduplicate SSH history cache fills --- .../runtime-owned-ssh-session-list.test.ts | 33 ++++++++ .../runtime-owned-ssh-session-list.ts | 54 ++++++++++--- src/main/ipc/ai-vault-host-leg-cache.test.ts | 80 +++++++++++++++++++ src/main/ipc/ai-vault-host-leg-cache.ts | 45 ++++++++++- 4 files changed, 197 insertions(+), 15 deletions(-) create mode 100644 src/main/ipc/ai-vault-host-leg-cache.test.ts 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 d4997c1ecd1..811afa5965f 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 @@ -82,6 +82,25 @@ describe('runtime-owned SSH AI Vault inventory', () => { expect(mocks.callRuntimeEnvironment).toHaveBeenCalledTimes(1) }) + it('deduplicates concurrent owner lookups', async () => { + let resolveInventory!: (value: { ok: true; result: { targets: { id: string }[] } }) => void + mocks.callRuntimeEnvironment.mockReturnValueOnce( + new Promise((resolve) => { + resolveInventory = resolve + }) + ) + + const first = findRuntimeOwningSshAiVaultHost('/user-data', 'hub-owned-host') + const second = findRuntimeOwningSshAiVaultHost('/user-data', 'hub-owned-host') + expect(mocks.callRuntimeEnvironment).toHaveBeenCalledTimes(1) + resolveInventory({ ok: true, result: { targets: [{ id: 'hub-owned-host' }] } }) + + await expect(Promise.all([first, second])).resolves.toEqual([ + expect.objectContaining({ environmentId: 'hub-runtime' }), + expect.objectContaining({ environmentId: 'hub-runtime' }) + ]) + }) + it('expires cached owner lookups after one minute', async () => { vi.useFakeTimers({ now: new Date('2026-08-14T00:00:00.000Z') }) try { @@ -106,6 +125,20 @@ describe('runtime-owned SSH AI Vault inventory', () => { expect(mocks.callRuntimeEnvironment).toHaveBeenCalledTimes(2) }) + it('does not share owner lookups across a re-pair of the same environment id', async () => { + mocks.listEnvironments.mockReturnValue([ + { id: 'hub-runtime', createdAt: 1, pairingRevision: 1, runtimeId: 'runtime-a' } + ]) + await findRuntimeOwningSshAiVaultHost('/user-data', 'hub-owned-host') + mocks.listEnvironments.mockReturnValue([ + { id: 'hub-runtime', createdAt: 1, pairingRevision: 2, runtimeId: 'runtime-b' } + ]) + + await findRuntimeOwningSshAiVaultHost('/user-data', 'hub-owned-host') + + 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 268d5c97f35..f3fde7d318f 100644 --- a/src/main/ai-vault/runtime-owned-ssh-session-list.ts +++ b/src/main/ai-vault/runtime-owned-ssh-session-list.ts @@ -38,7 +38,9 @@ type CachedRuntimeOwnedSshHost = { expiresAt: number } +// Missing-title refreshes run every 20s; cache by pairing identity so they do not fan out to every runtime. const cachedOwners = new Map() +const inFlightOwnerLookups = new Map>() export async function listRuntimeOwnedSshAiVaultTargets( userDataPath: string, @@ -90,25 +92,54 @@ export async function findRuntimeOwningSshAiVaultHost( const cacheKey = JSON.stringify({ userDataPath, targetId, - environmentIds: environments.map((environment) => environment.id).sort() + environments: environments + .map((environment) => ({ + id: environment.id, + pairingRevision: environment.pairingRevision ?? environment.createdAt, + runtimeId: environment.runtimeId + })) + .sort((a, b) => a.id.localeCompare(b.id)) }) 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) - ) + const inFlight = inFlightOwnerLookups.get(cacheKey) + if (inFlight) { + return inFlight + } + const lookup = findRuntimeOwnedSshHostInInventories( + userDataPath, + targetId, + environments.map((environment) => environment.id), + options ) - for (const hosts of inventories) { - const match = hosts.find((host) => host.targetId === targetId) - if (match) { - cacheRuntimeOwnedSshHost(cacheKey, match) - return match + inFlightOwnerLookups.set(cacheKey, lookup) + try { + const owner = await lookup + if (owner) { + cacheRuntimeOwnedSshHost(cacheKey, owner) + } + return owner + } finally { + if (inFlightOwnerLookups.get(cacheKey) === lookup) { + inFlightOwnerLookups.delete(cacheKey) } } - return null +} + +async function findRuntimeOwnedSshHostInInventories( + userDataPath: string, + targetId: string, + environmentIds: readonly string[], + options: RuntimeOwnedSshAiVaultScanOptions +): Promise { + const inventories = await Promise.all( + environmentIds.map((environmentId) => + listRuntimeOwnedSshAiVaultTargets(userDataPath, environmentId, options) + ) + ) + return inventories.flat().find((host) => host.targetId === targetId) ?? null } function cacheRuntimeOwnedSshHost(cacheKey: string, owner: RuntimeOwnedSshAiVaultHost): void { @@ -129,6 +160,7 @@ function cacheRuntimeOwnedSshHost(cacheKey: string, owner: RuntimeOwnedSshAiVaul export function resetRuntimeOwnedSshOwnerCacheForTests(): void { cachedOwners.clear() + inFlightOwnerLookups.clear() } export async function scanRuntimeOwnedSshAiVaultSessions( diff --git a/src/main/ipc/ai-vault-host-leg-cache.test.ts b/src/main/ipc/ai-vault-host-leg-cache.test.ts new file mode 100644 index 00000000000..10b5e17a0c7 --- /dev/null +++ b/src/main/ipc/ai-vault-host-leg-cache.test.ts @@ -0,0 +1,80 @@ +import { beforeEach, describe, expect, it, vi } from 'vitest' +import type { AiVaultListResult } from '../../shared/ai-vault-types' +import { + AI_VAULT_HOST_LEG_CACHE_MAX_ENTRIES, + resetAiVaultHostLegCacheForTests, + scanHostLegWithCache +} from './ai-vault-host-leg-cache' + +const RESULT: AiVaultListResult = { + sessions: [], + issues: [], + scannedAt: '2026-08-14T00:00:00.000Z' +} + +beforeEach(() => { + resetAiVaultHostLegCacheForTests() +}) + +describe('AI Vault host-leg cache', () => { + it('deduplicates concurrent non-force scans for the same host and depth', async () => { + let resolveScan!: (result: AiVaultListResult) => void + const scan = vi.fn( + () => + new Promise((resolve) => { + resolveScan = resolve + }) + ) + const args = { cacheKey: 'host-a', depth: 500 as const, force: false, scan } + + const first = scanHostLegWithCache(args) + const second = scanHostLegWithCache(args) + expect(scan).toHaveBeenCalledTimes(1) + resolveScan(RESULT) + + await expect(Promise.all([first, second])).resolves.toEqual([RESULT, RESULT]) + }) + + it('evicts an earlier healthy result when a forced refresh returns a host issue', async () => { + const hostFailure: AiVaultListResult = { + ...RESULT, + issues: [ + { + executionHostId: 'ssh:host-a', + agent: 'codex', + kind: 'host', + path: 'host-a', + message: 'offline' + } + ] + } + const scan = vi + .fn() + .mockResolvedValue(RESULT) + .mockResolvedValueOnce(RESULT) + .mockResolvedValueOnce(hostFailure) + const baseArgs = { cacheKey: 'host-a', depth: 500 as const, scan } + + await scanHostLegWithCache({ ...baseArgs, force: false }) + await scanHostLegWithCache({ ...baseArgs, force: true }) + await scanHostLegWithCache({ ...baseArgs, force: false }) + + expect(scan).toHaveBeenCalledTimes(3) + }) + + it('caps retained host legs', async () => { + const scan = vi.fn().mockResolvedValue(RESULT) + for (let index = 0; index <= AI_VAULT_HOST_LEG_CACHE_MAX_ENTRIES; index++) { + await scanHostLegWithCache({ + cacheKey: `host-${index}`, + depth: 500, + force: false, + scan + }) + } + + await scanHostLegWithCache({ cacheKey: 'host-0', depth: 500, force: false, scan }) + + expect(scan).toHaveBeenCalledTimes(AI_VAULT_HOST_LEG_CACHE_MAX_ENTRIES + 2) + }) +}) diff --git a/src/main/ipc/ai-vault-host-leg-cache.ts b/src/main/ipc/ai-vault-host-leg-cache.ts index 18b75517d6d..21d4f17a417 100644 --- a/src/main/ipc/ai-vault-host-leg-cache.ts +++ b/src/main/ipc/ai-vault-host-leg-cache.ts @@ -6,6 +6,8 @@ import { } from '../../shared/ai-vault-session-depth' export const AI_VAULT_CACHE_TTL_MS = 15_000 +// Scope paths are part of the key, so cap retained scans from long-lived workspace churn. +export const AI_VAULT_HOST_LEG_CACHE_MAX_ENTRIES = 256 type CachedHostLeg = { depth: AiVaultSessionDepth @@ -18,6 +20,10 @@ type CachedHostLeg = { // open. Legs are cached individually and the failing host is the only one // rescanned. const cachedHostLegs = new Map() +const inFlightHostLegs = new Map< + string, + { depth: AiVaultSessionDepth; promise: Promise } +>() // Bumped on every invalidation. A leg scan that started before an invalidation // carries the old generation and must not write its (now stale) result back — // otherwise a delete's invalidation is silently undone by a scan that resolves @@ -45,6 +51,30 @@ export async function scanHostLegWithCache(args: { ) { return truncateAiVaultListResult(cached.result, args.depth, args.scopePaths) } + const inFlight = inFlightHostLegs.get(args.cacheKey) + if (!args.force && inFlight && aiVaultSessionDepthCovers(inFlight.depth, args.depth)) { + return truncateAiVaultListResult(await inFlight.promise, args.depth, args.scopePaths) + } + const promise = scanAndCacheHostLeg(args, now, startGeneration) + if (args.force) { + return promise + } + const entry = { depth: args.depth, promise } + inFlightHostLegs.set(args.cacheKey, entry) + try { + return await promise + } finally { + if (inFlightHostLegs.get(args.cacheKey) === entry) { + inFlightHostLegs.delete(args.cacheKey) + } + } +} + +async function scanAndCacheHostLeg( + args: Parameters[0], + now: number, + startGeneration: number +): Promise { const result = await args.scan() // A delete's invalidation landed while this leg was scanning; caching a // pre-delete result would resurrect the deleted session for the TTL. @@ -52,10 +82,7 @@ export async function scanHostLegWithCache(args: { return result } if (result.issues.some((issue) => issue.kind === 'host')) { - const current = cachedHostLegs.get(args.cacheKey) - if (!current || current.expiresAt <= Date.now()) { - cachedHostLegs.delete(args.cacheKey) - } + cachedHostLegs.delete(args.cacheKey) return result } pruneExpiredHostLegs(now) @@ -68,6 +95,15 @@ export async function scanHostLegWithCache(args: { ) { return result } + if ( + !cachedHostLegs.has(args.cacheKey) && + cachedHostLegs.size >= AI_VAULT_HOST_LEG_CACHE_MAX_ENTRIES + ) { + const oldestKey = cachedHostLegs.keys().next().value + if (oldestKey) { + cachedHostLegs.delete(oldestKey) + } + } cachedHostLegs.set(args.cacheKey, { depth: args.depth, result, @@ -95,4 +131,5 @@ export function invalidateAiVaultHostLegCache(): void { export function resetAiVaultHostLegCacheForTests(): void { invalidateAiVaultHostLegCache() + inFlightHostLegs.clear() }