From d9b00956d271707dd5fabeb59b901c4912e6983a Mon Sep 17 00:00:00 2001
From: Brennan Benson <79079362+brennanb2025@users.noreply.github.com>
Date: Fri, 14 Aug 2026 14:29:44 -0700
Subject: [PATCH] fix(ai-vault): observe owned SSH cancellation

---
 .../ai-vault-runtime-ssh-scope.repro.test.ts  | 48 +++++++++++++++++++
 src/main/runtime/rpc/methods/ai-vault.ts      | 11 ++---
 2 files changed, 52 insertions(+), 7 deletions(-)

diff --git a/src/main/runtime/rpc/methods/ai-vault-runtime-ssh-scope.repro.test.ts b/src/main/runtime/rpc/methods/ai-vault-runtime-ssh-scope.repro.test.ts
index 4dd7e6f9bf1..036dedeed74 100644
--- a/src/main/runtime/rpc/methods/ai-vault-runtime-ssh-scope.repro.test.ts
+++ b/src/main/runtime/rpc/methods/ai-vault-runtime-ssh-scope.repro.test.ts
@@ -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 = {
diff --git a/src/main/runtime/rpc/methods/ai-vault.ts b/src/main/runtime/rpc/methods/ai-vault.ts
index 065dd5ed0cc..060dd4e34d8 100644
--- a/src/main/runtime/rpc/methods/ai-vault.ts
+++ b/src/main/runtime/rpc/methods/ai-vault.ts
@@ -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
       }