From 9787108532c89995daccc0c558de8706c794902c Mon Sep 17 00:00:00 2001 From: Tommaso Sciortino Date: Wed, 5 Nov 2025 15:38:44 -0800 Subject: [PATCH] List tools in a consistent order. (#12615) --- packages/core/src/agents/executor.ts | 1 + packages/core/src/config/config.test.ts | 1 + packages/core/src/config/config.ts | 1 + packages/core/src/tools/mcp-client.test.ts | 5 +++ packages/core/src/tools/mcp-client.ts | 2 + packages/core/src/tools/tool-registry.test.ts | 45 +++++++++++++++++++ packages/core/src/tools/tool-registry.ts | 37 +++++++++++++++ 7 files changed, 92 insertions(+) diff --git a/packages/core/src/agents/executor.ts b/packages/core/src/agents/executor.ts index ac0b9a27af..5982817660 100644 --- a/packages/core/src/agents/executor.ts +++ b/packages/core/src/agents/executor.ts @@ -128,6 +128,7 @@ export class AgentExecutor { // registered; their schemas are passed directly to the model later. } + agentToolRegistry.sortTools(); // Validate that all registered tools are safe for non-interactive // execution. await AgentExecutor.validateTools(agentToolRegistry, definition.name); diff --git a/packages/core/src/config/config.test.ts b/packages/core/src/config/config.test.ts index a9c3b04620..23cea2d1b8 100644 --- a/packages/core/src/config/config.test.ts +++ b/packages/core/src/config/config.test.ts @@ -49,6 +49,7 @@ vi.mock('../tools/tool-registry', () => { const ToolRegistryMock = vi.fn(); ToolRegistryMock.prototype.registerTool = vi.fn(); ToolRegistryMock.prototype.discoverAllTools = vi.fn(); + ToolRegistryMock.prototype.sortTools = vi.fn(); ToolRegistryMock.prototype.getAllTools = vi.fn(() => []); // Mock methods if needed ToolRegistryMock.prototype.getTool = vi.fn(); ToolRegistryMock.prototype.getFunctionDeclarations = vi.fn(() => []); diff --git a/packages/core/src/config/config.ts b/packages/core/src/config/config.ts index ef3598824b..676b4e32fa 100644 --- a/packages/core/src/config/config.ts +++ b/packages/core/src/config/config.ts @@ -1346,6 +1346,7 @@ export class Config { } await registry.discoverAllTools(); + registry.sortTools(); return registry; } diff --git a/packages/core/src/tools/mcp-client.test.ts b/packages/core/src/tools/mcp-client.test.ts index 7811888eec..14c134be0a 100644 --- a/packages/core/src/tools/mcp-client.test.ts +++ b/packages/core/src/tools/mcp-client.test.ts @@ -89,6 +89,7 @@ describe('mcp-client', () => { } as unknown as GenAiLib.CallableTool); const mockedToolRegistry = { registerTool: vi.fn(), + sortTools: vi.fn(), getMessageBus: vi.fn().mockReturnValue(undefined), } as unknown as ToolRegistry; const client = new McpClient( @@ -153,6 +154,7 @@ describe('mcp-client', () => { } as unknown as GenAiLib.CallableTool); const mockedToolRegistry = { registerTool: vi.fn(), + sortTools: vi.fn(), getMessageBus: vi.fn().mockReturnValue(undefined), } as unknown as ToolRegistry; const client = new McpClient( @@ -237,6 +239,7 @@ describe('mcp-client', () => { const mockedMcpToTool = vi.mocked(GenAiLib.mcpToTool); const mockedToolRegistry = { registerTool: vi.fn(), + sortTools: vi.fn(), getMessageBus: vi.fn().mockReturnValue(undefined), } as unknown as ToolRegistry; const client = new McpClient( @@ -286,6 +289,7 @@ describe('mcp-client', () => { } as unknown as GenAiLib.CallableTool); const mockedToolRegistry = { registerTool: vi.fn(), + sortTools: vi.fn(), getMessageBus: vi.fn().mockReturnValue(undefined), } as unknown as ToolRegistry; const client = new McpClient( @@ -340,6 +344,7 @@ describe('mcp-client', () => { unregisterTool: vi.fn(), getMessageBus: vi.fn().mockReturnValue(undefined), removeMcpToolsByServer: vi.fn(), + sortTools: vi.fn(), } as unknown as ToolRegistry; const mockedPromptRegistry = { registerPrompt: vi.fn(), diff --git a/packages/core/src/tools/mcp-client.ts b/packages/core/src/tools/mcp-client.ts index c5b1dc6caa..45e481390e 100644 --- a/packages/core/src/tools/mcp-client.ts +++ b/packages/core/src/tools/mcp-client.ts @@ -151,6 +151,7 @@ export class McpClient { for (const tool of tools) { this.toolRegistry.registerTool(tool); } + this.toolRegistry.sortTools(); } /** @@ -568,6 +569,7 @@ export async function connectAndDiscover( for (const tool of tools) { toolRegistry.registerTool(tool); } + toolRegistry.sortTools(); } catch (error) { if (mcpClient) { mcpClient.close(); diff --git a/packages/core/src/tools/tool-registry.test.ts b/packages/core/src/tools/tool-registry.test.ts index f002250910..80e9390cce 100644 --- a/packages/core/src/tools/tool-registry.test.ts +++ b/packages/core/src/tools/tool-registry.test.ts @@ -250,6 +250,51 @@ describe('ToolRegistry', () => { }); }); + describe('sortTools', () => { + it('should sort tools by priority: built-in, discovered, then MCP (by server name)', () => { + const builtIn1 = new MockTool({ name: 'builtin-1' }); + const builtIn2 = new MockTool({ name: 'builtin-2' }); + const discovered1 = new DiscoveredTool( + config, + 'discovered-1', + 'desc', + {}, + ); + const mockCallable = {} as CallableTool; + const mcpZebra = new DiscoveredMCPTool( + mockCallable, + 'zebra-server', + 'mcp-zebra', + 'desc', + {}, + ); + const mcpApple = new DiscoveredMCPTool( + mockCallable, + 'apple-server', + 'mcp-apple', + 'desc', + {}, + ); + + // Register in mixed order + toolRegistry.registerTool(mcpZebra); + toolRegistry.registerTool(discovered1); + toolRegistry.registerTool(builtIn1); + toolRegistry.registerTool(mcpApple); + toolRegistry.registerTool(builtIn2); + + toolRegistry.sortTools(); + + expect(toolRegistry.getAllToolNames()).toEqual([ + 'builtin-1', + 'builtin-2', + 'discovered-1', + 'mcp-apple', + 'mcp-zebra', + ]); + }); + }); + describe('discoverTools', () => { it('should will preserve tool parametersJsonSchema during discovery from command', async () => { const discoveryCommand = 'my-discovery-command'; diff --git a/packages/core/src/tools/tool-registry.ts b/packages/core/src/tools/tool-registry.ts index abb03d5329..c350abfbd2 100644 --- a/packages/core/src/tools/tool-registry.ts +++ b/packages/core/src/tools/tool-registry.ts @@ -205,6 +205,43 @@ export class ToolRegistry { this.tools.set(tool.name, tool); } + /** + * Sorts tools as: + * 1. Built in tools. + * 2. Discovered tools. + * 3. MCP tools ordered by server name. + * + * This is a stable sort in that ties preseve existing order. + */ + sortTools(): void { + const getPriority = (tool: AnyDeclarativeTool): number => { + if (tool instanceof DiscoveredMCPTool) return 2; + if (tool instanceof DiscoveredTool) return 1; + return 0; // Built-in + }; + + this.tools = new Map( + Array.from(this.tools.entries()).sort((a, b) => { + const toolA = a[1]; + const toolB = b[1]; + const priorityA = getPriority(toolA); + const priorityB = getPriority(toolB); + + if (priorityA !== priorityB) { + return priorityA - priorityB; + } + + if (priorityA === 2) { + const serverA = (toolA as DiscoveredMCPTool).serverName; + const serverB = (toolB as DiscoveredMCPTool).serverName; + return serverA.localeCompare(serverB); + } + + return 0; + }), + ); + } + private removeDiscoveredTools(): void { for (const tool of this.tools.values()) { if (tool instanceof DiscoveredTool || tool instanceof DiscoveredMCPTool) {