diff --git a/src/main/ai-vault/runtime-owned-ssh-session-list.test.ts b/src/main/ai-vault/runtime-owned-ssh-session-list.test.ts index 8dabe15833e..7c12537fb1f 100644 --- a/src/main/ai-vault/runtime-owned-ssh-session-list.test.ts +++ b/src/main/ai-vault/runtime-owned-ssh-session-list.test.ts @@ -60,6 +60,22 @@ describe('runtime-owned SSH AI Vault inventory', () => { ]) }) + it('skips malformed target rows from a paired runtime', async () => { + mocks.callRuntimeEnvironment.mockResolvedValueOnce({ + ok: true, + result: { targets: [null, 42, { id: 'hub-owned-host', connected: true }] } + }) + + await expect(listRuntimeOwnedSshAiVaultTargets('/user-data', 'hub-runtime')).resolves.toEqual([ + { + environmentId: 'hub-runtime', + targetId: 'hub-owned-host', + executionHostId: 'ssh:hub-owned-host', + connected: true + } + ]) + }) + it('finds which paired runtime registered an SSH target', async () => { await expect(findRuntimeOwningSshAiVaultHost('/user-data', 'hub-owned-host')).resolves.toEqual({ environmentId: 'hub-runtime', diff --git a/src/main/ai-vault/runtime-owned-ssh-session-list.ts b/src/main/ai-vault/runtime-owned-ssh-session-list.ts index 7e3bf2921f9..fefcab98e3c 100644 --- a/src/main/ai-vault/runtime-owned-ssh-session-list.ts +++ b/src/main/ai-vault/runtime-owned-ssh-session-list.ts @@ -75,7 +75,13 @@ async function loadRuntimeOwnedSshAiVaultTargets( return { hosts: [], authoritative: false } } const hosts = response.result.targets.flatMap((target) => { - if (typeof target.id !== 'string' || target.id.length === 0) { + if ( + typeof target !== 'object' || + target === null || + !('id' in target) || + typeof target.id !== 'string' || + target.id.length === 0 + ) { return [] } if (isRuntimeOwnedSshTargetId(target.id)) { @@ -86,7 +92,9 @@ async function loadRuntimeOwnedSshAiVaultTargets( environmentId, targetId: target.id, executionHostId: toSshExecutionHostId(target.id), - ...(typeof target.connected === 'boolean' ? { connected: target.connected } : {}) + ...('connected' in target && typeof target.connected === 'boolean' + ? { connected: target.connected } + : {}) } ] }) @@ -270,9 +278,7 @@ function unsupportedSshHostMessage(message: string): string { return message } -function isTargetSummaryList( - value: unknown -): value is { targets: { id?: unknown; connected?: unknown }[] } { +function isTargetSummaryList(value: unknown): value is { targets: unknown[] } { return ( typeof value === 'object' && value !== null && 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 381f5986509..4dd7e6f9bf1 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 @@ -149,6 +149,41 @@ describe('a runtime can scan an SSH host that it owns', () => { ]) }) + it('cancels a named SSH scan when its RPC request aborts', async () => { + let scanSignal: AbortSignal | undefined + scanSshAiVaultSessions.mockImplementationOnce( + ( + _targetId: string, + _args: unknown, + options: { signal?: AbortSignal } + ): Promise => + 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: SSH_HOST_OWNED_BY_RUNTIME }), + { signal: controller.signal } + ) + await vi.waitFor(() => expect(scanSignal).toBeDefined()) + + controller.abort() + + await expect(pending).resolves.toMatchObject({ ok: false }) + expect(scanSignal?.aborted).toBe(true) + }) + it('scans the identical host-local sources for the runtime ids it does accept', async () => { const dispatcher = makeDispatcher() diff --git a/src/main/runtime/rpc/methods/ai-vault.ts b/src/main/runtime/rpc/methods/ai-vault.ts index 3a9373ffd66..065dd5ed0cc 100644 --- a/src/main/runtime/rpc/methods/ai-vault.ts +++ b/src/main/runtime/rpc/methods/ai-vault.ts @@ -122,7 +122,7 @@ export const AI_VAULT_METHODS: RpcMethod[] = [ defineMethod({ name: 'aiVault.listSessions', params: AiVaultListSessionsParams, - handler: async (params, { runtime }) => { + handler: async (params, { runtime, signal }) => { const parsed = params.executionHostId ? parseExecutionHostId(params.executionHostId) : null const listArgs = { limit: params.unlimited ? undefined : params.limit, @@ -131,13 +131,13 @@ export const AI_VAULT_METHODS: RpcMethod[] = [ scopePaths: params.scopePaths } if (parsed?.kind === 'ssh') { - return scanRuntimeSshAiVaultSessions(parsed.targetId, listArgs) + return scanRuntimeSshAiVaultSessions(parsed.targetId, listArgs, signal) } const sshHosts = params.includeOwnedSshHosts === true ? getActiveSshAiVaultHostInfos() : [] const sshPromise = sshHosts.length > 0 ? Promise.all( - sshHosts.map((host) => scanRuntimeSshAiVaultSessions(host.targetId, listArgs)) + sshHosts.map((host) => scanRuntimeSshAiVaultSessions(host.targetId, listArgs, signal)) ) : Promise.resolve([]) let result: AiVaultListResult @@ -181,7 +181,8 @@ export const AI_VAULT_METHODS: RpcMethod[] = [ function scanRuntimeSshAiVaultSessions( targetId: string, - args: AiVaultListArgs + args: AiVaultListArgs, + signal?: AbortSignal ): Promise { const scopePaths = args.scopePaths ?? [] return scanHostLegWithCache({ @@ -195,6 +196,7 @@ function scanRuntimeSshAiVaultSessions( force: args.force === true, scan: () => scanSshAiVaultSessions(targetId, args, { + ...(signal ? { signal } : {}), timeoutMs: OWNED_SSH_SCAN_TIMEOUT_MS }) })