From ff58ea47fda468ee6acf95b7b747e0dd4fc48ad8 Mon Sep 17 00:00:00 2001 From: Jinwoo-H Date: Sun, 13 Sep 2026 16:18:19 -0400 Subject: [PATCH] feat(ai-vault-search): route desktop search by execution host scope, including runtimes --- .../ai-vault/runtime-session-search-call.ts | 24 +++ src/main/ipc/ai-vault-all-host-timeouts.ts | 14 ++ src/main/ipc/ai-vault-search-all-hosts.ts | 178 +++++++++++++++++ src/main/ipc/ai-vault-search.ts | 183 ++++++++++++++++-- src/main/ipc/ai-vault.ts | 16 +- .../register-core-handlers.ts | 8 +- 6 files changed, 392 insertions(+), 31 deletions(-) create mode 100644 src/main/ai-vault/runtime-session-search-call.ts create mode 100644 src/main/ipc/ai-vault-all-host-timeouts.ts create mode 100644 src/main/ipc/ai-vault-search-all-hosts.ts diff --git a/src/main/ai-vault/runtime-session-search-call.ts b/src/main/ai-vault/runtime-session-search-call.ts new file mode 100644 index 00000000000..6c700a2cb58 --- /dev/null +++ b/src/main/ai-vault/runtime-session-search-call.ts @@ -0,0 +1,24 @@ +import { callRuntimeEnvironment } from '../ipc/runtime-environment-transport-routing' + +// Why: runtime RPC failures resolve as ok:false, but the shared search client +// classifies thrown errors by code, so the refusal has to keep its code to be +// recognised as an old host that lacks the method. +export async function callRuntimeSessionSearch( + userDataPath: string, + environmentId: string, + method: string, + params: Record, + timeoutMs?: number +): Promise { + const response = await callRuntimeEnvironment( + userDataPath, + environmentId, + method, + params, + timeoutMs + ) + if (response.ok === true) { + return response.result + } + throw Object.assign(new Error(response.error.message), { code: response.error.code }) +} diff --git a/src/main/ipc/ai-vault-all-host-timeouts.ts b/src/main/ipc/ai-vault-all-host-timeouts.ts new file mode 100644 index 00000000000..d840ed8af6f --- /dev/null +++ b/src/main/ipc/ai-vault-all-host-timeouts.ts @@ -0,0 +1,14 @@ +// Per-leg bounds for the all-hosts fan-outs, so one slow host cannot hold a merge open. +export const AI_VAULT_ALL_HOST_TIMEOUT_MS = { + runtimeScan: 3_000, + // Why: a remote home with many agent roots routinely needs seconds to walk, + // stat and parse. The old shared 3s bound emptied healthy SSH hosts in the + // all-hosts view; the relay gets a real scan budget and the whole leg (relay + // attempt plus any legacy crawl) stays bounded. + sshScanRelay: 15_000, + sshScan: 20_000, + // Why: search answers from an index rather than a filesystem walk, so a leg + // needs the relay round trip plus the 5s `wait-until-current` ceiling, not + // the scan budget above. + search: 10_000 +} as const diff --git a/src/main/ipc/ai-vault-search-all-hosts.ts b/src/main/ipc/ai-vault-search-all-hosts.ts new file mode 100644 index 00000000000..7626134bca7 --- /dev/null +++ b/src/main/ipc/ai-vault-search-all-hosts.ts @@ -0,0 +1,178 @@ +import { resolveSessionSearchLimit } from '../../shared/ai-vault-search-limit' +import type { + AiVaultSearchHit, + AiVaultSearchHostOutcome, + AiVaultSearchRequest, + AiVaultSearchResponse +} from '../../shared/ai-vault-search-types' +import { LOCAL_EXECUTION_HOST_ID, type ExecutionHostId } from '../../shared/execution-host' + +export type SessionSearchHostLeg = { + executionHostId: ExecutionHostId + // Omitted for the in-process local leg, which has no transport to hang on. + timeoutMs?: number + search: (request: AiVaultSearchRequest) => Promise +} + +// Never trust a host id the far side returned; this parent owns which host it addressed. +export function withSearchExecutionHost( + response: AiVaultSearchResponse, + executionHostId: ExecutionHostId +): AiVaultSearchResponse { + return response.kind === 'results' + ? { ...response, hits: stampExecutionHost(response.hits, executionHostId) } + : response +} + +export function encodeMergedSearchCursor(byHost: Readonly>): string { + return Buffer.from(JSON.stringify(byHost), 'utf8').toString('base64url') +} + +export function decodeMergedSearchCursor(cursor: string): Record | null { + let parsed: unknown + try { + parsed = JSON.parse(Buffer.from(cursor, 'base64url').toString('utf8')) + } catch { + return null + } + if (!parsed || typeof parsed !== 'object' || Array.isArray(parsed)) { + return null + } + const entries = Object.entries(parsed) + return entries.every(([, value]) => typeof value === 'string') + ? (Object.fromEntries(entries) as Record) + : null +} + +/** + * Why: relevance scores come from independent indexes and are not comparable, so + * the merge orders by recency only and asks every leg for that same order. + */ +export async function searchAllExecutionHosts( + request: AiVaultSearchRequest, + legs: readonly SessionSearchHostLeg[] +): Promise { + const startedAt = Date.now() + const resumed = request.cursor === undefined ? null : decodeMergedSearchCursor(request.cursor) + if (request.cursor !== undefined && !resumed) { + return { kind: 'malformed-cursor' } + } + const legRequest: AiVaultSearchRequest = { + ...request, + filters: { ...request.filters, sort: 'newest' } + } + delete legRequest.cursor + const targeted = resumed + ? legs.filter((leg) => resumed[leg.executionHostId] !== undefined) + : [...legs] + const settled = await Promise.all( + targeted.map(async (leg) => { + const hostCursor = resumed?.[leg.executionHostId] + const hostRequest = + hostCursor === undefined ? legRequest : { ...legRequest, cursor: hostCursor } + try { + return { leg, response: await withLegTimeout(leg.search(hostRequest), leg.timeoutMs) } + } catch (error) { + console.error(`[ai-vault-search] ${leg.executionHostId} leg failed:`, error) + return { leg, response: null } + } + }) + ) + return mergeHostSearchResults(settled, { + limit: resolveSessionSearchLimit(request.limit), + durationMs: Date.now() - startedAt, + // A cursor may name a host that has since disconnected; report it, don't fail the merge. + unreachable: resumed + ? Object.keys(resumed).filter((id) => !legs.some((leg) => leg.executionHostId === id)) + : [] + }) +} + +type SettledHostLeg = { leg: SessionSearchHostLeg; response: AiVaultSearchResponse | null } + +function mergeHostSearchResults( + settled: readonly SettledHostLeg[], + merge: { limit: number; durationMs: number; unreachable: readonly string[] } +): AiVaultSearchResponse { + const hits: AiVaultSearchHit[] = [] + const hosts: AiVaultSearchHostOutcome[] = [] + const nextCursors: Record = {} + const truncated = { candidates: false, snippets: 0, query: false, freshness: false } + let generation = 0 + let hasMore = false + for (const { leg, response } of settled) { + const executionHostId = leg.executionHostId + if (!response) { + hosts.push({ executionHostId, outcome: 'unreachable' }) + continue + } + hosts.push({ executionHostId, outcome: response.kind }) + if (response.kind !== 'results') { + continue + } + hits.push(...stampExecutionHost(response.hits, executionHostId)) + hasMore ||= response.page.hasMore + if (response.page.hasMore && response.page.cursor !== null) { + nextCursors[executionHostId] = response.page.cursor + } + truncated.candidates ||= response.truncated.candidates + truncated.snippets += response.truncated.snippets + truncated.query ||= response.truncated.query + truncated.freshness ||= response.truncated.freshness + if (executionHostId === LOCAL_EXECUTION_HOST_ID) { + generation = response.generation + } + } + for (const executionHostId of merge.unreachable) { + hosts.push({ executionHostId, outcome: 'unreachable' }) + } + const hasNextCursors = Object.keys(nextCursors).length > 0 + return { + kind: 'results', + hits: hits.sort(byRecencyDescending).slice(0, merge.limit), + page: { cursor: hasNextCursors ? encodeMergedSearchCursor(nextCursors) : null, hasMore }, + // Per-host generations live inside the cursor; the merged fence is the local host's. + generation, + truncated, + durationMs: merge.durationMs, + hosts + } +} + +function stampExecutionHost( + hits: readonly AiVaultSearchHit[], + executionHostId: ExecutionHostId +): AiVaultSearchHit[] { + return hits.map((hit) => ({ ...hit, executionHostId })) +} + +function byRecencyDescending(left: AiVaultSearchHit, right: AiVaultSearchHit): number { + const leftMs = updatedAtMs(left) + const rightMs = updatedAtMs(right) + if (leftMs === rightMs) { + return 0 + } + return leftMs === null ? 1 : rightMs === null ? -1 : rightMs - leftMs +} + +function updatedAtMs(hit: AiVaultSearchHit): number | null { + const parsed = hit.updatedAt === null ? Number.NaN : Date.parse(hit.updatedAt) + return Number.isNaN(parsed) ? null : parsed +} + +async function withLegTimeout(pending: Promise, timeoutMs: number | undefined): Promise { + if (timeoutMs === undefined) { + return pending + } + let timer: ReturnType | undefined + try { + return await Promise.race([ + pending, + new Promise((_resolve, reject) => { + timer = setTimeout(() => reject(new Error('Session search host timed out.')), timeoutMs) + }) + ]) + } finally { + clearTimeout(timer) + } +} diff --git a/src/main/ipc/ai-vault-search.ts b/src/main/ipc/ai-vault-search.ts index d76a0143251..badafffad29 100644 --- a/src/main/ipc/ai-vault-search.ts +++ b/src/main/ipc/ai-vault-search.ts @@ -4,29 +4,176 @@ import { searchSessionService, sessionSearchServiceStatus } from '../ai-vault-search/session-search-service-registry' -import { createSessionSearchClient } from '../../shared/ai-vault-search-client' +import { + createSessionSearchClient, + unavailableSessionSearchStatus +} from '../../shared/ai-vault-search-client' import { AiVaultSearchRequestSchema } from '../../shared/ai-vault-search-contract' -import { requestActiveSshSessionSearch } from './ssh' +import type { + AiVaultSearchRequest, + AiVaultSearchResponse, + AiVaultSearchStatus +} from '../../shared/ai-vault-search-types' +import { + ALL_EXECUTION_HOSTS_SCOPE, + LOCAL_EXECUTION_HOST_ID, + parseExecutionHostId, + toSshExecutionHostId, + type ParsedExecutionHost +} from '../../shared/execution-host' +import { getActiveSshAiVaultHostInfos, requestActiveSshSessionSearch } from './ssh' +import type { RuntimeAiVaultHostInfo } from './ai-vault-runtime-scan' +import { AI_VAULT_ALL_HOST_TIMEOUT_MS } from './ai-vault-all-host-timeouts' +import { + searchAllExecutionHosts, + withSearchExecutionHost, + type SessionSearchHostLeg +} from './ai-vault-search-all-hosts' -const targetSchema = z.string().min(1).optional() +export type RuntimeSessionSearchCall = ( + environmentId: string, + method: string, + params: Record, + timeoutMs?: number +) => Promise -export function registerAiVaultSearchHandlers(): void { - ipcMain.handle('aiVault:searchSessions', (_event, raw: unknown, rawTarget?: unknown) => { - const targetId = targetSchema.parse(rawTarget) - const request = AiVaultSearchRequestSchema.parse(raw) - return targetId - ? remoteClient(targetId).searchSessions(request) - : searchSessionService(request, 'ipc') +export type AiVaultSearchHandlerOptions = { + getActiveRuntimeAiVaultHostInfos?: () => readonly RuntimeAiVaultHostInfo[] + callRuntimeSearch?: RuntimeSessionSearchCall +} + +// One wording with the session list, which refuses the same unroutable scope. +const UNROUTABLE_HOST_MESSAGE = 'Agent Session History is not available for this execution host.' +const scopeSchema = z.string().min(1).optional() + +type RequestedSearchScope = typeof ALL_EXECUTION_HOSTS_SCOPE | ParsedExecutionHost + +let handlerOptions: AiVaultSearchHandlerOptions = {} + +export function registerAiVaultSearchHandlers(options: AiVaultSearchHandlerOptions = {}): void { + handlerOptions = options + ipcMain.handle('aiVault:searchSessions', (_event, raw: unknown, rawScope?: unknown) => { + const scope = requestedSearchScope(rawScope) + return searchByExecutionHostScope(AiVaultSearchRequestSchema.parse(raw), scope) }) - ipcMain.handle('aiVault:searchStatus', (_event, rawTarget?: unknown) => { - const targetId = targetSchema.parse(rawTarget) - return targetId ? remoteClient(targetId).searchStatus() : sessionSearchServiceStatus({}, 'ipc') + ipcMain.handle('aiVault:searchStatus', (_event, rawScope?: unknown) => { + const scope = requestedSearchScope(rawScope) + // Status describes one index; there is nothing to merge across hosts. + if (scope === ALL_EXECUTION_HOSTS_SCOPE) { + throw new Error(UNROUTABLE_HOST_MESSAGE) + } + return statusByExecutionHost(scope) }) } -function remoteClient(targetId: string): ReturnType { - return createSessionSearchClient( - (method, params) => requestActiveSshSessionSearch(targetId, method, params), - 'relay' - ) +/** + * Why not the list's `requestedExecutionHostScope`: it normalizes an unparseable + * id to `all`, which would answer an unroutable request by searching every host. + * Same parser, same omitted-means-this-host rule, but garbage is refused. + */ +function requestedSearchScope(raw: unknown): RequestedSearchScope { + const value = scopeSchema.parse(raw) + if (value === undefined) { + return { kind: 'local', id: LOCAL_EXECUTION_HOST_ID } + } + if (value === ALL_EXECUTION_HOSTS_SCOPE) { + return ALL_EXECUTION_HOSTS_SCOPE + } + const parsed = parseExecutionHostId(value) + if (!parsed) { + throw new Error(UNROUTABLE_HOST_MESSAGE) + } + return parsed +} + +async function searchByExecutionHostScope( + request: AiVaultSearchRequest, + scope: RequestedSearchScope +): Promise { + if (scope === ALL_EXECUTION_HOSTS_SCOPE) { + return searchAllExecutionHosts(request, allExecutionHostLegs()) + } + if (scope.kind === 'local') { + return searchSessionService(request, 'ipc') + } + const client = remoteSearchClient(scope, handlerOptions.callRuntimeSearch) + if (!client) { + return { kind: 'unavailable', reason: 'no-service' } + } + return withSearchExecutionHost(await client.searchSessions(request), scope.id) +} + +function statusByExecutionHost(scope: ParsedExecutionHost): Promise { + if (scope.kind === 'local') { + return sessionSearchServiceStatus({}, 'ipc') + } + const client = remoteSearchClient(scope, handlerOptions.callRuntimeSearch) + return client ? client.searchStatus() : Promise.resolve(unavailableSessionSearchStatus()) +} + +// Null for the local host and for a runtime environment with no injected transport. +function remoteSearchClient( + host: ParsedExecutionHost, + call: RuntimeSessionSearchCall | undefined, + timeoutMs?: number +): ReturnType | null { + if (host.kind === 'ssh') { + const { targetId } = host + return createSessionSearchClient( + (method, params) => requestActiveSshSessionSearch(targetId, method, params), + 'relay' + ) + } + if (host.kind === 'runtime' && call) { + const { environmentId } = host + return createSessionSearchClient( + (method, params) => call(environmentId, method, params, timeoutMs), + 'relay' + ) + } + return null +} + +function allExecutionHostLegs(): SessionSearchHostLeg[] { + const call = handlerOptions.callRuntimeSearch + return [ + { + executionHostId: LOCAL_EXECUTION_HOST_ID, + search: (request) => searchSessionService(request, 'ipc') + }, + ...activeRemoteSearchHosts().flatMap((host) => { + const client = remoteSearchClient(host, call, AI_VAULT_ALL_HOST_TIMEOUT_MS.search) + return client + ? [ + { + executionHostId: host.id, + timeoutMs: AI_VAULT_ALL_HOST_TIMEOUT_MS.search, + search: (request: AiVaultSearchRequest) => client.searchSessions(request) + } + ] + : [] + }) + ] +} + +// Enumerating live SSH sessions can throw; that must cost those hosts, not the merge. +function activeRemoteSearchHosts(): ParsedExecutionHost[] { + let sshTargetIds: readonly string[] = [] + try { + sshTargetIds = getActiveSshAiVaultHostInfos().map((hostInfo) => hostInfo.targetId) + } catch (error) { + console.error('[ai-vault-search] SSH host enumeration failed:', error) + } + return [ + ...sshTargetIds.map((targetId) => ({ + kind: 'ssh' as const, + id: toSshExecutionHostId(targetId), + targetId + })), + ...(handlerOptions.getActiveRuntimeAiVaultHostInfos?.() ?? []).map((hostInfo) => ({ + kind: 'runtime' as const, + id: hostInfo.executionHostId, + environmentId: hostInfo.environmentId + })) + ] } diff --git a/src/main/ipc/ai-vault.ts b/src/main/ipc/ai-vault.ts index 746f469dba7..1d09043c28b 100644 --- a/src/main/ipc/ai-vault.ts +++ b/src/main/ipc/ai-vault.ts @@ -58,15 +58,7 @@ import { type RuntimeAiVaultSessionTitleResolver } from './ai-vault-session-title-routing' import { projectStructuredAiVaultSessions } from '../ai-vault/structured-session-ownership' - -const AI_VAULT_ALL_HOST_RUNTIME_TIMEOUT_MS = 3_000 -// Why: a remote home with many agent roots routinely needs seconds to walk, -// stat and parse. The old shared 3s bound emptied healthy SSH hosts in the -// all-hosts view; the relay gets a real scan budget and the whole leg (relay -// attempt plus any legacy crawl) stays bounded so one host can't hold the -// merge open. -const AI_VAULT_ALL_HOST_SSH_RELAY_TIMEOUT_MS = 15_000 -const AI_VAULT_ALL_HOST_SSH_TIMEOUT_MS = 20_000 +import { AI_VAULT_ALL_HOST_TIMEOUT_MS } from './ai-vault-all-host-timeouts' type AiVaultHandlerOptions = AiVaultSessionSources & AiVaultResumeHandlerOptions & { @@ -159,8 +151,8 @@ async function scanAiVaultSessionsByHostScope( scan: () => scanSshAiVaultSessions(hostInfo.targetId, args, { signal, - timeoutMs: AI_VAULT_ALL_HOST_SSH_TIMEOUT_MS, - relayTimeoutMs: AI_VAULT_ALL_HOST_SSH_RELAY_TIMEOUT_MS + timeoutMs: AI_VAULT_ALL_HOST_TIMEOUT_MS.sshScan, + relayTimeoutMs: AI_VAULT_ALL_HOST_TIMEOUT_MS.sshScanRelay }) }) ), @@ -175,7 +167,7 @@ async function scanAiVaultSessionsByHostScope( hostInfo, scanner: handlerOptions.scanRuntimeAiVaultSessions, listArgs: args, - options: { signal, timeoutMs: AI_VAULT_ALL_HOST_RUNTIME_TIMEOUT_MS } + options: { signal, timeoutMs: AI_VAULT_ALL_HOST_TIMEOUT_MS.runtimeScan } }) }) ) diff --git a/src/main/ipc/register-core-handlers/register-core-handlers.ts b/src/main/ipc/register-core-handlers/register-core-handlers.ts index 2ca4bedf314..d492bf5a9f4 100644 --- a/src/main/ipc/register-core-handlers/register-core-handlers.ts +++ b/src/main/ipc/register-core-handlers/register-core-handlers.ts @@ -90,6 +90,7 @@ import { resolveRuntimeAiVaultSessionTitles, scanRuntimeAiVaultSessions } from '../../ai-vault/runtime-session-scanner' +import { callRuntimeSessionSearch } from '../../ai-vault/runtime-session-search-call' import type { PluginService } from '../../plugins/plugin-service' import type { PluginMarketplaceHandlerServices } from '../plugin-marketplaces' @@ -215,7 +216,12 @@ export function registerCoreHandlers( registerRuntimeHandlers(runtime) registerRuntimeEnvironmentHandlers(store) registerEphemeralVmHandlers(store, pluginService) - registerAiVaultSearchHandlers() + registerAiVaultSearchHandlers({ + getActiveRuntimeAiVaultHostInfos: () => + getSavedRuntimeAiVaultHostInfos(app.getPath('userData')), + callRuntimeSearch: (environmentId, method, params, timeoutMs) => + callRuntimeSessionSearch(app.getPath('userData'), environmentId, method, params, timeoutMs) + }) registerAiVaultHandlers({ ensureStructuredSessionOwnership: () => runtime.ensureStructuredAgentSessionHost(), getAdditionalCodexHomePaths: lifecycleOptions.getAdditionalAiVaultCodexHomePaths,