fix(ai-vault): observe owned SSH cancellation

This commit is contained in:
Brennan Benson
2026-08-14 15:40:27 -07:00
parent 59f0e0850e
commit d9b00956d2
2 changed files with 52 additions and 7 deletions
@@ -267,6 +267,54 @@ describe('a runtime can scan an SSH host that it owns', () => {
])
})
it('observes an owned SSH cancellation while the local scan is pending', async () => {
let resolveLocal!: (result: AiVaultListResult) => void
let scanSignal: AbortSignal | undefined
scanAiVaultSessionsInWorker.mockReturnValueOnce(
new Promise<AiVaultListResult>((resolve) => {
resolveLocal = resolve
})
)
scanSshAiVaultSessions.mockImplementationOnce(
(
_targetId: string,
_args: unknown,
options: { signal?: AbortSignal }
): Promise<AiVaultListResult> =>
new Promise((_resolve, reject) => {
scanSignal = options.signal
options.signal?.addEventListener(
'abort',
() => {
const error = new Error('cancelled')
error.name = 'AbortError'
reject(error)
},
{ once: true }
)
})
)
const dispatcher = makeDispatcher()
const controller = new AbortController()
const pending = dispatcher.dispatch(
makeRequest('aiVault.listSessions', {
executionHostId: 'runtime:hub-runtime',
includeOwnedSshHosts: true
}),
{ signal: controller.signal }
)
await vi.waitFor(() => expect(scanSignal).toBeDefined())
controller.abort()
const settled = await Promise.race([
pending,
new Promise<'pending'>((resolve) => setTimeout(() => resolve('pending')))
])
resolveLocal({ sessions: [], issues: [], scannedAt: SCANNED_AT })
expect(settled).toMatchObject({ ok: false })
})
it('reuses each SSH host leg across repeated non-force all-host scans', async () => {
const dispatcher = makeDispatcher()
const params = {
+4 -7
View File
@@ -140,24 +140,21 @@ export const AI_VAULT_METHODS: RpcMethod[] = [
sshHosts.map((host) => scanRuntimeSshAiVaultSessions(host.targetId, listArgs, signal))
)
: Promise.resolve([])
let result: AiVaultListResult
try {
result = await runtime.listAiVaultSessions(listArgs)
} catch (error) {
const localPromise = runtime.listAiVaultSessions(listArgs).catch((error: unknown) => {
if (params.includeOwnedSshHosts !== true) {
throw error
}
result = aiVaultScanIssueResult({
return aiVaultScanIssueResult({
executionHostId: LOCAL_EXECUTION_HOST_ID,
path: 'this computer',
message: error instanceof Error ? error.message : 'Local session scan failed.'
})
}
})
const [result, sshResults] = await Promise.all([localPromise, sshPromise])
// Why: web clients consume this response directly (no parent-side retag),
// so host-local sessions must come back stamped as the runtime they addressed.
const stamped =
parsed?.kind === 'runtime' ? restampAiVaultListResult(result, parsed.id) : result
const sshResults = await sshPromise
if (sshResults.length === 0) {
return stamped
}