mirror of
https://github.com/stablyai/orca.git
synced 2026-09-22 16:02:32 +00:00
feat(session-search): merge all-computers search across hosts
The `all` scope on `aiVault:searchSessions` now fans out from the desktop to every host the session list enumerates and merges the pages into one. Legs run in parallel: the local index through the search service, SSH and runtime hosts through the existing remote search client. Two fixed orders, because relevance scores from independent indexes are not comparable. `newest` asks every leg for recency and k-way merges on `updatedAt`, nulls last, ties broken on execution host id. `relevance` rotates hosts in host-id order by their own rank. The merged cursor is an opaque base64url payload holding each host's cursor, how many of its current page were already emitted, and the generation that offset counts into, plus the page size and sort the cursor belongs to. A host whose index moved is fenced to `stale` and stops contributing; the rest keep paging. Per-host outcomes ride back on one new optional `hosts` field on the results response. `aiVault:searchStatus` with `all` stays refused, and neither the runtime RPC nor the CLI gains the scope, so a fan-out is never two hops.
This commit is contained in:
@@ -6,5 +6,8 @@ export const AI_VAULT_ALL_HOST_TIMEOUT_MS = {
|
||||
// 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
|
||||
sshScan: 20_000,
|
||||
// A search reads an index rather than walking a home, but it shares the relay
|
||||
// with the scans, so it gets the relay budget rather than one of its own.
|
||||
search: 15_000
|
||||
} as const
|
||||
|
||||
@@ -0,0 +1,292 @@
|
||||
import { mkdir, rm, writeFile } from 'node:fs/promises'
|
||||
import { join } from 'node:path'
|
||||
import { afterEach, beforeEach, expect, it } from 'vitest'
|
||||
import { resetSessionParseCacheForTests } from '../ai-vault/session-scanner-parse-cache'
|
||||
import { resetTranscriptConsumersForTests } from '../ai-vault/session-transcript-consumers'
|
||||
import { SessionSearchInstance } from '../ai-vault-search/session-search-instance'
|
||||
import {
|
||||
claudeLines,
|
||||
openSessionSearchIndexerHarness,
|
||||
type SessionSearchIndexerHarness
|
||||
} from '../ai-vault-search/session-search-indexer-test-fixture'
|
||||
import type { AiVaultSearchResponse } from '../../shared/ai-vault-search-types'
|
||||
import type { ExecutionHostId } from '../../shared/execution-host'
|
||||
import { searchAllExecutionHosts, type SessionSearchHostLeg } from './ai-vault-search-all-hosts'
|
||||
import { encodeMergedSearchCursor } from './ai-vault-search-merged-cursor'
|
||||
|
||||
/**
|
||||
* Three real indexes over real transcripts, wired as three legs. Everything a
|
||||
* merged page claims — every hit exactly once, a purge fencing one host, an
|
||||
* unreachable host retried — is checked against the hit sets the indexes hold.
|
||||
*/
|
||||
|
||||
const HOST_IDS = ['local', 'ssh:alpha', 'ssh:beta'] as const
|
||||
const SESSIONS_PER_HOST = [14, 13, 13] as const
|
||||
|
||||
type Host = {
|
||||
executionHostId: ExecutionHostId
|
||||
harness: SessionSearchIndexerHarness
|
||||
instance: SessionSearchInstance
|
||||
sessionIds: string[]
|
||||
}
|
||||
|
||||
let hosts: Host[]
|
||||
|
||||
function sessionIdFor(index: number): string {
|
||||
return `aaaaaaaa-bbbb-4ccc-8ddd-${String(index).padStart(12, '0')}`
|
||||
}
|
||||
|
||||
async function writeSession(harness: SessionSearchIndexerHarness, index: number): Promise<string> {
|
||||
const sessionId = sessionIdFor(index)
|
||||
const path = join(harness.claudeProjectDir, `${sessionId}.jsonl`)
|
||||
await mkdir(harness.claudeProjectDir, { recursive: true })
|
||||
// A distinct start index per session gives every hit in the fixture its own
|
||||
// `updatedAt`, so the newest-first order across hosts is total.
|
||||
const lines = claudeLines([`needle transcript ${index}`], sessionId, index * 2)
|
||||
await writeFile(path, `${lines.join('\n')}\n`)
|
||||
return sessionId
|
||||
}
|
||||
|
||||
beforeEach(async () => {
|
||||
resetSessionParseCacheForTests()
|
||||
resetTranscriptConsumersForTests()
|
||||
hosts = []
|
||||
let nextIndex = 0
|
||||
for (const [position, executionHostId] of HOST_IDS.entries()) {
|
||||
const harness = await openSessionSearchIndexerHarness(`ss-all-hosts-${position}`)
|
||||
const sessionIds: string[] = []
|
||||
for (let n = 0; n < SESSIONS_PER_HOST[position]!; n++) {
|
||||
sessionIds.push(await writeSession(harness, nextIndex++))
|
||||
}
|
||||
const instance = new SessionSearchInstance({
|
||||
databasePath: harness.databasePath,
|
||||
roots: harness.roots,
|
||||
onError: (error) => {
|
||||
throw error
|
||||
}
|
||||
})
|
||||
// One process holds one index, so the transcript reader publishes every read
|
||||
// to every live consumer. Three machines means three indexes built alone.
|
||||
instance.apply({ enabled: true, historyDays: null })
|
||||
await instance.settled()
|
||||
instance.close()
|
||||
hosts.push({ executionHostId, harness, instance, sessionIds })
|
||||
}
|
||||
for (const host of hosts) {
|
||||
host.instance.apply({ enabled: true, historyDays: null })
|
||||
await host.instance.settled()
|
||||
const own = resultsOf(await host.instance.search({ query: 'needle', limit: 100 }))
|
||||
expect(own.hits.map((hit) => hit.sessionId).sort()).toEqual([...host.sessionIds].sort())
|
||||
}
|
||||
})
|
||||
|
||||
afterEach(async () => {
|
||||
for (const host of hosts) {
|
||||
host.instance.close()
|
||||
await host.harness.cleanup()
|
||||
}
|
||||
resetTranscriptConsumersForTests()
|
||||
resetSessionParseCacheForTests()
|
||||
})
|
||||
|
||||
function legs(overrides: Partial<Record<string, SessionSearchHostLeg['search']>> = {}) {
|
||||
return hosts.map((host) => ({
|
||||
executionHostId: host.executionHostId,
|
||||
search: overrides[host.executionHostId] ?? ((request) => host.instance.search(request))
|
||||
})) satisfies SessionSearchHostLeg[]
|
||||
}
|
||||
|
||||
function resultsOf(response: AiVaultSearchResponse) {
|
||||
if (response.kind !== 'results') {
|
||||
throw new Error(`expected results, got ${response.kind}`)
|
||||
}
|
||||
return response
|
||||
}
|
||||
|
||||
function keysOf(response: AiVaultSearchResponse): string[] {
|
||||
return resultsOf(response).hits.map((hit) => `${hit.executionHostId}/${hit.sessionId}`)
|
||||
}
|
||||
|
||||
function everyKey(): string[] {
|
||||
return hosts.flatMap((host) =>
|
||||
host.sessionIds.map((sessionId) => `${host.executionHostId}/${sessionId}`)
|
||||
)
|
||||
}
|
||||
|
||||
async function paginate(
|
||||
limit: number,
|
||||
sort: 'relevance' | 'newest',
|
||||
hostLegs = legs()
|
||||
): Promise<{ keys: string[]; pages: number }> {
|
||||
const request = { query: 'needle', limit, filters: { sort } }
|
||||
const keys: string[] = []
|
||||
let cursor: string | null = null
|
||||
let pages = 0
|
||||
do {
|
||||
const response = resultsOf(
|
||||
await searchAllExecutionHosts(cursor === null ? request : { ...request, cursor }, hostLegs)
|
||||
)
|
||||
keys.push(...keysOf(response))
|
||||
cursor = response.page.cursor
|
||||
pages++
|
||||
expect(pages).toBeLessThan(40)
|
||||
} while (cursor !== null)
|
||||
return { keys, pages }
|
||||
}
|
||||
|
||||
it('hands out every hit on every host exactly once, at limit 5 and limit 20', async () => {
|
||||
const expected = everyKey().sort()
|
||||
expect(expected).toHaveLength(40)
|
||||
for (const limit of [5, 20]) {
|
||||
for (const sort of ['relevance', 'newest'] as const) {
|
||||
const { keys } = await paginate(limit, sort)
|
||||
expect(new Set(keys).size, `${sort} at ${limit} repeated a hit`).toBe(keys.length)
|
||||
expect([...keys].sort(), `${sort} at ${limit} lost a hit`).toEqual(expected)
|
||||
}
|
||||
}
|
||||
})
|
||||
|
||||
it('orders a newest merge by recency across hosts, newest first', async () => {
|
||||
const response = resultsOf(
|
||||
await searchAllExecutionHosts(
|
||||
{ query: 'needle', limit: 20, filters: { sort: 'newest' } },
|
||||
legs()
|
||||
)
|
||||
)
|
||||
const updated = response.hits.map((hit) => hit.updatedAt)
|
||||
expect(updated).toEqual([...updated].sort().toReversed())
|
||||
// The 20 newest of the 40 are the 20 highest session indexes, which span hosts.
|
||||
expect(new Set(response.hits.map((hit) => hit.executionHostId)).size).toBeGreaterThan(1)
|
||||
})
|
||||
|
||||
it('rotates hosts in host-id order when merging by relevance', async () => {
|
||||
const response = resultsOf(await searchAllExecutionHosts({ query: 'needle', limit: 6 }, legs()))
|
||||
expect(response.hits.map((hit) => hit.executionHostId)).toEqual([
|
||||
'local',
|
||||
'ssh:alpha',
|
||||
'ssh:beta',
|
||||
'local',
|
||||
'ssh:alpha',
|
||||
'ssh:beta'
|
||||
])
|
||||
})
|
||||
|
||||
it('fences the purged host and keeps the other two paginating', async () => {
|
||||
const request = { query: 'needle', limit: 5 }
|
||||
const first = resultsOf(await searchAllExecutionHosts(request, legs()))
|
||||
expect(first.hosts?.every((host) => host.outcome === 'searched')).toBe(true)
|
||||
|
||||
// A real purge: the transcript is gone and a full reconcile publishes that.
|
||||
const purged = hosts[2]!
|
||||
await rm(join(purged.harness.claudeProjectDir, `${purged.sessionIds[0]!}.jsonl`))
|
||||
await purged.instance.reconcile()
|
||||
|
||||
const second = resultsOf(
|
||||
await searchAllExecutionHosts({ ...request, cursor: first.page.cursor! }, legs())
|
||||
)
|
||||
expect(second.hosts).toContainEqual({
|
||||
executionHostId: 'ssh:beta',
|
||||
outcome: 'stale'
|
||||
})
|
||||
expect(second.hits.some((hit) => hit.executionHostId === 'ssh:beta')).toBe(false)
|
||||
|
||||
const seen = [...keysOf(first), ...keysOf(second)]
|
||||
let cursor = second.page.cursor
|
||||
let pages = 0
|
||||
while (cursor !== null) {
|
||||
const page = resultsOf(await searchAllExecutionHosts({ ...request, cursor }, legs()))
|
||||
seen.push(...keysOf(page))
|
||||
cursor = page.page.cursor
|
||||
// A merge that never retires a host would page for ever; fail instead.
|
||||
expect((pages += 1)).toBeLessThan(40)
|
||||
}
|
||||
// Beta contributed only what it handed out before the purge; nothing repeats,
|
||||
// and both healthy hosts finished their own hit sets.
|
||||
expect(new Set(seen).size).toBe(seen.length)
|
||||
const betaEmitted = keysOf(first).filter((key) => key.startsWith('ssh:beta/'))
|
||||
expect([...seen].sort()).toEqual(
|
||||
[
|
||||
...hosts[0]!.sessionIds.map((id) => `local/${id}`),
|
||||
...hosts[1]!.sessionIds.map((id) => `ssh:alpha/${id}`),
|
||||
...betaEmitted
|
||||
].sort()
|
||||
)
|
||||
})
|
||||
|
||||
it('reports an unreachable host, keeps hasMore, and picks it up on the retry', async () => {
|
||||
let reject = true
|
||||
const flaky = () =>
|
||||
reject
|
||||
? Promise.reject(new Error('relay down'))
|
||||
: hosts[1]!.instance.search({
|
||||
query: 'needle',
|
||||
limit: 5,
|
||||
filters: { sort: 'relevance' }
|
||||
})
|
||||
const request = { query: 'needle', limit: 5 }
|
||||
const first = resultsOf(await searchAllExecutionHosts(request, legs({ 'ssh:alpha': flaky })))
|
||||
expect(first.hosts).toContainEqual({
|
||||
executionHostId: 'ssh:alpha',
|
||||
outcome: 'unreachable'
|
||||
})
|
||||
expect(first.page.hasMore).toBe(true)
|
||||
expect(first.hits.some((hit) => hit.executionHostId === 'ssh:alpha')).toBe(false)
|
||||
|
||||
reject = false
|
||||
const second = resultsOf(
|
||||
await searchAllExecutionHosts({ ...request, cursor: first.page.cursor! }, legs())
|
||||
)
|
||||
expect(second.hosts).toContainEqual({
|
||||
executionHostId: 'ssh:alpha',
|
||||
outcome: 'searched'
|
||||
})
|
||||
expect(second.hits.some((hit) => hit.executionHostId === 'ssh:alpha')).toBe(true)
|
||||
|
||||
const seen = [...keysOf(first), ...keysOf(second)]
|
||||
let cursor = second.page.cursor
|
||||
let pages = 0
|
||||
while (cursor !== null) {
|
||||
const page = resultsOf(await searchAllExecutionHosts({ ...request, cursor }, legs()))
|
||||
seen.push(...keysOf(page))
|
||||
cursor = page.page.cursor
|
||||
// A merge that never retires a host would page for ever; fail instead.
|
||||
expect((pages += 1)).toBeLessThan(40)
|
||||
}
|
||||
expect(new Set(seen).size).toBe(seen.length)
|
||||
expect([...seen].sort()).toEqual(everyKey().sort())
|
||||
})
|
||||
|
||||
it('refuses a cursor whose page size or host set no longer matches the request', async () => {
|
||||
const first = resultsOf(await searchAllExecutionHosts({ query: 'needle', limit: 5 }, legs()))
|
||||
expect(
|
||||
await searchAllExecutionHosts(
|
||||
{ query: 'needle', limit: 20, cursor: first.page.cursor! },
|
||||
legs()
|
||||
)
|
||||
).toEqual({ kind: 'malformed-cursor' })
|
||||
const nonHost = encodeMergedSearchCursor({
|
||||
limit: 5,
|
||||
sort: 'relevance',
|
||||
hosts: { 'not-a-host': { c: null, e: 0, g: 0 } }
|
||||
})
|
||||
expect(
|
||||
await searchAllExecutionHosts({ query: 'needle', limit: 5, cursor: nonHost }, legs())
|
||||
).toEqual({ kind: 'malformed-cursor' })
|
||||
})
|
||||
|
||||
it('reports a disabled host without aborting the merge', async () => {
|
||||
hosts[2]!.instance.apply({ enabled: false, historyDays: null })
|
||||
const { keys } = await paginate(5, 'relevance')
|
||||
expect(new Set(keys).size).toBe(keys.length)
|
||||
expect([...keys].sort()).toEqual(
|
||||
[
|
||||
...hosts[0]!.sessionIds.map((id) => `local/${id}`),
|
||||
...hosts[1]!.sessionIds.map((id) => `ssh:alpha/${id}`)
|
||||
].sort()
|
||||
)
|
||||
const first = resultsOf(await searchAllExecutionHosts({ query: 'needle', limit: 5 }, legs()))
|
||||
expect(first.hosts).toContainEqual({
|
||||
executionHostId: 'ssh:beta',
|
||||
outcome: 'disabled'
|
||||
})
|
||||
})
|
||||
@@ -0,0 +1,340 @@
|
||||
import { expect, it, vi } from 'vitest'
|
||||
import type {
|
||||
AiVaultSearchHit,
|
||||
AiVaultSearchRequest,
|
||||
AiVaultSearchResponse
|
||||
} from '../../shared/ai-vault-search-types'
|
||||
import type { ExecutionHostId } from '../../shared/execution-host'
|
||||
import { searchAllExecutionHosts, type SessionSearchHostLeg } from './ai-vault-search-all-hosts'
|
||||
import { encodeMergedSearchCursor } from './ai-vault-search-merged-cursor'
|
||||
|
||||
/** A host that ranks a fixed list and fences its own cursors on a generation change. */
|
||||
class StubHost {
|
||||
generation = 1
|
||||
pages = 0
|
||||
lastRequests: AiVaultSearchRequest[] = []
|
||||
|
||||
constructor(
|
||||
readonly executionHostId: ExecutionHostId,
|
||||
private sessions: readonly { id: string; updatedAt: string | null }[],
|
||||
/** Caps this host's own page, the way a smaller remote page size would. */
|
||||
private readonly pageSize = Number.POSITIVE_INFINITY
|
||||
) {}
|
||||
|
||||
purge(): void {
|
||||
this.sessions = this.sessions.slice(1)
|
||||
this.generation += 1
|
||||
}
|
||||
|
||||
leg(timeoutMs?: number): SessionSearchHostLeg {
|
||||
return {
|
||||
executionHostId: this.executionHostId,
|
||||
...(timeoutMs === undefined ? {} : { timeoutMs }),
|
||||
search: (request) => Promise.resolve(this.search(request))
|
||||
}
|
||||
}
|
||||
|
||||
private search(request: AiVaultSearchRequest): AiVaultSearchResponse {
|
||||
this.pages += 1
|
||||
this.lastRequests.push(request)
|
||||
const limit = Math.min(request.limit ?? 20, this.pageSize)
|
||||
let offset = 0
|
||||
if (request.cursor !== undefined) {
|
||||
const parsed = JSON.parse(Buffer.from(request.cursor, 'base64url').toString('utf8'))
|
||||
if (parsed.g !== this.generation) {
|
||||
return {
|
||||
kind: 'stale-cursor',
|
||||
generation: this.generation,
|
||||
expectedGeneration: parsed.g
|
||||
}
|
||||
}
|
||||
offset = parsed.o
|
||||
}
|
||||
const page = this.sessions.slice(offset, offset + limit)
|
||||
const hasMore = offset + page.length < this.sessions.length
|
||||
return {
|
||||
kind: 'results',
|
||||
hits: page.map((session) => stubHit(session.id, session.updatedAt)),
|
||||
page: {
|
||||
cursor: hasMore
|
||||
? Buffer.from(
|
||||
JSON.stringify({ g: this.generation, o: offset + page.length }),
|
||||
'utf8'
|
||||
).toString('base64url')
|
||||
: null,
|
||||
hasMore
|
||||
},
|
||||
generation: this.generation,
|
||||
truncated: {
|
||||
candidates: false,
|
||||
snippets: 1,
|
||||
query: false,
|
||||
freshness: false
|
||||
},
|
||||
durationMs: 1
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
function stubHit(sessionId: string, updatedAt: string | null): AiVaultSearchHit {
|
||||
return {
|
||||
agent: 'claude',
|
||||
// A host may claim any id; the merge overwrites it with the one it addressed.
|
||||
executionHostId: 'ssh:impostor',
|
||||
sessionId,
|
||||
title: sessionId,
|
||||
cwd: null,
|
||||
branch: null,
|
||||
updatedAt,
|
||||
messageCount: 1,
|
||||
score: 1,
|
||||
source: { presence: 'unverifiable' },
|
||||
evidence: null
|
||||
}
|
||||
}
|
||||
|
||||
function sessions(prefix: string, count: number, day = 1) {
|
||||
return Array.from({ length: count }, (_unused, index) => ({
|
||||
id: `${prefix}-${index}`,
|
||||
updatedAt: `2026-09-${String(day + index).padStart(2, '0')}T00:00:00.000Z`
|
||||
}))
|
||||
}
|
||||
|
||||
function resultsOf(response: AiVaultSearchResponse) {
|
||||
if (response.kind !== 'results') {
|
||||
throw new Error(`expected results, got ${response.kind}`)
|
||||
}
|
||||
return response
|
||||
}
|
||||
|
||||
function unavailableLeg(
|
||||
executionHostId: ExecutionHostId,
|
||||
reason: 'disabled' | 'not-ready' | 'no-service'
|
||||
): SessionSearchHostLeg {
|
||||
return {
|
||||
executionHostId,
|
||||
search: () => Promise.resolve({ kind: 'unavailable', reason })
|
||||
}
|
||||
}
|
||||
|
||||
it('answers an empty merge when no host is reachable to ask', async () => {
|
||||
const response = resultsOf(await searchAllExecutionHosts({ query: 'needle' }, []))
|
||||
expect(response).toMatchObject({
|
||||
hits: [],
|
||||
page: { cursor: null, hasMore: false },
|
||||
generation: 0,
|
||||
hosts: []
|
||||
})
|
||||
})
|
||||
|
||||
it('asks every leg for the merged order and passes freshness through unchanged', async () => {
|
||||
const host = new StubHost('local', sessions('a', 2))
|
||||
await searchAllExecutionHosts(
|
||||
{
|
||||
query: 'needle',
|
||||
freshness: 'wait-until-current',
|
||||
filters: { sort: 'newest' },
|
||||
debug: true
|
||||
},
|
||||
[host.leg()]
|
||||
)
|
||||
expect(host.lastRequests[0]).toEqual({
|
||||
query: 'needle',
|
||||
freshness: 'wait-until-current',
|
||||
filters: { sort: 'newest' }
|
||||
})
|
||||
})
|
||||
|
||||
it('forces the merged order onto a leg the caller left to default', async () => {
|
||||
const host = new StubHost('local', sessions('a', 2))
|
||||
await searchAllExecutionHosts({ query: 'needle' }, [host.leg()])
|
||||
expect(host.lastRequests[0]?.filters).toEqual({ sort: 'relevance' })
|
||||
})
|
||||
|
||||
it('stamps every hit with the host the desktop addressed', async () => {
|
||||
const host = new StubHost('ssh:box', sessions('a', 2))
|
||||
const response = resultsOf(await searchAllExecutionHosts({ query: 'needle' }, [host.leg()]))
|
||||
expect(response.hits.map((hit) => hit.executionHostId)).toEqual(['ssh:box', 'ssh:box'])
|
||||
})
|
||||
|
||||
it('merges newest first across hosts, with undated hits last', async () => {
|
||||
const left = new StubHost('local', [
|
||||
{ id: 'l-new', updatedAt: '2026-09-09T00:00:00.000Z' },
|
||||
{ id: 'l-none', updatedAt: null }
|
||||
])
|
||||
const right = new StubHost('ssh:box', [
|
||||
{ id: 'r-mid', updatedAt: '2026-09-05T00:00:00.000Z' },
|
||||
{ id: 'r-old', updatedAt: '2026-09-01T00:00:00.000Z' }
|
||||
])
|
||||
const response = resultsOf(
|
||||
await searchAllExecutionHosts({ query: 'needle', filters: { sort: 'newest' } }, [
|
||||
left.leg(),
|
||||
right.leg()
|
||||
])
|
||||
)
|
||||
expect(response.hits.map((hit) => hit.sessionId)).toEqual(['l-new', 'r-mid', 'r-old', 'l-none'])
|
||||
})
|
||||
|
||||
it('breaks a recency tie on execution host id', async () => {
|
||||
const at = '2026-09-05T00:00:00.000Z'
|
||||
const response = resultsOf(
|
||||
await searchAllExecutionHosts({ query: 'needle', filters: { sort: 'newest' } }, [
|
||||
new StubHost('ssh:box', [{ id: 'later-host', updatedAt: at }]).leg(),
|
||||
new StubHost('local', [{ id: 'earlier-host', updatedAt: at }]).leg()
|
||||
])
|
||||
)
|
||||
expect(response.hits.map((hit) => hit.sessionId)).toEqual(['earlier-host', 'later-host'])
|
||||
})
|
||||
|
||||
it('rotates hosts by host-id order when merging by relevance', async () => {
|
||||
const response = resultsOf(
|
||||
await searchAllExecutionHosts({ query: 'needle', limit: 4 }, [
|
||||
new StubHost('ssh:box', sessions('b', 3)).leg(),
|
||||
new StubHost('local', sessions('a', 3)).leg()
|
||||
])
|
||||
)
|
||||
expect(response.hits.map((hit) => hit.sessionId)).toEqual(['a-0', 'b-0', 'a-1', 'b-1'])
|
||||
})
|
||||
|
||||
it('reports a host that refuses its own cursor as stale and keeps the merge going', async () => {
|
||||
const healthy = new StubHost('local', sessions('a', 6))
|
||||
const purged = new StubHost('ssh:box', sessions('b', 6))
|
||||
const legs = [healthy.leg(), purged.leg()]
|
||||
const first = resultsOf(await searchAllExecutionHosts({ query: 'needle', limit: 4 }, legs))
|
||||
purged.purge()
|
||||
const second = resultsOf(
|
||||
await searchAllExecutionHosts({ query: 'needle', limit: 4, cursor: first.page.cursor! }, legs)
|
||||
)
|
||||
expect(second.hosts).toEqual([
|
||||
{ executionHostId: 'local', outcome: 'searched' },
|
||||
{ executionHostId: 'ssh:box', outcome: 'stale' }
|
||||
])
|
||||
expect(second.hits.every((hit) => hit.executionHostId === 'local')).toBe(true)
|
||||
})
|
||||
|
||||
it('names an unavailable host by its reason without losing the other hosts', async () => {
|
||||
const healthy = new StubHost('local', sessions('a', 2))
|
||||
const response = resultsOf(
|
||||
await searchAllExecutionHosts({ query: 'needle' }, [
|
||||
healthy.leg(),
|
||||
unavailableLeg('ssh:off', 'disabled'),
|
||||
unavailableLeg('ssh:cold', 'not-ready'),
|
||||
unavailableLeg('runtime:old', 'no-service')
|
||||
])
|
||||
)
|
||||
expect(response.hosts).toEqual([
|
||||
{ executionHostId: 'local', outcome: 'searched' },
|
||||
{ executionHostId: 'runtime:old', outcome: 'no-service' },
|
||||
{ executionHostId: 'ssh:cold', outcome: 'not-ready' },
|
||||
{ executionHostId: 'ssh:off', outcome: 'disabled' }
|
||||
])
|
||||
expect(response.hits).toHaveLength(2)
|
||||
// Nothing more is owed, so a disabled host does not hold the page open.
|
||||
expect(response.page.hasMore).toBe(false)
|
||||
})
|
||||
|
||||
it('calls a leg that times out unreachable and retries it on the next page', async () => {
|
||||
vi.useFakeTimers()
|
||||
try {
|
||||
const stalled: SessionSearchHostLeg = {
|
||||
executionHostId: 'ssh:slow',
|
||||
timeoutMs: 50,
|
||||
search: () => new Promise(() => undefined)
|
||||
}
|
||||
const healthy = new StubHost('local', sessions('a', 2))
|
||||
const pending = searchAllExecutionHosts({ query: 'needle' }, [healthy.leg(), stalled])
|
||||
await vi.advanceTimersByTimeAsync(60)
|
||||
const first = resultsOf(await pending)
|
||||
expect(first.hosts).toContainEqual({
|
||||
executionHostId: 'ssh:slow',
|
||||
outcome: 'unreachable'
|
||||
})
|
||||
expect(first.page.hasMore).toBe(true)
|
||||
|
||||
const recovered = new StubHost('ssh:slow', sessions('s', 2))
|
||||
const second = resultsOf(
|
||||
await searchAllExecutionHosts({ query: 'needle', cursor: first.page.cursor! }, [
|
||||
healthy.leg(),
|
||||
recovered.leg()
|
||||
])
|
||||
)
|
||||
// The healthy host finished, so only the retried host is still in the walk.
|
||||
expect(second.hosts).toEqual([{ executionHostId: 'ssh:slow', outcome: 'searched' }])
|
||||
expect(second.hits.map((hit) => hit.sessionId)).toEqual(['s-0', 's-1'])
|
||||
} finally {
|
||||
vi.useRealTimers()
|
||||
}
|
||||
})
|
||||
|
||||
it('reads at most three pages from one host per merged request', async () => {
|
||||
// Twelve hits behind pages of two is six host pages; the bound stops at three
|
||||
// and the cursor carries the unread page so nothing is lost.
|
||||
const host = new StubHost('local', sessions('a', 12), 2)
|
||||
const first = resultsOf(
|
||||
await searchAllExecutionHosts({ query: 'needle', limit: 10 }, [host.leg()])
|
||||
)
|
||||
expect(host.pages).toBe(3)
|
||||
expect(first.hits.map((hit) => hit.sessionId)).toEqual(['a-0', 'a-1', 'a-2', 'a-3', 'a-4', 'a-5'])
|
||||
expect(first.page.hasMore).toBe(true)
|
||||
|
||||
const second = resultsOf(
|
||||
await searchAllExecutionHosts({ query: 'needle', limit: 10, cursor: first.page.cursor! }, [
|
||||
host.leg()
|
||||
])
|
||||
)
|
||||
expect(second.hits.map((hit) => hit.sessionId)).toEqual([
|
||||
'a-6',
|
||||
'a-7',
|
||||
'a-8',
|
||||
'a-9',
|
||||
'a-10',
|
||||
'a-11'
|
||||
])
|
||||
})
|
||||
|
||||
it('refuses a cursor that belongs to a different query or host set', async () => {
|
||||
const host = new StubHost('local', sessions('a', 6))
|
||||
const first = resultsOf(
|
||||
await searchAllExecutionHosts({ query: 'needle', limit: 2 }, [host.leg()])
|
||||
)
|
||||
const refused = [
|
||||
{ query: 'needle', limit: 4, cursor: first.page.cursor! },
|
||||
{
|
||||
query: 'needle',
|
||||
limit: 2,
|
||||
filters: { sort: 'newest' as const },
|
||||
cursor: first.page.cursor!
|
||||
},
|
||||
{ query: 'needle', limit: 2, cursor: 'not a cursor' },
|
||||
{
|
||||
query: 'needle',
|
||||
limit: 2,
|
||||
cursor: encodeMergedSearchCursor({
|
||||
limit: 2,
|
||||
sort: 'relevance',
|
||||
hosts: { 'ssh:gone': { c: null, e: 0, g: 1 } }
|
||||
})
|
||||
}
|
||||
]
|
||||
for (const request of refused) {
|
||||
expect(await searchAllExecutionHosts(request, [host.leg()])).toEqual({
|
||||
kind: 'malformed-cursor'
|
||||
})
|
||||
}
|
||||
})
|
||||
|
||||
it('sums truncation across the hosts it searched', async () => {
|
||||
const response = resultsOf(
|
||||
await searchAllExecutionHosts({ query: 'needle' }, [
|
||||
new StubHost('local', sessions('a', 2)).leg(),
|
||||
new StubHost('ssh:box', sessions('b', 2)).leg(),
|
||||
unavailableLeg('ssh:off', 'disabled')
|
||||
])
|
||||
)
|
||||
expect(response.truncated).toEqual({
|
||||
candidates: false,
|
||||
snippets: 2,
|
||||
query: false,
|
||||
freshness: false
|
||||
})
|
||||
})
|
||||
@@ -0,0 +1,318 @@
|
||||
import { resolveSessionSearchLimit } from '../../shared/ai-vault-search-limit'
|
||||
import type {
|
||||
AiVaultSearchHit,
|
||||
AiVaultSearchHostOutcome,
|
||||
AiVaultSearchRequest,
|
||||
AiVaultSearchResponse
|
||||
} from '../../shared/ai-vault-search-types'
|
||||
import type { ExecutionHostId } from '../../shared/execution-host'
|
||||
import {
|
||||
decodeMergedSearchCursor,
|
||||
encodeMergedSearchCursor,
|
||||
type MergedSearchCursorEntry
|
||||
} from './ai-vault-search-merged-cursor'
|
||||
|
||||
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>
|
||||
}
|
||||
|
||||
type MergedSort = 'relevance' | 'newest'
|
||||
type HostOutcome = AiVaultSearchHostOutcome['outcome']
|
||||
type Truncation = {
|
||||
candidates: boolean
|
||||
snippets: number
|
||||
query: boolean
|
||||
freshness: boolean
|
||||
}
|
||||
|
||||
// One merged request reads at most this many pages from any single host.
|
||||
const MAX_HOST_PAGES_PER_REQUEST = 3
|
||||
|
||||
type HostWalk = {
|
||||
executionHostId: ExecutionHostId
|
||||
leg: SessionSearchHostLeg
|
||||
request: AiVaultSearchRequest
|
||||
outcome: HostOutcome
|
||||
cursor: string | null
|
||||
emitted: number
|
||||
generation: number
|
||||
pending: AiVaultSearchHit[]
|
||||
nextCursor: string | null
|
||||
pages: number
|
||||
/** This host still owes hits this request could not read; keep its entry. */
|
||||
carry: boolean
|
||||
truncated: Truncation
|
||||
}
|
||||
|
||||
/**
|
||||
* Fans one query out to every execution host and merges the pages into one.
|
||||
*
|
||||
* Relevance scores come from independent indexes and are not comparable, so the
|
||||
* two orders are the only two that mean anything across hosts: recency, which
|
||||
* every host can be asked for directly, and round-robin over each host's own
|
||||
* ranking. Legs are walked page by page, so a hit that lost the cut on one page
|
||||
* is emitted on the next instead of being dropped.
|
||||
*/
|
||||
export async function searchAllExecutionHosts(
|
||||
request: AiVaultSearchRequest,
|
||||
legs: readonly SessionSearchHostLeg[]
|
||||
): Promise<AiVaultSearchResponse> {
|
||||
const startedAt = Date.now()
|
||||
const limit = resolveSessionSearchLimit(request.limit)
|
||||
const sort: MergedSort = request.filters?.sort ?? 'relevance'
|
||||
const resumed = request.cursor === undefined ? null : decodeMergedSearchCursor(request.cursor)
|
||||
if (request.cursor !== undefined) {
|
||||
const known = new Set<string>(legs.map((leg) => leg.executionHostId))
|
||||
// A cursor belongs to one query over one host set; anything else is not ours.
|
||||
if (
|
||||
!resumed ||
|
||||
resumed.limit !== limit ||
|
||||
resumed.sort !== sort ||
|
||||
Object.keys(resumed.hosts).some((executionHostId) => !known.has(executionHostId))
|
||||
) {
|
||||
return { kind: 'malformed-cursor' }
|
||||
}
|
||||
}
|
||||
// A host absent from the cursor either finished or joined mid-walk; either way
|
||||
// it contributes nothing to this page. Host-id order fixes every tiebreak.
|
||||
const walks = legs
|
||||
.filter((leg) => !resumed || resumed.hosts[leg.executionHostId] !== undefined)
|
||||
.map((leg) => newHostWalk(leg, legRequest(request, sort)))
|
||||
.sort((left, right) => left.executionHostId.localeCompare(right.executionHostId))
|
||||
await Promise.all(
|
||||
walks.map((walk) => fetchHostPage(walk, resumed?.hosts[walk.executionHostId] ?? null))
|
||||
)
|
||||
const hits = await drainMergedPage(walks, limit, sort)
|
||||
return mergedSearchResponse(walks, { limit, sort }, hits, Date.now() - startedAt)
|
||||
}
|
||||
|
||||
/** Every leg answers in the merged order; the cursor and debug are this merge's own. */
|
||||
function legRequest(request: AiVaultSearchRequest, sort: MergedSort): AiVaultSearchRequest {
|
||||
const { cursor: _cursor, debug: _debug, ...rest } = request
|
||||
return { ...rest, filters: { ...request.filters, sort } }
|
||||
}
|
||||
|
||||
function newHostWalk(leg: SessionSearchHostLeg, request: AiVaultSearchRequest): HostWalk {
|
||||
return {
|
||||
executionHostId: leg.executionHostId,
|
||||
leg,
|
||||
request,
|
||||
outcome: 'unreachable',
|
||||
cursor: null,
|
||||
emitted: 0,
|
||||
generation: 0,
|
||||
pending: [],
|
||||
nextCursor: null,
|
||||
pages: 0,
|
||||
carry: false,
|
||||
truncated: {
|
||||
candidates: false,
|
||||
snippets: 0,
|
||||
query: false,
|
||||
freshness: false
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
async function fetchHostPage(walk: HostWalk, entry: MergedSearchCursorEntry | null): Promise<void> {
|
||||
walk.cursor = entry?.c ?? null
|
||||
walk.emitted = entry?.e ?? 0
|
||||
walk.generation = entry?.g ?? 0
|
||||
walk.pages += 1
|
||||
let response: AiVaultSearchResponse
|
||||
try {
|
||||
const { cursor } = walk
|
||||
response = await withLegTimeout(
|
||||
walk.leg.search(cursor === null ? walk.request : { ...walk.request, cursor }),
|
||||
walk.leg.timeoutMs
|
||||
)
|
||||
} catch (error) {
|
||||
console.error(`[ai-vault-search] ${walk.executionHostId} leg failed:`, error)
|
||||
// Keep its place: the next merged page retries this host from here.
|
||||
endHostWalk(walk, 'unreachable', true)
|
||||
return
|
||||
}
|
||||
if (response.kind === 'unavailable') {
|
||||
endHostWalk(walk, response.reason, false)
|
||||
return
|
||||
}
|
||||
// A cursor this merge minted can only be refused because the host's index
|
||||
// moved, so both refusals mean the same thing: this host is done for now.
|
||||
if (response.kind !== 'results') {
|
||||
endHostWalk(walk, 'stale', false)
|
||||
return
|
||||
}
|
||||
// `e` is an offset into one generation's ranked page, so it is only
|
||||
// meaningful while that generation stands. Nothing emitted, nothing to fence.
|
||||
if (walk.emitted > 0 && walk.generation !== response.generation) {
|
||||
endHostWalk(walk, 'stale', false)
|
||||
return
|
||||
}
|
||||
walk.outcome = 'searched'
|
||||
walk.generation = response.generation
|
||||
walk.pending = stampExecutionHost(response.hits, walk.executionHostId).slice(walk.emitted)
|
||||
walk.nextCursor = response.page.hasMore ? response.page.cursor : null
|
||||
walk.carry = false
|
||||
walk.truncated.candidates ||= response.truncated.candidates
|
||||
walk.truncated.snippets += response.truncated.snippets
|
||||
walk.truncated.query ||= response.truncated.query
|
||||
walk.truncated.freshness ||= response.truncated.freshness
|
||||
}
|
||||
|
||||
function endHostWalk(walk: HostWalk, outcome: HostOutcome, carry: boolean): void {
|
||||
walk.outcome = outcome
|
||||
walk.pending = []
|
||||
walk.nextCursor = null
|
||||
walk.carry = carry
|
||||
}
|
||||
|
||||
async function advanceHostWalk(walk: HostWalk): Promise<void> {
|
||||
while (walk.pending.length === 0 && walk.nextCursor !== null) {
|
||||
if (walk.pages >= MAX_HOST_PAGES_PER_REQUEST) {
|
||||
// Budget spent; the unread page's cursor is already this walk's nextCursor.
|
||||
walk.carry = true
|
||||
return
|
||||
}
|
||||
await fetchHostPage(walk, { c: walk.nextCursor, e: 0, g: walk.generation })
|
||||
}
|
||||
}
|
||||
|
||||
async function drainMergedPage(
|
||||
walks: readonly HostWalk[],
|
||||
limit: number,
|
||||
sort: MergedSort
|
||||
): Promise<AiVaultSearchHit[]> {
|
||||
const hits: AiVaultSearchHit[] = []
|
||||
let turn = 0
|
||||
while (hits.length < limit) {
|
||||
// Every head must be known before picking, so a lagging host is never skipped.
|
||||
for (const walk of walks) {
|
||||
await advanceHostWalk(walk)
|
||||
}
|
||||
const next = sort === 'newest' ? mostRecentWalk(walks) : walkWithTurn(walks, turn)
|
||||
if (!next) {
|
||||
return hits
|
||||
}
|
||||
turn = walks.indexOf(next) + 1
|
||||
hits.push(next.pending.shift()!)
|
||||
next.emitted += 1
|
||||
}
|
||||
return hits
|
||||
}
|
||||
|
||||
/** Round-robin in host-id order, resuming after whichever host answered last. */
|
||||
function walkWithTurn(walks: readonly HostWalk[], turn: number): HostWalk | null {
|
||||
for (let step = 0; step < walks.length; step++) {
|
||||
const walk = walks[(turn + step) % walks.length]
|
||||
if (walk && walk.pending.length > 0) {
|
||||
return walk
|
||||
}
|
||||
}
|
||||
return null
|
||||
}
|
||||
|
||||
function mostRecentWalk(walks: readonly HostWalk[]): HostWalk | null {
|
||||
let best: HostWalk | null = null
|
||||
for (const walk of walks) {
|
||||
const head = walk.pending[0]
|
||||
// Walks are in host-id order, so a strict comparison keeps the first host on a tie.
|
||||
if (head && (!best || byRecencyDescending(head, best.pending[0]!) < 0)) {
|
||||
best = walk
|
||||
}
|
||||
}
|
||||
return best
|
||||
}
|
||||
|
||||
function mergedSearchResponse(
|
||||
walks: readonly HostWalk[],
|
||||
query: { limit: number; sort: MergedSort },
|
||||
hits: AiVaultSearchHit[],
|
||||
durationMs: number
|
||||
): AiVaultSearchResponse {
|
||||
const hosts: AiVaultSearchHostOutcome[] = []
|
||||
const nextHosts: Record<string, MergedSearchCursorEntry> = {}
|
||||
const truncated: Truncation = {
|
||||
candidates: false,
|
||||
snippets: 0,
|
||||
query: false,
|
||||
freshness: false
|
||||
}
|
||||
for (const walk of walks) {
|
||||
hosts.push({
|
||||
executionHostId: walk.executionHostId,
|
||||
outcome: walk.outcome
|
||||
})
|
||||
const entry = nextCursorEntry(walk)
|
||||
if (entry) {
|
||||
nextHosts[walk.executionHostId] = entry
|
||||
}
|
||||
truncated.candidates ||= walk.truncated.candidates
|
||||
truncated.snippets += walk.truncated.snippets
|
||||
truncated.query ||= walk.truncated.query
|
||||
truncated.freshness ||= walk.truncated.freshness
|
||||
}
|
||||
const hasMore = Object.keys(nextHosts).length > 0
|
||||
const cursor = hasMore ? encodeMergedSearchCursor({ ...query, hosts: nextHosts }) : null
|
||||
return {
|
||||
kind: 'results',
|
||||
hits,
|
||||
page: { cursor, hasMore },
|
||||
generation: 0,
|
||||
truncated,
|
||||
durationMs,
|
||||
hosts
|
||||
}
|
||||
}
|
||||
|
||||
/** Resume where this request stopped: mid-page by skip count, else the unread page. */
|
||||
function nextCursorEntry(walk: HostWalk): MergedSearchCursorEntry | null {
|
||||
if (walk.pending.length > 0) {
|
||||
return { c: walk.cursor, e: walk.emitted, g: walk.generation }
|
||||
}
|
||||
if (walk.nextCursor !== null) {
|
||||
return { c: walk.nextCursor, e: 0, g: walk.generation }
|
||||
}
|
||||
return walk.carry ? { c: walk.cursor, e: walk.emitted, g: walk.generation } : null
|
||||
}
|
||||
|
||||
// This desktop owns which host it addressed; never trust an id the far side returned.
|
||||
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)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,38 @@
|
||||
import { expect, it } from 'vitest'
|
||||
import { decodeMergedSearchCursor, encodeMergedSearchCursor } from './ai-vault-search-merged-cursor'
|
||||
|
||||
const cursor = {
|
||||
limit: 20,
|
||||
sort: 'newest',
|
||||
hosts: {
|
||||
local: { c: null, e: 3, g: 7 },
|
||||
'ssh:box': { c: 'opaque', e: 0, g: 2 }
|
||||
}
|
||||
} as const
|
||||
|
||||
it('round-trips a merged cursor through base64url', () => {
|
||||
expect(decodeMergedSearchCursor(encodeMergedSearchCursor(cursor))).toEqual(cursor)
|
||||
})
|
||||
|
||||
it('refuses anything that is not a cursor this module minted', () => {
|
||||
for (const raw of ['', 'not-base64url!!', Buffer.from('{]').toString('base64url')]) {
|
||||
expect(decodeMergedSearchCursor(raw)).toBeNull()
|
||||
}
|
||||
})
|
||||
|
||||
it('refuses a payload whose shape would change what a skip count means', () => {
|
||||
const refused = [
|
||||
{ l: 20, s: 'newest', h: { local: { c: null, e: 3 } } },
|
||||
{ l: 20, s: 'newest', h: { local: { c: null, e: -1, g: 7 } } },
|
||||
{ l: 20, s: 'newest', h: { local: { c: null, e: 1.5, g: 7 } } },
|
||||
{ l: 20, s: 'sideways', h: {} },
|
||||
{ l: 0, s: 'newest', h: {} },
|
||||
{ s: 'newest', h: {} },
|
||||
// A host cursor is an opaque string; a decoded object is a forged one.
|
||||
{ l: 20, s: 'newest', h: { local: { c: { o: 1 }, e: 0, g: 7 } } }
|
||||
]
|
||||
for (const payload of refused) {
|
||||
const raw = Buffer.from(JSON.stringify(payload), 'utf8').toString('base64url')
|
||||
expect(decodeMergedSearchCursor(raw), JSON.stringify(payload)).toBeNull()
|
||||
}
|
||||
})
|
||||
@@ -0,0 +1,53 @@
|
||||
import { z } from 'zod'
|
||||
|
||||
/**
|
||||
* One host's place in a merged walk. `c` is the host cursor that produced the
|
||||
* page being consumed (null for that host's first page), `e` is how many of that
|
||||
* page's hits the merge already emitted, and `g` is the host generation `e`
|
||||
* counts into. Refetch with `c`, skip `e`, and no hit is skipped or repeated.
|
||||
*/
|
||||
export type MergedSearchCursorEntry = {
|
||||
c: string | null
|
||||
e: number
|
||||
g: number
|
||||
}
|
||||
|
||||
export type MergedSearchCursor = {
|
||||
/** Page size and sort the cursor was minted for; a cursor belongs to one query. */
|
||||
limit: number
|
||||
sort: 'relevance' | 'newest'
|
||||
hosts: Record<string, MergedSearchCursorEntry>
|
||||
}
|
||||
|
||||
const mergedCursorPayloadSchema = z.object({
|
||||
l: z.number().int().positive(),
|
||||
s: z.enum(['relevance', 'newest']),
|
||||
h: z.record(
|
||||
z.string().min(1),
|
||||
z.object({
|
||||
c: z.string().nullable(),
|
||||
e: z.number().int().nonnegative(),
|
||||
g: z.number().int().nonnegative()
|
||||
})
|
||||
)
|
||||
})
|
||||
|
||||
export function encodeMergedSearchCursor(cursor: MergedSearchCursor): string {
|
||||
const payload = { l: cursor.limit, s: cursor.sort, h: cursor.hosts }
|
||||
return Buffer.from(JSON.stringify(payload), 'utf8').toString('base64url')
|
||||
}
|
||||
|
||||
/** Null for anything that is not a cursor this module minted. */
|
||||
export function decodeMergedSearchCursor(raw: string): MergedSearchCursor | null {
|
||||
let parsed: unknown
|
||||
try {
|
||||
parsed = JSON.parse(Buffer.from(raw, 'base64url').toString('utf8'))
|
||||
} catch {
|
||||
return null
|
||||
}
|
||||
const result = mergedCursorPayloadSchema.safeParse(parsed)
|
||||
if (!result.success) {
|
||||
return null
|
||||
}
|
||||
return { limit: result.data.l, sort: result.data.s, hosts: result.data.h }
|
||||
}
|
||||
@@ -1,9 +1,10 @@
|
||||
import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'
|
||||
|
||||
const { clearSearch, handlers, sshSearch, runtimeSearch } = vi.hoisted(() => ({
|
||||
const { clearSearch, handlers, sshSearch, sshHostInfos, runtimeSearch } = vi.hoisted(() => ({
|
||||
clearSearch: vi.fn(),
|
||||
handlers: new Map<string, (...args: unknown[]) => Promise<unknown>>(),
|
||||
sshSearch: vi.fn(),
|
||||
sshHostInfos: vi.fn<() => { targetId: string }[]>(() => []),
|
||||
runtimeSearch: vi.fn()
|
||||
}))
|
||||
vi.mock('../ai-vault/session-scanner-service-spawn', () => ({
|
||||
@@ -17,17 +18,24 @@ vi.mock('electron', () => ({
|
||||
ipcRenderer: { invoke: (name: string, ...args: unknown[]) => handlers.get(name)!(null, ...args) }
|
||||
}))
|
||||
vi.mock('./ssh', () => ({
|
||||
requestActiveSshSessionSearch: sshSearch
|
||||
requestActiveSshSessionSearch: sshSearch,
|
||||
getActiveSshAiVaultHostInfos: sshHostInfos
|
||||
}))
|
||||
|
||||
import { registerAiVaultSearchHandlers } from './ai-vault-search'
|
||||
import { aiVaultApi } from '../../preload/api/ai-vault-bridge'
|
||||
import { setSessionSearchService } from '../ai-vault-search/session-search-service-registry'
|
||||
import { unavailableSessionSearchStatus } from '../../shared/ai-vault-search-client'
|
||||
import { fakeSearchService, searchResults } from '../../shared/ai-vault-search-test-fixture'
|
||||
import {
|
||||
fakeSearchService,
|
||||
searchHit,
|
||||
searchResults
|
||||
} from '../../shared/ai-vault-search-test-fixture'
|
||||
beforeEach(() => {
|
||||
handlers.clear()
|
||||
sshSearch.mockReset()
|
||||
sshHostInfos.mockReset()
|
||||
sshHostInfos.mockReturnValue([])
|
||||
runtimeSearch.mockReset()
|
||||
clearSearch.mockReset()
|
||||
registerAiVaultSearchHandlers({
|
||||
@@ -138,7 +146,7 @@ describe('desktop IPC and preload search boundary', () => {
|
||||
it('refuses an unroutable host instead of widening it to every host', async () => {
|
||||
const local = fakeSearchService()
|
||||
setSessionSearchService(local)
|
||||
for (const scope of ['all', 'nope', 'ssh:', 'runtime:a|b']) {
|
||||
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')
|
||||
@@ -151,4 +159,32 @@ describe('desktop IPC and preload search boundary', () => {
|
||||
expect(sshSearch).not.toHaveBeenCalled()
|
||||
expect(runtimeSearch).not.toHaveBeenCalled()
|
||||
})
|
||||
it('merges every enumerated host under the all scope, and still refuses an all status', async () => {
|
||||
setSessionSearchService(fakeSearchService())
|
||||
sshHostInfos.mockReturnValue([{ targetId: 'box' }])
|
||||
sshSearch.mockResolvedValue({
|
||||
...searchResults(),
|
||||
hits: [{ ...searchHit(), sessionId: 'far' }]
|
||||
})
|
||||
const merged = await aiVaultApi.searchSessions({ query: 'needle' }, 'all')
|
||||
expect(merged).toMatchObject({
|
||||
kind: 'results',
|
||||
hosts: [
|
||||
{ executionHostId: 'local', outcome: 'searched' },
|
||||
{ executionHostId: 'ssh:box', outcome: 'searched' }
|
||||
]
|
||||
})
|
||||
expect(
|
||||
merged.kind === 'results'
|
||||
? merged.hits.map((hit) => [hit.executionHostId, hit.sessionId])
|
||||
: null
|
||||
).toEqual([
|
||||
['local', 'host-session'],
|
||||
['ssh:box', 'far']
|
||||
])
|
||||
// A merged status would have to reconcile six phases into one; it stays refused.
|
||||
await expect(handlers.get('aiVault:searchStatus')!(null, 'all')).rejects.toThrow(
|
||||
'not available for this execution host'
|
||||
)
|
||||
})
|
||||
})
|
||||
|
||||
@@ -15,12 +15,20 @@ import type {
|
||||
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 { requestActiveSshSessionSearch } from './ssh'
|
||||
import { clearSessionSearchInService } from '../ai-vault/session-scanner-service-spawn'
|
||||
import { searchAllExecutionHosts, type SessionSearchHostLeg } from './ai-vault-search-all-hosts'
|
||||
import {
|
||||
getActiveRuntimeAiVaultHostInfosResult,
|
||||
getActiveSshAiVaultHostInfosResult
|
||||
} from './ai-vault'
|
||||
import { AI_VAULT_ALL_HOST_TIMEOUT_MS } from './ai-vault-all-host-timeouts'
|
||||
|
||||
export type RuntimeSessionSearchCall = (
|
||||
environmentId: string,
|
||||
@@ -42,8 +50,12 @@ export function registerAiVaultSearchHandlers(options: AiVaultSearchHandlerOptio
|
||||
handlerOptions = options
|
||||
// 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)
|
||||
const request = AiVaultSearchRequestSchema.parse(raw)
|
||||
// Only the desktop fans out: a runtime or CLI caller would make it two hops.
|
||||
if (scopeSchema.parse(rawScope) === ALL_EXECUTION_HOSTS_SCOPE) {
|
||||
return searchAllExecutionHosts(request, allExecutionHostLegs())
|
||||
}
|
||||
return searchByExecutionHostScope(request, requestedSearchScope(rawScope))
|
||||
})
|
||||
ipcMain.handle('aiVault:searchStatus', async (_event, rawScope?: unknown) => {
|
||||
const scope = requestedSearchScope(rawScope)
|
||||
@@ -87,6 +99,41 @@ async function searchByExecutionHostScope(
|
||||
: response
|
||||
}
|
||||
|
||||
/**
|
||||
* Every host the session list's `all` scope would enumerate, in one leg each.
|
||||
* A broken enumerator already degrades to an empty list rather than throwing,
|
||||
* so one unusable host class costs its own rows and not the merge.
|
||||
*/
|
||||
function allExecutionHostLegs(): SessionSearchHostLeg[] {
|
||||
const localLeg: SessionSearchHostLeg = {
|
||||
executionHostId: LOCAL_EXECUTION_HOST_ID,
|
||||
search: (request) => searchSessionService(request, 'ipc')
|
||||
}
|
||||
const sshLegs = getActiveSshAiVaultHostInfosResult().hostInfos.map(({ targetId }) =>
|
||||
remoteHostLeg({ kind: 'ssh', id: toSshExecutionHostId(targetId), targetId })
|
||||
)
|
||||
const runtimeLegs = getActiveRuntimeAiVaultHostInfosResult().hostInfos.map((hostInfo) =>
|
||||
remoteHostLeg({
|
||||
kind: 'runtime',
|
||||
id: hostInfo.executionHostId,
|
||||
environmentId: hostInfo.environmentId
|
||||
})
|
||||
)
|
||||
return [localLeg, ...sshLegs, ...runtimeLegs]
|
||||
}
|
||||
|
||||
function remoteHostLeg(host: ParsedExecutionHost): SessionSearchHostLeg {
|
||||
const client = remoteSearchClient(host, handlerOptions.callRuntimeSearch)
|
||||
return {
|
||||
executionHostId: host.id,
|
||||
timeoutMs: AI_VAULT_ALL_HOST_TIMEOUT_MS.search,
|
||||
search: (request) =>
|
||||
client
|
||||
? client.searchSessions(request)
|
||||
: Promise.resolve({ kind: 'unavailable', reason: 'no-service' })
|
||||
}
|
||||
}
|
||||
|
||||
function statusByExecutionHost(scope: ParsedExecutionHost): Promise<AiVaultSearchStatus> {
|
||||
if (scope.kind === 'local') {
|
||||
return sessionSearchServiceStatus({}, 'ipc')
|
||||
|
||||
@@ -202,14 +202,16 @@ async function scanAiVaultSessionsByHostScope(
|
||||
})
|
||||
}
|
||||
|
||||
function getActiveRuntimeAiVaultHostInfosResult(): AiVaultHostDiscoveryResult<RuntimeAiVaultHostInfo> {
|
||||
export function getActiveRuntimeAiVaultHostInfosResult(): AiVaultHostDiscoveryResult<RuntimeAiVaultHostInfo> {
|
||||
return discoverAiVaultHosts(() => handlerOptions.getActiveRuntimeAiVaultHostInfos?.() ?? [], {
|
||||
path: 'runtime environments',
|
||||
fallbackMessage: 'Runtime hosts are unavailable.'
|
||||
})
|
||||
}
|
||||
|
||||
function getActiveSshAiVaultHostInfosResult(): AiVaultHostDiscoveryResult<{ targetId: string }> {
|
||||
export function getActiveSshAiVaultHostInfosResult(): AiVaultHostDiscoveryResult<{
|
||||
targetId: string
|
||||
}> {
|
||||
return discoverAiVaultHosts(getActiveSshAiVaultHostInfos, {
|
||||
path: 'SSH hosts',
|
||||
fallbackMessage: 'SSH hosts are unavailable.'
|
||||
|
||||
@@ -61,6 +61,14 @@ export const AiVaultSearchTruncationSchema = z.object({
|
||||
query: z.boolean(),
|
||||
freshness: z.boolean()
|
||||
})
|
||||
/**
|
||||
* Per-host outcomes of an all-computers merge. Additive and desktop-only: no
|
||||
* host publishes this, and a reader that does not know it simply drops it.
|
||||
*/
|
||||
export const AiVaultSearchHostOutcomeSchema = z.object({
|
||||
executionHostId: executionHostIdSchema,
|
||||
outcome: z.enum(['searched', 'stale', 'disabled', 'not-ready', 'no-service', 'unreachable'])
|
||||
})
|
||||
const routeSchema = z.enum(['phrase', 'and', 'or', 'typo+phrase', 'typo+and', 'typo+or'])
|
||||
export const AiVaultSearchPlannerReportSchema = z.object({
|
||||
route: routeSchema,
|
||||
@@ -80,7 +88,8 @@ export const AiVaultSearchResponseSchema = z.discriminatedUnion('kind', [
|
||||
generation: z.number().int().nonnegative(),
|
||||
truncated: AiVaultSearchTruncationSchema,
|
||||
durationMs: z.number().nonnegative(),
|
||||
debug: AiVaultSearchDebugSchema.optional()
|
||||
debug: AiVaultSearchDebugSchema.optional(),
|
||||
hosts: z.array(AiVaultSearchHostOutcomeSchema).optional()
|
||||
}),
|
||||
z.object({
|
||||
kind: z.literal('stale-cursor'),
|
||||
|
||||
@@ -3,6 +3,7 @@ import type {
|
||||
AiVaultSearchRequestSchema,
|
||||
AiVaultSearchResponseSchema,
|
||||
AiVaultSearchHitSchema,
|
||||
AiVaultSearchHostOutcomeSchema,
|
||||
AiVaultSearchStatusSchema
|
||||
} from './ai-vault-search-contract'
|
||||
|
||||
@@ -20,3 +21,5 @@ export type AiVaultSearchResponse = z.infer<typeof AiVaultSearchResponseSchema>
|
||||
*/
|
||||
export type AiVaultSearchHit = z.infer<typeof AiVaultSearchHitSchema>
|
||||
export type AiVaultSearchStatus = z.infer<typeof AiVaultSearchStatusSchema>
|
||||
/** Only an all-computers merge reports these; a single-host answer omits them. */
|
||||
export type AiVaultSearchHostOutcome = z.infer<typeof AiVaultSearchHostOutcomeSchema>
|
||||
|
||||
Reference in New Issue
Block a user