abbyswag/qwen-code
1
0
Fork 0
mirror of https://github.com/QwenLM/qwen-code.git synced 2025-12-21 01:07:46 +00:00
Code Issues Packages Projects Releases Wiki Activity

refactor: Extract MCP discovery from ToolRegistry

Browse Source 
- Moves MCP tool discovery logic from ToolRegistry into a new, dedicated MCP client (mcp-client.ts and mcp-tool.ts).
- Updates ToolRegistry to utilize the new MCP client.
- Adds comprehensive tests for the new MCP client and its integration with ToolRegistry.

Part of https://github.com/google-gemini/gemini-cli/issues/577
This commit is contained in:
Taylor Mullen
2025-05-28 00:43:23 -07:00
committed by N. Taylor Mullen
parent c413988ae0
commit d74c0f581b
6 changed files with 737 additions and 162 deletions
Download Patch File Download Diff File Expand all files Collapse all files

26
packages/server/src/tools/tool-registry.test.ts
View File

@@ -14,11 +14,8 @@ import {
afterEach,
Mocked,
} from 'vitest';
import {
ToolRegistry,
DiscoveredTool,
DiscoveredMCPTool,
} from './tool-registry.js';
import { ToolRegistry, DiscoveredTool } from './tool-registry.js';
import { DiscoveredMCPTool } from './mcp-tool.js';
import { Config } from '../config/config.js';
import { BaseTool, ToolResult } from './tools.js';
import { FunctionDeclaration } from '@google/genai';
@@ -347,7 +344,7 @@ describe('ToolRegistry', () => {
toolRegistry = new ToolRegistry(config);
});
it('should discover tools using discovery command', () => {
it('should discover tools using discovery command', async () => {
const discoveryCommand = 'my-discovery-command';
mockConfigGetToolDiscoveryCommand.mockReturnValue(discoveryCommand);
const mockToolDeclarations: FunctionDeclaration[] = [
@@ -366,7 +363,7 @@ describe('ToolRegistry', () => {
),
);
toolRegistry.discoverTools();
await toolRegistry.discoverTools();
expect(execSync).toHaveBeenCalledWith(discoveryCommand);
const discoveredTool = toolRegistry.getTool('discovered-tool-1');
@@ -376,7 +373,7 @@ describe('ToolRegistry', () => {
expect(discoveredTool?.description).toContain(discoveryCommand);
});
it('should remove previously discovered tools before discovering new ones', () => {
it('should remove previously discovered tools before discovering new ones', async () => {
const discoveryCommand = 'my-discovery-command';
mockConfigGetToolDiscoveryCommand.mockReturnValue(discoveryCommand);
mockExecSync.mockReturnValueOnce(
@@ -394,7 +391,7 @@ describe('ToolRegistry', () => {
]),
),
);
toolRegistry.discoverTools();
await toolRegistry.discoverTools();
expect(toolRegistry.getTool('old-discovered-tool')).toBeInstanceOf(
DiscoveredTool,
);
@@ -414,7 +411,7 @@ describe('ToolRegistry', () => {
]),
),
);
toolRegistry.discoverTools();
await toolRegistry.discoverTools();
expect(toolRegistry.getTool('old-discovered-tool')).toBeUndefined();
expect(toolRegistry.getTool('new-discovered-tool')).toBeInstanceOf(
DiscoveredTool,
@@ -457,8 +454,7 @@ describe('ToolRegistry', () => {
});
mockMcpClientInstance.connect.mockResolvedValue(undefined);
toolRegistry.discoverTools();
await new Promise((resolve) => setTimeout(resolve, 100)); // Wait for async operations
await toolRegistry.discoverTools();
expect(Client).toHaveBeenCalledTimes(1);
expect(StdioClientTransport).toHaveBeenCalledWith({
@@ -511,8 +507,7 @@ describe('ToolRegistry', () => {
});
mockMcpClientInstance.connect.mockResolvedValue(undefined);
toolRegistry.discoverTools();
await new Promise((resolve) => setTimeout(resolve, 100));
await toolRegistry.discoverTools();
expect(Client).toHaveBeenCalledTimes(1);
expect(StdioClientTransport).toHaveBeenCalledWith({
@@ -544,8 +539,7 @@ describe('ToolRegistry', () => {
// Need to await the async IIFE within discoverTools.
// Since discoverTools itself isn't async, we can't directly await it.
// We'll check the console.error mock.
toolRegistry.discoverTools();
await new Promise((resolve) => setTimeout(resolve, 100)); // Wait for async operations
await toolRegistry.discoverTools();
expect(console.error).toHaveBeenCalledWith(
`failed to start or connect to MCP server 'failing-mcp' ${JSON.stringify({ command: 'fail-cmd' })}; \nError: Connection failed`,