diff --git a/ARCHITECTURE.md b/ARCHITECTURE.md index 63c4f58..80d6176 100644 --- a/ARCHITECTURE.md +++ b/ARCHITECTURE.md @@ -276,6 +276,7 @@ Tooling is deliberately split into two levels: - **dynamic runtime tools** reported by the Copilot CLI, with a fallback catalog for startup/offline cases - **global definitions** for MCP servers and LSP profiles - **MCP tool discovery** — when MCP server configs declare wildcard tools (empty `tools` array), the main process probes each server directly via the MCP protocol `tools/list` method to discover available tools, using the same auth credentials Aryx manages for OAuth-protected servers +- **incremental probe progress** — MCP probing runs concurrently and publishes per-server progress through the pushed workspace snapshot, using the runtime-only `mcpProbingServerIds` field so the renderer can reflect in-flight discovery without persisting transient UI state - **pattern defaults** where tool-call approval is enabled by default, plus which known runtime tools can bypass manual approval - **per-session overrides** for both tool enablement and tool auto-approval diff --git a/src/main/AryxAppService.ts b/src/main/AryxAppService.ts index 7bbb30b..38444cc 100644 --- a/src/main/AryxAppService.ts +++ b/src/main/AryxAppService.ts @@ -181,12 +181,14 @@ export class AryxAppService extends EventEmitter { private readonly gitService = new GitService(); private readonly configScanner = new ConfigScannerRegistry(); private readonly customizationScanner = new ProjectCustomizationScanner(); + private readonly probeMcpServers = probeServers; private readonly pendingApprovalHandles = new Map(); private readonly pendingUserInputHandles = new Map(); private workspace?: WorkspaceState; private sidecarCapabilities?: SidecarCapabilities; private sidecarCapabilitiesPromise?: Promise; private didScheduleInitialProjectGitRefresh = false; + private mcpProbeUpdateQueue = Promise.resolve(); async describeSidecarCapabilities(): Promise { return this.loadSidecarCapabilities(); @@ -346,7 +348,7 @@ export class AryxAppService extends EventEmitter { const result = await this.persistAndBroadcast(workspace); if (resolution === 'accept') { - void this.probeDiscoveredMcpServers(workspace.settings.discoveredUserTooling, serverIds).catch((error) => { + void this.probeDiscoveredMcpServers(workspace, workspace.settings.discoveredUserTooling, serverIds).catch((error) => { console.error('[aryx mcp-probe]', error); }); } @@ -362,7 +364,7 @@ export class AryxAppService extends EventEmitter { await this.pruneUnavailableApprovalTools(workspace); const result = await this.persistAndBroadcast(workspace); - void this.probeDiscoveredMcpServersFromState(project.discoveredTooling).catch((error) => { + void this.probeDiscoveredMcpServersFromState(workspace, project.discoveredTooling).catch((error) => { console.error('[aryx mcp-probe]', error); }); @@ -409,7 +411,7 @@ export class AryxAppService extends EventEmitter { const result = await this.persistAndBroadcast(workspace); if (resolution === 'accept') { - void this.probeDiscoveredMcpServers(project.discoveredTooling, serverIds).catch((error) => { + void this.probeDiscoveredMcpServers(workspace, project.discoveredTooling, serverIds).catch((error) => { console.error('[aryx mcp-probe]', error); }); } @@ -2138,86 +2140,73 @@ export class AryxAppService extends EventEmitter { } private async probeAllAcceptedMcpServers(workspace: WorkspaceState): Promise { - // Probe discovered MCP servers (from config files) - await this.probeDiscoveredMcpServersFromState(workspace.settings.discoveredUserTooling); - for (const project of workspace.projects) { - await this.probeDiscoveredMcpServersFromState(project.discoveredTooling); - } + const targets = [ + ...this.listAcceptedDiscoveredServerDefinitions( + workspace, + (server) => !server.probedTools || server.probedTools.length === 0, + ), + ...workspace.settings.tooling.mcpServers.filter( + (server) => server.tools.length === 0 && (!server.probedTools || server.probedTools.length === 0), + ), + ]; - // Probe manually configured MCP servers that have empty tools arrays - await this.probeManualMcpServers(workspace); + await this.probeWorkspaceMcpServers(workspace, targets); } private async probeDiscoveredMcpServersFromState( + workspace: WorkspaceState, state?: DiscoveredToolingState, ): Promise { - const accepted = listAcceptedDiscoveredMcpServers(state); - const serversNeedingProbe = accepted.filter( - (server) => !server.probedTools || server.probedTools.length === 0, - ); - if (serversNeedingProbe.length === 0) return; - - await this.probeAndApplyResults(state, serversNeedingProbe); + const targets = listAcceptedDiscoveredMcpServers(state) + .filter((server) => !server.probedTools || server.probedTools.length === 0) + .map((server) => this.discoveredServerToDefinition(server)); + await this.probeWorkspaceMcpServers(workspace, targets); } private async probeDiscoveredMcpServers( + workspace: WorkspaceState, state: DiscoveredToolingState | undefined, serverIds: ReadonlyArray, ): Promise { - const accepted = listAcceptedDiscoveredMcpServers(state); - const targets = accepted.filter((server) => serverIds.includes(server.id)); - if (targets.length === 0) return; - - await this.probeAndApplyResults(state, targets); + const targets = listAcceptedDiscoveredMcpServers(state) + .filter((server) => serverIds.includes(server.id)) + .map((server) => this.discoveredServerToDefinition(server)); + await this.probeWorkspaceMcpServers(workspace, targets); } - private async probeAndApplyResults( - state: DiscoveredToolingState | undefined, - targets: ReadonlyArray, + private async probeWorkspaceMcpServers( + workspace: WorkspaceState, + targets: ReadonlyArray, ): Promise { + const uniqueTargets = [...new Map(targets.map((server) => [server.id, server])).values()]; + if (uniqueTargets.length === 0) { + return; + } + + const targetIds = uniqueTargets.map((server) => server.id); const tokenLookup = (url: string) => getStoredToken(url)?.accessToken; - const serverDefs = targets.map((server) => this.discoveredServerToDefinition(server)); - const results = await probeServers(serverDefs, tokenLookup); - - const resultsById = new Map( - results.map((r) => [r.serverId, r]), - ); - - let changed = false; - for (const server of state?.mcpServers ?? []) { - const result = resultsById.get(server.id); - if (result?.status === 'success' && result.tools.length > 0) { - server.probedTools = result.tools; - changed = true; + await this.enqueueMcpProbeUpdate(async () => { + if (this.addMcpProbingServerIds(workspace, targetIds)) { + await this.persistAndBroadcast(workspace); } - } + }); - if (changed && this.workspace) { - await this.persistAndBroadcast(this.workspace); - } - } - - private async probeManualMcpServers(workspace: WorkspaceState): Promise { - const manualServers = workspace.settings.tooling.mcpServers.filter( - (server) => server.tools.length === 0 && (!server.probedTools || server.probedTools.length === 0), - ); - if (manualServers.length === 0) return; - - const tokenLookup = (url: string) => getStoredToken(url)?.accessToken; - const results = await probeServers(manualServers, tokenLookup); - - let changed = false; - for (const result of results) { - if (result.status !== 'success' || result.tools.length === 0) continue; - const server = workspace.settings.tooling.mcpServers.find((s) => s.id === result.serverId); - if (server) { - server.probedTools = result.tools; - changed = true; - } - } - - if (changed) { - await this.persistAndBroadcast(workspace); + try { + await this.probeMcpServers(uniqueTargets, tokenLookup, (result) => + this.enqueueMcpProbeUpdate(async () => { + const didUpdateProbing = this.removeMcpProbingServerIds(workspace, [result.serverId]); + const didApplyResult = this.applyMcpProbeResult(workspace, result); + if (didUpdateProbing || didApplyResult) { + await this.persistAndBroadcast(workspace); + } + }), + ); + } finally { + await this.enqueueMcpProbeUpdate(async () => { + if (this.removeMcpProbingServerIds(workspace, targetIds)) { + await this.persistAndBroadcast(workspace); + } + }); } } @@ -2227,7 +2216,6 @@ export class AryxAppService extends EventEmitter { */ private async reprobeServerByUrl(serverUrl: string): Promise { const workspace = await this.loadWorkspace(); - const tokenLookup = (url: string) => getStoredToken(url)?.accessToken; // Collect matching servers from manual config and discovered tooling const targets: McpServerDefinition[] = []; @@ -2248,36 +2236,103 @@ export class AryxAppService extends EventEmitter { } } - if (targets.length === 0) return; + await this.probeWorkspaceMcpServers(workspace, targets); + } - const results = await probeServers(targets, tokenLookup); - const resultsById = new Map(results.map((r) => [r.serverId, r])); + private listAcceptedDiscoveredServerDefinitions( + workspace: WorkspaceState, + predicate?: (server: DiscoveredMcpServer) => boolean, + ): McpServerDefinition[] { + const definitions: McpServerDefinition[] = []; + + for (const state of this.listDiscoveredToolingStates(workspace)) { + for (const server of listAcceptedDiscoveredMcpServers(state)) { + if (predicate && !predicate(server)) { + continue; + } + definitions.push(this.discoveredServerToDefinition(server)); + } + } + + return definitions; + } + + private listDiscoveredToolingStates(workspace: WorkspaceState): Array { + return [ + workspace.settings.discoveredUserTooling, + ...workspace.projects.map((project) => project.discoveredTooling), + ]; + } + + private addMcpProbingServerIds(workspace: WorkspaceState, serverIds: ReadonlyArray): boolean { + return this.updateMcpProbingServerIds(workspace, serverIds, 'add'); + } + + private removeMcpProbingServerIds(workspace: WorkspaceState, serverIds: ReadonlyArray): boolean { + return this.updateMcpProbingServerIds(workspace, serverIds, 'remove'); + } + + private updateMcpProbingServerIds( + workspace: WorkspaceState, + serverIds: ReadonlyArray, + operation: 'add' | 'remove', + ): boolean { + const next = new Set(workspace.mcpProbingServerIds ?? []); + const before = next.size; + + for (const serverId of serverIds) { + if (operation === 'add') { + next.add(serverId); + } else { + next.delete(serverId); + } + } + + if (next.size === before) { + return false; + } + + if (next.size === 0) { + delete workspace.mcpProbingServerIds; + } else { + workspace.mcpProbingServerIds = [...next]; + } + + return true; + } + + private applyMcpProbeResult(workspace: WorkspaceState, result: McpProbeResult): boolean { + if (result.status !== 'success' || result.tools.length === 0) { + return false; + } let changed = false; - // Apply to manual servers for (const server of workspace.settings.tooling.mcpServers) { - const result = resultsById.get(server.id); - if (result?.status === 'success' && result.tools.length > 0) { + if (server.id !== result.serverId) { + continue; + } + server.probedTools = result.tools; + changed = true; + } + + for (const state of this.listDiscoveredToolingStates(workspace)) { + for (const server of state?.mcpServers ?? []) { + if (server.id !== result.serverId) { + continue; + } server.probedTools = result.tools; changed = true; } } - // Apply to discovered servers - for (const state of [workspace.settings.discoveredUserTooling, ...workspace.projects.map((p) => p.discoveredTooling)]) { - for (const server of state?.mcpServers ?? []) { - const result = resultsById.get(server.id); - if (result?.status === 'success' && result.tools.length > 0) { - server.probedTools = result.tools; - changed = true; - } - } - } + return changed; + } - if (changed) { - await this.persistAndBroadcast(workspace); - } + private async enqueueMcpProbeUpdate(update: () => Promise): Promise { + const next = this.mcpProbeUpdateQueue.then(update, update); + this.mcpProbeUpdateQueue = next.catch(() => undefined); + await next; } private discoveredServerToDefinition(server: DiscoveredMcpServer): McpServerDefinition { diff --git a/src/main/persistence/workspaceRepository.ts b/src/main/persistence/workspaceRepository.ts index 82af88e..c20c627 100644 --- a/src/main/persistence/workspaceRepository.ts +++ b/src/main/persistence/workspaceRepository.ts @@ -123,8 +123,9 @@ export class WorkspaceRepository { } async save(workspace: WorkspaceState): Promise { + const { mcpProbingServerIds: _mcpProbingServerIds, ...persistedWorkspace } = workspace; await writeJsonFile(this.filePath, { - ...workspace, + ...persistedWorkspace, lastUpdatedAt: nowIso(), }); } diff --git a/src/main/services/mcpToolProber.ts b/src/main/services/mcpToolProber.ts index 1029fca..a075f45 100644 --- a/src/main/services/mcpToolProber.ts +++ b/src/main/services/mcpToolProber.ts @@ -25,30 +25,32 @@ const MAX_CONCURRENCY = 5; export async function probeServers( servers: ReadonlyArray, tokenLookup?: (serverUrl: string) => string | undefined, + onResult?: (result: McpProbeResult) => void | Promise, ): Promise { - const pending = [...servers]; - const results: McpProbeResult[] = []; - const active: Promise[] = []; + if (servers.length === 0) { + return []; + } - for (const server of pending) { - const task = probeServer(server, tokenLookup).then((result) => { - results.push(result); - }); - active.push(task); + const results = new Array(servers.length); + let nextIndex = 0; + const workerCount = Math.min(MAX_CONCURRENCY, servers.length); - if (active.length >= MAX_CONCURRENCY) { - await Promise.race(active); - // Remove settled promises - for (let i = active.length - 1; i >= 0; i--) { - const status = await Promise.race([active[i].then(() => 'done'), Promise.resolve('pending')]); - if (status === 'done') { - active.splice(i, 1); - } + async function worker(): Promise { + while (true) { + const index = nextIndex; + nextIndex += 1; + const server = servers[index]; + if (!server) { + return; } + + const result = await probeServer(server, tokenLookup); + results[index] = result; + await onResult?.(result); } } - await Promise.allSettled(active); + await Promise.all(Array.from({ length: workerCount }, () => worker())); return results; } diff --git a/src/shared/domain/workspace.ts b/src/shared/domain/workspace.ts index fb985df..4cf5323 100644 --- a/src/shared/domain/workspace.ts +++ b/src/shared/domain/workspace.ts @@ -10,6 +10,8 @@ export interface WorkspaceState { patterns: PatternDefinition[]; sessions: SessionRecord[]; settings: WorkspaceSettings; + /** Runtime-only MCP probe progress for live UI updates. */ + mcpProbingServerIds?: string[]; selectedProjectId?: string; selectedPatternId?: string; selectedSessionId?: string; diff --git a/tests/main/appServiceMcpProbing.test.ts b/tests/main/appServiceMcpProbing.test.ts new file mode 100644 index 0000000..b3415d5 --- /dev/null +++ b/tests/main/appServiceMcpProbing.test.ts @@ -0,0 +1,278 @@ +import { beforeEach, describe, expect, mock, test } from 'bun:test'; + +import type { ProjectRecord } from '@shared/domain/project'; +import { createWorkspaceSeed, type WorkspaceState } from '@shared/domain/workspace'; + +const TIMESTAMP = '2026-03-28T00:00:00.000Z'; + +type MockProbeResult = { + serverId: string; + serverName: string; + tools: Array<{ name: string; description?: string }>; + status: 'success' | 'failed'; + error?: string; +}; + +const probeCalls: string[][] = []; +let probeResults: MockProbeResult[] = []; + +mock.module('electron', () => { + const electronMock = { + app: { + isPackaged: false, + getAppPath: () => 'C:\\workspace\\personal\\repositories\\aryx', + getPath: () => 'C:\\workspace\\personal\\repositories\\aryx\\tests\\fixtures', + }, + dialog: { + showOpenDialog: async () => ({ canceled: true, filePaths: [] }), + }, + shell: { + openPath: async () => '', + }, + }; + + return { + ...electronMock, + default: electronMock, + }; +}); + +mock.module('keytar', () => ({ + default: { + getPassword: async () => null, + setPassword: async () => undefined, + deletePassword: async () => false, + }, +})); + +const { AryxAppService } = await import('@main/AryxAppService'); + +beforeEach(() => { + probeCalls.length = 0; + probeResults = []; +}); + +function createProject(overrides?: Partial): ProjectRecord { + return { + id: 'project-alpha', + name: 'alpha', + path: 'C:\\workspace\\alpha', + addedAt: TIMESTAMP, + ...overrides, + }; +} + +function cloneWorkspaceState(workspace: WorkspaceState): WorkspaceState { + return JSON.parse(JSON.stringify(workspace)) as WorkspaceState; +} + +function createService(workspace: WorkspaceState): { + service: InstanceType; + snapshots: WorkspaceState[]; +} { + const service = new AryxAppService(); + const internals = service as unknown as Record; + const snapshots: WorkspaceState[] = []; + + internals.loadWorkspace = async () => workspace; + internals.persistAndBroadcast = async (nextWorkspace: WorkspaceState) => { + snapshots.push(cloneWorkspaceState(nextWorkspace)); + return nextWorkspace; + }; + internals.probeMcpServers = async ( + servers: Array<{ id: string }>, + _tokenLookup?: (serverUrl: string) => string | undefined, + onResult?: (result: MockProbeResult) => void | Promise, + ) => { + probeCalls.push(servers.map((server) => server.id)); + for (const result of probeResults) { + await onResult?.(result); + } + return probeResults; + }; + + return { service, snapshots }; +} + +describe('AryxAppService MCP probing', () => { + test('probes accepted discovered and manual MCP servers in one batch with incremental workspace updates', async () => { + const workspace = createWorkspaceSeed(); + const project = createProject({ + discoveredTooling: { + mcpServers: [ + { + id: 'project-server', + name: 'Project MCP', + transport: 'local', + command: 'project-mcp', + args: [], + tools: [], + scope: 'project', + scannerId: 'vscode-mcp', + sourcePath: 'C:\\workspace\\alpha\\.vscode\\mcp.json', + sourceLabel: '.vscode\\mcp.json', + fingerprint: 'project-fingerprint', + status: 'accepted', + }, + ], + lastScannedAt: TIMESTAMP, + }, + }); + + workspace.projects = [project]; + workspace.settings.discoveredUserTooling = { + mcpServers: [ + { + id: 'user-server', + name: 'User MCP', + transport: 'local', + command: 'user-mcp', + args: [], + tools: [], + scope: 'user', + scannerId: 'copilot-user-mcp', + sourcePath: 'C:\\Users\\tester\\.copilot\\mcp.json', + sourceLabel: '~\\.copilot\\mcp.json', + fingerprint: 'user-fingerprint', + status: 'accepted', + }, + ], + lastScannedAt: TIMESTAMP, + }; + workspace.settings.tooling.mcpServers = [ + { + id: 'manual-server', + name: 'Manual MCP', + transport: 'local', + command: 'manual-mcp', + args: [], + tools: [], + createdAt: TIMESTAMP, + updatedAt: TIMESTAMP, + }, + ]; + + probeResults = [ + { + serverId: 'project-server', + serverName: 'Project MCP', + status: 'success', + tools: [{ name: 'project.status' }], + }, + { + serverId: 'manual-server', + serverName: 'Manual MCP', + status: 'success', + tools: [{ name: 'manual.status' }], + }, + { + serverId: 'user-server', + serverName: 'User MCP', + status: 'failed', + tools: [], + error: 'boom', + }, + ]; + + const { service, snapshots } = createService(workspace); + + await ( + service as unknown as { + probeAllAcceptedMcpServers: (nextWorkspace: WorkspaceState) => Promise; + } + ).probeAllAcceptedMcpServers(workspace); + + expect(probeCalls).toEqual([[ + 'user-server', + 'project-server', + 'manual-server', + ]]); + expect(snapshots.map((snapshot) => snapshot.mcpProbingServerIds)).toEqual([ + ['user-server', 'project-server', 'manual-server'], + ['user-server', 'manual-server'], + ['user-server'], + undefined, + ]); + expect(workspace.projects[0]?.discoveredTooling?.mcpServers[0]?.probedTools).toEqual([ + { name: 'project.status' }, + ]); + expect(workspace.settings.tooling.mcpServers[0]?.probedTools).toEqual([ + { name: 'manual.status' }, + ]); + expect(workspace.settings.discoveredUserTooling.mcpServers[0]?.probedTools).toBeUndefined(); + }); + + test('tracks probing progress while re-probing matching remote MCP servers after OAuth', async () => { + const workspace = createWorkspaceSeed(); + workspace.settings.discoveredUserTooling = { + mcpServers: [ + { + id: 'discovered-remote', + name: 'Discovered Remote', + transport: 'http', + url: 'https://example.com/mcp', + headers: { 'X-Test': '1' }, + tools: [], + scope: 'user', + scannerId: 'copilot-user-mcp', + sourcePath: 'C:\\Users\\tester\\.copilot\\mcp.json', + sourceLabel: '~\\.copilot\\mcp.json', + fingerprint: 'remote-fingerprint', + status: 'accepted', + }, + ], + lastScannedAt: TIMESTAMP, + }; + workspace.settings.tooling.mcpServers = [ + { + id: 'manual-remote', + name: 'Manual Remote', + transport: 'http', + url: 'https://example.com/mcp', + headers: { Authorization: 'Bearer token' }, + tools: [], + createdAt: TIMESTAMP, + updatedAt: TIMESTAMP, + }, + ]; + + probeResults = [ + { + serverId: 'manual-remote', + serverName: 'Manual Remote', + status: 'success', + tools: [{ name: 'manual.remote' }], + }, + { + serverId: 'discovered-remote', + serverName: 'Discovered Remote', + status: 'success', + tools: [{ name: 'discovered.remote' }], + }, + ]; + + const { service, snapshots } = createService(workspace); + + await ( + service as unknown as { + reprobeServerByUrl: (serverUrl: string) => Promise; + } + ).reprobeServerByUrl('https://example.com/mcp'); + + expect(probeCalls).toEqual([[ + 'manual-remote', + 'discovered-remote', + ]]); + expect(snapshots.map((snapshot) => snapshot.mcpProbingServerIds)).toEqual([ + ['manual-remote', 'discovered-remote'], + ['discovered-remote'], + undefined, + ]); + expect(workspace.settings.tooling.mcpServers[0]?.probedTools).toEqual([ + { name: 'manual.remote' }, + ]); + expect(workspace.settings.discoveredUserTooling.mcpServers[0]?.probedTools).toEqual([ + { name: 'discovered.remote' }, + ]); + }); +}); diff --git a/tests/main/mcpToolProber.test.ts b/tests/main/mcpToolProber.test.ts new file mode 100644 index 0000000..fbb26ba --- /dev/null +++ b/tests/main/mcpToolProber.test.ts @@ -0,0 +1,133 @@ +import { beforeEach, describe, expect, mock, test } from 'bun:test'; + +const commandDelaysMs = new Map(); +let activeListTools = 0; +let maxActiveListTools = 0; + +class FakeStdioClientTransport { + readonly command: string; + readonly args?: string[]; + readonly env?: Record; + readonly cwd?: string; + readonly stderr?: 'ignore'; + onmessage?: (message: unknown) => void; + + constructor(options: { + command: string; + args?: string[]; + env?: Record; + cwd?: string; + stderr?: 'ignore'; + }) { + this.command = options.command; + this.args = options.args; + this.env = options.env; + this.cwd = options.cwd; + this.stderr = options.stderr; + } + + async send(_message: unknown): Promise {} +} + +class FakeSSEClientTransport { + onmessage?: (message: unknown) => void; + + constructor(_url: URL, _options?: unknown) {} + + async send(_message: unknown): Promise {} +} + +class FakeStreamableHTTPClientTransport { + onmessage?: (message: unknown) => void; + + constructor(_url: URL, _options?: unknown) {} + + async send(_message: unknown): Promise {} +} + +class FakeClient { + private transport?: FakeStdioClientTransport; + + constructor(_clientInfo: unknown, _options: unknown) {} + + async connect(transport: FakeStdioClientTransport): Promise { + this.transport = transport; + } + + async listTools(): Promise<{ tools: Array<{ name: string }> }> { + const command = this.transport?.command; + if (!command) { + throw new Error('Expected a command for the fake MCP transport.'); + } + + activeListTools += 1; + maxActiveListTools = Math.max(maxActiveListTools, activeListTools); + await new Promise((resolve) => setTimeout(resolve, commandDelaysMs.get(command) ?? 0)); + activeListTools -= 1; + + return { + tools: [{ name: `${command}.tool` }], + }; + } + + async close(): Promise {} +} + +mock.module('@modelcontextprotocol/sdk/client/index.js', () => ({ + Client: FakeClient, +})); + +mock.module('@modelcontextprotocol/sdk/client/stdio.js', () => ({ + StdioClientTransport: FakeStdioClientTransport, +})); + +mock.module('@modelcontextprotocol/sdk/client/sse.js', () => ({ + SSEClientTransport: FakeSSEClientTransport, +})); + +mock.module('@modelcontextprotocol/sdk/client/streamableHttp.js', () => ({ + StreamableHTTPClientTransport: FakeStreamableHTTPClientTransport, +})); + +const { probeServers } = await import('../../src/main/services/mcpToolProber'); + +const TIMESTAMP = '2026-03-28T00:00:00.000Z'; + +beforeEach(() => { + commandDelaysMs.clear(); + activeListTools = 0; + maxActiveListTools = 0; +}); + +describe('probeServers', () => { + test('fires onResult as probes finish while preserving input order and concurrency limit', async () => { + const servers = Array.from({ length: 7 }, (_, index) => ({ + id: `server-${index}`, + name: `Server ${index}`, + transport: 'local' as const, + command: `cmd-${index}`, + args: [] as string[], + tools: [] as string[], + createdAt: TIMESTAMP, + updatedAt: TIMESTAMP, + })); + + const delays = [70, 60, 50, 40, 30, 20, 10]; + for (const [index, server] of servers.entries()) { + commandDelaysMs.set(server.command, delays[index] ?? 0); + } + + const callbackOrder: string[] = []; + const results = await probeServers(servers, undefined, (result) => { + callbackOrder.push(result.serverId); + }); + + expect(results.map((result) => result.serverId)).toEqual(servers.map((server) => server.id)); + expect(results.map((result) => result.tools[0]?.name)).toEqual( + servers.map((server) => `${server.command}.tool`), + ); + expect(callbackOrder).not.toEqual(servers.map((server) => server.id)); + expect([...callbackOrder].sort()).toEqual(servers.map((server) => server.id).sort()); + expect(maxActiveListTools).toBe(5); + }); +}); diff --git a/tests/main/workspaceRepositoryScratchpad.test.ts b/tests/main/workspaceRepositoryScratchpad.test.ts index 20757a5..3212b29 100644 --- a/tests/main/workspaceRepositoryScratchpad.test.ts +++ b/tests/main/workspaceRepositoryScratchpad.test.ts @@ -87,4 +87,18 @@ describe('WorkspaceRepository scratchpad migration', () => { const persisted = JSON.parse(await readFile(workspaceFilePath, 'utf8')) as WorkspaceState; expect(persisted.sessions[0]?.cwd).toBe(expectedSessionPath); }); + + test('strips runtime MCP probing state before persisting workspace data', async () => { + const workspaceFilePath = join(USER_DATA_PATH, 'workspace.json'); + await mkdir(USER_DATA_PATH, { recursive: true }); + + const repository = new WorkspaceRepository(); + const workspace = createStoredWorkspace(); + workspace.mcpProbingServerIds = ['server-a', 'server-b']; + + await repository.save(workspace); + + const persisted = JSON.parse(await readFile(workspaceFilePath, 'utf8')) as WorkspaceState; + expect('mcpProbingServerIds' in persisted).toBe(false); + }); });