diff --git a/mobile/src/agent-history/use-mobile-agent-history-state.ts b/mobile/src/agent-history/use-mobile-agent-history-state.ts index a0a82bee019..f97cf832c77 100644 --- a/mobile/src/agent-history/use-mobile-agent-history-state.ts +++ b/mobile/src/agent-history/use-mobile-agent-history-state.ts @@ -117,7 +117,8 @@ export function useMobileAgentHistoryState(params: MobileAgentHistoryStateParams const response = await client.sendRequest('aiVault.listSessions', { limit: MOBILE_AI_VAULT_SESSION_LIMIT, force: options.force, - scopePaths + scopePaths, + includeOwnedSshHosts: true }) if (!isCurrent()) { return 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 new file mode 100644 index 00000000000..48faeef19ef --- /dev/null +++ b/src/main/ai-vault/runtime-owned-ssh-session-list.test.ts @@ -0,0 +1,162 @@ +import { beforeEach, describe, expect, it, vi } from 'vitest' +import type { AiVaultListResult, AiVaultSession } from '../../shared/ai-vault-types' + +const mocks = vi.hoisted(() => ({ + callRuntimeEnvironment: vi.fn(), + listEnvironments: vi.fn() +})) + +vi.mock('../../shared/runtime-environment-store', () => ({ + listEnvironments: mocks.listEnvironments +})) + +vi.mock('../ipc/runtime-environment-transport-routing', () => ({ + callRuntimeEnvironment: mocks.callRuntimeEnvironment +})) + +const { + findRuntimeOwningSshAiVaultHost, + listRuntimeOwnedSshAiVaultTargets, + scanRuntimeOwnedSshAiVaultSessions +} = await import('./runtime-owned-ssh-session-list') + +beforeEach(() => { + vi.clearAllMocks() + mocks.listEnvironments.mockReturnValue([{ id: 'hub-runtime' }]) + mocks.callRuntimeEnvironment.mockResolvedValue({ + ok: true, + result: { targets: [{ id: 'hub-owned-host', label: 'Hub host' }] } + }) +}) + +describe('runtime-owned SSH AI Vault inventory', () => { + it('lists SSH targets owned by a paired runtime and skips recipe-VM ids', async () => { + mocks.callRuntimeEnvironment.mockResolvedValueOnce({ + ok: true, + result: { + targets: [ + { id: 'hub-owned-host', label: 'Hub host' }, + { id: 'runtime-ssh-recipe', label: 'Recipe VM' } + ] + } + }) + + await expect(listRuntimeOwnedSshAiVaultTargets('/user-data', 'hub-runtime')).resolves.toEqual([ + { + environmentId: 'hub-runtime', + targetId: 'hub-owned-host', + executionHostId: 'ssh:hub-owned-host' + } + ]) + }) + + it('finds which paired runtime registered an SSH target', async () => { + await expect(findRuntimeOwningSshAiVaultHost('/user-data', 'hub-owned-host')).resolves.toEqual({ + environmentId: 'hub-runtime', + targetId: 'hub-owned-host', + executionHostId: 'ssh:hub-owned-host' + }) + expect(mocks.callRuntimeEnvironment).toHaveBeenCalledWith( + '/user-data', + 'hub-runtime', + 'ssh.listTargetSummaries', + undefined, + undefined + ) + }) + + it('does not treat recipe-VM targets as runtime-owned SSH history hosts', async () => { + await expect(findRuntimeOwningSshAiVaultHost('/user-data', 'runtime-ssh-recipe')).resolves.toBe( + null + ) + expect(mocks.callRuntimeEnvironment).not.toHaveBeenCalled() + }) + + it('asks the owning runtime to scan a named SSH host and keeps the SSH stamp', async () => { + mocks.callRuntimeEnvironment.mockResolvedValueOnce({ + ok: true, + result: result([session('local', 'ssh-session')]) + }) + + const scanned = await scanRuntimeOwnedSshAiVaultSessions( + '/user-data', + 'hub-runtime', + 'hub-owned-host', + { limit: 25, force: true } + ) + + expect(mocks.callRuntimeEnvironment).toHaveBeenCalledWith( + '/user-data', + 'hub-runtime', + 'aiVault.listSessions', + { + limit: 25, + force: true, + executionHostId: 'ssh:hub-owned-host' + }, + undefined + ) + expect(scanned.sessions).toEqual([ + expect.objectContaining({ + executionHostId: 'ssh:hub-owned-host', + sessionId: 'ssh-session' + }) + ]) + }) + + it('turns an old-runtime host-id rejection into an SSH-host issue', async () => { + mocks.callRuntimeEnvironment.mockResolvedValueOnce({ + ok: false, + error: { message: 'Invalid runtime execution host id' } + }) + + const scanned = await scanRuntimeOwnedSshAiVaultSessions( + '/user-data', + 'hub-runtime', + 'hub-owned-host', + {} + ) + + expect(scanned.sessions).toEqual([]) + expect(scanned.issues).toEqual([ + expect.objectContaining({ + executionHostId: 'ssh:hub-owned-host', + path: 'hub-owned-host', + message: expect.stringContaining('cannot scan Agent Session History on its SSH hosts') + }) + ]) + }) +}) + +function result(sessions: AiVaultSession[]): AiVaultListResult { + return { sessions, issues: [], scannedAt: '2026-08-12T00:00:00.000Z' } +} + +function session( + executionHostId: AiVaultSession['executionHostId'], + sessionId: string +): AiVaultSession { + return { + id: `${executionHostId}:codex:${sessionId}:/sessions/${sessionId}.jsonl`, + executionHostId, + executionHostPlatform: 'linux', + agent: 'codex', + sessionId, + title: sessionId, + cwd: '/srv/app', + branch: null, + model: null, + filePath: `/sessions/${sessionId}.jsonl`, + codexHome: null, + createdAt: null, + updatedAt: '2026-08-12T03:00:00.000Z', + modifiedAt: '2026-08-12T00:00:00.000Z', + messageCount: 1, + totalTokens: 0, + previewMessages: [], + queuedMessageCount: 0, + subagentTranscriptCount: 0, + resumeCommand: `codex resume ${sessionId}`, + subagent: null + } +} diff --git a/src/main/ai-vault/runtime-owned-ssh-session-list.ts b/src/main/ai-vault/runtime-owned-ssh-session-list.ts new file mode 100644 index 00000000000..f73c8d7c9eb --- /dev/null +++ b/src/main/ai-vault/runtime-owned-ssh-session-list.ts @@ -0,0 +1,160 @@ +import { + AI_VAULT_SCOPE_PATHS_MAX_COUNT, + type AiVaultListArgs, + type AiVaultListResult +} from '../../shared/ai-vault-types' +import type { + AiVaultSessionTitlesArgs, + AiVaultSessionTitlesResult +} from '../../shared/ai-vault-session-title' +import { + isRuntimeOwnedSshTargetId, + toSshExecutionHostId, + type ExecutionHostId +} from '../../shared/execution-host' +import { listEnvironments } from '../../shared/runtime-environment-store' +import { callRuntimeEnvironment } from '../ipc/runtime-environment-transport-routing' +import { parseAiVaultListResult } from './session-list-result-validation' +import { parseAiVaultSessionTitlesResult } from './session-title-result-validation' +import { aiVaultScanIssueResult, restampAiVaultListResult } from './session-list-results' + +export type RuntimeOwnedSshAiVaultHost = { + environmentId: string + targetId: string + executionHostId: `ssh:${string}` +} + +export type RuntimeOwnedSshAiVaultScanOptions = { + timeoutMs?: number +} + +const UNSUPPORTED_SSH_HOST_PARAM = /invalid (runtime )?execution host id/i + +export async function listRuntimeOwnedSshAiVaultTargets( + userDataPath: string, + environmentId: string, + options: RuntimeOwnedSshAiVaultScanOptions = {} +): Promise { + const response = await callRuntimeEnvironment( + userDataPath, + environmentId, + 'ssh.listTargetSummaries', + undefined, + options.timeoutMs + ) + if (response.ok !== true || !isTargetSummaryList(response.result)) { + return [] + } + return response.result.targets.flatMap((target) => { + if (typeof target.id !== 'string' || target.id.length === 0) { + return [] + } + if (isRuntimeOwnedSshTargetId(target.id)) { + return [] + } + return [ + { + environmentId, + targetId: target.id, + executionHostId: toSshExecutionHostId(target.id) + } + ] + }) +} + +export async function findRuntimeOwningSshAiVaultHost( + userDataPath: string, + targetId: string, + options: RuntimeOwnedSshAiVaultScanOptions = {} +): Promise { + if (isRuntimeOwnedSshTargetId(targetId)) { + return null + } + for (const environment of listEnvironments(userDataPath)) { + const hosts = await listRuntimeOwnedSshAiVaultTargets(userDataPath, environment.id, options) + const match = hosts.find((host) => host.targetId === targetId) + if (match) { + return match + } + } + return null +} + +export async function scanRuntimeOwnedSshAiVaultSessions( + userDataPath: string, + environmentId: string, + targetId: string, + args: AiVaultListArgs, + options: RuntimeOwnedSshAiVaultScanOptions = {} +): Promise { + const executionHostId = toSshExecutionHostId(targetId) + const response = await callRuntimeEnvironment( + userDataPath, + environmentId, + 'aiVault.listSessions', + { + limit: args.limit, + unlimited: args.unlimited, + force: args.force, + scopePaths: args.scopePaths?.slice(0, AI_VAULT_SCOPE_PATHS_MAX_COUNT), + executionHostId + }, + options.timeoutMs + ) + if (response.ok !== true) { + return aiVaultScanIssueResult({ + executionHostId, + path: targetId, + message: unsupportedSshHostMessage(response.error.message) + }) + } + try { + return restampAiVaultListResult(parseAiVaultListResult(response.result), executionHostId) + } catch (error) { + return aiVaultScanIssueResult({ + executionHostId, + path: targetId, + message: `Invalid aiVault.listSessions response: ${ + error instanceof Error ? error.message : 'unexpected result shape' + }` + }) + } +} + +export async function resolveRuntimeOwnedSshAiVaultSessionTitles( + userDataPath: string, + environmentId: string, + targetId: string, + args: AiVaultSessionTitlesArgs +): Promise { + const executionHostId: ExecutionHostId = toSshExecutionHostId(targetId) + const response = await callRuntimeEnvironment( + userDataPath, + environmentId, + 'aiVault.resolveSessionTitles', + { requests: args.requests, executionHostId } + ) + if (response.ok !== true) { + return { titles: [] } + } + try { + return parseAiVaultSessionTitlesResult(response.result) + } catch { + return { titles: [] } + } +} + +function unsupportedSshHostMessage(message: string): string { + if (UNSUPPORTED_SSH_HOST_PARAM.test(message)) { + return 'This Orca server cannot scan Agent Session History on its SSH hosts. Update the server and try again.' + } + return message +} + +function isTargetSummaryList(value: unknown): value is { targets: { id?: unknown }[] } { + return ( + typeof value === 'object' && + value !== null && + Array.isArray((value as { targets?: unknown }).targets) + ) +} diff --git a/src/main/ai-vault/runtime-session-scanner.test.ts b/src/main/ai-vault/runtime-session-scanner.test.ts index 922a6af96d7..3c79b911162 100644 --- a/src/main/ai-vault/runtime-session-scanner.test.ts +++ b/src/main/ai-vault/runtime-session-scanner.test.ts @@ -68,6 +68,27 @@ describe('runtime AI Vault session scanner', () => { ) }) + it('asks the runtime to include SSH hosts it owns when the all-hosts scan opts in', async () => { + await scanRuntimeAiVaultSessions( + '/user-data', + 'env-1', + { limit: 25 }, + { includeOwnedSshHosts: true } + ) + + expect(mocks.callRuntimeEnvironment).toHaveBeenCalledWith( + '/user-data', + 'env-1', + 'aiVault.listSessions', + { + limit: 25, + executionHostId: 'runtime:env-1', + includeOwnedSshHosts: true + }, + undefined + ) + }) + it('surfaces project scope truncation from the runtime transport bound', async () => { const scopePaths = Array.from({ length: 80 }, (_, index) => `/srv/repo-${index}`) @@ -109,13 +130,42 @@ describe('runtime AI Vault session scanner', () => { ]) expect(scanResult.issues).toEqual([ expect.objectContaining({ - executionHostId: 'runtime:env-1', + executionHostId: 'ssh:dev-box', agent: 'codex', path: '/sessions/session-1.jsonl' }) ]) }) + it('keeps SSH-stamped rows from a runtime that scanned hosts it owns', async () => { + mocks.callRuntimeEnvironment.mockResolvedValueOnce({ + ok: true, + result: result( + [session('local', 'session-1'), session('ssh:hub-owned-host', 'session-ssh')], + [] + ) + }) + + const scanResult = await scanRuntimeAiVaultSessions( + '/user-data', + 'env-1', + {}, + { includeOwnedSshHosts: true } + ) + + expect(scanResult.sessions).toEqual([ + expect.objectContaining({ + executionHostId: 'runtime:env-1', + sessionId: 'session-1' + }), + expect.objectContaining({ + id: 'ssh:hub-owned-host:codex:session-ssh:/sessions/session-ssh.jsonl', + executionHostId: 'ssh:hub-owned-host', + sessionId: 'session-ssh' + }) + ]) + }) + it('stamps accepted sessions with the requested runtime host', async () => { const scanResult = await scanRuntimeAiVaultSessions('/user-data', 'env-1', {}) diff --git a/src/main/ai-vault/runtime-session-scanner.ts b/src/main/ai-vault/runtime-session-scanner.ts index b876c029b45..bf210253956 100644 --- a/src/main/ai-vault/runtime-session-scanner.ts +++ b/src/main/ai-vault/runtime-session-scanner.ts @@ -5,7 +5,7 @@ import { type AiVaultListResult, type AiVaultSession } from '../../shared/ai-vault-types' -import { toRuntimeExecutionHostId } from '../../shared/execution-host' +import { parseExecutionHostId, toRuntimeExecutionHostId } from '../../shared/execution-host' import { listEnvironments } from '../../shared/runtime-environment-store' import { callRuntimeEnvironment } from '../ipc/runtime-environment-transport-routing' import type { @@ -26,6 +26,7 @@ export type RuntimeAiVaultHostInfo = { export type RuntimeAiVaultScanOptions = { timeoutMs?: number + includeOwnedSshHosts?: boolean } // Why: zod strips unknown keys, so the repin home must be declared or the @@ -64,7 +65,8 @@ export async function scanRuntimeAiVaultSessions( // cap). Dropped paths only lose the older-than-recency-cap guarantee, // never the recent sessions themselves. scopePaths: args.scopePaths?.slice(0, AI_VAULT_SCOPE_PATHS_MAX_COUNT), - executionHostId + executionHostId, + ...(options.includeOwnedSshHosts === true ? { includeOwnedSshHosts: true } : {}) }, options.timeoutMs ) @@ -156,8 +158,16 @@ function withRuntimeExecutionHost( executionHostId: `runtime:${string}` ): AiVaultListResult { return { - sessions: result.sessions.map((session) => retagRuntimeSession(session, executionHostId)), - issues: result.issues.map((issue) => ({ ...issue, executionHostId })), + sessions: result.sessions.map((session) => + parseExecutionHostId(session.executionHostId)?.kind === 'ssh' + ? session + : retagRuntimeSession(session, executionHostId) + ), + issues: result.issues.map((issue) => + issue.executionHostId && parseExecutionHostId(issue.executionHostId)?.kind === 'ssh' + ? issue + : { ...issue, executionHostId } + ), scannedAt: result.scannedAt } } diff --git a/src/main/ipc/ai-vault-host-scope-scan.ts b/src/main/ipc/ai-vault-host-scope-scan.ts new file mode 100644 index 00000000000..ef0577d5c65 --- /dev/null +++ b/src/main/ipc/ai-vault-host-scope-scan.ts @@ -0,0 +1,171 @@ +import { scanSshAiVaultSessions } from '../ai-vault/ssh-session-list' +import { aiVaultScanIssueResult, mergeAiVaultListResults } from '../ai-vault/session-list-results' +import { requestedAiVaultSessionDepth } from '../../shared/ai-vault-session-depth' +import { + isAiVaultScanCancelledError, + type AiVaultListArgs, + type AiVaultListResult +} from '../../shared/ai-vault-types' +import { + LOCAL_EXECUTION_HOST_ID, + parseExecutionHostId, + toRuntimeExecutionHostId, + toSshExecutionHostId, + type ExecutionHostScope +} from '../../shared/execution-host' +import { getActiveSshAiVaultHostInfos } from './ssh' +import { discoverAiVaultHosts } from './ai-vault-host-discovery' +import { + scanRuntimeAiVaultSessions, + type RuntimeAiVaultHostInfo, + type RuntimeAiVaultScanner +} from './ai-vault-runtime-scan' +import { scanHostLegWithCache } from './ai-vault-host-leg-cache' +import { + scanSshAiVaultSessionsByOwner, + type RuntimeOwnedSshAiVaultScanner +} from './ai-vault-runtime-owned-ssh' +import type { RuntimeOwnedSshAiVaultHost } from '../ai-vault/runtime-owned-ssh-session-list' + +export const AI_VAULT_ALL_HOST_RUNTIME_TIMEOUT_MS = 3_000 +export const AI_VAULT_ALL_HOST_SSH_RELAY_TIMEOUT_MS = 15_000 +export const AI_VAULT_ALL_HOST_SSH_TIMEOUT_MS = 20_000 + +export type AiVaultHostScopeScanOptions = { + getActiveRuntimeAiVaultHostInfos?: () => readonly RuntimeAiVaultHostInfo[] + scanRuntimeAiVaultSessions?: RuntimeAiVaultScanner + findRuntimeOwningSshAiVaultHost?: (targetId: string) => Promise + scanRuntimeOwnedSshAiVaultSessions?: RuntimeOwnedSshAiVaultScanner + scanLocal: ( + args: AiVaultListArgs | undefined, + signal: AbortSignal | undefined + ) => Promise +} + +export async function scanAiVaultSessionsByHostScope( + args: AiVaultListArgs | undefined, + executionHostScope: ExecutionHostScope, + signal: AbortSignal | undefined, + cacheKey: string, + options: AiVaultHostScopeScanOptions +): Promise { + const depth = requestedAiVaultSessionDepth(args) + const scopePaths = args?.scopePaths ?? [] + if (executionHostScope === LOCAL_EXECUTION_HOST_ID) { + return scanLocalAiVaultSessionsAsIssue(options.scanLocal, args, signal) + } + if (executionHostScope === 'all') { + return scanAllAiVaultHosts(args, signal, cacheKey, depth, scopePaths, options) + } + + const parsed = parseExecutionHostId(executionHostScope) + if (parsed?.kind === 'ssh') { + return scanSshAiVaultSessionsByOwner({ + targetId: parsed.targetId, + listArgs: args, + signal, + timeoutMs: AI_VAULT_ALL_HOST_SSH_TIMEOUT_MS, + findOwner: options.findRuntimeOwningSshAiVaultHost, + scanOwned: options.scanRuntimeOwnedSshAiVaultSessions + }) + } + if (parsed?.kind === 'runtime') { + return scanRuntimeAiVaultSessions({ + hostInfo: { + environmentId: parsed.environmentId, + executionHostId: toRuntimeExecutionHostId(parsed.environmentId) + }, + scanner: options.scanRuntimeAiVaultSessions, + listArgs: args, + options: { signal } + }) + } + + return aiVaultScanIssueResult({ + executionHostId: executionHostScope, + path: executionHostScope, + message: 'Agent Session History is not available for this execution host.' + }) +} + +async function scanAllAiVaultHosts( + args: AiVaultListArgs | undefined, + signal: AbortSignal | undefined, + cacheKey: string, + depth: ReturnType, + scopePaths: readonly string[], + options: AiVaultHostScopeScanOptions +): Promise { + const runtimeHosts = discoverAiVaultHosts( + () => options.getActiveRuntimeAiVaultHostInfos?.() ?? [], + { path: 'runtime environments', fallbackMessage: 'Runtime hosts are unavailable.' } + ) + const sshHosts = discoverAiVaultHosts(getActiveSshAiVaultHostInfos, { + path: 'SSH hosts', + fallbackMessage: 'SSH hosts are unavailable.' + }) + const scannedResults = await Promise.all([ + scanLocalAiVaultSessionsAsIssue(options.scanLocal, args, signal), + ...sshHosts.hostInfos.map((hostInfo) => + scanHostLegWithCache({ + cacheKey: `${cacheKey}|${toSshExecutionHostId(hostInfo.targetId)}`, + depth, + scopePaths, + force: args?.force === true, + scan: () => + scanSshAiVaultSessions(hostInfo.targetId, args, { + signal, + timeoutMs: AI_VAULT_ALL_HOST_SSH_TIMEOUT_MS, + relayTimeoutMs: AI_VAULT_ALL_HOST_SSH_RELAY_TIMEOUT_MS + }) + }) + ), + ...runtimeHosts.hostInfos.map((hostInfo) => + scanHostLegWithCache({ + cacheKey: `${cacheKey}|${hostInfo.executionHostId}`, + depth, + scopePaths, + force: args?.force === true, + scan: () => + scanRuntimeAiVaultSessions({ + hostInfo, + scanner: options.scanRuntimeAiVaultSessions, + listArgs: args, + options: { + signal, + timeoutMs: AI_VAULT_ALL_HOST_RUNTIME_TIMEOUT_MS, + includeOwnedSshHosts: true + } + }) + }) + ) + ]) + return mergeAiVaultListResults( + [ + ...scannedResults, + ...(runtimeHosts.issue ? [runtimeHosts.issue] : []), + ...(sshHosts.issue ? [sshHosts.issue] : []) + ], + args?.limit, + args?.unlimited + ) +} + +async function scanLocalAiVaultSessionsAsIssue( + scanLocal: AiVaultHostScopeScanOptions['scanLocal'], + args: AiVaultListArgs | undefined, + signal: AbortSignal | undefined +): Promise { + try { + return await scanLocal(args, signal) + } catch (error) { + if (isAiVaultScanCancelledError(error)) { + throw error + } + return aiVaultScanIssueResult({ + executionHostId: LOCAL_EXECUTION_HOST_ID, + path: 'this computer', + message: error instanceof Error ? error.message : 'Local session scan failed.' + }) + } +} diff --git a/src/main/ipc/ai-vault-runtime-owned-ssh.test.ts b/src/main/ipc/ai-vault-runtime-owned-ssh.test.ts new file mode 100644 index 00000000000..ba1e45b6365 --- /dev/null +++ b/src/main/ipc/ai-vault-runtime-owned-ssh.test.ts @@ -0,0 +1,185 @@ +import { beforeEach, describe, expect, it, vi } from 'vitest' +import type { AiVaultListResult, AiVaultSession } from '../../shared/ai-vault-types' +import type { IFilesystemProvider } from '../providers/types' + +const mocks = vi.hoisted(() => ({ + scanAiVaultSessionsInWorker: vi.fn(), + resolveAiVaultSessionTitlesInWorker: vi.fn(), + scanRemoteAiVaultSessions: vi.fn(), + getSshFilesystemProvider: vi.fn(), + getActiveSshAiVaultHostInfo: vi.fn(), + getActiveSshAiVaultHostInfos: vi.fn(), + requestActiveSshAiVaultSessionList: vi.fn(), + requestActiveSshAiVaultSessionTitles: vi.fn(), + ipcHandle: vi.fn() +})) + +vi.mock('electron', () => ({ + app: { on: vi.fn() }, + ipcMain: { handle: mocks.ipcHandle } +})) + +vi.mock('../ai-vault/session-scanner-worker-spawn', () => ({ + scanAiVaultSessionsInWorker: mocks.scanAiVaultSessionsInWorker, + resolveAiVaultSessionTitlesInWorker: mocks.resolveAiVaultSessionTitlesInWorker, + resetAiVaultScannerWorkerForTests: vi.fn() +})) + +vi.mock('../ai-vault/remote-session-scanner', () => ({ + scanRemoteAiVaultSessions: mocks.scanRemoteAiVaultSessions +})) + +vi.mock('../wsl', () => ({ + getWslHomeAsync: vi.fn(), + listWslDistrosAsync: vi.fn().mockResolvedValue([]) +})) + +vi.mock('../providers/ssh-filesystem-dispatch', () => ({ + SSH_FILESYSTEM_PROVIDER_UNAVAILABLE_MESSAGE: + 'Remote connection dropped. Click Reconnect on the SSH target before retrying.', + getSshFilesystemProvider: mocks.getSshFilesystemProvider +})) + +vi.mock('./ssh', () => ({ + getActiveSshAiVaultHostInfo: mocks.getActiveSshAiVaultHostInfo, + getActiveSshAiVaultHostInfos: mocks.getActiveSshAiVaultHostInfos, + requestActiveSshAiVaultSessionList: mocks.requestActiveSshAiVaultSessionList, + requestActiveSshAiVaultSessionTitles: mocks.requestActiveSshAiVaultSessionTitles +})) + +const { _internals, registerAiVaultHandlers } = await import('./ai-vault') + +beforeEach(() => { + vi.clearAllMocks() + _internals.resetAiVaultCacheForTests() + mocks.scanAiVaultSessionsInWorker.mockResolvedValue(result([])) + mocks.resolveAiVaultSessionTitlesInWorker.mockResolvedValue({ titles: [] }) + mocks.scanRemoteAiVaultSessions.mockResolvedValue(result([])) + mocks.getSshFilesystemProvider.mockReturnValue({} as IFilesystemProvider) + mocks.requestActiveSshAiVaultSessionList.mockResolvedValue(null) + mocks.requestActiveSshAiVaultSessionTitles.mockResolvedValue(null) + mocks.getActiveSshAiVaultHostInfo.mockReturnValue(null) + mocks.getActiveSshAiVaultHostInfos.mockReturnValue([]) +}) + +describe('runtime-owned SSH AI Vault routing', () => { + it('routes an SSH host the local process does not own through its paired runtime', async () => { + const scanRuntimeOwnedSshAiVaultSessions = vi + .fn() + .mockResolvedValue(result([session('ssh:hub-owned-host', 'hub-ssh-session')])) + registerAiVaultHandlers({ + findRuntimeOwningSshAiVaultHost: async () => ({ + environmentId: 'hub-runtime', + targetId: 'hub-owned-host', + executionHostId: 'ssh:hub-owned-host' + }), + scanRuntimeOwnedSshAiVaultSessions + }) + + const scanned = await _internals.listAiVaultSessions({ + executionHostScope: 'ssh:hub-owned-host' + }) + + expect(scanRuntimeOwnedSshAiVaultSessions).toHaveBeenCalledWith( + 'hub-runtime', + 'hub-owned-host', + { executionHostScope: 'ssh:hub-owned-host' }, + { timeoutMs: 20_000 } + ) + expect(mocks.scanRemoteAiVaultSessions).not.toHaveBeenCalled() + expect(scanned.sessions).toEqual([ + expect.objectContaining({ + executionHostId: 'ssh:hub-owned-host', + sessionId: 'hub-ssh-session' + }) + ]) + }) + + it('does not send recipe-VM SSH targets through a paired runtime', async () => { + const findRuntimeOwningSshAiVaultHost = vi.fn() + registerAiVaultHandlers({ findRuntimeOwningSshAiVaultHost }) + + const scanned = await _internals.listAiVaultSessions({ + executionHostScope: 'ssh:runtime-ssh-recipe' + }) + + expect(findRuntimeOwningSshAiVaultHost).not.toHaveBeenCalled() + expect(scanned.sessions).toEqual([]) + expect(scanned.issues).toMatchObject([ + { + executionHostId: 'ssh:runtime-ssh-recipe', + agent: 'codex', + path: 'runtime-ssh-recipe' + } + ]) + }) + + it('resolves SSH titles through the owning runtime when this process has no session', async () => { + mocks.requestActiveSshAiVaultSessionTitles.mockRejectedValue( + new Error('SSH relay is not ready') + ) + const titles = { + titles: [{ agent: 'codex' as const, sessionId: 'session-1', title: 'Exact title' }] + } + const requests = [ + { agent: 'codex' as const, sessionId: 'session-1', transcriptPath: '/tmp/session.jsonl' } + ] + const resolveRuntimeOwnedSshAiVaultSessionTitles = vi.fn().mockResolvedValue(titles) + registerAiVaultHandlers({ + findRuntimeOwningSshAiVaultHost: async () => ({ + environmentId: 'hub-runtime', + targetId: 'hub-owned-host', + executionHostId: 'ssh:hub-owned-host' + }), + resolveRuntimeOwnedSshAiVaultSessionTitles + }) + + await expect( + _internals.resolveAiVaultSessionTitles({ + executionHostScope: 'ssh:hub-owned-host', + requests + }) + ).resolves.toEqual(titles) + + expect(resolveRuntimeOwnedSshAiVaultSessionTitles).toHaveBeenCalledWith( + 'hub-runtime', + 'hub-owned-host', + { + executionHostScope: 'ssh:hub-owned-host', + requests + } + ) + }) +}) + +function result(sessions: AiVaultSession[]): AiVaultListResult { + return { sessions, issues: [], scannedAt: new Date().toISOString() } +} + +function session( + executionHostId: AiVaultSession['executionHostId'], + sessionId: string +): AiVaultSession { + return { + id: `${executionHostId}:codex:${sessionId}:/tmp/${sessionId}.jsonl`, + executionHostId, + agent: 'codex', + sessionId, + title: sessionId, + cwd: '/repo', + branch: null, + model: null, + filePath: `/tmp/${sessionId}.jsonl`, + codexHome: null, + createdAt: null, + updatedAt: null, + modifiedAt: new Date().toISOString(), + messageCount: 1, + totalTokens: 0, + previewMessages: [], + queuedMessageCount: 0, + subagentTranscriptCount: 0, + resumeCommand: `codex resume ${sessionId}`, + subagent: null + } +} diff --git a/src/main/ipc/ai-vault-runtime-owned-ssh.ts b/src/main/ipc/ai-vault-runtime-owned-ssh.ts new file mode 100644 index 00000000000..1b73a6c3454 --- /dev/null +++ b/src/main/ipc/ai-vault-runtime-owned-ssh.ts @@ -0,0 +1,36 @@ +import type { AiVaultListArgs, AiVaultListResult } from '../../shared/ai-vault-types' +import { isRuntimeOwnedSshTargetId } from '../../shared/execution-host' +import { scanSshAiVaultSessions } from '../ai-vault/ssh-session-list' +import type { RuntimeOwnedSshAiVaultHost } from '../ai-vault/runtime-owned-ssh-session-list' +import { getActiveSshAiVaultHostInfo } from './ssh' + +export type RuntimeOwnedSshAiVaultScanner = ( + environmentId: string, + targetId: string, + args: AiVaultListArgs, + options?: { timeoutMs?: number } +) => Promise + +export async function scanSshAiVaultSessionsByOwner(args: { + targetId: string + listArgs?: AiVaultListArgs + signal?: AbortSignal + timeoutMs: number + findOwner?: (targetId: string) => Promise + scanOwned?: RuntimeOwnedSshAiVaultScanner +}): Promise { + if ( + !isRuntimeOwnedSshTargetId(args.targetId) && + !getActiveSshAiVaultHostInfo(args.targetId) && + args.findOwner && + args.scanOwned + ) { + const owner = await args.findOwner(args.targetId) + if (owner) { + return args.scanOwned(owner.environmentId, owner.targetId, args.listArgs ?? {}, { + timeoutMs: args.timeoutMs + }) + } + } + return scanSshAiVaultSessions(args.targetId, args.listArgs, { signal: args.signal }) +} diff --git a/src/main/ipc/ai-vault-runtime-scan.ts b/src/main/ipc/ai-vault-runtime-scan.ts index f9432f143a8..75a29472e20 100644 --- a/src/main/ipc/ai-vault-runtime-scan.ts +++ b/src/main/ipc/ai-vault-runtime-scan.ts @@ -16,6 +16,7 @@ export type RuntimeAiVaultHostInfo = { export type RuntimeAiVaultScanOptions = { timeoutMs?: number + includeOwnedSshHosts?: boolean } export type RuntimeAiVaultScanner = ( diff --git a/src/main/ipc/ai-vault-session-title-routing.ts b/src/main/ipc/ai-vault-session-title-routing.ts index a62fbcda04d..e01e76c7fb2 100644 --- a/src/main/ipc/ai-vault-session-title-routing.ts +++ b/src/main/ipc/ai-vault-session-title-routing.ts @@ -3,10 +3,12 @@ import type { AiVaultSessionTitlesResult } from '../../shared/ai-vault-session-title' import { + isRuntimeOwnedSshTargetId, LOCAL_EXECUTION_HOST_ID, normalizeExecutionHostScope, parseExecutionHostId } from '../../shared/execution-host' +import type { RuntimeOwnedSshAiVaultHost } from '../ai-vault/runtime-owned-ssh-session-list' import { resolveLocalAiVaultSessionTitles } from '../ai-vault/session-title-resolver' import { parseAiVaultSessionTitlesResult } from '../ai-vault/session-title-result-validation' import { requestActiveSshAiVaultSessionTitles } from './ssh' @@ -16,9 +18,21 @@ export type RuntimeAiVaultSessionTitleResolver = ( args: AiVaultSessionTitlesArgs ) => Promise +export type RuntimeOwnedSshAiVaultSessionTitleResolver = ( + environmentId: string, + targetId: string, + args: AiVaultSessionTitlesArgs +) => Promise + export async function resolveAiVaultSessionTitlesByHost( args: AiVaultSessionTitlesArgs, - resolveRuntime?: RuntimeAiVaultSessionTitleResolver + options: { + resolveRuntime?: RuntimeAiVaultSessionTitleResolver + findRuntimeOwningSshAiVaultHost?: ( + targetId: string + ) => Promise + resolveRuntimeOwnedSsh?: RuntimeOwnedSshAiVaultSessionTitleResolver + } = {} ): Promise { const executionHostScope = normalizeExecutionHostScope( args.executionHostScope ?? LOCAL_EXECUTION_HOST_ID @@ -28,17 +42,44 @@ export async function resolveAiVaultSessionTitlesByHost( } const parsed = parseExecutionHostId(executionHostScope) if (parsed?.kind === 'ssh') { - try { - const result = await requestActiveSshAiVaultSessionTitles(parsed.targetId, { - requests: args.requests - }) - return result === null ? { titles: [] } : parseAiVaultSessionTitlesResult(result) - } catch { - return { titles: [] } - } + return resolveSshAiVaultSessionTitles(parsed.targetId, args, options) } - if (parsed?.kind === 'runtime' && resolveRuntime) { - return resolveRuntime(parsed.environmentId, args).catch(() => ({ titles: [] })) + if (parsed?.kind === 'runtime' && options.resolveRuntime) { + return options.resolveRuntime(parsed.environmentId, args).catch(() => ({ titles: [] })) } return { titles: [] } } + +async function resolveSshAiVaultSessionTitles( + targetId: string, + args: AiVaultSessionTitlesArgs, + options: { + findRuntimeOwningSshAiVaultHost?: ( + targetId: string + ) => Promise + resolveRuntimeOwnedSsh?: RuntimeOwnedSshAiVaultSessionTitleResolver + } +): Promise { + try { + const result = await requestActiveSshAiVaultSessionTitles(targetId, { + requests: args.requests + }) + if (result !== null) { + return parseAiVaultSessionTitlesResult(result) + } + } catch { + // Fall through to the owning runtime when this process does not hold the SSH session. + } + if ( + isRuntimeOwnedSshTargetId(targetId) || + !options.findRuntimeOwningSshAiVaultHost || + !options.resolveRuntimeOwnedSsh + ) { + return { titles: [] } + } + const owner = await options.findRuntimeOwningSshAiVaultHost(targetId) + if (!owner) { + return { titles: [] } + } + return options.resolveRuntimeOwnedSsh(owner.environmentId, owner.targetId, args) +} diff --git a/src/main/ipc/ai-vault.test.ts b/src/main/ipc/ai-vault.test.ts index 13b2bc51926..b0a2910f564 100644 --- a/src/main/ipc/ai-vault.test.ts +++ b/src/main/ipc/ai-vault.test.ts @@ -324,7 +324,7 @@ describe('listAiVaultSessions host routing', () => { { executionHostScope: 'runtime:remote-server' }, - expect.objectContaining({ timeoutMs: expect.any(Number) }) + expect.objectContaining({ timeoutMs: expect.any(Number), includeOwnedSshHosts: true }) ) expect(result.sessions.map((entry) => entry.executionHostId)).toEqual([ 'runtime:remote-server', diff --git a/src/main/ipc/ai-vault.ts b/src/main/ipc/ai-vault.ts index 8f4b411cb6f..207408ee3b8 100644 --- a/src/main/ipc/ai-vault.ts +++ b/src/main/ipc/ai-vault.ts @@ -7,12 +7,7 @@ import { } from '../ai-vault/cached-session-list' import { deleteAiVaultSession, registerAiVaultDeleteHandler } from './ai-vault-delete' import { listAiVaultSubagentSessions } from './ai-vault-subagent-list' -import { - aiVaultScanIssueResult, - cancelledAiVaultListResult, - mergeAiVaultListResults -} from '../ai-vault/session-list-results' -import { scanSshAiVaultSessions } from '../ai-vault/ssh-session-list' +import { cancelledAiVaultListResult } from '../ai-vault/session-list-results' import { AiVaultScanCoordinator } from '../ai-vault/ai-vault-scan-coordinator' import type { AiVaultDeleteSessionArgs } from '../../shared/ai-vault-session-deletion' import { @@ -26,27 +21,15 @@ import { } from '../../shared/ai-vault-types' import { handleAiVaultGetFirstUserPrompt } from '../ai-vault/session-first-user-prompt-handler' import { registerAiVaultResumeHandler, type AiVaultResumeHandlerOptions } from './ai-vault-resume' -import { - LOCAL_EXECUTION_HOST_ID, - normalizeExecutionHostScope, - parseExecutionHostId, - toRuntimeExecutionHostId, - toSshExecutionHostId, - type ExecutionHostScope -} from '../../shared/execution-host' -import { getActiveSshAiVaultHostInfos } from './ssh' +import { LOCAL_EXECUTION_HOST_ID, normalizeExecutionHostScope } from '../../shared/execution-host' import { createSenderScopedRequestCancellations } from './sender-scoped-request-cancellation' -import { discoverAiVaultHosts, type AiVaultHostDiscoveryResult } from './ai-vault-host-discovery' -import { - scanRuntimeAiVaultSessions, - type RuntimeAiVaultHostInfo, - type RuntimeAiVaultScanner -} from './ai-vault-runtime-scan' +import type { RuntimeAiVaultHostInfo, RuntimeAiVaultScanner } from './ai-vault-runtime-scan' import { invalidateAiVaultHostLegCache, resetAiVaultHostLegCacheForTests, scanHostLegWithCache } from './ai-vault-host-leg-cache' +import { scanAiVaultSessionsByHostScope } from './ai-vault-host-scope-scan' import { requestedAiVaultSessionDepth } from '../../shared/ai-vault-session-depth' import type { AiVaultSessionTitlesArgs, @@ -54,23 +37,22 @@ import type { } from '../../shared/ai-vault-session-title' import { resolveAiVaultSessionTitlesByHost, - type RuntimeAiVaultSessionTitleResolver + type RuntimeAiVaultSessionTitleResolver, + type RuntimeOwnedSshAiVaultSessionTitleResolver } from './ai-vault-session-title-routing' - -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 type { RuntimeOwnedSshAiVaultHost } from '../ai-vault/runtime-owned-ssh-session-list' +import type { RuntimeOwnedSshAiVaultScanner } from './ai-vault-runtime-owned-ssh' type AiVaultHandlerOptions = AiVaultSessionSources & AiVaultResumeHandlerOptions & { getActiveRuntimeAiVaultHostInfos?: () => readonly RuntimeAiVaultHostInfo[] scanRuntimeAiVaultSessions?: RuntimeAiVaultScanner resolveRuntimeAiVaultSessionTitles?: RuntimeAiVaultSessionTitleResolver + findRuntimeOwningSshAiVaultHost?: ( + targetId: string + ) => Promise + scanRuntimeOwnedSshAiVaultSessions?: RuntimeOwnedSshAiVaultScanner + resolveRuntimeOwnedSshAiVaultSessionTitles?: RuntimeOwnedSshAiVaultSessionTitleResolver } let scanCoordinator = new AiVaultScanCoordinator() @@ -85,7 +67,11 @@ const aiVaultDeleteDeps = { const resolveAiVaultSessionTitles = ( args: AiVaultSessionTitlesArgs ): Promise => - resolveAiVaultSessionTitlesByHost(args, handlerOptions.resolveRuntimeAiVaultSessionTitles) + resolveAiVaultSessionTitlesByHost(args, { + resolveRuntime: handlerOptions.resolveRuntimeAiVaultSessionTitles, + findRuntimeOwningSshAiVaultHost: handlerOptions.findRuntimeOwningSshAiVaultHost, + resolveRuntimeOwnedSsh: handlerOptions.resolveRuntimeOwnedSshAiVaultSessionTitles + }) async function listAiVaultSessions( args?: AiVaultListArgs, @@ -115,7 +101,14 @@ async function listAiVaultSessions( force: args?.force, signal: options.signal, start: (scanSignal) => { - const scan = () => scanAiVaultSessionsByHostScope(args, executionHostScope, scanSignal, key) + const scan = () => + scanAiVaultSessionsByHostScope(args, executionHostScope, scanSignal, key, { + getActiveRuntimeAiVaultHostInfos: handlerOptions.getActiveRuntimeAiVaultHostInfos, + scanRuntimeAiVaultSessions: handlerOptions.scanRuntimeAiVaultSessions, + findRuntimeOwningSshAiVaultHost: handlerOptions.findRuntimeOwningSshAiVaultHost, + scanRuntimeOwnedSshAiVaultSessions: handlerOptions.scanRuntimeOwnedSshAiVaultSessions, + scanLocal: scanLocalAiVaultSessions + }) if (executionHostScope === LOCAL_EXECUTION_HOST_ID) { return scan() } @@ -130,123 +123,6 @@ async function listAiVaultSessions( }) } -async function scanAiVaultSessionsByHostScope( - args: AiVaultListArgs | undefined, - executionHostScope: ExecutionHostScope, - signal?: AbortSignal, - cacheKey = '' -): Promise { - const depth = requestedAiVaultSessionDepth(args) - const scopePaths = args?.scopePaths ?? [] - if (executionHostScope === LOCAL_EXECUTION_HOST_ID) { - return scanLocalAiVaultSessionsAsIssue(args, signal) - } - if (executionHostScope === 'all') { - const runtimeHosts = getActiveRuntimeAiVaultHostInfosResult() - const sshHosts = getActiveSshAiVaultHostInfosResult() - const runtimeResults = [ - ...(runtimeHosts.issue ? [runtimeHosts.issue] : []), - ...(sshHosts.issue ? [sshHosts.issue] : []) - ] - const scannedResults = await Promise.all([ - scanLocalAiVaultSessionsAsIssue(args, signal), - ...sshHosts.hostInfos.map((hostInfo) => - scanHostLegWithCache({ - cacheKey: `${cacheKey}|${toSshExecutionHostId(hostInfo.targetId)}`, - depth, - scopePaths, - force: args?.force === true, - scan: () => - scanSshAiVaultSessions(hostInfo.targetId, args, { - signal, - timeoutMs: AI_VAULT_ALL_HOST_SSH_TIMEOUT_MS, - relayTimeoutMs: AI_VAULT_ALL_HOST_SSH_RELAY_TIMEOUT_MS - }) - }) - ), - ...runtimeHosts.hostInfos.map((hostInfo) => - scanHostLegWithCache({ - cacheKey: `${cacheKey}|${hostInfo.executionHostId}`, - depth, - scopePaths, - force: args?.force === true, - scan: () => - scanRuntimeAiVaultSessions({ - hostInfo, - scanner: handlerOptions.scanRuntimeAiVaultSessions, - listArgs: args, - options: { signal, timeoutMs: AI_VAULT_ALL_HOST_RUNTIME_TIMEOUT_MS } - }) - }) - ) - ]) - return mergeAiVaultListResults( - [...scannedResults, ...runtimeResults], - args?.limit, - args?.unlimited - ) - } - - const parsed = parseExecutionHostId(executionHostScope) - if (parsed?.kind === 'ssh') { - return scanSshAiVaultSessions(parsed.targetId, args, { signal }) - } - if (parsed?.kind === 'runtime') { - return scanRuntimeAiVaultSessions({ - hostInfo: { - environmentId: parsed.environmentId, - executionHostId: toRuntimeExecutionHostId(parsed.environmentId) - }, - scanner: handlerOptions.scanRuntimeAiVaultSessions, - listArgs: args, - options: { signal } - }) - } - - return aiVaultScanIssueResult({ - executionHostId: executionHostScope, - path: executionHostScope, - message: 'Agent Session History is not available for this execution host.' - }) -} - -function getActiveRuntimeAiVaultHostInfosResult(): AiVaultHostDiscoveryResult { - return discoverAiVaultHosts(() => handlerOptions.getActiveRuntimeAiVaultHostInfos?.() ?? [], { - path: 'runtime environments', - fallbackMessage: 'Runtime hosts are unavailable.' - }) -} - -function getActiveSshAiVaultHostInfosResult(): AiVaultHostDiscoveryResult<{ targetId: string }> { - return discoverAiVaultHosts(getActiveSshAiVaultHostInfos, { - path: 'SSH hosts', - fallbackMessage: 'SSH hosts are unavailable.' - }) -} - -// Why: the SSH legs already degrade to an issue row so one bad host can't take -// the shared Promise.all down; the local leg can throw too (parse-cache load, -// WSL home resolution, scanner service supervision) and would otherwise discard -// every host's sessions under 'all', or replace the list with a raw error string -// under single-host scope. -async function scanLocalAiVaultSessionsAsIssue( - args: AiVaultListArgs | undefined, - signal: AbortSignal | undefined -): Promise { - try { - return await scanLocalAiVaultSessions(args, signal) - } catch (error) { - if (isAiVaultScanCancelledError(error)) { - throw error - } - return aiVaultScanIssueResult({ - executionHostId: LOCAL_EXECUTION_HOST_ID, - path: 'this computer', - message: error instanceof Error ? error.message : 'Local session scan failed.' - }) - } -} - async function scanLocalAiVaultSessions( args?: AiVaultListArgs, signal?: AbortSignal diff --git a/src/main/ipc/register-core-handlers.ts b/src/main/ipc/register-core-handlers.ts index 7825e73c57e..13b3a813aab 100644 --- a/src/main/ipc/register-core-handlers.ts +++ b/src/main/ipc/register-core-handlers.ts @@ -91,6 +91,11 @@ import { resolveRuntimeAiVaultSessionTitles, scanRuntimeAiVaultSessions } from '../ai-vault/runtime-session-scanner' +import { + findRuntimeOwningSshAiVaultHost, + resolveRuntimeOwnedSshAiVaultSessionTitles, + scanRuntimeOwnedSshAiVaultSessions +} from '../ai-vault/runtime-owned-ssh-session-list' import type { PluginService } from '../plugins/plugin-service' import type { PluginMarketplaceHandlerServices } from './plugin-marketplaces' @@ -224,6 +229,23 @@ export function registerCoreHandlers( scanRuntimeAiVaultSessions(app.getPath('userData'), environmentId, args, options), resolveRuntimeAiVaultSessionTitles: async (environmentId, args) => resolveRuntimeAiVaultSessionTitles(app.getPath('userData'), environmentId, args), + findRuntimeOwningSshAiVaultHost: async (targetId) => + findRuntimeOwningSshAiVaultHost(app.getPath('userData'), targetId), + scanRuntimeOwnedSshAiVaultSessions: async (environmentId, targetId, args, options) => + scanRuntimeOwnedSshAiVaultSessions( + app.getPath('userData'), + environmentId, + targetId, + args, + options + ), + resolveRuntimeOwnedSshAiVaultSessionTitles: async (environmentId, targetId, args) => + resolveRuntimeOwnedSshAiVaultSessionTitles( + app.getPath('userData'), + environmentId, + targetId, + args + ), prepareRuntimeSessionResume: async (environmentId, args) => prepareRuntimeAiVaultSessionResume(app.getPath('userData'), environmentId, args) }) 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 new file mode 100644 index 00000000000..f76d52cee9d --- /dev/null +++ b/src/main/runtime/rpc/methods/ai-vault-runtime-ssh-scope.repro.test.ts @@ -0,0 +1,228 @@ +import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest' +import { RpcDispatcher } from '../dispatcher' +import type { RpcRequest } from '../core' +import type { OrcaRuntimeService } from '../../orca-runtime' +import type { AiVaultListResult, AiVaultSession } from '../../../../shared/ai-vault-types' + +const { + scanAiVaultSessionsInWorker, + resolveAiVaultSessionTitlesInWorker, + scanSshAiVaultSessions, + getActiveSshAiVaultHostInfos, + requestActiveSshAiVaultSessionTitles +} = vi.hoisted(() => ({ + scanAiVaultSessionsInWorker: vi.fn(), + resolveAiVaultSessionTitlesInWorker: vi.fn(), + scanSshAiVaultSessions: vi.fn(), + getActiveSshAiVaultHostInfos: vi.fn(), + requestActiveSshAiVaultSessionTitles: vi.fn() +})) + +vi.mock('../../../ai-vault/session-scanner-worker-spawn', () => ({ + scanAiVaultSessionsInWorker, + resolveAiVaultSessionTitlesInWorker, + resetAiVaultScannerWorkerForTests: vi.fn() +})) + +vi.mock('../../../ai-vault/ssh-session-list', () => ({ + scanSshAiVaultSessions +})) + +vi.mock('../../../ipc/ssh', () => ({ + getActiveSshAiVaultHostInfos, + requestActiveSshAiVaultSessionTitles +})) + +import { AI_VAULT_METHODS } from './ai-vault' +import { + listAiVaultSessions, + resetAiVaultSessionListCacheForTests +} from '../../../ai-vault/cached-session-list' + +const SCANNED_AT = '2026-08-12T00:00:00.000Z' +const SSH_HOST_OWNED_BY_RUNTIME = 'ssh:hub-owned-host' + +function makeRequest(method: string, params?: unknown): RpcRequest { + return { id: 'req-1', authToken: 'tok', method, params } +} + +function makeRuntimeLocalSession(): AiVaultSession { + return { + id: 'local:claude:sess-hub:/hub/home/.claude/projects/p/t.jsonl', + executionHostId: 'local', + agent: 'claude', + sessionId: 'sess-hub', + title: 'Session on the runtime host itself', + cwd: '/hub/home/projects/p', + branch: null, + model: null, + filePath: '/hub/home/.claude/projects/p/t.jsonl', + codexHome: null, + createdAt: null, + updatedAt: null, + modifiedAt: SCANNED_AT, + messageCount: 2, + totalTokens: 0, + previewMessages: [], + queuedMessageCount: 0, + subagentTranscriptCount: 0, + resumeCommand: 'claude --resume sess-hub', + subagent: null + } +} + +function makeSshSession(): AiVaultSession { + return { + ...makeRuntimeLocalSession(), + id: 'ssh:hub-owned-host:claude:sess-ssh:/remote/home/.claude/projects/p/t.jsonl', + executionHostId: 'ssh:hub-owned-host', + sessionId: 'sess-ssh', + title: 'Session on an SSH host this runtime owns', + cwd: '/remote/home/projects/p', + filePath: '/remote/home/.claude/projects/p/t.jsonl', + resumeCommand: 'claude --resume sess-ssh' + } +} + +function makeDispatcher(): RpcDispatcher { + const runtime = { + getRuntimeId: () => 'hub-runtime', + listAiVaultSessions: (args?: Parameters[0]) => + listAiVaultSessions(args), + resolveAiVaultSessionTitles: (requests: unknown[], signal?: AbortSignal) => + resolveAiVaultSessionTitlesInWorker(requests, signal) + } as unknown as OrcaRuntimeService + return new RpcDispatcher({ runtime, methods: AI_VAULT_METHODS }) +} + +describe('a runtime can scan an SSH host that it owns', () => { + beforeEach(() => { + resetAiVaultSessionListCacheForTests() + scanAiVaultSessionsInWorker.mockReset() + scanAiVaultSessionsInWorker.mockResolvedValue({ + sessions: [makeRuntimeLocalSession()], + issues: [], + scannedAt: SCANNED_AT + } satisfies AiVaultListResult) + resolveAiVaultSessionTitlesInWorker.mockReset() + resolveAiVaultSessionTitlesInWorker.mockResolvedValue({ titles: [] }) + scanSshAiVaultSessions.mockReset() + scanSshAiVaultSessions.mockResolvedValue({ + sessions: [makeSshSession()], + issues: [], + scannedAt: SCANNED_AT + } satisfies AiVaultListResult) + getActiveSshAiVaultHostInfos.mockReset() + getActiveSshAiVaultHostInfos.mockReturnValue([{ targetId: 'hub-owned-host' }]) + requestActiveSshAiVaultSessionTitles.mockReset() + requestActiveSshAiVaultSessionTitles.mockResolvedValue({ + titles: [{ agent: 'claude', sessionId: 'sess-ssh', title: 'SSH title' }] + }) + }) + + afterEach(() => { + resetAiVaultSessionListCacheForTests() + }) + + it('routes an SSH execution host to that host instead of rejecting it', async () => { + const dispatcher = makeDispatcher() + + const scoped = (await dispatcher.dispatch( + makeRequest('aiVault.listSessions', { + limit: 500, + executionHostId: SSH_HOST_OWNED_BY_RUNTIME + }) + )) as { ok: boolean; result: AiVaultListResult } + + expect(scoped.ok).toBe(true) + expect(scanAiVaultSessionsInWorker).not.toHaveBeenCalled() + expect(scanSshAiVaultSessions).toHaveBeenCalledWith( + 'hub-owned-host', + expect.objectContaining({ limit: 500 }) + ) + expect(scoped.result.sessions.map((session) => session.executionHostId)).toEqual([ + 'ssh:hub-owned-host' + ]) + }) + + it('scans the identical host-local sources for the runtime ids it does accept', async () => { + const dispatcher = makeDispatcher() + + await dispatcher.dispatch(makeRequest('aiVault.listSessions', { limit: 500 })) + const baselineScanArgs = scanAiVaultSessionsInWorker.mock.calls[0]?.[0] + + resetAiVaultSessionListCacheForTests() + const scoped = (await dispatcher.dispatch( + makeRequest('aiVault.listSessions', { + limit: 500, + executionHostId: 'runtime:some-other-env' + }) + )) as { ok: boolean; result: AiVaultListResult } + const scopedScanArgs = scanAiVaultSessionsInWorker.mock.calls[1]?.[0] + + expect(scoped.ok).toBe(true) + expect(scopedScanArgs).toEqual(baselineScanArgs) + expect(scanSshAiVaultSessions).not.toHaveBeenCalled() + expect(scoped.result.sessions.map((session) => session.filePath)).toEqual([ + '/hub/home/.claude/projects/p/t.jsonl' + ]) + }) + + it('keeps unscoped scans host-local so old clients do not receive SSH rows', async () => { + const dispatcher = makeDispatcher() + + const unscoped = (await dispatcher.dispatch( + makeRequest('aiVault.listSessions', { limit: 500 }) + )) as { ok: boolean; result: AiVaultListResult } + + expect(unscoped.ok).toBe(true) + expect(scanSshAiVaultSessions).not.toHaveBeenCalled() + expect(unscoped.result.sessions.map((session) => session.executionHostId)).toEqual(['local']) + }) + + it('merges this runtime’s connected SSH hosts only when the caller opts in', async () => { + const dispatcher = makeDispatcher() + + const optedIn = (await dispatcher.dispatch( + makeRequest('aiVault.listSessions', { + limit: 500, + executionHostId: 'runtime:hub-runtime', + includeOwnedSshHosts: true + }) + )) as { ok: boolean; result: AiVaultListResult } + + expect(optedIn.ok).toBe(true) + expect(scanAiVaultSessionsInWorker).toHaveBeenCalledTimes(1) + expect(scanSshAiVaultSessions).toHaveBeenCalledWith( + 'hub-owned-host', + expect.objectContaining({ limit: 500 }) + ) + expect(optedIn.result.sessions.map((session) => session.executionHostId).sort()).toEqual([ + 'runtime:hub-runtime', + 'ssh:hub-owned-host' + ]) + }) + + it('resolves SSH session titles on the transcript-owning host', async () => { + const dispatcher = makeDispatcher() + const requests = [{ agent: 'claude', sessionId: 'sess-ssh' }] + + const response = await dispatcher.dispatch( + makeRequest('aiVault.resolveSessionTitles', { + requests, + executionHostId: SSH_HOST_OWNED_BY_RUNTIME + }) + ) + + expect(response).toMatchObject({ + ok: true, + result: { titles: [{ sessionId: 'sess-ssh', title: 'SSH title' }] } + }) + expect(requestActiveSshAiVaultSessionTitles).toHaveBeenCalledWith( + 'hub-owned-host', + { requests }, + expect.objectContaining({ signal: undefined }) + ) + expect(resolveAiVaultSessionTitlesInWorker).not.toHaveBeenCalled() + }) +}) diff --git a/src/main/runtime/rpc/methods/ai-vault.test.ts b/src/main/runtime/rpc/methods/ai-vault.test.ts index 0eb9535a097..2154c058e7e 100644 --- a/src/main/runtime/rpc/methods/ai-vault.test.ts +++ b/src/main/runtime/rpc/methods/ai-vault.test.ts @@ -164,14 +164,17 @@ describe('aiVault.listSessions params schema', () => { expect(parsed.success).toBe(false) }) - it('rejects non-runtime execution host ids before dispatch', () => { + it('rejects host ids that are neither a runtime nor an SSH target', () => { expect( AiVaultListSessionsParams.safeParse({ executionHostId: 'not-a-runtime-host' }).success ).toBe(false) expect(AiVaultListSessionsParams.safeParse({ executionHostId: 'local' }).success).toBe(false) expect(AiVaultListSessionsParams.safeParse({ executionHostId: 'ssh:dev-box' }).success).toBe( - false + true ) + expect( + AiVaultListSessionsParams.safeParse({ executionHostId: 'runtime:hub-runtime' }).success + ).toBe(true) }) }) diff --git a/src/main/runtime/rpc/methods/ai-vault.ts b/src/main/runtime/rpc/methods/ai-vault.ts index 9578c903eb2..005d6717192 100644 --- a/src/main/runtime/rpc/methods/ai-vault.ts +++ b/src/main/runtime/rpc/methods/ai-vault.ts @@ -1,7 +1,16 @@ import { z } from 'zod' import { defineMethod, type RpcMethod } from '../core' import { OptionalBoolean } from '../schemas' -import { restampAiVaultListResult } from '../../../ai-vault/session-list-results' +import { + mergeAiVaultListResults, + restampAiVaultListResult +} from '../../../ai-vault/session-list-results' +import { parseAiVaultSessionTitlesResult } from '../../../ai-vault/session-title-result-validation' +import { scanSshAiVaultSessions } from '../../../ai-vault/ssh-session-list' +import { + getActiveSshAiVaultHostInfos, + requestActiveSshAiVaultSessionTitles +} from '../../../ipc/ssh' import { AI_VAULT_AGENTS, AI_VAULT_SCOPE_PATHS_MAX_COUNT } from '../../../../shared/ai-vault-types' import { AI_VAULT_SESSION_TITLE_REQUEST_MAX_COUNT } from '../../../../shared/ai-vault-session-title' import { LOCAL_EXECUTION_HOST_ID, parseExecutionHostId } from '../../../../shared/execution-host' @@ -12,17 +21,19 @@ import { LOCAL_EXECUTION_HOST_ID, parseExecutionHostId } from '../../../../share const AI_VAULT_SCOPE_PATH_MAX_LENGTH = 4096 const AI_VAULT_LIMIT_MAX = 2000 -const executionHostIdSchema = z.string().transform((value, ctx): `runtime:${string}` => { - const parsed = parseExecutionHostId(value) - if (parsed?.kind === 'runtime') { - return parsed.id - } - ctx.addIssue({ - code: 'custom', - message: 'Invalid runtime execution host id' +const executionHostIdSchema = z + .string() + .transform((value, ctx): `runtime:${string}` | `ssh:${string}` => { + const parsed = parseExecutionHostId(value) + if (parsed?.kind === 'runtime' || parsed?.kind === 'ssh') { + return parsed.id + } + ctx.addIssue({ + code: 'custom', + message: 'Invalid execution host id' + }) + return z.NEVER }) - return z.NEVER -}) export const AiVaultListSessionsParams = z .object({ @@ -42,10 +53,13 @@ export const AiVaultListSessionsParams = z // desktop parents) that send more than the bound. .transform((paths) => paths.slice(0, AI_VAULT_SCOPE_PATHS_MAX_COUNT)) .optional(), - // Why: desktop/web callers name the runtime host they are addressing; mobile - // omits it. The scan itself is host-local either way, so the id must never - // change what is scanned — it only restamps the shared cached result. - executionHostId: executionHostIdSchema.optional() + // Why: desktop/web callers name the runtime host they are addressing; a + // paired client can also name an SSH host this runtime owns. Runtime ids + // still only restamp the host-local cache. SSH ids route to that host. + executionHostId: executionHostIdSchema.optional(), + // Why: omitted on old clients so mixed-version peers keep the host-local + // projection. New all-hosts callers opt in to this runtime's SSH inventory. + includeOwnedSshHosts: OptionalBoolean }) .superRefine((params, ctx) => { if (params.unlimited !== true && params.limit && params.limit > AI_VAULT_LIMIT_MAX) { @@ -69,31 +83,61 @@ export const AiVaultSessionTitlesParams = z.object({ transcriptPath: z.string().min(1).max(32_768).optional() }) ) - .max(AI_VAULT_SESSION_TITLE_REQUEST_MAX_COUNT) + .max(AI_VAULT_SESSION_TITLE_REQUEST_MAX_COUNT), + executionHostId: executionHostIdSchema.optional() }) export const AI_VAULT_METHODS: RpcMethod[] = [ defineMethod({ name: 'aiVault.resolveSessionTitles', params: AiVaultSessionTitlesParams, - handler: (params, { runtime, signal }) => - runtime.resolveAiVaultSessionTitles(params.requests, signal) + handler: async (params, { runtime, signal }) => { + const parsed = params.executionHostId ? parseExecutionHostId(params.executionHostId) : null + if (parsed?.kind === 'ssh') { + try { + const result = await requestActiveSshAiVaultSessionTitles( + parsed.targetId, + { requests: params.requests }, + { signal } + ) + return result === null ? { titles: [] } : parseAiVaultSessionTitlesResult(result) + } catch { + return { titles: [] } + } + } + return runtime.resolveAiVaultSessionTitles(params.requests, signal) + } }), defineMethod({ name: 'aiVault.listSessions', params: AiVaultListSessionsParams, handler: async (params, { runtime }) => { - const result = await runtime.listAiVaultSessions({ + const parsed = params.executionHostId ? parseExecutionHostId(params.executionHostId) : null + const listArgs = { limit: params.unlimited ? undefined : params.limit, unlimited: params.unlimited, force: params.force, scopePaths: params.scopePaths - }) + } + if (parsed?.kind === 'ssh') { + return scanSshAiVaultSessions(parsed.targetId, listArgs) + } + const result = await runtime.listAiVaultSessions(listArgs) // Why: web clients consume this response directly (no parent-side retag), - // so sessions must come back stamped as the runtime host they addressed. - return params.executionHostId - ? restampAiVaultListResult(result, params.executionHostId) - : result + // so host-local sessions must come back stamped as the runtime they addressed. + const stamped = + parsed?.kind === 'runtime' ? restampAiVaultListResult(result, parsed.id) : result + if (params.includeOwnedSshHosts !== true) { + return stamped + } + const sshHosts = getActiveSshAiVaultHostInfos() + if (sshHosts.length === 0) { + return stamped + } + const sshResults = await Promise.all( + sshHosts.map((host) => scanSshAiVaultSessions(host.targetId, listArgs)) + ) + return mergeAiVaultListResults([stamped, ...sshResults], listArgs.limit, listArgs.unlimited) } }), defineMethod({ diff --git a/src/renderer/src/web/web-preload-api.test.ts b/src/renderer/src/web/web-preload-api.test.ts index 9fcec640925..b6fa003e2b4 100644 --- a/src/renderer/src/web/web-preload-api.test.ts +++ b/src/renderer/src/web/web-preload-api.test.ts @@ -1609,7 +1609,8 @@ describe('web AI Vault preload API', () => { limit: 25, force: true, scopePaths: ['/srv/app'], - executionHostId: 'runtime:web-env-1' + executionHostId: 'runtime:web-env-1', + includeOwnedSshHosts: true } } ]) @@ -1652,6 +1653,51 @@ describe('web AI Vault preload API', () => { }) expect(runtimeCalls).toEqual([]) }) + + it('forwards SSH host scope to the paired runtime that owns the target', async () => { + const runtimeCalls: { method: string; params: unknown }[] = [] + const scanResult = { + sessions: [], + issues: [], + scannedAt: '2026-08-12T00:00:00.000Z' + } + vi.doMock('./web-runtime-client', () => ({ + WebRuntimeClient: class { + call(method: string, params?: unknown): Promise> { + runtimeCalls.push({ method, params }) + return Promise.resolve({ + id: `call-${runtimeCalls.length}`, + ok: true, + result: scanResult, + _meta: { runtimeId: 'runtime-1' } + }) + } + + close(): void {} + } + })) + + const globals = installBrowserGlobals('Linux') + writeStoredRuntimeEnvironment(globals.storage) + const { installWebPreloadApi } = await import('./web-preload-api') + installWebPreloadApi() + + await expect( + globals.window.api.aiVault.listSessions({ + executionHostScope: 'ssh:hub-owned-host', + limit: 25 + }) + ).resolves.toEqual(scanResult) + expect(runtimeCalls).toEqual([ + { + method: 'aiVault.listSessions', + params: { + limit: 25, + executionHostId: 'ssh:hub-owned-host' + } + } + ]) + }) }) describe('web UI preload API', () => { diff --git a/src/renderer/src/web/web-preload-api.ts b/src/renderer/src/web/web-preload-api.ts index 2baa23a7b28..0a6eb9ef11e 100644 --- a/src/renderer/src/web/web-preload-api.ts +++ b/src/renderer/src/web/web-preload-api.ts @@ -1578,24 +1578,42 @@ function createAiVaultApi(): NonNullable['aiVault']> { const requestedScope = normalizeExecutionHostScope( args?.executionHostScope ?? executionHostId ) + const parsedScope = parseExecutionHostId(requestedScope) + if (parsedScope?.kind === 'ssh') { + return callRuntimeResult('aiVault.listSessions', { + limit: args?.limit, + force: args?.force, + scopePaths: args?.scopePaths, + executionHostId: parsedScope.id + }) + } if (requestedScope !== 'all' && requestedScope !== executionHostId) { return Promise.resolve(webAiVaultUnavailableResult(requestedScope)) } - // Why: no local filesystem in the browser, so every history scan runs on and is stamped as the paired runtime host. + // Why: no local filesystem in the browser, so every history scan runs on + // the paired runtime. `all` also asks that runtime for SSH hosts it owns. return callRuntimeResult('aiVault.listSessions', { limit: args?.limit, force: args?.force, scopePaths: args?.scopePaths, - executionHostId + executionHostId, + ...(requestedScope === 'all' ? { includeOwnedSshHosts: true } : {}) }) }, resolveSessionTitles: (args: AiVaultSessionTitlesArgs) => { const environment = requireActiveEnvironment() const executionHostId = toRuntimeExecutionHostId(environment.id) - if ( - args.executionHostScope && - normalizeExecutionHostScope(args.executionHostScope) !== executionHostId - ) { + const requestedScope = args.executionHostScope + ? normalizeExecutionHostScope(args.executionHostScope) + : executionHostId + const parsedScope = parseExecutionHostId(requestedScope) + if (parsedScope?.kind === 'ssh') { + return callRuntimeResult('aiVault.resolveSessionTitles', { + requests: args.requests, + executionHostId: parsedScope.id + }).catch(() => ({ titles: [] })) + } + if (requestedScope !== executionHostId) { return Promise.resolve({ titles: [] }) } return callRuntimeResult('aiVault.resolveSessionTitles', {