feat(ai-vault-search): route desktop search by execution host scope, including runtimes

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<string, unknown>,
+  timeoutMs?: number
+): Promise<unknown> {
+  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 })
+}

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

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<AiVaultSearchResponse>
+}
+
+// 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<Record<string, string>>): string {
+  return Buffer.from(JSON.stringify(byHost), 'utf8').toString('base64url')
+}
+
+export function decodeMergedSearchCursor(cursor: string): Record<string, string> | 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<string, string>)
+    : 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<AiVaultSearchResponse> {
+  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<string, string> = {}
+  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<T>(pending: Promise<T>, timeoutMs: number | undefined): Promise<T> {
+  if (timeoutMs === undefined) {
+    return pending
+  }
+  let timer: ReturnType<typeof setTimeout> | undefined
+  try {
+    return await Promise.race([
+      pending,
+      new Promise<never>((_resolve, reject) => {
+        timer = setTimeout(() => reject(new Error('Session search host timed out.')), timeoutMs)
+      })
+    ])
+  } finally {
+    clearTimeout(timer)
+  }
+}

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<string, unknown>,
+  timeoutMs?: number
+) => Promise<unknown>
 
-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<typeof createSessionSearchClient> {
-  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<AiVaultSearchResponse> {
+  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<AiVaultSearchStatus> {
+  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<typeof createSessionSearchClient> | 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
+    }))
+  ]
 }

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 }
             })
         })
       )

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,