mirror of
https://github.com/davidkaya/aryx.git
synced 2026-07-28 23:48:39 +02:00
feat: stream MCP probe progress
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
This commit is contained in:
David Kaya
@@ -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>): 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<typeof AryxAppService>;
|
||||
snapshots: WorkspaceState[];
|
||||
} {
|
||||
const service = new AryxAppService();
|
||||
const internals = service as unknown as Record<string, unknown>;
|
||||
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<void>,
|
||||
) => {
|
||||
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<void>;
|
||||
}
|
||||
).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<void>;
|
||||
}
|
||||
).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' },
|
||||
]);
|
||||
});
|
||||
});
|
||||
@@ -0,0 +1,133 @@
|
||||
import { beforeEach, describe, expect, mock, test } from 'bun:test';
|
||||
|
||||
const commandDelaysMs = new Map<string, number>();
|
||||
let activeListTools = 0;
|
||||
let maxActiveListTools = 0;
|
||||
|
||||
class FakeStdioClientTransport {
|
||||
readonly command: string;
|
||||
readonly args?: string[];
|
||||
readonly env?: Record<string, string>;
|
||||
readonly cwd?: string;
|
||||
readonly stderr?: 'ignore';
|
||||
onmessage?: (message: unknown) => void;
|
||||
|
||||
constructor(options: {
|
||||
command: string;
|
||||
args?: string[];
|
||||
env?: Record<string, string>;
|
||||
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<void> {}
|
||||
}
|
||||
|
||||
class FakeSSEClientTransport {
|
||||
onmessage?: (message: unknown) => void;
|
||||
|
||||
constructor(_url: URL, _options?: unknown) {}
|
||||
|
||||
async send(_message: unknown): Promise<void> {}
|
||||
}
|
||||
|
||||
class FakeStreamableHTTPClientTransport {
|
||||
onmessage?: (message: unknown) => void;
|
||||
|
||||
constructor(_url: URL, _options?: unknown) {}
|
||||
|
||||
async send(_message: unknown): Promise<void> {}
|
||||
}
|
||||
|
||||
class FakeClient {
|
||||
private transport?: FakeStdioClientTransport;
|
||||
|
||||
constructor(_clientInfo: unknown, _options: unknown) {}
|
||||
|
||||
async connect(transport: FakeStdioClientTransport): Promise<void> {
|
||||
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<void> {}
|
||||
}
|
||||
|
||||
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);
|
||||
});
|
||||
});
|
||||
@@ -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);
|
||||
});
|
||||
});
|
||||
|
||||
Reference in New Issue
Block a user