From 65f6f95e48eefd2b9934c04dc84a2f0bd2b7584b Mon Sep 17 00:00:00 2001 From: Jinwoo-H Date: Sun, 13 Sep 2026 16:23:46 -0400 Subject: [PATCH] test(ai-vault-search): cover every host scope, the all-hosts merge and wire compat --- .../runtime-session-search-call.test.ts | 40 +++ src/main/ipc/ai-vault-search.test.ts | 257 +++++++++++++++++- src/main/ipc/ai-vault-search.ts | 5 +- .../web-ai-vault-search-compat.test.ts | 23 +- src/shared/ai-vault-search-contract.test.ts | 27 ++ 5 files changed, 331 insertions(+), 21 deletions(-) create mode 100644 src/main/ai-vault/runtime-session-search-call.test.ts diff --git a/src/main/ai-vault/runtime-session-search-call.test.ts b/src/main/ai-vault/runtime-session-search-call.test.ts new file mode 100644 index 00000000000..aee3ae8c116 --- /dev/null +++ b/src/main/ai-vault/runtime-session-search-call.test.ts @@ -0,0 +1,40 @@ +import { beforeEach, describe, expect, it, vi } from 'vitest' + +const callRuntimeEnvironment = vi.hoisted(() => vi.fn()) +vi.mock('../ipc/runtime-environment-transport-routing', () => ({ callRuntimeEnvironment })) + +import { callRuntimeSessionSearch } from './runtime-session-search-call' + +beforeEach(() => callRuntimeEnvironment.mockReset()) + +describe('runtime session search transport', () => { + it('addresses the environment by method and params and returns its result', async () => { + callRuntimeEnvironment.mockResolvedValue({ id: '1', ok: true, result: { kind: 'ok' } }) + expect( + await callRuntimeSessionSearch( + '/user/data', + 'env-1', + 'aiVault.searchSessions', + { query: 'needle' }, + 10_000 + ) + ).toEqual({ kind: 'ok' }) + expect(callRuntimeEnvironment).toHaveBeenCalledWith( + '/user/data', + 'env-1', + 'aiVault.searchSessions', + { query: 'needle' }, + 10_000 + ) + }) + it('rethrows a refusal with its code so an old host reads as an absent method', async () => { + callRuntimeEnvironment.mockResolvedValue({ + id: '1', + ok: false, + error: { code: 'method_not_found', message: 'unknown method' } + }) + await expect( + callRuntimeSessionSearch('/user/data', 'env-1', 'aiVault.searchSessions', {}) + ).rejects.toMatchObject({ code: 'method_not_found', message: 'unknown method' }) + }) +}) diff --git a/src/main/ipc/ai-vault-search.test.ts b/src/main/ipc/ai-vault-search.test.ts index be9ed253bcb..e89484bb99e 100644 --- a/src/main/ipc/ai-vault-search.test.ts +++ b/src/main/ipc/ai-vault-search.test.ts @@ -1,8 +1,13 @@ import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest' -const { handlers, remote } = vi.hoisted(() => ({ +const { handlers, sshSearch, sshHostInfos, runtimeSearch, runtimeHostInfos } = vi.hoisted(() => ({ handlers: new Map Promise>(), - remote: vi.fn() + sshSearch: vi.fn(), + sshHostInfos: vi.fn<() => { targetId: string }[]>(() => []), + runtimeSearch: vi.fn(), + runtimeHostInfos: vi.fn<() => { environmentId: string; executionHostId: `runtime:${string}` }[]>( + () => [] + ) })) vi.mock('electron', () => ({ ipcMain: { @@ -11,17 +16,48 @@ vi.mock('electron', () => ({ }, ipcRenderer: { invoke: (name: string, ...args: unknown[]) => handlers.get(name)!(null, ...args) } })) -vi.mock('./ssh', () => ({ requestActiveSshSessionSearch: remote })) +vi.mock('./ssh', () => ({ + requestActiveSshSessionSearch: sshSearch, + getActiveSshAiVaultHostInfos: sshHostInfos +})) import { registerAiVaultSearchHandlers } from './ai-vault-search' +import { decodeMergedSearchCursor } from './ai-vault-search-all-hosts' import { aiVaultApi } from '../../preload/api/ai-vault-bridge' import { setSessionSearchService } from '../ai-vault-search/session-search-service-registry' -import { fakeSearchService, searchResults } from '../../shared/ai-vault-search-test-fixture' +import { unavailableSessionSearchStatus } from '../../shared/ai-vault-search-client' +import { + fakeSearchService, + searchHit, + searchResults +} from '../../shared/ai-vault-search-test-fixture' +import type { AiVaultSearchResponse } from '../../shared/ai-vault-search-types' + +function resultsAt( + updatedAt: string | null, + page: { cursor: string | null; hasMore: boolean } = { cursor: null, hasMore: false } +) { + return { ...searchResults(), hits: [{ ...searchHit(), updatedAt }], page } +} + +function serviceReturning(response: AiVaultSearchResponse) { + return { ...fakeSearchService(), search: vi.fn(async () => response) } +} + +function updatedAtOf(response: AiVaultSearchResponse): (string | null)[] { + return response.kind === 'results' ? response.hits.map((hit) => hit.updatedAt) : [] +} beforeEach(() => { handlers.clear() - remote.mockReset() - registerAiVaultSearchHandlers() + sshSearch.mockReset() + runtimeSearch.mockReset() + sshHostInfos.mockReset().mockReturnValue([]) + runtimeHostInfos.mockReset().mockReturnValue([]) + registerAiVaultSearchHandlers({ + getActiveRuntimeAiVaultHostInfos: runtimeHostInfos, + callRuntimeSearch: runtimeSearch + }) }) afterEach(() => setSessionSearchService(null)) @@ -37,32 +73,225 @@ describe('desktop IPC and preload search boundary', () => { } ] }) + expect(await aiVaultApi.searchSessions({ query: 'needle' }, 'local')).toMatchObject({ + hits: [{ source: { filePath: '/host/transcript.jsonl' } }] + }) expect(await aiVaultApi.searchStatus()).toMatchObject({ enabled: true, generation: 7 }) + expect(sshSearch).not.toHaveBeenCalled() + expect(runtimeSearch).not.toHaveBeenCalled() }) it('rejects malformed renderer input and uses typed unavailable', async () => { expect(await aiVaultApi.searchSessions({ query: 'needle' })).toEqual({ kind: 'unavailable', reason: 'no-service' }) - expect(() => handlers.get('aiVault:searchSessions')!(null, { query: 1 })).toThrow() - expect(() => handlers.get('aiVault:searchStatus')!(null, 42)).toThrow() + await expect(handlers.get('aiVault:searchSessions')!(null, { query: 1 })).rejects.toThrow() + await expect(handlers.get('aiVault:searchStatus')!(null, 42)).rejects.toThrow() }) it('routes one SSH target without touching the local index and redacts received paths', async () => { const local = fakeSearchService() setSessionSearchService(local) - remote.mockResolvedValue(searchResults()) - const result = await aiVaultApi.searchSessions({ query: 'needle' }, 'ssh-host') - expect(remote).toHaveBeenCalledWith('ssh-host', 'aiVault.searchSessions', { + sshSearch.mockResolvedValue(searchResults()) + const result = await aiVaultApi.searchSessions({ query: 'needle' }, 'ssh:ssh-host') + expect(sshSearch).toHaveBeenCalledWith('ssh-host', 'aiVault.searchSessions', { query: 'needle', limit: 20 }) - expect(result).toMatchObject({ hits: [{ source: { presence: 'present' } }] }) + expect(result).toMatchObject({ + hits: [{ executionHostId: 'ssh:ssh-host', source: { presence: 'present' } }] + }) expect(JSON.stringify(result)).not.toContain('resumeCommand') expect(local.search).not.toHaveBeenCalled() - remote.mockRejectedValue(new Error('SSH relay is not ready')) - await expect(aiVaultApi.searchSessions({ query: 'needle' }, 'ssh-host')).rejects.toThrow( + sshSearch.mockRejectedValue(new Error('SSH relay is not ready')) + await expect(aiVaultApi.searchSessions({ query: 'needle' }, 'ssh:ssh-host')).rejects.toThrow( 'SSH relay is not ready' ) expect(local.search).not.toHaveBeenCalled() }) + it('routes one runtime environment over its RPC and stamps the answering host', async () => { + const local = fakeSearchService() + setSessionSearchService(local) + runtimeSearch.mockResolvedValue(searchResults()) + const result = await aiVaultApi.searchSessions({ query: 'needle' }, 'runtime:env-1') + expect(runtimeSearch).toHaveBeenCalledWith( + 'env-1', + 'aiVault.searchSessions', + { query: 'needle', limit: 20 }, + undefined + ) + expect(result).toMatchObject({ + hits: [{ executionHostId: 'runtime:env-1', source: { presence: 'present' } }] + }) + expect(JSON.stringify(result)).not.toContain('/host/transcript.jsonl') + expect(local.search).not.toHaveBeenCalled() + runtimeSearch.mockResolvedValue(unavailableSessionSearchStatus()) + expect(await aiVaultApi.searchStatus('runtime:env-1')).toEqual(unavailableSessionSearchStatus()) + expect(runtimeSearch).toHaveBeenLastCalledWith('env-1', 'aiVault.searchStatus', {}, undefined) + }) + it('maps a runtime unknown-method refusal to unavailable and keeps transport errors', async () => { + runtimeSearch.mockRejectedValue( + Object.assign(new Error('unknown method'), { code: 'method_not_found' }) + ) + expect(await aiVaultApi.searchSessions({ query: 'needle' }, 'runtime:env-1')).toEqual({ + kind: 'unavailable', + reason: 'no-service' + }) + runtimeSearch.mockRejectedValue( + Object.assign(new Error('runtime disconnected'), { code: 'connection_lost' }) + ) + await expect(aiVaultApi.searchSessions({ query: 'needle' }, 'runtime:env-1')).rejects.toThrow( + 'runtime disconnected' + ) + }) + it('reports unavailable when no runtime transport is injected', async () => { + handlers.clear() + registerAiVaultSearchHandlers() + expect(await aiVaultApi.searchSessions({ query: 'needle' }, 'runtime:env-1')).toEqual({ + kind: 'unavailable', + reason: 'no-service' + }) + expect(await aiVaultApi.searchStatus('runtime:env-1')).toMatchObject({ + enabled: false, + phase: 'idle' + }) + }) + it('refuses an unroutable host instead of widening it to every host', async () => { + setSessionSearchService(fakeSearchService()) + for (const scope of ['nope', 'ssh:', 'runtime:a|b']) { + await expect( + handlers.get('aiVault:searchSessions')!(null, { query: 'needle' }, scope) + ).rejects.toThrow('not available for this execution host') + await expect(handlers.get('aiVault:searchStatus')!(null, scope)).rejects.toThrow( + 'not available for this execution host' + ) + } + // Status describes one index, so the everything-scope is not routable either. + await expect(handlers.get('aiVault:searchStatus')!(null, 'all')).rejects.toThrow( + 'not available for this execution host' + ) + }) +}) + +describe('all-hosts search fan-out', () => { + beforeEach(() => { + sshHostInfos.mockReturnValue([{ targetId: 'ssh-host' }]) + runtimeHostInfos.mockReturnValue([{ environmentId: 'env-1', executionHostId: 'runtime:env-1' }]) + }) + + it('merges every host by recency, keeps local paths and withholds remote ones', async () => { + setSessionSearchService(serviceReturning(resultsAt('2026-01-02T00:00:00.000Z'))) + sshSearch.mockResolvedValue(resultsAt('2026-01-03T00:00:00.000Z')) + runtimeSearch.mockResolvedValue(resultsAt(null)) + const result = await aiVaultApi.searchSessions({ query: 'needle' }, 'all') + expect(result.kind).toBe('results') + if (result.kind !== 'results') { + return + } + expect(result.hits.map((hit) => hit.executionHostId)).toEqual([ + 'ssh:ssh-host', + 'local', + 'runtime:env-1' + ]) + expect(result.hosts).toEqual([ + { executionHostId: 'local', outcome: 'results' }, + { executionHostId: 'ssh:ssh-host', outcome: 'results' }, + { executionHostId: 'runtime:env-1', outcome: 'results' } + ]) + expect(result.generation).toBe(7) + expect(result.page).toEqual({ cursor: null, hasMore: false }) + const localHit = result.hits.find((hit) => hit.executionHostId === 'local') + expect(localHit?.source).toEqual({ + presence: 'present', + filePath: '/host/transcript.jsonl', + codexHome: '/host/codex' + }) + for (const hit of result.hits.filter((candidate) => candidate.executionHostId !== 'local')) { + expect(hit.source).toEqual({ presence: 'present' }) + } + expect( + JSON.stringify(result.hits.filter((hit) => hit.executionHostId !== 'local')) + ).not.toContain('/host/transcript.jsonl') + }) + it('asks every leg for newest order and cuts the merge to the requested limit', async () => { + setSessionSearchService(serviceReturning(resultsAt('2026-01-02T00:00:00.000Z'))) + sshSearch.mockResolvedValue(resultsAt('2026-01-03T00:00:00.000Z')) + runtimeSearch.mockResolvedValue(resultsAt('2026-01-01T00:00:00.000Z')) + const result = await aiVaultApi.searchSessions( + { query: 'needle', limit: 2, filters: { sort: 'relevance' } }, + 'all' + ) + expect(updatedAtOf(result)).toEqual(['2026-01-03T00:00:00.000Z', '2026-01-02T00:00:00.000Z']) + expect(sshSearch).toHaveBeenCalledWith('ssh-host', 'aiVault.searchSessions', { + query: 'needle', + limit: 2, + filters: { sort: 'newest' } + }) + }) + it('round-trips a per-host cursor map and re-asks only the hosts with more', async () => { + setSessionSearchService( + serviceReturning( + resultsAt('2026-01-02T00:00:00.000Z', { cursor: 'local-page-2', hasMore: true }) + ) + ) + sshSearch.mockResolvedValue(resultsAt('2026-01-03T00:00:00.000Z')) + runtimeSearch.mockResolvedValue( + resultsAt('2026-01-01T00:00:00.000Z', { cursor: 'runtime-page-2', hasMore: true }) + ) + const first = await aiVaultApi.searchSessions({ query: 'needle' }, 'all') + const cursor = first.kind === 'results' ? first.page.cursor : null + expect(first.kind === 'results' && first.page.hasMore).toBe(true) + expect(decodeMergedSearchCursor(cursor!)).toEqual({ + local: 'local-page-2', + 'runtime:env-1': 'runtime-page-2' + }) + + sshSearch.mockClear() + runtimeSearch.mockClear() + const second = await aiVaultApi.searchSessions({ query: 'needle', cursor: cursor! }, 'all') + expect(sshSearch).not.toHaveBeenCalled() + expect(runtimeSearch).toHaveBeenCalledWith( + 'env-1', + 'aiVault.searchSessions', + { query: 'needle', limit: 20, cursor: 'runtime-page-2', filters: { sort: 'newest' } }, + 10_000 + ) + expect(second.kind === 'results' && second.hosts).toEqual([ + { executionHostId: 'local', outcome: 'results' }, + { executionHostId: 'runtime:env-1', outcome: 'results' } + ]) + }) + it('reports a cursor for a host that has gone away without failing the merge', async () => { + sshHostInfos.mockReturnValue([]) + runtimeHostInfos.mockReturnValue([]) + setSessionSearchService(serviceReturning(resultsAt('2026-01-02T00:00:00.000Z'))) + const cursor = Buffer.from( + JSON.stringify({ local: 'local-page-2', 'ssh:gone': 'gone-page-2' }), + 'utf8' + ).toString('base64url') + const result = await aiVaultApi.searchSessions({ query: 'needle', cursor }, 'all') + expect(result.kind === 'results' && result.hosts).toEqual([ + { executionHostId: 'local', outcome: 'results' }, + { executionHostId: 'ssh:gone', outcome: 'unreachable' } + ]) + expect(updatedAtOf(result)).toEqual(['2026-01-02T00:00:00.000Z']) + }) + it('reports a stale, unavailable or unreachable leg without failing the merge', async () => { + setSessionSearchService(serviceReturning(resultsAt('2026-01-02T00:00:00.000Z'))) + sshSearch.mockResolvedValue({ kind: 'stale-cursor', generation: 3 }) + runtimeSearch.mockRejectedValue(new Error('runtime disconnected')) + const result = await aiVaultApi.searchSessions({ query: 'needle' }, 'all') + expect(result.kind === 'results' && result.hosts).toEqual([ + { executionHostId: 'local', outcome: 'results' }, + { executionHostId: 'ssh:ssh-host', outcome: 'stale-cursor' }, + { executionHostId: 'runtime:env-1', outcome: 'unreachable' } + ]) + expect(updatedAtOf(result)).toEqual(['2026-01-02T00:00:00.000Z']) + }) + it('refuses a merged cursor it did not mint', async () => { + setSessionSearchService(serviceReturning(resultsAt(null))) + expect( + await aiVaultApi.searchSessions({ query: 'needle', cursor: 'not-a-map' }, 'all') + ).toEqual({ kind: 'malformed-cursor' }) + expect(sshSearch).not.toHaveBeenCalled() + }) }) diff --git a/src/main/ipc/ai-vault-search.ts b/src/main/ipc/ai-vault-search.ts index badafffad29..24447a76542 100644 --- a/src/main/ipc/ai-vault-search.ts +++ b/src/main/ipc/ai-vault-search.ts @@ -52,11 +52,12 @@ let handlerOptions: AiVaultSearchHandlerOptions = {} export function registerAiVaultSearchHandlers(options: AiVaultSearchHandlerOptions = {}): void { handlerOptions = options - ipcMain.handle('aiVault:searchSessions', (_event, raw: unknown, rawScope?: unknown) => { + // Async so a refused scope reaches the renderer as a rejection, like every other parse failure. + ipcMain.handle('aiVault:searchSessions', async (_event, raw: unknown, rawScope?: unknown) => { const scope = requestedSearchScope(rawScope) return searchByExecutionHostScope(AiVaultSearchRequestSchema.parse(raw), scope) }) - ipcMain.handle('aiVault:searchStatus', (_event, rawScope?: unknown) => { + ipcMain.handle('aiVault:searchStatus', async (_event, rawScope?: unknown) => { const scope = requestedSearchScope(rawScope) // Status describes one index; there is nothing to merge across hosts. if (scope === ALL_EXECUTION_HOSTS_SCOPE) { diff --git a/src/renderer/src/web/preload-api/web-ai-vault-search-compat.test.ts b/src/renderer/src/web/preload-api/web-ai-vault-search-compat.test.ts index 5af48f7ed6b..17053ce41f6 100644 --- a/src/renderer/src/web/preload-api/web-ai-vault-search-compat.test.ts +++ b/src/renderer/src/web/preload-api/web-ai-vault-search-compat.test.ts @@ -45,16 +45,29 @@ describe('web session search preload compatibility', () => { await expect(api.searchSessions({ query: 'needle' })).rejects.toThrow('disconnected') await expect(api.searchStatus()).rejects.toThrow('disconnected') }) - it('rejects invalid host responses and unsupported SSH selection', async () => { + it('rejects invalid host responses', async () => { const api = createWebAiVaultApi() callRuntimeResult.mockResolvedValue({ kind: 'results' }) await expect(api.searchSessions({ query: 'needle' })).rejects.toThrow() await expect(api.searchStatus()).rejects.toThrow() + }) + it('answers for its own runtime and for `all`, and reports any other host unavailable', async () => { + const api = createWebAiVaultApi() + callRuntimeResult.mockResolvedValue(searchResults()) + for (const scope of ['runtime:owning-host', 'all'] as const) { + expect(await api.searchSessions({ query: 'needle' }, scope)).toMatchObject({ + kind: 'results' + }) + } + expect(callRuntimeResult).toHaveBeenCalledTimes(2) callRuntimeResult.mockClear() - await expect(api.searchSessions({ query: 'needle' }, 'other-host')).rejects.toThrow( - 'transcript-owning runtime' - ) - await expect(api.searchStatus('other-host')).rejects.toThrow('transcript-owning runtime') + for (const scope of ['runtime:other-host', 'ssh:box', 'local'] as const) { + expect(await api.searchSessions({ query: 'needle' }, scope)).toEqual({ + kind: 'unavailable', + reason: 'no-service' + }) + expect(await api.searchStatus(scope)).toEqual(unavailableSessionSearchStatus()) + } expect(callRuntimeResult).not.toHaveBeenCalled() }) }) diff --git a/src/shared/ai-vault-search-contract.test.ts b/src/shared/ai-vault-search-contract.test.ts index f5c092736a5..83c8b503e8d 100644 --- a/src/shared/ai-vault-search-contract.test.ts +++ b/src/shared/ai-vault-search-contract.test.ts @@ -41,6 +41,33 @@ describe('session search public contract', () => { ) expect(AiVaultSearchResponseSchema.safeParse({ kind: 'results', hits: [] }).success).toBe(false) }) + it('keeps host attribution optional in both wire directions', () => { + const legacy = searchResults() + expect(AiVaultSearchResponseSchema.parse(legacy)).toEqual(legacy) + expect(legacy.hits[0]).not.toHaveProperty('executionHostId') + expect(legacy).not.toHaveProperty('hosts') + const merged = { + ...searchResults(), + hits: [{ ...searchHit(), executionHostId: 'runtime:env-1' }], + hosts: [ + { executionHostId: 'local', outcome: 'results' }, + { executionHostId: 'ssh:box', outcome: 'unreachable' } + ] + } + expect(AiVaultSearchResponseSchema.parse(merged)).toEqual(merged) + expect( + AiVaultSearchResponseSchema.safeParse({ + ...searchResults(), + hosts: [{ executionHostId: 'local', outcome: 'exploded' }] + }).success + ).toBe(false) + expect( + AiVaultSearchResponseSchema.safeParse({ + ...searchResults(), + hits: [{ ...searchHit(), executionHostId: '' }] + }).success + ).toBe(false) + }) it('never accepts resume commands for an unverified or missing source', () => { for (const presence of ['unverifiable', 'missing'] as const) { const response = searchResults()