From 38b32a95e72aec1beb4172e7404a51bf5168f560 Mon Sep 17 00:00:00 2001 From: Brennan Benson <79079362+brennanb2025@users.noreply.github.com> Date: Mon, 3 Aug 2026 18:24:19 -0700 Subject: [PATCH] fix(skills): harden SSH pane discovery --- src/main/runtime/orca-runtime.ts | 37 +-- .../skills/discovery-wire-contract.test.ts | 26 ++ src/main/skills/discovery.ts | 39 ++- .../pane-skill-discovery-workspace.test.ts | 150 +++++++++++ .../skills/pane-skill-discovery-workspace.ts | 245 ++++++++++++++++++ .../native-chat-pane-discovery-cache.ts | 31 +++ .../native-chat-skill-discovery-context.ts | 26 +- .../use-native-chat-skills.react.test.tsx | 32 ++- .../use-native-chat-skills.test.ts | 41 ++- .../native-chat/use-native-chat-skills.ts | 23 +- src/shared/skills.ts | 30 ++- 11 files changed, 629 insertions(+), 51 deletions(-) create mode 100644 src/main/skills/discovery-wire-contract.test.ts create mode 100644 src/main/skills/pane-skill-discovery-workspace.test.ts create mode 100644 src/main/skills/pane-skill-discovery-workspace.ts create mode 100644 src/renderer/src/components/native-chat/native-chat-pane-discovery-cache.ts diff --git a/src/main/runtime/orca-runtime.ts b/src/main/runtime/orca-runtime.ts index d33bf4aec94..5ce1e615422 100644 --- a/src/main/runtime/orca-runtime.ts +++ b/src/main/runtime/orca-runtime.ts @@ -668,6 +668,7 @@ import { resolveLocalProjectRuntimeForWorktreeId } from '../local-project-runtim import type { ProjectExecutionRuntimeResolution } from '../../shared/project-execution-runtime' import type { SkillDiscoveryForPaneResponse } from '../../shared/skills' import { discoverPaneSkills } from '../skills/pane-skill-discovery' +import { resolvePaneSkillDiscoveryWorkspace } from '../skills/pane-skill-discovery-workspace' import { resolveTerminalOrchestrationCliCommand } from './orchestration/cli-command' import { getLocalWorktreePathAccess, @@ -12281,26 +12282,26 @@ export class OrcaRuntimeService { args: { worktreeId: string; terminalTabId?: string }, signal?: AbortSignal ): Promise { - const target = await this.resolveRuntimeFileTarget(args.worktreeId) - const session = this.getWorkspaceSessionForWorktree(args.worktreeId) - const sessionWorktreeId = session - ? resolveTerminalSessionWorktreeId(session, args.worktreeId) - : null - const tab = - args.terminalTabId && session && sessionWorktreeId - ? session.tabsByWorktree[sessionWorktreeId]?.find( - (candidate) => candidate.id === args.terminalTabId - ) - : undefined + const store = this.requireStore() + const scope = parseWorkspaceKey(args.worktreeId) + const metadataWorktreeId = scope?.type === 'worktree' ? scope.worktreeId : args.worktreeId + const workspace = resolvePaneSkillDiscoveryWorkspace({ + worktreeId: args.worktreeId, + terminalTabId: args.terminalTabId, + repos: store.getRepos(), + projectGroups: store.getProjectGroups?.() ?? [], + folderWorkspaces: store.getFolderWorkspaces?.() ?? [], + worktreeMeta: + scope?.type === 'folder' ? undefined : store.getWorktreeMeta(metadataWorktreeId), + sessions: (store.getWorkspaceSessionHostIds?.() ?? [LOCAL_EXECUTION_HOST_ID]).map( + (hostId) => ({ hostId, session: store.getWorkspaceSession(hostId) }) + ) + }) return discoverPaneSkills({ worktreeId: args.worktreeId, - // A missing/unpersisted tab still scans a runtime-owned path: the workspace root. - cwd: tab?.startupCwd?.trim() || target.worktree.path, - connectionId: target.connectionId, - projectRuntime: target.connectionId - ? undefined - : this.resolveProjectRuntimeForWorktree(args.worktreeId), - repos: this.listRepos(), + cwd: workspace.cwd, + connectionId: workspace.connectionId, + repos: [], signal }) } diff --git a/src/main/skills/discovery-wire-contract.test.ts b/src/main/skills/discovery-wire-contract.test.ts new file mode 100644 index 00000000000..4dc646d201c --- /dev/null +++ b/src/main/skills/discovery-wire-contract.test.ts @@ -0,0 +1,26 @@ +import { mkdir, mkdtemp, writeFile } from 'node:fs/promises' +import { tmpdir } from 'node:os' +import { join } from 'node:path' +import { describe, expect, it } from 'vitest' +import { parseSkillDiscoveryResult, SKILL_DISCOVERY_LIMITS } from '../../shared/skills' +import { discoverSkills } from './discovery' + +describe('skill discovery wire contract', () => { + it('bounds scanner metadata before returning it to an SSH client', async () => { + const root = await mkdtemp(join(tmpdir(), 'orca-skills-')) + const home = join(root, 'home') + const skillDir = join(home, '.agents', 'skills', 'oversized') + await mkdir(skillDir, { recursive: true }) + await writeFile( + join(skillDir, 'SKILL.md'), + `---\nname: ${'n'.repeat(SKILL_DISCOVERY_LIMITS.nameLength + 1)}\ndescription: ${'d'.repeat(SKILL_DISCOVERY_LIMITS.descriptionLength + 1)}\n---\n` + ) + + const result = await discoverSkills({ homeDir: home, repos: [], includeCwd: false }) + const discovered = result.skills.find((entry) => entry.directoryPath === skillDir) + + expect(discovered?.name).toHaveLength(SKILL_DISCOVERY_LIMITS.nameLength) + expect(discovered?.description).toHaveLength(SKILL_DISCOVERY_LIMITS.descriptionLength) + expect(() => parseSkillDiscoveryResult(result)).not.toThrow() + }) +}) diff --git a/src/main/skills/discovery.ts b/src/main/skills/discovery.ts index 0d19c7320b3..7977fa4e996 100644 --- a/src/main/skills/discovery.ts +++ b/src/main/skills/discovery.ts @@ -4,10 +4,11 @@ import { homedir } from 'node:os' import { basename, dirname, isAbsolute, join, relative, sep } from 'node:path' import { summarizeSkillMarkdown } from '../../shared/skill-metadata' import type { Repo } from '../../shared/types' -import type { - DiscoveredSkill, - SkillDiscoveryResult, - SkillDiscoverySource +import { + SKILL_DISCOVERY_LIMITS, + type DiscoveredSkill, + type SkillDiscoveryResult, + type SkillDiscoverySource } from '../../shared/skills' import { buildSkillDiscoverySources, @@ -18,6 +19,7 @@ import { type SkillScanRoot } from './skill-discovery-sources' import { discoverClaudePluginSkillSources } from './claude-plugin-skill-sources' +import { runSkillCandidateTasks } from './skill-candidate-concurrency' export { buildSkillDiscoverySources } from './skill-discovery-sources' @@ -65,6 +67,9 @@ async function findSkillFiles( const visitedDirectoryPaths = new Set() async function visit(dirPath: string): Promise { throwIfDiscoveryAborted(signal) + if (out.length >= SKILL_DISCOVERY_LIMITS.skills) { + return + } if (!isWithinDepth(rootPath, dirPath, maxDepth)) { return } @@ -87,6 +92,9 @@ async function findSkillFiles( } for (const entry of entries) { throwIfDiscoveryAborted(signal) + if (out.length >= SKILL_DISCOVERY_LIMITS.skills) { + return + } const entryPath = join(dirPath, entry.name) if (entry.name === SKILL_FILE_NAME) { if (entry.isFile()) { @@ -204,8 +212,8 @@ type ScannedSkill = DiscoveredSkill & { canonicalSkillFilePath: string } async function scanRoot(root: SkillScanRoot, signal?: AbortSignal): Promise { const maxDepth = root.sourceKind === 'plugin' ? 9 : 4 const skillFiles = await findSkillFiles(root.path, maxDepth, signal) - const skills = await Promise.all( - skillFiles.map(async (skillFilePath): Promise => { + const skills = await runSkillCandidateTasks( + skillFiles.map((skillFilePath) => async (): Promise => { throwIfDiscoveryAborted(signal) // Why: path identity belongs to the scanning host; canonicalizing before // returning prevents symlinked roots from becoming duplicate picker rows. @@ -218,8 +226,9 @@ async function scanRoot(root: SkillScanRoot, signal?: AbortSignal): Promise root.path.length <= SKILL_DISCOVERY_LIMITS.pathLength) + .slice(0, SKILL_DISCOVERY_LIMITS.sources) const sources: SkillDiscoverySource[] = [] - const skillGroups = await Promise.all( - roots.map(async (root) => { + const skillGroups = await runSkillCandidateTasks( + roots.map((root) => async () => { throwIfDiscoveryAborted(signal) const exists = await pathExists(root.path) sources.push({ @@ -281,7 +292,11 @@ export async function discoverSkills(args: { // record every contributing root so per-agent visibility survives dedup. const existing = seen.get(skill.canonicalSkillFilePath) if (existing) { - if (existing.rootPaths && !existing.rootPaths.includes(skill.rootPath)) { + if ( + existing.rootPaths && + existing.rootPaths.length < SKILL_DISCOVERY_LIMITS.rootPaths && + !existing.rootPaths.includes(skill.rootPath) + ) { existing.rootPaths.push(skill.rootPath) } // Why: providers is per-agent visibility just like rootPaths; keeping only @@ -302,7 +317,7 @@ export async function discoverSkills(args: { seen.set(canonicalSkillFilePath, { ...publicSkill, rootPaths: [skill.rootPath] }) } return { - skills: Array.from(seen.values()).sort(compareSkills), + skills: Array.from(seen.values()).sort(compareSkills).slice(0, SKILL_DISCOVERY_LIMITS.skills), sources: sources.sort((a, b) => a.label.localeCompare(b.label, undefined, { sensitivity: 'base' }) ), diff --git a/src/main/skills/pane-skill-discovery-workspace.test.ts b/src/main/skills/pane-skill-discovery-workspace.test.ts new file mode 100644 index 00000000000..5c4f48d9d50 --- /dev/null +++ b/src/main/skills/pane-skill-discovery-workspace.test.ts @@ -0,0 +1,150 @@ +import { describe, expect, it } from 'vitest' +import type { FolderWorkspace, ProjectGroup, Repo, WorkspaceSessionState } from '../../shared/types' +import { resolvePaneSkillDiscoveryWorkspace } from './pane-skill-discovery-workspace' + +const EMPTY_SESSION = { tabsByWorktree: {} } as WorkspaceSessionState + +function repo(overrides: Partial = {}): Repo { + return { + id: 'repo-1', + path: '/remote/repo', + displayName: 'repo', + badgeColor: 'blue', + addedAt: 1, + connectionId: 'target-1', + ...overrides + } +} + +function session(worktreeId: string, startupCwd?: string): WorkspaceSessionState { + return { + tabsByWorktree: { + [worktreeId]: [{ id: 'tab-1', startupCwd }] + } + } as WorkspaceSessionState +} + +describe('resolvePaneSkillDiscoveryWorkspace', () => { + it('uses persisted pane ownership without scanning the repo catalog', () => { + const worktreeId = 'repo-1::/remote/repo/worktree' + expect( + resolvePaneSkillDiscoveryWorkspace({ + worktreeId, + terminalTabId: 'tab-1', + repos: [repo()], + projectGroups: [], + folderWorkspaces: [], + sessions: [ + { hostId: 'local', session: EMPTY_SESSION }, + { hostId: 'ssh:target-1', session: session(worktreeId, '/remote/repo/packages/app') } + ] + }) + ).toEqual({ connectionId: 'target-1', cwd: '/remote/repo/packages/app' }) + }) + + it('routes an executionHostId-only SSH repo instead of falling back locally', () => { + expect( + resolvePaneSkillDiscoveryWorkspace({ + worktreeId: 'repo-1::/remote/repo', + repos: [repo({ connectionId: null, executionHostId: 'ssh:target-1' })], + projectGroups: [], + folderWorkspaces: [], + sessions: [] + }) + ).toEqual({ connectionId: 'target-1', cwd: '/remote/repo' }) + }) + + it('uses the pane session to isolate duplicate repo identities across hosts', () => { + const worktreeId = 'repo-1::/same/path' + expect( + resolvePaneSkillDiscoveryWorkspace({ + worktreeId, + terminalTabId: 'tab-1', + repos: [ + repo({ path: '/same/path', connectionId: 'target-1' }), + repo({ path: '/same/path', connectionId: 'target-2' }) + ], + projectGroups: [], + folderWorkspaces: [], + sessions: [{ hostId: 'ssh:target-2', session: session(worktreeId) }] + }) + ).toEqual({ connectionId: 'target-2', cwd: '/same/path' }) + }) + + it('rejects contradictory repo ownership instead of selecting either host', () => { + expect(() => + resolvePaneSkillDiscoveryWorkspace({ + worktreeId: 'repo-1::/remote/repo', + repos: [repo({ connectionId: 'target-1', executionHostId: 'ssh:target-2' })], + projectGroups: [], + folderWorkspaces: [], + sessions: [] + }) + ).toThrow('pane_skill_discovery_owner_invalid') + }) + + it('rejects an unpersisted renderer-claimed linked-worktree path', () => { + expect(() => + resolvePaneSkillDiscoveryWorkspace({ + worktreeId: 'repo-1::/remote/other', + repos: [repo()], + projectGroups: [], + folderWorkspaces: [], + sessions: [] + }) + ).toThrow('pane_skill_discovery_workspace_not_found') + }) + + it('resolves SSH folder workspaces from persisted folder ownership', () => { + const group = { + id: 'group-1', + connectionId: 'target-1', + parentPath: '/remote/folder' + } as ProjectGroup + const workspace = { + id: 'folder-1', + projectGroupId: group.id, + folderPath: '/remote/folder', + connectionId: 'target-1' + } as FolderWorkspace + expect( + resolvePaneSkillDiscoveryWorkspace({ + worktreeId: 'folder:folder-1', + repos: [], + projectGroups: [group], + folderWorkspaces: [workspace], + sessions: [] + }) + ).toEqual({ connectionId: 'target-1', cwd: '/remote/folder' }) + }) + + it('routes an executionHostId-only SSH folder workspace', () => { + expect( + resolvePaneSkillDiscoveryWorkspace({ + worktreeId: 'folder:folder-1', + repos: [], + projectGroups: [], + folderWorkspaces: [ + { + id: 'folder-1', + projectGroupId: 'group-1', + name: 'Remote folder', + folderPath: '/remote/folder', + connectionId: null, + executionHostId: 'ssh:target-1', + linkedTask: null, + comment: '', + isArchived: false, + isUnread: false, + isPinned: false, + sortOrder: 0, + lastActivityAt: 1, + createdAt: 1, + updatedAt: 1 + } + ], + sessions: [] + }) + ).toEqual({ connectionId: 'target-1', cwd: '/remote/folder' }) + }) +}) diff --git a/src/main/skills/pane-skill-discovery-workspace.ts b/src/main/skills/pane-skill-discovery-workspace.ts new file mode 100644 index 00000000000..881a4b078c0 --- /dev/null +++ b/src/main/skills/pane-skill-discovery-workspace.ts @@ -0,0 +1,245 @@ +import { normalizeRuntimePathForComparison } from '../../shared/cross-platform-path' +import { + LOCAL_EXECUTION_HOST_ID, + parseExecutionHostId, + toSshExecutionHostId, + type ExecutionHostId +} from '../../shared/execution-host' +import type { + FolderWorkspace, + ProjectGroup, + Repo, + WorktreeMeta, + WorkspaceSessionState +} from '../../shared/types' +import { splitWorktreeIdForFilesystem } from '../../shared/worktree-id' +import { parseWorkspaceKey } from '../../shared/workspace-scope' +import { inferFolderWorkspacePathConnection } from '../project-groups/folder-workspace-path-status' + +export type PaneSkillDiscoveryWorkspace = { + connectionId: string + cwd: string +} + +type HostSession = { hostId: ExecutionHostId; session: WorkspaceSessionState } + +type PaneSkillDiscoveryWorkspaceArgs = { + worktreeId: string + terminalTabId?: string + repos: readonly Repo[] + projectGroups: readonly ProjectGroup[] + folderWorkspaces: readonly FolderWorkspace[] + worktreeMeta?: WorktreeMeta + sessions: readonly HostSession[] +} + +export function resolvePaneSkillDiscoveryWorkspace( + args: PaneSkillDiscoveryWorkspaceArgs +): PaneSkillDiscoveryWorkspace { + const scope = parseWorkspaceKey(args.worktreeId) + const worktreeId = scope?.type === 'worktree' ? scope.worktreeId : args.worktreeId + const sessionOwner = findPaneSession(args.sessions, worktreeId, args.terminalTabId) + + if (scope?.type === 'folder') { + const workspace = selectFolderWorkspace( + args.folderWorkspaces.filter((candidate) => candidate.id === scope.folderWorkspaceId), + sessionOwner?.hostId + ) + if (!workspace) { + throw new Error('pane_skill_discovery_workspace_not_found') + } + const group = args.projectGroups.find((candidate) => candidate.id === workspace.projectGroupId) + const executionHostId = parseOptionalHostId(workspace.executionHostId ?? group?.executionHostId) + const connectionId = parseOptionalConnectionId(workspace.connectionId ?? group?.connectionId) + const connectionHostId = connectionId ? toSshExecutionHostId(connectionId) : null + assertMatchingHost(executionHostId, connectionHostId) + const hostId = + executionHostId ?? + connectionHostId ?? + inferLegacyFolderHostId(workspace, args.projectGroups, args.repos) + assertMatchingHost(sessionOwner?.hostId, hostId) + return sshWorkspace(hostId, sessionOwner?.tab.startupCwd?.trim() || workspace.folderPath) + } + + const parsed = splitWorktreeIdForFilesystem(worktreeId) + if (!parsed?.repoId || !parsed.worktreePath) { + throw new Error('pane_skill_discovery_workspace_not_found') + } + const candidates = args.repos.filter((repo) => repo.id === parsed.repoId) + if (candidates.length === 0) { + throw new Error('pane_skill_discovery_workspace_not_found') + } + const metaHostId = parseOptionalHostId(args.worktreeMeta?.hostId) + assertMatchingHost(sessionOwner?.hostId, metaHostId) + const knownHostId = sessionOwner?.hostId ?? metaHostId + const repo = selectPaneRepo(candidates, knownHostId) + const repoHostId = repoExecutionHostId(repo) + const hostId = knownHostId ?? repoHostId + if (knownHostId && hasExplicitRepoOwner(repo)) { + assertMatchingHost(knownHostId, repoHostId) + } + const isPersistedWorktree = Boolean(sessionOwner || args.worktreeMeta) + if ( + !isPersistedWorktree && + normalizeRuntimePathForComparison(parsed.worktreePath) !== + normalizeRuntimePathForComparison(repo.path) + ) { + throw new Error('pane_skill_discovery_workspace_not_found') + } + return sshWorkspace(hostId, sessionOwner?.tab.startupCwd?.trim() || parsed.worktreePath) +} + +function inferLegacyFolderHostId( + workspace: FolderWorkspace, + projectGroups: readonly ProjectGroup[], + repos: readonly Repo[] +): ExecutionHostId { + const connection = inferFolderWorkspacePathConnection({ + folderPath: workspace.folderPath, + projectGroupId: workspace.projectGroupId, + connectionId: null, + projectGroups, + repos + }) + if (connection.kind === 'ambiguous') { + throw new Error('pane_skill_discovery_owner_ambiguous') + } + return connection.kind === 'ssh' + ? toSshExecutionHostId(connection.connectionId) + : LOCAL_EXECUTION_HOST_ID +} + +function selectFolderWorkspace( + candidates: readonly FolderWorkspace[], + hostId: ExecutionHostId | undefined +): FolderWorkspace | null { + if (candidates.length === 0) { + return null + } + if (candidates.length === 1) { + return candidates[0] + } + if (!hostId) { + throw new Error('pane_skill_discovery_owner_ambiguous') + } + const exact = candidates.filter((workspace) => { + const explicitHostId = parseOptionalHostId(workspace.executionHostId) + const connectionId = parseOptionalConnectionId(workspace.connectionId) + const connectionHostId = connectionId ? toSshExecutionHostId(connectionId) : null + assertMatchingHost(explicitHostId, connectionHostId) + return (explicitHostId ?? connectionHostId) === hostId + }) + if (exact.length !== 1) { + throw new Error('pane_skill_discovery_owner_ambiguous') + } + return exact[0] +} + +function findPaneSession( + sessions: readonly HostSession[], + worktreeId: string, + terminalTabId: string | undefined +): { hostId: ExecutionHostId; tab: { startupCwd?: string } } | null { + if (!terminalTabId) { + return null + } + const matches = sessions.flatMap(({ hostId, session }) => + Object.entries(session.tabsByWorktree ?? {}).flatMap(([candidateWorktreeId, tabs]) => + worktreeIdsEqual(candidateWorktreeId, worktreeId) + ? tabs + .filter((tab) => tab.id === terminalTabId) + .map((tab) => ({ hostId, tab: { startupCwd: tab.startupCwd } })) + : [] + ) + ) + if (matches.length > 1) { + throw new Error('pane_skill_discovery_owner_ambiguous') + } + return matches[0] ?? null +} + +function selectPaneRepo(candidates: readonly Repo[], hostId: ExecutionHostId | null): Repo { + if (!hostId) { + if (candidates.length !== 1) { + throw new Error('pane_skill_discovery_owner_ambiguous') + } + repoExecutionHostId(candidates[0]) + return candidates[0] + } + const exact = candidates.filter((repo) => repoExecutionHostId(repo) === hostId) + if (exact.length === 1) { + return exact[0] + } + const legacy = candidates.filter((repo) => !hasExplicitRepoOwner(repo)) + if (exact.length === 0 && legacy.length === 1) { + return legacy[0] + } + throw new Error('pane_skill_discovery_owner_ambiguous') +} + +function repoExecutionHostId(repo: Repo): ExecutionHostId { + const hasHost = repo.executionHostId !== null && repo.executionHostId !== undefined + const parsedHost = hasHost ? parseExecutionHostId(repo.executionHostId) : null + if (hasHost && !parsedHost) { + throw new Error('pane_skill_discovery_owner_invalid') + } + const connectionId = parseOptionalConnectionId(repo.connectionId) + const connectionHostId = connectionId ? toSshExecutionHostId(connectionId) : null + if (parsedHost && connectionHostId && parsedHost.id !== connectionHostId) { + throw new Error('pane_skill_discovery_owner_invalid') + } + return parsedHost?.id ?? connectionHostId ?? LOCAL_EXECUTION_HOST_ID +} + +function parseOptionalHostId(hostId: string | null | undefined): ExecutionHostId | null { + if (hostId === null || hostId === undefined) { + return null + } + const parsed = parseExecutionHostId(hostId) + if (!parsed) { + throw new Error('pane_skill_discovery_owner_invalid') + } + return parsed.id +} + +function parseOptionalConnectionId(connectionId: string | null | undefined): string | null { + if (connectionId === null || connectionId === undefined) { + return null + } + const normalized = connectionId.trim() + if (!normalized) { + throw new Error('pane_skill_discovery_owner_invalid') + } + return normalized +} + +function hasExplicitRepoOwner(repo: Repo): boolean { + return Boolean(repo.executionHostId?.trim() || repo.connectionId?.trim()) +} + +function assertMatchingHost( + left: ExecutionHostId | null | undefined, + right: ExecutionHostId | null | undefined +): void { + if (left && right && left !== right) { + throw new Error('pane_skill_discovery_owner_ambiguous') + } +} + +function worktreeIdsEqual(left: string, right: string): boolean { + const parsedLeft = splitWorktreeIdForFilesystem(left) + const parsedRight = splitWorktreeIdForFilesystem(right) + return parsedLeft && parsedRight + ? parsedLeft.repoId === parsedRight.repoId && + normalizeRuntimePathForComparison(parsedLeft.worktreePath) === + normalizeRuntimePathForComparison(parsedRight.worktreePath) + : left === right +} + +function sshWorkspace(hostId: ExecutionHostId, cwd: string): PaneSkillDiscoveryWorkspace { + const host = parseExecutionHostId(hostId) + if (host?.kind !== 'ssh') { + throw new Error('pane_skill_discovery_requires_ssh_workspace') + } + return { connectionId: host.targetId, cwd } +} diff --git a/src/renderer/src/components/native-chat/native-chat-pane-discovery-cache.ts b/src/renderer/src/components/native-chat/native-chat-pane-discovery-cache.ts new file mode 100644 index 00000000000..52b146cfb65 --- /dev/null +++ b/src/renderer/src/components/native-chat/native-chat-pane-discovery-cache.ts @@ -0,0 +1,31 @@ +import type { SkillDiscoveryResult } from '../../../../shared/skills' + +const PANE_DISCOVERY_CACHE_LIMIT = 8 + +export function readPaneDiscoveryCache( + cache: Map, + key: string +): SkillDiscoveryResult | undefined { + const result = cache.get(key) + if (result) { + cache.delete(key) + cache.set(key, result) + } + return result +} + +export function writePaneDiscoveryCache( + cache: Map, + key: string, + result: SkillDiscoveryResult +): void { + cache.delete(key) + cache.set(key, result) + while (cache.size > PANE_DISCOVERY_CACHE_LIMIT) { + const oldestKey = cache.keys().next().value + if (oldestKey === undefined) { + return + } + cache.delete(oldestKey) + } +} diff --git a/src/renderer/src/components/native-chat/native-chat-skill-discovery-context.ts b/src/renderer/src/components/native-chat/native-chat-skill-discovery-context.ts index 8f3ecbee8e6..4af023e78d7 100644 --- a/src/renderer/src/components/native-chat/native-chat-skill-discovery-context.ts +++ b/src/renderer/src/components/native-chat/native-chat-skill-discovery-context.ts @@ -134,7 +134,9 @@ export function resolveNativeChatSkillDiscoveryContext( runtimeEnvironmentId ?? null, hostId, connectionState?.connectionGeneration ?? 0, - cwd + cwd, + worktreeId, + terminalTabId ]), cwd, executionHostKind: 'ssh', @@ -180,6 +182,28 @@ export function resolveNativeChatSkillDiscoveryContext( } } +export function resolveNativeChatSkillDiscoverySubscriptionKey( + state: AppState, + terminalTabId: string +): string | null { + const context = resolveNativeChatSkillDiscoveryContext( + selectNativeChatSkillStateInputs(state), + terminalTabId + ) + return getNativeChatSkillDiscoverySubscriptionKey(context) +} + +export function getNativeChatSkillDiscoverySubscriptionKey( + context: NativeChatSkillDiscoveryContext | null +): string | null { + if (!context) { + return null + } + return context.executionHostKind === 'ssh' + ? JSON.stringify([context.key, context.sshDisconnected]) + : context.key +} + function findTerminalTab( tabsByWorktree: Record, terminalTabId: string diff --git a/src/renderer/src/components/native-chat/use-native-chat-skills.react.test.tsx b/src/renderer/src/components/native-chat/use-native-chat-skills.react.test.tsx index 368bfc35df6..a38d8df0a85 100644 --- a/src/renderer/src/components/native-chat/use-native-chat-skills.react.test.tsx +++ b/src/renderer/src/components/native-chat/use-native-chat-skills.react.test.tsx @@ -11,7 +11,10 @@ const mocks = vi.hoisted(() => ({ })) vi.mock('../../store', () => ({ - useAppStore: (selector: (state: Record) => unknown) => selector(mocks.state) + useAppStore: Object.assign( + (selector: (state: Record) => unknown) => selector(mocks.state), + { getState: () => mocks.state } + ) })) vi.mock('@/runtime/runtime-rpc-client', () => ({ callRuntimeRpc: (...args: unknown[]) => mocks.callRuntimeRpc(...args) @@ -53,7 +56,7 @@ function stateForHost(hostId: string) { } } -function connectedSshState() { +function connectedSshState(connectionGeneration = 3) { return { ...stateForHost('ssh:target-1'), sshConnectionStates: new Map([ @@ -64,7 +67,7 @@ function connectedSshState() { status: 'connected', error: null, reconnectAttempt: 0, - connectionGeneration: 3 + connectionGeneration } ] ]) @@ -198,4 +201,27 @@ describe('useNativeChatSkills', () => { { timeoutMs: 10_000 } ) }) + + it('bounds cached SSH reconnect generations per composer', async () => { + mocks.state = connectedSshState(1) + mocks.callRuntimeRpc.mockResolvedValue({ status: 'ok', result: DISCOVERY_RESULT }) + const view = render() + await waitFor(() => expect(mocks.snapshots.at(-1)?.status).toBe('ready')) + + for (let generation = 2; generation <= 9; generation += 1) { + mocks.state = connectedSshState(generation) + view.rerender() + await waitFor(() => { + expect(mocks.callRuntimeRpc).toHaveBeenCalledTimes(generation) + expect(mocks.snapshots.at(-1)?.status).toBe('ready') + }) + } + + mocks.state = connectedSshState(1) + view.rerender() + await waitFor(() => { + expect(mocks.callRuntimeRpc).toHaveBeenCalledTimes(10) + expect(mocks.snapshots.at(-1)?.status).toBe('ready') + }) + }) }) diff --git a/src/renderer/src/components/native-chat/use-native-chat-skills.test.ts b/src/renderer/src/components/native-chat/use-native-chat-skills.test.ts index e6b329ea266..df6fb3d2e5f 100644 --- a/src/renderer/src/components/native-chat/use-native-chat-skills.test.ts +++ b/src/renderer/src/components/native-chat/use-native-chat-skills.test.ts @@ -5,7 +5,10 @@ import { resolveNativeChatSkillDiscoveryContext, resolveNativeChatSkillDiscoveryCwd } from './use-native-chat-skills' -import type { NativeChatSkillStateInputs } from './native-chat-skill-discovery-context' +import { + getNativeChatSkillDiscoverySubscriptionKey, + type NativeChatSkillStateInputs +} from './native-chat-skill-discovery-context' function skill(overrides: Partial): DiscoveredSkill { return { @@ -341,4 +344,40 @@ describe('resolveNativeChatSkillDiscoveryContext for SSH panes', () => { expect(onTargetOne?.cwd).toBe(onTargetTwo?.cwd) expect(onTargetOne?.key).not.toBe(onTargetTwo?.key) }) + + it('keeps pane identity in the key when renderer paths are stale or equal', () => { + const firstPane = resolveNativeChatSkillDiscoveryContext(sshInputs(), 'tab-1') + const secondPane = resolveNativeChatSkillDiscoveryContext( + sshInputs({ + tabsByWorktree: { 'worktree-2': [{ id: 'tab-2' }] }, + worktreesByRepo: { + 'repo-1': [ + { id: 'worktree-2', repoId: 'repo-1', path: '/repo/worktree', hostId: 'ssh:target-1' } + ] + } + }), + 'tab-2' + ) + expect(firstPane?.cwd).toBe(secondPane?.cwd) + expect(firstPane?.key).not.toBe(secondPane?.key) + }) + + it('ignores unrelated SSH state updates in the hook subscription key', () => { + const initial = sshInputs() + const unrelatedUpdate = sshInputs({ + sshConnectionStates: new Map([ + ['target-1', connectionState()], + ['target-2', connectionState({ targetId: 'target-2', connectionGeneration: 99 })] + ]) + }) + expect( + getNativeChatSkillDiscoverySubscriptionKey( + resolveNativeChatSkillDiscoveryContext(unrelatedUpdate, 'tab-1') + ) + ).toBe( + getNativeChatSkillDiscoverySubscriptionKey( + resolveNativeChatSkillDiscoveryContext(initial, 'tab-1') + ) + ) + }) }) diff --git a/src/renderer/src/components/native-chat/use-native-chat-skills.ts b/src/renderer/src/components/native-chat/use-native-chat-skills.ts index a461501630d..44c9f26d271 100644 --- a/src/renderer/src/components/native-chat/use-native-chat-skills.ts +++ b/src/renderer/src/components/native-chat/use-native-chat-skills.ts @@ -1,5 +1,4 @@ import { useCallback, useEffect, useMemo, useRef, useState } from 'react' -import { useShallow } from 'zustand/react/shallow' import { useAppStore } from '../../store' import type { AgentType } from '../../../../shared/agent-status-types' import type { @@ -13,11 +12,14 @@ import { RuntimeRpcCallError } from '@/runtime/runtime-rpc-result' import { isRuntimeCompatBlockError } from '@/runtime/runtime-protocol-compat' import { emitNativeChatSkillDiscovery } from '@/lib/native-chat-telemetry' import { + getNativeChatSkillDiscoverySubscriptionKey, resolveNativeChatSkillDiscoveryContext, + resolveNativeChatSkillDiscoverySubscriptionKey, selectNativeChatSkillStateInputs, type NativeChatSkillDiscoveryContext } from './native-chat-skill-discovery-context' import type { NativeChatSkillDiscoveryErrorKind } from './native-chat-picker-items' +import { readPaneDiscoveryCache, writePaneDiscoveryCache } from './native-chat-pane-discovery-cache' export { resolveNativeChatSkillDiscoveryContext, @@ -90,11 +92,18 @@ export function useNativeChatSkills( terminalTabId: string, enabled = false ): NativeChatSkillDiscovery { - const inputs = useAppStore(useShallow(selectNativeChatSkillStateInputs)) - const context = useMemo( - () => resolveNativeChatSkillDiscoveryContext(inputs, terminalTabId), - [inputs, terminalTabId] + const contextSubscriptionKey = useAppStore((state) => + resolveNativeChatSkillDiscoverySubscriptionKey(state, terminalTabId) ) + const context = useMemo(() => { + const currentContext = resolveNativeChatSkillDiscoveryContext( + selectNativeChatSkillStateInputs(useAppStore.getState()), + terminalTabId + ) + return getNativeChatSkillDiscoverySubscriptionKey(currentContext) === contextSubscriptionKey + ? currentContext + : null + }, [contextSubscriptionKey, terminalTabId]) const [state, setState] = useState(IDLE_STATE) const [retryGeneration, setRetryGeneration] = useState(0) const paneDiscoveryCache = useRef(new Map()) @@ -123,7 +132,7 @@ export function useNativeChatSkills( } const paneCacheKey = context.key - const cached = paneDiscoveryCache.current.get(paneCacheKey) + const cached = readPaneDiscoveryCache(paneDiscoveryCache.current, paneCacheKey) if (cached) { emitNativeChatSkillDiscovery({ agent, @@ -137,7 +146,7 @@ export function useNativeChatSkills( const request = getOrStartDiscovery(context) void request.then( (result) => { - paneDiscoveryCache.current.set(paneCacheKey, result) + writePaneDiscoveryCache(paneDiscoveryCache.current, paneCacheKey, result) if (cancelled) { return } diff --git a/src/shared/skills.ts b/src/shared/skills.ts index fb60c53e60c..c76b4728108 100644 --- a/src/shared/skills.ts +++ b/src/shared/skills.ts @@ -127,6 +127,15 @@ export type SkillDiscoveryForPaneResponse = const SKILL_PROVIDER_VALUES = ['codex', 'claude', 'agent-skills'] as const const SKILL_SOURCE_KIND_VALUES = ['home', 'repo', 'bundled', 'plugin'] as const +export const SKILL_DISCOVERY_LIMITS = { + descriptionLength: 8192, + nameLength: 512, + pathLength: 4096, + rootPaths: 64, + skills: 5000, + sources: 1000 +} as const + const boundedString = (max: number): z.ZodString => z.string().max(max) /** Validates untrusted skill metadata at the SSH relay boundary before it can @@ -136,27 +145,30 @@ export const SkillDiscoveryResultSchema = z.object({ .array( z.object({ id: boundedString(512), - name: boundedString(512), - description: boundedString(8192).nullable(), + name: boundedString(SKILL_DISCOVERY_LIMITS.nameLength), + description: boundedString(SKILL_DISCOVERY_LIMITS.descriptionLength).nullable(), providers: z.array(z.enum(SKILL_PROVIDER_VALUES)).max(8), sourceKind: z.enum(SKILL_SOURCE_KIND_VALUES), sourceLabel: boundedString(1024), - rootPath: boundedString(4096), - rootPaths: z.array(boundedString(4096)).max(64).optional(), - directoryPath: boundedString(4096), - skillFilePath: boundedString(4096), + rootPath: boundedString(SKILL_DISCOVERY_LIMITS.pathLength), + rootPaths: z + .array(boundedString(SKILL_DISCOVERY_LIMITS.pathLength)) + .max(SKILL_DISCOVERY_LIMITS.rootPaths) + .optional(), + directoryPath: boundedString(SKILL_DISCOVERY_LIMITS.pathLength), + skillFilePath: boundedString(SKILL_DISCOVERY_LIMITS.pathLength), installed: z.boolean(), fileCount: z.number().int().min(0).max(1_000_000), updatedAt: z.number().finite().nullable() }) ) - .max(5000), + .max(SKILL_DISCOVERY_LIMITS.skills), sources: z .array( z.object({ id: boundedString(512), label: boundedString(1024), - path: boundedString(4096), + path: boundedString(SKILL_DISCOVERY_LIMITS.pathLength), sourceKind: z.enum(SKILL_SOURCE_KIND_VALUES), providers: z.array(z.enum(SKILL_PROVIDER_VALUES)).max(8), owner: boundedString(64).nullable(), @@ -164,7 +176,7 @@ export const SkillDiscoveryResultSchema = z.object({ skippedReason: z.enum(['missing', 'remote-repo']).optional() }) ) - .max(1000), + .max(SKILL_DISCOVERY_LIMITS.sources), scannedAt: z.number().finite() })