From 32c0b47afe406d0048218f69c443bf97de73f285 Mon Sep 17 00:00:00 2001 From: Jinjing <6427696+AmethystLiang@users.noreply.github.com> Date: Thu, 25 Jun 2026 21:58:08 -0700 Subject: [PATCH] Restrict AI Vault resume actions to local execution hosts (#6406) Previously, the AI Vault only blocked resume actions for workspaces with an active SSH connection (checking `connectionId`). This allowed resume actions to run on runtime-owned workspaces, which are non-local but do not use SSH. To resolve this: - Introduce `getAiVaultResumeWorkspaceTargetStatus` to classify targets based on both `connectionId` and `executionHostId`. - Restrict AI Vault panel actions, session resume checks, and drop layers to local workspaces. - Preserve the `executionHostId` on project groups during normalization to ensure runtime-owned groups retain their execution host status across persistence reloads. --- .../components/right-sidebar/AiVaultPanel.tsx | 12 +- .../ai-vault-session-launch-actions.ts | 3 +- .../ai-vault-session-resume.test.ts | 90 ++++++++++++ .../right-sidebar/ai-vault-session-resume.ts | 5 +- .../tab-group/AiVaultSessionDropLayer.tsx | 9 +- .../src/lib/ai-vault-resume-target.test.ts | 134 ++++++++++++++++++ .../src/lib/ai-vault-resume-target.ts | 114 +++++++++++++++ src/shared/project-groups.test.ts | 12 ++ src/shared/project-groups.ts | 6 +- 9 files changed, 372 insertions(+), 13 deletions(-) create mode 100644 src/renderer/src/lib/ai-vault-resume-target.test.ts create mode 100644 src/renderer/src/lib/ai-vault-resume-target.ts diff --git a/src/renderer/src/components/right-sidebar/AiVaultPanel.tsx b/src/renderer/src/components/right-sidebar/AiVaultPanel.tsx index 5b00eaace..a52bbf276 100644 --- a/src/renderer/src/components/right-sidebar/AiVaultPanel.tsx +++ b/src/renderer/src/components/right-sidebar/AiVaultPanel.tsx @@ -6,7 +6,6 @@ import { useActiveWorktree, useAllWorktrees, useProjectHostSetupProjection, - useRepoById, useRepos } from '@/store/selectors' import { filterAiVaultSessions, groupAiVaultSessions } from './ai-vault-session-filters' @@ -27,6 +26,7 @@ import { import { useAiVaultSessionLaunchActions } from './ai-vault-session-launch-actions' import { useAiVaultSessionWorktreeMap } from './ai-vault-session-worktree' import { useAiVaultOriginalPaneActions } from './ai-vault-original-pane-actions' +import { getAiVaultResumeWorkspaceTargetStatus } from '@/lib/ai-vault-resume-target' import { AI_VAULT_AGENTS, type AiVaultAgent, @@ -44,7 +44,6 @@ import { useAiVaultSessionRefresh } from './ai-vault-session-refresh' export default function AiVaultPanel(): React.JSX.Element { const activeWorktree = useActiveWorktree() const activeRepo = useActiveRepo() - const activeWorktreeRepo = useRepoById(activeWorktree?.repoId ?? null) const repos = useRepos() const allWorktrees = useAllWorktrees() const projectHostSetupProjection = useProjectHostSetupProjection() @@ -62,7 +61,10 @@ export default function AiVaultPanel(): React.JSX.Element { const userChangedScopeRef = useRef(false) const preferredScopeRef = useRef(DEFAULT_AI_VAULT_SCOPE) - const isRemoteWorktree = Boolean(activeWorktreeRepo?.connectionId) + const isNonLocalWorktree = useAppStore( + (state) => + getAiVaultResumeWorkspaceTargetStatus(state, activeWorktree?.id ?? null) === 'non-local' + ) const activeWorktreePath = activeWorktree?.path ?? null // Why: AI Vault ownership is cwd-based, so we must consider live worktrees across all repos. const activeWorktreePaths = useMemo( @@ -276,11 +278,11 @@ export default function AiVaultPanel(): React.JSX.Element { onRefresh={() => void refresh({ force: true })} /> - {isRemoteWorktree ? ( + {isNonLocalWorktree ? (
{translate( 'auto.components.right.sidebar.AiVaultPanel.remoteBrowseLocalHistory', - 'SSH-host workspaces can browse local history. Resume actions run from {{value0}} workspaces.', + 'Non-local workspaces can browse local history. Resume actions run from {{value0}} workspaces.', { value0: getLocalExecutionHostLabel() } )}
diff --git a/src/renderer/src/components/right-sidebar/ai-vault-session-launch-actions.ts b/src/renderer/src/components/right-sidebar/ai-vault-session-launch-actions.ts index a2d074b6b..031f3c81f 100644 --- a/src/renderer/src/components/right-sidebar/ai-vault-session-launch-actions.ts +++ b/src/renderer/src/components/right-sidebar/ai-vault-session-launch-actions.ts @@ -8,6 +8,7 @@ import { import { launchAiVaultSessionInNewTab } from '@/lib/launch-ai-vault-session' import { activateAndRevealWorktree } from '@/lib/worktree-activation' import { useAppStore } from '@/store' +import { isNonLocalAiVaultResumeRepo } from '@/lib/ai-vault-resume-target' import type { AiVaultAgent, AiVaultSession } from '../../../../shared/ai-vault-types' import type { Repo, Worktree } from '../../../../shared/types' import { translate } from '@/i18n/i18n' @@ -80,7 +81,7 @@ export function useAiVaultSessionLaunchActions({ } const worktreeRepo = repos.find((repo) => repo.id === worktree.repoId) - if (worktreeRepo?.connectionId) { + if (isNonLocalAiVaultResumeRepo(worktreeRepo)) { toast.error( translate( 'auto.components.right.sidebar.AiVaultPanel.localWorkspacesOnly', diff --git a/src/renderer/src/components/right-sidebar/ai-vault-session-resume.test.ts b/src/renderer/src/components/right-sidebar/ai-vault-session-resume.test.ts index 9c9e86b29..b6cce33d3 100644 --- a/src/renderer/src/components/right-sidebar/ai-vault-session-resume.test.ts +++ b/src/renderer/src/components/right-sidebar/ai-vault-session-resume.test.ts @@ -30,6 +30,17 @@ function makeWorktree(overrides: Partial = {}): Worktree { } } +function makeRepo(overrides: Partial = {}): Repo { + return { + id: 'repo-1', + path: '/repo/orca', + displayName: 'orca', + badgeColor: '#000000', + addedAt: 1, + ...overrides + } +} + function makeWorktreeInfo( status: AiVaultSessionWorktreeInfo['status'] ): AiVaultSessionWorktreeInfo { @@ -90,6 +101,26 @@ describe('resolveAiVaultSessionResumeState', () => { }) }) + it('blocks runtime-owned targets even when they have no SSH connection', () => { + expect( + resolveAiVaultSessionResumeState({ + worktreeInfo: makeWorktreeInfo('active'), + activeWorktreeId: null, + worktrees: [makeWorktree()], + repos: [ + makeRepo({ + connectionId: null, + executionHostId: 'runtime:env-1' + }) + ] + }) + ).toEqual({ + blocked: true, + worktreeId: null, + usesSessionWorktree: false + }) + }) + it('uses a local session worktree when the active workspace is remote', () => { expect( resolveAiVaultSessionResumeState({ @@ -111,6 +142,36 @@ describe('resolveAiVaultSessionResumeState', () => { usesSessionWorktree: true }) }) + + it('falls back to the active workspace when the session worktree is runtime-owned', () => { + expect( + resolveAiVaultSessionResumeState({ + worktreeInfo: makeWorktreeInfo('active'), + activeWorktreeId: 'repo-2::/repo/other', + worktrees: [ + makeWorktree(), + makeWorktree({ + id: 'repo-2::/repo/other', + repoId: 'repo-2', + path: '/repo/other' + }) + ], + repos: [ + makeRepo({ connectionId: null, executionHostId: 'runtime:env-1' }), + makeRepo({ + id: 'repo-2', + path: '/repo/other', + connectionId: null, + executionHostId: 'local' + }) + ] + }) + ).toEqual({ + blocked: false, + worktreeId: 'repo-2::/repo/other', + usesSessionWorktree: false + }) + }) }) describe('resolveAiVaultSessionResumeActions', () => { @@ -152,6 +213,35 @@ describe('resolveAiVaultSessionResumeActions', () => { }) }) + it('disables runtime-owned targets without disabling local targets', () => { + expect( + resolveAiVaultSessionResumeActions({ + worktreeInfo: makeWorktreeInfo('active'), + activeWorktreeId: 'repo-2::/repo/other', + worktrees: [ + makeWorktree(), + makeWorktree({ + id: 'repo-2::/repo/other', + repoId: 'repo-2', + path: '/repo/other' + }) + ], + repos: [ + makeRepo({ connectionId: null, executionHostId: 'runtime:env-1' }), + makeRepo({ + id: 'repo-2', + path: '/repo/other', + connectionId: null, + executionHostId: 'local' + }) + ] + }) + ).toEqual({ + worktree: { worktreeId: 'repo-1::/repo/orca', disabled: true }, + newTab: { worktreeId: 'repo-2::/repo/other', disabled: false } + }) + }) + it('does not expose the active workspace as a duplicate new-tab target', () => { expect( resolveAiVaultSessionResumeActions({ diff --git a/src/renderer/src/components/right-sidebar/ai-vault-session-resume.ts b/src/renderer/src/components/right-sidebar/ai-vault-session-resume.ts index 32481998a..67b540d5b 100644 --- a/src/renderer/src/components/right-sidebar/ai-vault-session-resume.ts +++ b/src/renderer/src/components/right-sidebar/ai-vault-session-resume.ts @@ -1,4 +1,5 @@ import type { Repo, Worktree } from '../../../../shared/types' +import { isLocalAiVaultResumeRepo } from '@/lib/ai-vault-resume-target' import { translate } from '@/i18n/i18n' import { canJumpToAiVaultSessionWorktree, @@ -45,7 +46,7 @@ export function resolveAiVaultSessionResumeState(args: { continue } const repo = args.repos.find((candidate) => candidate.id === worktree.repoId) - if (repo?.connectionId) { + if (!isLocalAiVaultResumeRepo(repo)) { continue } return { @@ -117,7 +118,7 @@ function resolveLocalResumeWorktreeId(args: { } const repo = args.repos.find((candidate) => candidate.id === worktree.repoId) - if (repo?.connectionId) { + if (!isLocalAiVaultResumeRepo(repo)) { return null } diff --git a/src/renderer/src/components/tab-group/AiVaultSessionDropLayer.tsx b/src/renderer/src/components/tab-group/AiVaultSessionDropLayer.tsx index f9d791d3c..302c89997 100644 --- a/src/renderer/src/components/tab-group/AiVaultSessionDropLayer.tsx +++ b/src/renderer/src/components/tab-group/AiVaultSessionDropLayer.tsx @@ -1,6 +1,6 @@ import { useCallback, useEffect, useRef, useState, type CSSProperties } from 'react' import { toast } from 'sonner' -import { getConnectionId } from '@/lib/connection-context' +import { getAiVaultResumeWorkspaceTargetStatus } from '@/lib/ai-vault-resume-target' import { AI_VAULT_SESSION_DRAG_END_EVENT, AI_VAULT_SESSION_DRAG_START_EVENT, @@ -9,6 +9,7 @@ import { readAiVaultSessionDragData } from '@/lib/ai-vault-session-drag' import { launchAiVaultSessionInNewTab } from '@/lib/launch-ai-vault-session' +import { useAppStore } from '@/store' import { resolveDropZone } from './tab-drop-zone' import type { TabDropZone } from './useTabDragSplit' import { translate } from '@/i18n/i18n' @@ -165,8 +166,8 @@ export default function AiVaultSessionDropLayer({ return true } - const connectionId = getConnectionId(worktreeId) - if (connectionId) { + const targetStatus = getAiVaultResumeWorkspaceTargetStatus(useAppStore.getState(), worktreeId) + if (targetStatus === 'non-local') { toast.error( translate( 'auto.components.tab.group.AiVaultSessionDropLayer.localWorkspacesOnly', @@ -175,7 +176,7 @@ export default function AiVaultSessionDropLayer({ ) return true } - if (connectionId === undefined) { + if (targetStatus === 'unknown') { toast.error( translate( 'auto.components.tab.group.AiVaultSessionDropLayer.openLocalWorkspace', diff --git a/src/renderer/src/lib/ai-vault-resume-target.test.ts b/src/renderer/src/lib/ai-vault-resume-target.test.ts new file mode 100644 index 000000000..151745801 --- /dev/null +++ b/src/renderer/src/lib/ai-vault-resume-target.test.ts @@ -0,0 +1,134 @@ +import { describe, expect, it } from 'vitest' +import type { AppState } from '@/store/types' +import { + getAiVaultResumeRepoTargetStatus, + getAiVaultResumeWorktreeTargetStatus, + getAiVaultResumeWorkspaceTargetStatus, + isLocalAiVaultResumeRepo, + isNonLocalAiVaultResumeRepo +} from './ai-vault-resume-target' +import { folderWorkspaceKey } from '../../../shared/workspace-scope' + +type ResumeTargetState = Pick< + AppState, + 'folderWorkspaces' | 'projectGroups' | 'repos' | 'worktreesByRepo' +> + +function makeState( + overrides: Partial> +): ResumeTargetState { + return { + folderWorkspaces: [], + projectGroups: [], + repos: [], + worktreesByRepo: {}, + ...overrides + } as unknown as ResumeTargetState +} + +describe('ai vault resume target ownership', () => { + it('classifies local, SSH, runtime, and unknown repo owners', () => { + expect(getAiVaultResumeRepoTargetStatus({ connectionId: null, executionHostId: 'local' })).toBe( + 'local' + ) + expect(getAiVaultResumeRepoTargetStatus({ connectionId: 'ssh-1', executionHostId: null })).toBe( + 'non-local' + ) + expect( + getAiVaultResumeRepoTargetStatus({ + connectionId: null, + executionHostId: 'runtime:env-1' + }) + ).toBe('non-local') + expect(getAiVaultResumeRepoTargetStatus(null)).toBe('unknown') + }) + + it('exposes boolean predicates for resume gates', () => { + expect(isLocalAiVaultResumeRepo({ connectionId: null, executionHostId: 'local' })).toBe(true) + expect( + isNonLocalAiVaultResumeRepo({ connectionId: null, executionHostId: 'runtime:env-1' }) + ).toBe(true) + }) + + it('resolves runtime-owned worktree targets through their repo owner', () => { + expect( + getAiVaultResumeWorktreeTargetStatus({ + worktreeId: 'repo-1::/repo/orca', + worktrees: [{ id: 'repo-1::/repo/orca', repoId: 'repo-1' }], + repos: [{ id: 'repo-1', connectionId: null, executionHostId: 'runtime:env-1' }] + }) + ).toBe('non-local') + }) + + it('uses the composite worktree repo id when worktree discovery is incomplete', () => { + expect( + getAiVaultResumeWorkspaceTargetStatus( + makeState({ + repos: [{ id: 'repo-1', connectionId: null, executionHostId: 'runtime:env-1' }] + }), + 'repo-1::/repo/orca' + ) + ).toBe('non-local') + }) + + it('resolves explicit workspace keys through the target worktree owner', () => { + expect( + getAiVaultResumeWorkspaceTargetStatus( + makeState({ + worktreesByRepo: { + 'repo-1': [{ id: 'repo-1::/repo/orca', repoId: 'repo-1' }] + }, + repos: [{ id: 'repo-1', connectionId: null, executionHostId: 'runtime:env-1' }] + }), + 'worktree:repo-1::/repo/orca' + ) + ).toBe('non-local') + }) + + it('blocks folder workspaces owned by runtime project groups', () => { + expect( + getAiVaultResumeWorkspaceTargetStatus( + makeState({ + folderWorkspaces: [ + { + id: 'folder-1', + projectGroupId: 'group-1', + name: 'Platform', + folderPath: '/repo/platform' + } + ], + projectGroups: [{ id: 'group-1', executionHostId: 'runtime:env-1' }] + }), + folderWorkspaceKey('folder-1') + ) + ).toBe('non-local') + }) + + it('blocks mixed local and runtime folder workspace targets', () => { + expect( + getAiVaultResumeWorkspaceTargetStatus( + makeState({ + folderWorkspaces: [ + { + id: 'folder-1', + projectGroupId: 'group-1', + name: 'Platform', + folderPath: '/repo/platform' + } + ], + projectGroups: [{ id: 'group-1' }], + repos: [ + { id: 'repo-local', path: '/repo/platform/web', connectionId: null }, + { + id: 'repo-runtime', + path: '/repo/platform/api', + connectionId: null, + executionHostId: 'runtime:env-1' + } + ] + }), + folderWorkspaceKey('folder-1') + ) + ).toBe('non-local') + }) +}) diff --git a/src/renderer/src/lib/ai-vault-resume-target.ts b/src/renderer/src/lib/ai-vault-resume-target.ts new file mode 100644 index 000000000..036796a3d --- /dev/null +++ b/src/renderer/src/lib/ai-vault-resume-target.ts @@ -0,0 +1,114 @@ +import { + getRepoExecutionHostId, + LOCAL_EXECUTION_HOST_ID, + normalizeExecutionHostId +} from '../../../shared/execution-host' +import type { Repo } from '../../../shared/types' +import { getRepoIdFromWorktreeId } from '../../../shared/worktree-id' +import { parseWorkspaceKey } from '../../../shared/workspace-scope' +import type { AppState } from '@/store/types' +import { getFolderWorkspaceCandidateRepos } from './folder-workspace-connection' + +export type AiVaultResumeTargetStatus = 'local' | 'non-local' | 'unknown' + +type AiVaultResumeRepoOwner = Pick + +export function getAiVaultResumeRepoTargetStatus( + repo: AiVaultResumeRepoOwner | null | undefined +): AiVaultResumeTargetStatus { + if (!repo) { + return 'unknown' + } + // Why: runtime-owned repos intentionally keep connectionId null, so resume + // availability must follow the execution owner instead of SSH state alone. + return repo.connectionId || getRepoExecutionHostId(repo) !== LOCAL_EXECUTION_HOST_ID + ? 'non-local' + : 'local' +} + +export function isLocalAiVaultResumeRepo(repo: AiVaultResumeRepoOwner | null | undefined): boolean { + return getAiVaultResumeRepoTargetStatus(repo) === 'local' +} + +export function isNonLocalAiVaultResumeRepo( + repo: AiVaultResumeRepoOwner | null | undefined +): boolean { + return getAiVaultResumeRepoTargetStatus(repo) === 'non-local' +} + +export function getAiVaultResumeWorktreeTargetStatus(args: { + worktreeId: string | null + worktrees: readonly { id: string; repoId: string }[] + repos: readonly AiVaultResumeRepoOwnerWithId[] +}): AiVaultResumeTargetStatus { + if (!args.worktreeId) { + return 'unknown' + } + const worktree = args.worktrees.find((candidate) => candidate.id === args.worktreeId) + if (!worktree) { + return 'unknown' + } + return getAiVaultResumeRepoTargetStatus( + args.repos.find((candidate) => candidate.id === worktree.repoId) + ) +} + +export function getAiVaultResumeWorkspaceTargetStatus( + state: Pick, + workspaceId: string | null +): AiVaultResumeTargetStatus { + if (!workspaceId) { + return 'unknown' + } + + const workspaceKey = parseWorkspaceKey(workspaceId) + if (workspaceKey?.type === 'folder') { + return getAiVaultResumeFolderTargetStatus(state, workspaceKey.folderWorkspaceId) + } + + const worktreeId = workspaceKey?.type === 'worktree' ? workspaceKey.worktreeId : workspaceId + const worktree = Object.values(state.worktreesByRepo ?? {}) + .flat() + .find((candidate) => candidate.id === worktreeId) + const repoId = worktree?.repoId ?? getRepoIdFromWorktreeId(worktreeId) + return getAiVaultResumeRepoTargetStatus(state.repos.find((repo) => repo.id === repoId)) +} + +type AiVaultResumeRepoOwnerWithId = AiVaultResumeRepoOwner & { id: string } + +function getAiVaultResumeFolderTargetStatus( + state: Pick, + folderWorkspaceId: string +): AiVaultResumeTargetStatus { + const workspace = state.folderWorkspaces.find((entry) => entry.id === folderWorkspaceId) + if (!workspace) { + return 'unknown' + } + + const group = state.projectGroups.find((entry) => entry.id === workspace.projectGroupId) + const groupHostId = normalizeExecutionHostId(group?.executionHostId) + if ( + workspace.connectionId || + group?.connectionId || + (groupHostId && groupHostId !== LOCAL_EXECUTION_HOST_ID) + ) { + return 'non-local' + } + + return mergeAiVaultResumeTargetStatuses( + getFolderWorkspaceCandidateRepos(state, folderWorkspaceId).map(getAiVaultResumeRepoTargetStatus) + ) +} + +function mergeAiVaultResumeTargetStatuses( + statuses: readonly AiVaultResumeTargetStatus[] +): AiVaultResumeTargetStatus { + if (statuses.length === 0) { + return 'local' + } + const uniqueStatuses = new Set(statuses) + if (uniqueStatuses.size === 1) { + return statuses[0] ?? 'unknown' + } + return uniqueStatuses.has('unknown') ? 'unknown' : 'non-local' +} diff --git a/src/shared/project-groups.test.ts b/src/shared/project-groups.test.ts index 9517d9934..4c11e8d59 100644 --- a/src/shared/project-groups.test.ts +++ b/src/shared/project-groups.test.ts @@ -72,6 +72,18 @@ describe('project-groups', () => { }) }) + it('preserves normalized execution ownership for persisted groups', () => { + const groups = normalizeProjectGroups([ + { id: 'runtime', name: 'Runtime', tabOrder: 1, executionHostId: 'runtime:env-1' }, + { id: 'local', name: 'Local', tabOrder: 2, executionHostId: 'local' }, + { id: 'invalid', name: 'Invalid', tabOrder: 3, executionHostId: 'runtime:' } + ]) + + expect(groups.find((group) => group.id === 'runtime')?.executionHostId).toBe('runtime:env-1') + expect(groups.find((group) => group.id === 'local')?.executionHostId).toBe('local') + expect(groups.find((group) => group.id === 'invalid')?.executionHostId).toBeUndefined() + }) + it('clears repo memberships whose group no longer exists', () => { const groups = [createProjectGroup({ name: 'Known', createdFrom: 'manual', tabOrder: 0 })] const repos = clearMissingProjectGroupMemberships( diff --git a/src/shared/project-groups.ts b/src/shared/project-groups.ts index 1badd1f69..6d9b02d83 100644 --- a/src/shared/project-groups.ts +++ b/src/shared/project-groups.ts @@ -1,3 +1,4 @@ +import { normalizeExecutionHostId } from './execution-host' import type { Repo, ProjectGroup, ProjectGroupCreatedFrom } from './types' export const UNGROUPED_PROJECT_GROUP_KEY = 'project-group:ungrouped' @@ -56,6 +57,7 @@ export function normalizeProjectGroups(value: unknown): ProjectGroup[] { } seen.add(raw.id) const now = Date.now() + const executionHostId = normalizeExecutionHostId(raw.executionHostId) groups.push({ id: raw.id, name: normalizeProjectGroupName(typeof raw.name === 'string' ? raw.name : ''), @@ -80,7 +82,9 @@ export function normalizeProjectGroups(value: unknown): ProjectGroup[] { createdAt: typeof raw.createdAt === 'number' && Number.isFinite(raw.createdAt) ? raw.createdAt : now, updatedAt: - typeof raw.updatedAt === 'number' && Number.isFinite(raw.updatedAt) ? raw.updatedAt : now + typeof raw.updatedAt === 'number' && Number.isFinite(raw.updatedAt) ? raw.updatedAt : now, + // Why: runtime-owned groups otherwise look local after persistence reload. + ...(executionHostId ? { executionHostId } : {}) }) } groups.sort(