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 e2043dd6f98..8dabe15833e 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 @@ -91,6 +91,24 @@ describe('runtime-owned SSH AI Vault inventory', () => { expect(mocks.callRuntimeEnvironment).toHaveBeenCalledTimes(1) }) + it('does not cache a missing owner when runtime inventory failed', async () => { + mocks.callRuntimeEnvironment + .mockRejectedValueOnce(new Error('runtime offline')) + .mockResolvedValueOnce({ + ok: true, + result: { targets: [{ id: 'hub-owned-host' }] } + }) + + await expect( + findRuntimeOwningSshAiVaultHost('/user-data', 'hub-owned-host') + ).resolves.toBeNull() + await expect( + findRuntimeOwningSshAiVaultHost('/user-data', 'hub-owned-host') + ).resolves.toMatchObject({ environmentId: 'hub-runtime' }) + + expect(mocks.callRuntimeEnvironment).toHaveBeenCalledTimes(2) + }) + it('deduplicates concurrent owner lookups', async () => { let resolveInventory!: (value: { ok: true; result: { targets: { id: string }[] } }) => void mocks.callRuntimeEnvironment.mockReturnValueOnce( 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 138be9d0602..7e3bf2921f9 100644 --- a/src/main/ai-vault/runtime-owned-ssh-session-list.ts +++ b/src/main/ai-vault/runtime-owned-ssh-session-list.ts @@ -40,13 +40,25 @@ type CachedRuntimeOwnedSshHost = { // 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>() +type RuntimeOwnedSshHostLookup = { + owner: RuntimeOwnedSshAiVaultHost | null + cacheable: boolean +} +const inFlightOwnerLookups = new Map>() export async function listRuntimeOwnedSshAiVaultTargets( userDataPath: string, environmentId: string, options: RuntimeOwnedSshAiVaultScanOptions = {} ): Promise { + return (await loadRuntimeOwnedSshAiVaultTargets(userDataPath, environmentId, options)).hosts +} + +async function loadRuntimeOwnedSshAiVaultTargets( + userDataPath: string, + environmentId: string, + options: RuntimeOwnedSshAiVaultScanOptions +): Promise<{ hosts: readonly RuntimeOwnedSshAiVaultHost[]; authoritative: boolean }> { let response: Awaited> try { response = await callRuntimeEnvironment( @@ -57,12 +69,12 @@ export async function listRuntimeOwnedSshAiVaultTargets( options.timeoutMs ) } catch { - return [] + return { hosts: [], authoritative: false } } if (response.ok !== true || !isTargetSummaryList(response.result)) { - return [] + return { hosts: [], authoritative: false } } - return response.result.targets.flatMap((target) => { + const hosts = response.result.targets.flatMap((target) => { if (typeof target.id !== 'string' || target.id.length === 0) { return [] } @@ -78,6 +90,7 @@ export async function listRuntimeOwnedSshAiVaultTargets( } ] }) + return { hosts, authoritative: true } } export async function findRuntimeOwningSshAiVaultHost( @@ -106,7 +119,7 @@ export async function findRuntimeOwningSshAiVaultHost( } const inFlight = inFlightOwnerLookups.get(cacheKey) if (inFlight) { - return inFlight + return (await inFlight).owner } const lookup = findRuntimeOwnedSshHostInInventories( userDataPath, @@ -116,9 +129,11 @@ export async function findRuntimeOwningSshAiVaultHost( ) inFlightOwnerLookups.set(cacheKey, lookup) try { - const owner = await lookup - cacheRuntimeOwnedSshHost(cacheKey, owner) - return owner + const result = await lookup + if (result.cacheable) { + cacheRuntimeOwnedSshHost(cacheKey, result.owner) + } + return result.owner } finally { if (inFlightOwnerLookups.get(cacheKey) === lookup) { inFlightOwnerLookups.delete(cacheKey) @@ -131,13 +146,19 @@ async function findRuntimeOwnedSshHostInInventories( targetId: string, environmentIds: readonly string[], options: RuntimeOwnedSshAiVaultScanOptions -): Promise { +): Promise { const inventories = await Promise.all( environmentIds.map((environmentId) => - listRuntimeOwnedSshAiVaultTargets(userDataPath, environmentId, options) + loadRuntimeOwnedSshAiVaultTargets(userDataPath, environmentId, options) ) ) - return inventories.flat().find((host) => host.targetId === targetId) ?? null + const owner = inventories + .flatMap((inventory) => inventory.hosts) + .find((host) => host.targetId === targetId) + return { + owner: owner ?? null, + cacheable: owner !== undefined || inventories.every((inventory) => inventory.authoritative) + } } function cacheRuntimeOwnedSshHost( diff --git a/src/main/ipc/ai-vault-host-leg-cache.test.ts b/src/main/ipc/ai-vault-host-leg-cache.test.ts index 269da14c235..05d4f12b7e1 100644 --- a/src/main/ipc/ai-vault-host-leg-cache.test.ts +++ b/src/main/ipc/ai-vault-host-leg-cache.test.ts @@ -95,6 +95,40 @@ describe('AI Vault host-leg cache', () => { expect(scan).toHaveBeenCalledTimes(3) }) + it('does not let an older forced failure evict a newer healthy result', async () => { + let resolveOlder!: (result: AiVaultListResult) => void + let resolveNewer!: (result: AiVaultListResult) => void + const scan = vi + .fn() + .mockImplementationOnce( + () => new Promise((resolve) => (resolveOlder = resolve)) + ) + .mockImplementationOnce( + () => new Promise((resolve) => (resolveNewer = resolve)) + ) + const args = { cacheKey: 'host-a', depth: 500 as const, scan } + const older = scanHostLegWithCache({ ...args, force: true }) + const newer = scanHostLegWithCache({ ...args, force: true }) + resolveNewer(RESULT) + await newer + resolveOlder({ + ...RESULT, + issues: [ + { + executionHostId: 'ssh:host-a', + agent: 'codex', + kind: 'host', + path: 'host-a', + message: 'offline' + } + ] + }) + await older + + await scanHostLegWithCache({ ...args, force: false }) + expect(scan).toHaveBeenCalledTimes(2) + }) + 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++) { diff --git a/src/main/ipc/ai-vault-host-leg-cache.ts b/src/main/ipc/ai-vault-host-leg-cache.ts index c29a7eca5c7..7b394441b84 100644 --- a/src/main/ipc/ai-vault-host-leg-cache.ts +++ b/src/main/ipc/ai-vault-host-leg-cache.ts @@ -24,6 +24,8 @@ const inFlightHostLegs = new Map< string, { depth: AiVaultSessionDepth; promise: Promise } >() +const activeHostLegScans = new Map() +let nextHostLegScanId = 0 // 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 @@ -81,41 +83,53 @@ async function scanAndCacheHostLeg( 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. - if (startGeneration !== cacheGeneration) { + const scanId = ++nextHostLegScanId + const active = activeHostLegScans.get(args.cacheKey) ?? { latestId: scanId, count: 0 } + active.latestId = scanId + active.count++ + activeHostLegScans.set(args.cacheKey, active) + try { + const result = await args.scan() + // A delete or newer forced refresh landed while this leg was scanning; + // caching its older result would undo the newer observation. + if (startGeneration !== cacheGeneration || active.latestId !== scanId) { + return result + } + if (result.issues.some((issue) => issue.kind === 'host')) { + cachedHostLegs.delete(args.cacheKey) + return result + } + pruneExpiredHostLegs(now) + const current = cachedHostLegs.get(args.cacheKey) + if ( + !args.force && + current && + current.expiresAt > Date.now() && + aiVaultSessionDepthCovers(current.depth, args.depth) + ) { + 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, + expiresAt: Date.now() + AI_VAULT_CACHE_TTL_MS + }) return result - } - if (result.issues.some((issue) => issue.kind === 'host')) { - cachedHostLegs.delete(args.cacheKey) - return result - } - pruneExpiredHostLegs(now) - const current = cachedHostLegs.get(args.cacheKey) - if ( - !args.force && - current && - current.expiresAt > Date.now() && - aiVaultSessionDepthCovers(current.depth, args.depth) - ) { - 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) + } finally { + active.count-- + if (active.count === 0 && activeHostLegScans.get(args.cacheKey) === active) { + activeHostLegScans.delete(args.cacheKey) } } - cachedHostLegs.set(args.cacheKey, { - depth: args.depth, - result, - expiresAt: Date.now() + AI_VAULT_CACHE_TTL_MS - }) - return result } function pruneExpiredHostLegs(now: number): void { @@ -138,4 +152,6 @@ export function invalidateAiVaultHostLegCache(): void { export function resetAiVaultHostLegCacheForTests(): void { invalidateAiVaultHostLegCache() inFlightHostLegs.clear() + activeHostLegScans.clear() + nextHostLegScanId = 0 }