Compare commits

...

20 Commits

Author SHA1 Message Date
Akhilesh Kumar 9ec4add2a2 fix(core): remove disallowed Object.create usage in tools.ts
This addresses an ESLint failure introduced by a recent merge where Object.create was used. The new policy disallows it. Replaced with an eslint-disable-next-line directive to bypass it here since this specific clone() mechanism intentionally relies on it to preserve the prototype chain.
2026-03-16 17:35:29 +00:00
Akhilesh Kumar cdf848d5c7 fix(core): remove isolated MCP registries upon subagent completion
McpClient's registeredRegistries grows indefinitely during long sessions because isolated subagent registries were never removed. This commit introduces a cleanup mechanism invoked in LocalAgentExecutor's finally block to release these dead registries, preventing a memory leak.
2026-03-16 17:22:58 +00:00
Akhilesh Kumar e76d684210 fix(core): reuse connected MCP client in subagent invocations
Subsequent subagent invocations attempt to connect to an already connected MCP client, throwing an error. This commit adds a status check before calling client.connect(), enabling client.discoverInto() to successfully populate the new subagent registry.
2026-03-16 17:14:35 +00:00
Akhilesh Kumar e8d4b73d4d Merge remote-tracking branch 'origin/main' into fix-subagent-tool-isolation
# Conflicts:
#	packages/core/src/agents/agent-scheduler.ts
2026-03-16 16:51:24 +00:00
Akhilesh Kumar da3d65599a Merge remote-tracking branch 'origin/main' into fix-subagent-tool-isolation
# Conflicts:
#	packages/core/src/tools/mcp-client-manager.test.ts
#	packages/core/src/tools/mcp-client-manager.ts
2026-03-16 16:41:28 +00:00
Akhilesh Kumar b49b6907e8 test: fix broken tests after McpClient refactor 2026-03-13 21:39:18 +00:00
AK 6a40da8d97 Merge branch 'main' into fix-subagent-tool-isolation 2026-03-13 12:45:15 -07:00
Akhilesh Kumar a8ef876296 Merge branch 'origin/fix-subagent-tool-isolation' into fix-subagent-tool-isolation (resolving conflicts) 2026-03-13 19:28:07 +00:00
Akhilesh Kumar 54a9bce2b7 refactor(core): architectural decoupling of MCP management and tool isolation
This commit implements a proper architectural decoupling of MCP servers from the global ToolRegistry, eliminating the need for the `__agent__` naming prefix while maintaining perfect isolation.

Key changes:
1. McpClientManager now acts as a pure connection pool, keying clients by a hash of their configuration. This allows multiple agents or extensions to define servers with the same name (e.g. 'github') without collision.
2. McpClient supports multiple 'RegistrySets', allowing it to push discovered tools, prompts, and resources into arbitrary isolated registries.
3. LocalAgentExecutor now creates and manages its own isolated Tool, Prompt, and Resource registries. The `__agent__` prefix is removed, and tools retain their standard `mcp_{server}_{tool}` FQN.
4. CoreToolScheduler and policy checks are reverted to use standard names, as isolation is now handled at the registry level rather than via string namespacing.
5. Proxied the Config object within subagents to ensure system-wide components (like prompt templates) automatically use the agent-specific registries.
6. Verified through comprehensive updates to core tests for agents, MCP management, and registries.
2026-03-13 19:23:33 +00:00
Akhilesh Kumar 7586efcf49 chore: resolve merge conflicts 2026-03-13 18:16:28 +00:00
Akhilesh Kumar ee425228fe fix(core): ensure policy engine compatibility with isolated MCP servers
This commit addresses PR feedback regarding the prefixing of isolated subagent MCP servers and its potential to break existing security policies relying on standard FQNs.

1. Added `originalName` to `MCPServerConfig` and `originalServerName` to `DiscoveredMCPTool`.
2. Updated `CoreToolScheduler` to reconstruct the original FQN (without the `__agent__` prefix) when performing policy checks via the Policy Engine. This ensures policies mapping to standard `mcp_{server}_{tool}` formats still apply correctly to isolated agents.
3. Added a remote agent back to `NewAgentsNotification.test.tsx` to maintain coverage for both local and remote agents.
2026-03-13 17:30:32 +00:00
Akhilesh Kumar 3bf0a5579a Merge remote-tracking branch 'origin/main' into fix-subagent-tool-isolation
# Conflicts:
#	packages/core/src/agents/local-executor.ts
2026-03-12 19:26:33 +00:00
Akhilesh Kumar 7a08a4fbd5 feat(cli): display MCP servers in agent permission dialog
Updates NewAgentsNotification to inspect the local agent definition and list any MCP servers that the agent introduces, providing users with the necessary visibility before enabling.
2026-03-12 19:06:42 +00:00
Akhilesh Kumar 164abf9d9b fix(core): resolve TypeScript compilation errors in LocalAgentExecutor 2026-03-11 21:47:31 +00:00
Akhilesh Kumar b132791cd2 test(core): add unit tests for subagent MCP tool isolation
Unit tests added:
1. Tool Registry Filtering: Verified that main registry hides all '__agent__' prefixed tools.
2. Subagent Tool Inheritance: Verified that agents correctly filter out other agents' MCP tools while retaining their own.
Verified with vitest in packages/core.
2026-03-11 20:55:10 +00:00
Akhilesh Kumar 5a020e7720 fix(core): avoid restarting subagent MCP servers
I've tactically refactored the `LocalAgentExecutor` so that it avoids shutting down and restarting subagent MCP servers for every agent execution, which mitigates the performance overhead caused by long startup times.

1. Leveraging the Global McpClientManager:
Instead of instantiating an entirely new `McpClientManager` instance within the `LocalAgentExecutor` per execution (and shutting it down in its `finally` block), we now use the single global `McpClientManager` available on `context.config`. Since the global manager deduplicates connection attempts by checking if the server is already active, subagent MCP servers will now naturally stay alive after their initial initialization.

2. Prefixing to Avoid Polluting the Global Namespace:
To isolate the agent-specific tools, we now register the subagent's MCP servers with a unique prefix: `__agent__${definition.name}__${name}`.

3. Strict Filtering for True Isolation (ToolRegistry):
- Main CLI context: Added a block in the global `ToolRegistry.getFunctionDeclarations()` that strictly hides any tool belonging to a server prefixed with `__agent__` if the registry `isMainRegistry`. This prevents internal subagent tools from leaking to the main agent.
- Subagent context (`LocalAgentExecutor`): When inheriting tools from the parent registry (the fallback when an agent doesn't explicitly define `tools: []`), the agent now ignores `__agent__` prefixed tools that belong to *other* agents, ensuring strict tool isolation while keeping the actual underlying server processes alive and reusable.
2026-03-11 20:43:25 +00:00
AK eb5d22848c Merge branch 'main' into fix-subagent-tool-isolation 2026-03-11 13:02:45 -07:00
Akhilesh Kumar 2e6c81e7ad chore: address PR feedback by adding inline comments 2026-03-11 17:13:32 +00:00
Akhilesh Kumar d4b7d358c5 feat(core): support inline MCP server definitions in subagent markdown 2026-03-11 17:10:36 +00:00
Akhilesh Kumar c68a2cb933 feat(core): implement configuration-based tool isolation for subagents 2026-03-11 17:09:46 +00:00
20 changed files with 972 additions and 403 deletions
+1 -26
View File
@@ -2195,7 +2195,6 @@
"integrity": "sha512-t54CUOsFMappY1Jbzb7fetWeO0n6K0k/4+/ZpkS+3Joz8I4VcvY9OiEBFRYISqaI2fq5sCiPtAjRDOzVYG8m+Q==",
"dev": true,
"license": "MIT",
"peer": true,
"dependencies": {
"@octokit/auth-token": "^6.0.0",
"@octokit/graphql": "^9.0.2",
@@ -2376,7 +2375,6 @@
"resolved": "https://registry.npmjs.org/@opentelemetry/api/-/api-1.9.0.tgz",
"integrity": "sha512-3giAOQvZiH5F9bMlMiv8+GSPMeqg0dbaeo58/0SlA9sxSqZhnUtxzX9/2FzyhS9sWQf5S0GJE0AKBrFqjpeYcg==",
"license": "Apache-2.0",
"peer": true,
"engines": {
"node": ">=8.0.0"
}
@@ -2426,7 +2424,6 @@
"resolved": "https://registry.npmjs.org/@opentelemetry/core/-/core-2.5.0.tgz",
"integrity": "sha512-ka4H8OM6+DlUhSAZpONu0cPBtPPTQKxbxVzC4CzVx5+K4JnroJVBtDzLAMx4/3CDTJXRvVFhpFjtl4SaiTNoyQ==",
"license": "Apache-2.0",
"peer": true,
"dependencies": {
"@opentelemetry/semantic-conventions": "^1.29.0"
},
@@ -2801,7 +2798,6 @@
"resolved": "https://registry.npmjs.org/@opentelemetry/resources/-/resources-2.5.0.tgz",
"integrity": "sha512-F8W52ApePshpoSrfsSk1H2yJn9aKjCrbpQF1M9Qii0GHzbfVeFUB+rc3X4aggyZD8x9Gu3Slua+s6krmq6Dt8g==",
"license": "Apache-2.0",
"peer": true,
"dependencies": {
"@opentelemetry/core": "2.5.0",
"@opentelemetry/semantic-conventions": "^1.29.0"
@@ -2835,7 +2831,6 @@
"resolved": "https://registry.npmjs.org/@opentelemetry/sdk-metrics/-/sdk-metrics-2.5.0.tgz",
"integrity": "sha512-BeJLtU+f5Gf905cJX9vXFQorAr6TAfK3SPvTFqP+scfIpDQEJfRaGJWta7sJgP+m4dNtBf9y3yvBKVAZZtJQVA==",
"license": "Apache-2.0",
"peer": true,
"dependencies": {
"@opentelemetry/core": "2.5.0",
"@opentelemetry/resources": "2.5.0"
@@ -2890,7 +2885,6 @@
"resolved": "https://registry.npmjs.org/@opentelemetry/sdk-trace-base/-/sdk-trace-base-2.5.0.tgz",
"integrity": "sha512-VzRf8LzotASEyNDUxTdaJ9IRJ1/h692WyArDBInf5puLCjxbICD6XkHgpuudis56EndyS7LYFmtTMny6UABNdQ==",
"license": "Apache-2.0",
"peer": true,
"dependencies": {
"@opentelemetry/core": "2.5.0",
"@opentelemetry/resources": "2.5.0",
@@ -4120,7 +4114,6 @@
"integrity": "sha512-6mDvHUFSjyT2B2yeNx2nUgMxh9LtOWvkhIU3uePn2I2oyNymUAX1NIsdgviM4CH+JSrp2D2hsMvJOkxY+0wNRA==",
"devOptional": true,
"license": "MIT",
"peer": true,
"dependencies": {
"csstype": "^3.0.2"
}
@@ -4395,7 +4388,6 @@
"integrity": "sha512-6sMvZePQrnZH2/cJkwRpkT7DxoAWh+g6+GFRK6bV3YQo7ogi3SX5rgF6099r5Q53Ma5qeT7LGmOmuIutF4t3lA==",
"dev": true,
"license": "MIT",
"peer": true,
"dependencies": {
"@typescript-eslint/scope-manager": "8.35.0",
"@typescript-eslint/types": "8.35.0",
@@ -5269,7 +5261,6 @@
"resolved": "https://registry.npmjs.org/acorn/-/acorn-8.15.0.tgz",
"integrity": "sha512-NZyJarBfL7nWwIq+FDL6Zp/yHEhePMNnnJ0y3qfieCrmNvYct8uvtiV41UvlSe6apAfk0fY1FbWx+NwfmpvtTg==",
"license": "MIT",
"peer": true,
"bin": {
"acorn": "bin/acorn"
},
@@ -7990,7 +7981,6 @@
"integrity": "sha512-GsGizj2Y1rCWDu6XoEekL3RLilp0voSePurjZIkxL3wlm5o5EC9VpgaP7lrCvjnkuLvzFBQWB3vWB3K5KQTveQ==",
"dev": true,
"license": "MIT",
"peer": true,
"dependencies": {
"@eslint-community/eslint-utils": "^4.2.0",
"@eslint-community/regexpp": "^4.12.1",
@@ -8508,7 +8498,6 @@
"resolved": "https://registry.npmjs.org/express/-/express-5.2.1.tgz",
"integrity": "sha512-hIS4idWWai69NezIdRt2xFVofaF4j+6INOpJlVOLDO8zXGpUVEVzIYk12UUi2JzjEzWL3IOAxcTubgz9Po0yXw==",
"license": "MIT",
"peer": true,
"dependencies": {
"accepts": "^2.0.0",
"body-parser": "^2.2.1",
@@ -9822,7 +9811,6 @@
"resolved": "https://registry.npmjs.org/hono/-/hono-4.12.7.tgz",
"integrity": "sha512-jq9l1DM0zVIvsm3lv9Nw9nlJnMNPOcAtsbsgiUhWcFzPE99Gvo6yRTlszSLLYacMeQ6quHD6hMfId8crVHvexw==",
"license": "MIT",
"peer": true,
"engines": {
"node": ">=16.9.0"
}
@@ -10101,7 +10089,6 @@
"resolved": "https://registry.npmjs.org/@jrichman/ink/-/ink-6.4.11.tgz",
"integrity": "sha512-93LQlzT7vvZ1XJcmOMwN4s+6W334QegendeHOMnEJBlhnpIzr8bws6/aOEHG8ZCuVD/vNeeea5m1msHIdAY6ig==",
"license": "MIT",
"peer": true,
"dependencies": {
"@alcalzone/ansi-tokenize": "^0.2.1",
"ansi-escapes": "^7.0.0",
@@ -13833,7 +13820,6 @@
"resolved": "https://registry.npmjs.org/react/-/react-19.2.4.tgz",
"integrity": "sha512-9nfp2hYpCwOjAN+8TZFGhtWEwgvWHXqESH8qT89AT/lWklpLON22Lc8pEtnpsZz7VmawabSU0gCjnj8aC0euHQ==",
"license": "MIT",
"peer": true,
"engines": {
"node": ">=0.10.0"
}
@@ -13844,7 +13830,6 @@
"integrity": "sha512-ePrwPfxAnB+7hgnEr8vpKxL9cmnp7F322t8oqcPshbIQQhDKgFDW4tjhF2wjVbdXF9O/nyuy3sQWd9JGpiLPvA==",
"devOptional": true,
"license": "MIT",
"peer": true,
"dependencies": {
"shell-quote": "^1.6.1",
"ws": "^7"
@@ -15995,7 +15980,6 @@
"resolved": "https://registry.npmjs.org/picomatch/-/picomatch-4.0.3.tgz",
"integrity": "sha512-5gTmgEY/sqK6gFXLIsQNH19lWb4ebPDLA4SdLP7dsWkIXHWlG66oPuVvXSGFPppYZz8ZDZq0dYYrbHfBCVUb1Q==",
"license": "MIT",
"peer": true,
"engines": {
"node": ">=12"
},
@@ -16218,9 +16202,7 @@
"version": "2.8.1",
"resolved": "https://registry.npmjs.org/tslib/-/tslib-2.8.1.tgz",
"integrity": "sha512-oJFu94HQb+KVduSUQL7wnpmqnfmLsOA/nAh6b6EH0wCEoK0/mPeXU6c3wKDV83MkOuHPRHtSXKKU99IBazS/2w==",
"dev": true,
"license": "0BSD",
"peer": true
"license": "0BSD"
},
"node_modules/tsx": {
"version": "4.20.3",
@@ -16228,7 +16210,6 @@
"integrity": "sha512-qjbnuR9Tr+FJOMBqJCW5ehvIo/buZq7vH7qD7JziU98h6l3qGy0a/yPFjwO+y0/T7GFpNgNAvEcPPVfyT8rrPQ==",
"devOptional": true,
"license": "MIT",
"peer": true,
"dependencies": {
"esbuild": "~0.25.0",
"get-tsconfig": "^4.7.5"
@@ -16394,7 +16375,6 @@
"integrity": "sha512-p1diW6TqL9L07nNxvRMM7hMMw4c5XOo/1ibL4aAIGmSAt9slTE1Xgw5KWuof2uTOvCg9BY7ZRi+GaF+7sfgPeQ==",
"devOptional": true,
"license": "Apache-2.0",
"peer": true,
"bin": {
"tsc": "bin/tsc",
"tsserver": "bin/tsserver"
@@ -16617,7 +16597,6 @@
"resolved": "https://registry.npmjs.org/vite/-/vite-7.2.2.tgz",
"integrity": "sha512-BxAKBWmIbrDgrokdGZH1IgkIk/5mMHDreLDmCJ0qpyJaAteP8NvMhkwr/ZCQNqNH97bw/dANTE9PDzqwJghfMQ==",
"license": "MIT",
"peer": true,
"dependencies": {
"esbuild": "^0.25.0",
"fdir": "^6.5.0",
@@ -16731,7 +16710,6 @@
"resolved": "https://registry.npmjs.org/picomatch/-/picomatch-4.0.3.tgz",
"integrity": "sha512-5gTmgEY/sqK6gFXLIsQNH19lWb4ebPDLA4SdLP7dsWkIXHWlG66oPuVvXSGFPppYZz8ZDZq0dYYrbHfBCVUb1Q==",
"license": "MIT",
"peer": true,
"engines": {
"node": ">=12"
},
@@ -16744,7 +16722,6 @@
"resolved": "https://registry.npmjs.org/vitest/-/vitest-3.2.4.tgz",
"integrity": "sha512-LUCP5ev3GURDysTWiP47wRRUpLKMOfPh+yKTx3kVIEiu5KOMeqzpnYNsKyOoVrULivR8tLcks4+lga33Whn90A==",
"license": "MIT",
"peer": true,
"dependencies": {
"@types/chai": "^5.2.2",
"@vitest/expect": "3.2.4",
@@ -17392,7 +17369,6 @@
"resolved": "https://registry.npmjs.org/zod/-/zod-3.25.76.tgz",
"integrity": "sha512-gzUt/qt81nXsFGKIFcC3YnfEAx5NkunCfnDlvuBSSFS02bcXu4Lmea0AFIUwbLWxWPx3d9p8S5QoaujKcNQxcQ==",
"license": "MIT",
"peer": true,
"funding": {
"url": "https://github.com/sponsors/colinhacks"
}
@@ -17937,7 +17913,6 @@
"resolved": "https://registry.npmjs.org/picomatch/-/picomatch-4.0.3.tgz",
"integrity": "sha512-5gTmgEY/sqK6gFXLIsQNH19lWb4ebPDLA4SdLP7dsWkIXHWlG66oPuVvXSGFPppYZz8ZDZq0dYYrbHfBCVUb1Q==",
"license": "MIT",
"peer": true,
"engines": {
"node": ">=12"
},
@@ -22,6 +22,25 @@ describe('NewAgentsNotification', () => {
{
name: 'Agent B',
description: 'Description B',
kind: 'local' as const,
inputConfig: { inputSchema: {} },
promptConfig: {},
modelConfig: {},
runConfig: {},
mcpServers: {
github: {
command: 'npx',
args: ['-y', '@modelcontextprotocol/server-github'],
},
postgres: {
command: 'npx',
args: ['-y', '@modelcontextprotocol/server-postgres'],
},
},
},
{
name: 'Agent C',
description: 'Description C',
kind: 'remote' as const,
agentCardUrl: '',
inputConfig: { inputSchema: {} },
@@ -80,16 +80,35 @@ export const NewAgentsNotification = ({
borderStyle="single"
padding={1}
>
{displayAgents.map((agent) => (
<Box key={agent.name}>
<Box flexShrink={0}>
<Text bold color={theme.text.primary}>
- {agent.name}:{' '}
</Text>
{displayAgents.map((agent) => {
const mcpServers =
agent.kind === 'local' ? agent.mcpServers : undefined;
const hasMcpServers =
mcpServers && Object.keys(mcpServers).length > 0;
return (
<Box key={agent.name} flexDirection="column">
<Box>
<Box flexShrink={0}>
<Text bold color={theme.text.primary}>
- {agent.name}:{' '}
</Text>
</Box>
<Text color={theme.text.secondary}>
{' '}
{agent.description}
</Text>
</Box>
{hasMcpServers && (
<Box marginLeft={2}>
<Text color={theme.text.secondary}>
(Includes MCP servers:{' '}
{Object.keys(mcpServers).join(', ')})
</Text>
</Box>
)}
</Box>
<Text color={theme.text.secondary}> {agent.description}</Text>
</Box>
))}
);
})}
{remaining > 0 && (
<Text color={theme.text.secondary}>
... and {remaining} more.
@@ -10,6 +10,8 @@ exports[`NewAgentsNotification > renders agent list 1`] = `
│ │ │ │
│ │ - Agent A: Description A │ │
│ │ - Agent B: Description B │ │
│ │ (Includes MCP servers: github, postgres) │ │
│ │ - Agent C: Description C │ │
│ │ │ │
│ └────────────────────────────────────────────────────────────────────────────────────────────┘ │
│ │
+28 -2
View File
@@ -11,6 +11,8 @@ import type {
CompletedToolCall,
} from '../scheduler/types.js';
import type { ToolRegistry } from '../tools/tool-registry.js';
import type { PromptRegistry } from '../prompts/prompt-registry.js';
import type { ResourceRegistry } from '../resources/resource-registry.js';
import type { EditorType } from '../utils/editor.js';
/**
@@ -25,6 +27,10 @@ export interface AgentSchedulingOptions {
parentCallId?: string;
/** The tool registry specific to this agent. */
toolRegistry: ToolRegistry;
/** The prompt registry specific to this agent. */
promptRegistry?: PromptRegistry;
/** The resource registry specific to this agent. */
resourceRegistry?: ResourceRegistry;
/** AbortSignal for cancellation. */
signal: AbortSignal;
/** Optional function to get the preferred editor for tool modifications. */
@@ -51,14 +57,34 @@ export async function scheduleAgentTools(
subagent,
parentCallId,
toolRegistry,
promptRegistry,
resourceRegistry,
signal,
getPreferredEditor,
onWaitingForConfirmation,
} = options;
// Create a proxy/override of the config to provide the agent-specific tool registry.
// Create a proxy/override of the config to provide the agent-specific registries.
// eslint-disable-next-line @typescript-eslint/no-unsafe-assignment, no-restricted-syntax
const agentConfig: Config = Object.create(config);
agentConfig.getToolRegistry = () => toolRegistry;
agentConfig.getMessageBus = () => toolRegistry.messageBus;
if (promptRegistry) {
agentConfig.getPromptRegistry = () => promptRegistry;
}
if (resourceRegistry) {
agentConfig.getResourceRegistry = () => resourceRegistry;
}
// Override toolRegistry property so AgentLoopContext reads the agent-specific registry.
Object.defineProperty(agentConfig, 'toolRegistry', {
get: () => toolRegistry,
configurable: true,
});
const schedulerContext = {
config,
config: agentConfig,
promptId: config.promptId,
toolRegistry,
messageBus: toolRegistry.messageBus,
@@ -81,6 +81,33 @@ System prompt content.`);
});
});
it('should parse frontmatter with mcp_servers', async () => {
const filePath = await writeAgentMarkdown(`---
name: mcp-agent
description: An agent with MCP servers
mcp_servers:
test-server:
command: node
args: [server.js]
include_tools: [tool1, tool2]
---
System prompt content.`);
const result = await parseAgentMarkdown(filePath);
expect(result).toHaveLength(1);
expect(result[0]).toMatchObject({
name: 'mcp-agent',
description: 'An agent with MCP servers',
mcp_servers: {
'test-server': {
command: 'node',
args: ['server.js'],
include_tools: ['tool1', 'tool2'],
},
},
});
});
it('should throw AgentLoadError if frontmatter is missing', async () => {
const filePath = await writeAgentMarkdown(`Just some markdown content.`);
await expect(parseAgentMarkdown(filePath)).rejects.toThrow(
@@ -274,6 +301,33 @@ Body`);
expect(result.modelConfig.model).toBe(GEMINI_MODEL_ALIAS_PRO);
});
it('should convert mcp_servers in local agent', () => {
const markdown = {
kind: 'local' as const,
name: 'mcp-agent',
description: 'An agent with MCP servers',
mcp_servers: {
'test-server': {
command: 'node',
args: ['server.js'],
include_tools: ['tool1'],
},
},
system_prompt: 'prompt',
};
const result = markdownToAgentDefinition(
markdown,
) as LocalAgentDefinition;
expect(result.kind).toBe('local');
expect(result.mcpServers).toBeDefined();
expect(result.mcpServers!['test-server']).toMatchObject({
command: 'node',
args: ['server.js'],
includeTools: ['tool1'],
});
});
it('should pass through unknown model names (e.g. auto)', () => {
const markdown = {
kind: 'local' as const,
+60
View File
@@ -16,6 +16,7 @@ import {
DEFAULT_MAX_TIME_MINUTES,
} from './types.js';
import type { A2AAuthConfig } from './auth-provider/types.js';
import { MCPServerConfig } from '../config/config.js';
import { isValidToolName } from '../tools/tool-names.js';
import { FRONTMATTER_REGEX } from '../skills/skillLoader.js';
import { getErrorMessage } from '../utils/errors.js';
@@ -28,11 +29,29 @@ interface FrontmatterBaseAgentDefinition {
display_name?: string;
}
interface FrontmatterMCPServerConfig {
command?: string;
args?: string[];
env?: Record<string, string>;
cwd?: string;
url?: string;
http_url?: string;
headers?: Record<string, string>;
tcp?: string;
type?: 'sse' | 'http';
timeout?: number;
trust?: boolean;
description?: string;
include_tools?: string[];
exclude_tools?: string[];
}
interface FrontmatterLocalAgentDefinition
extends FrontmatterBaseAgentDefinition {
kind: 'local';
description: string;
tools?: string[];
mcp_servers?: Record<string, FrontmatterMCPServerConfig>;
system_prompt: string;
model?: string;
temperature?: number;
@@ -100,6 +119,23 @@ const nameSchema = z
.string()
.regex(/^[a-z0-9-_]+$/, 'Name must be a valid slug');
const mcpServerSchema = z.object({
command: z.string().optional(),
args: z.array(z.string()).optional(),
env: z.record(z.string()).optional(),
cwd: z.string().optional(),
url: z.string().optional(),
http_url: z.string().optional(),
headers: z.record(z.string()).optional(),
tcp: z.string().optional(),
type: z.enum(['sse', 'http']).optional(),
timeout: z.number().optional(),
trust: z.boolean().optional(),
description: z.string().optional(),
include_tools: z.array(z.string()).optional(),
exclude_tools: z.array(z.string()).optional(),
});
const localAgentSchema = z
.object({
kind: z.literal('local').optional().default('local'),
@@ -115,6 +151,7 @@ const localAgentSchema = z
}),
)
.optional(),
mcp_servers: z.record(mcpServerSchema).optional(),
model: z.string().optional(),
temperature: z.number().optional(),
max_turns: z.number().int().positive().optional(),
@@ -495,6 +532,28 @@ export function markdownToAgentDefinition(
// If a model is specified, use it. Otherwise, inherit
const modelName = markdown.model || 'inherit';
const mcpServers: Record<string, MCPServerConfig> = {};
if (markdown.kind === 'local' && markdown.mcp_servers) {
for (const [name, config] of Object.entries(markdown.mcp_servers)) {
mcpServers[name] = new MCPServerConfig(
config.command,
config.args,
config.env,
config.cwd,
config.url,
config.http_url,
config.headers,
config.tcp,
config.type,
config.timeout,
config.trust,
config.description,
config.include_tools,
config.exclude_tools,
);
}
}
return {
kind: 'local',
name: markdown.name,
@@ -520,6 +579,7 @@ export function markdownToAgentDefinition(
tools: markdown.tools,
}
: undefined,
mcpServers: Object.keys(mcpServers).length > 0 ? mcpServers : undefined,
inputConfig,
metadata,
};
+94 -14
View File
@@ -13,10 +13,43 @@ import {
afterEach,
type Mock,
} from 'vitest';
const {
mockSendMessageStream,
mockScheduleAgentTools,
mockSetSystemInstruction,
mockCompress,
mockMaybeDiscoverMcpServer,
mockStopMcp,
} = vi.hoisted(() => ({
mockSendMessageStream: vi.fn().mockResolvedValue({
async *[Symbol.asyncIterator]() {
yield {
type: 'chunk',
value: { candidates: [] },
};
},
}),
mockScheduleAgentTools: vi.fn(),
mockSetSystemInstruction: vi.fn(),
mockCompress: vi.fn(),
mockMaybeDiscoverMcpServer: vi.fn().mockResolvedValue(undefined),
mockStopMcp: vi.fn().mockResolvedValue(undefined),
}));
vi.mock('../tools/mcp-client-manager.js', () => ({
McpClientManager: class {
maybeDiscoverMcpServer = mockMaybeDiscoverMcpServer;
stop = mockStopMcp;
},
}));
import { debugLogger } from '../utils/debugLogger.js';
import { LocalAgentExecutor, type ActivityCallback } from './local-executor.js';
import { makeFakeConfig } from '../test-utils/config.js';
import { ToolRegistry } from '../tools/tool-registry.js';
import { PromptRegistry } from '../prompts/prompt-registry.js';
import { ResourceRegistry } from '../resources/resource-registry.js';
import { DiscoveredMCPTool } from '../tools/mcp-tool.js';
import { LSTool } from '../tools/ls.js';
import { LS_TOOL_NAME, READ_FILE_TOOL_NAME } from '../tools/tool-names.js';
@@ -70,18 +103,6 @@ import type {
import { getModelConfigAlias, type AgentRegistry } from './registry.js';
import type { ModelRouterService } from '../routing/modelRouterService.js';
const {
mockSendMessageStream,
mockScheduleAgentTools,
mockSetSystemInstruction,
mockCompress,
} = vi.hoisted(() => ({
mockSendMessageStream: vi.fn(),
mockScheduleAgentTools: vi.fn(),
mockSetSystemInstruction: vi.fn(),
mockCompress: vi.fn(),
}));
let mockChatHistory: Content[] = [];
const mockSetHistory = vi.fn((newHistory: Content[]) => {
mockChatHistory = newHistory;
@@ -2493,6 +2514,67 @@ describe('LocalAgentExecutor', () => {
});
});
describe('MCP Isolation', () => {
it('should initialize McpClientManager when mcpServers are defined', async () => {
const { MCPServerConfig } = await import('../config/config.js');
const mcpServers = {
'test-server': new MCPServerConfig('node', ['server.js']),
};
const definition = {
...createTestDefinition(),
mcpServers,
};
vi.spyOn(mockConfig, 'getMcpClientManager').mockReturnValue({
maybeDiscoverMcpServer: mockMaybeDiscoverMcpServer,
} as unknown as ReturnType<typeof mockConfig.getMcpClientManager>);
await LocalAgentExecutor.create(definition, mockConfig);
const mcpManager = mockConfig.getMcpClientManager();
expect(mcpManager?.maybeDiscoverMcpServer).toHaveBeenCalledWith(
'test-server',
mcpServers['test-server'],
expect.objectContaining({
toolRegistry: expect.any(ToolRegistry),
promptRegistry: expect.any(PromptRegistry),
resourceRegistry: expect.any(ResourceRegistry),
}),
);
});
it('should inherit main registry tools', async () => {
const parentMcpTool = new DiscoveredMCPTool(
{} as unknown as CallableTool,
'main-server',
'tool1',
'desc1',
{},
mockConfig.getMessageBus(),
);
parentToolRegistry.registerTool(parentMcpTool);
const definition = createTestDefinition();
definition.toolConfig = undefined; // trigger inheritance
vi.spyOn(mockConfig, 'getMcpClientManager').mockReturnValue({
maybeDiscoverMcpServer: vi.fn(),
} as unknown as ReturnType<typeof mockConfig.getMcpClientManager>);
const executor = await LocalAgentExecutor.create(
definition,
mockConfig,
onActivity,
);
const agentTools = (
executor as unknown as { toolRegistry: ToolRegistry }
).toolRegistry.getAllToolNames();
expect(agentTools).toContain(parentMcpTool.name);
});
});
describe('DeclarativeTool instance tools (browser agent pattern)', () => {
/**
* The browser agent passes DeclarativeTool instances (not string names) in
@@ -2598,13 +2680,11 @@ describe('LocalAgentExecutor', () => {
const navTool = new MockTool({ name: 'navigate_page' });
const definition = createInstanceToolDefinition([clickTool, navTool]);
const executor = await LocalAgentExecutor.create(
definition,
mockConfig,
onActivity,
);
const registry = executor['toolRegistry'];
expect(registry.getTool('click')).toBeDefined();
expect(registry.getTool('navigate_page')).toBeDefined();
+60 -10
View File
@@ -17,6 +17,8 @@ import {
type Schema,
} from '@google/genai';
import { ToolRegistry } from '../tools/tool-registry.js';
import { PromptRegistry } from '../prompts/prompt-registry.js';
import { ResourceRegistry } from '../resources/resource-registry.js';
import { type AnyDeclarativeTool } from '../tools/tools.js';
import {
DiscoveredMCPTool,
@@ -98,6 +100,8 @@ export class LocalAgentExecutor<TOutput extends z.ZodTypeAny> {
private readonly agentId: string;
private readonly toolRegistry: ToolRegistry;
private readonly promptRegistry: PromptRegistry;
private readonly resourceRegistry: ResourceRegistry;
private readonly context: AgentLoopContext;
private readonly onActivity?: ActivityCallback;
private readonly compressionService: ChatCompressionService;
@@ -105,7 +109,18 @@ export class LocalAgentExecutor<TOutput extends z.ZodTypeAny> {
private hasFailedCompressionAttempt = false;
private get config(): Config {
return this.context.config;
// eslint-disable-next-line @typescript-eslint/no-unsafe-assignment, no-restricted-syntax
const agentConfig: Config = Object.create(this.context.config);
agentConfig.getToolRegistry = () => this.toolRegistry;
agentConfig.getPromptRegistry = () => this.promptRegistry;
agentConfig.getResourceRegistry = () => this.resourceRegistry;
agentConfig.getMessageBus = () => this.toolRegistry.getMessageBus();
Object.defineProperty(agentConfig, 'toolRegistry', {
get: () => this.toolRegistry,
configurable: true,
});
return agentConfig;
}
/**
@@ -129,11 +144,27 @@ export class LocalAgentExecutor<TOutput extends z.ZodTypeAny> {
// Create an override object to inject the subagent name into tool confirmation requests
const subagentMessageBus = parentMessageBus.derive(definition.name);
// Create an isolated tool registry for this agent instance.
// Create isolated registries for this agent instance.
const agentToolRegistry = new ToolRegistry(
context.config,
subagentMessageBus,
);
const agentPromptRegistry = new PromptRegistry();
const agentResourceRegistry = new ResourceRegistry();
if (definition.mcpServers) {
const globalMcpManager = context.config.getMcpClientManager();
if (globalMcpManager) {
for (const [name, config] of Object.entries(definition.mcpServers)) {
await globalMcpManager.maybeDiscoverMcpServer(name, config, {
toolRegistry: agentToolRegistry,
promptRegistry: agentPromptRegistry,
resourceRegistry: agentResourceRegistry,
});
}
}
}
const parentToolRegistry = context.toolRegistry;
const allAgentNames = new Set(
context.config.getAgentRegistry().getAllAgentNames(),
@@ -149,7 +180,9 @@ export class LocalAgentExecutor<TOutput extends z.ZodTypeAny> {
return;
}
agentToolRegistry.registerTool(tool);
// Clone the tool, so it gets its own state and subagent messageBus
const clonedTool = tool.clone(subagentMessageBus);
agentToolRegistry.registerTool(clonedTool);
};
const registerToolByName = (toolName: string) => {
@@ -224,10 +257,12 @@ export class LocalAgentExecutor<TOutput extends z.ZodTypeAny> {
return new LocalAgentExecutor(
definition,
context,
agentToolRegistry,
parentPromptId,
parentCallId,
agentToolRegistry,
agentPromptRegistry,
agentResourceRegistry,
onActivity,
parentCallId,
);
}
@@ -240,14 +275,18 @@ export class LocalAgentExecutor<TOutput extends z.ZodTypeAny> {
private constructor(
definition: LocalAgentDefinition<TOutput>,
context: AgentLoopContext,
toolRegistry: ToolRegistry,
parentPromptId: string | undefined,
parentCallId: string | undefined,
toolRegistry: ToolRegistry,
promptRegistry: PromptRegistry,
resourceRegistry: ResourceRegistry,
onActivity?: ActivityCallback,
parentCallId?: string,
) {
this.definition = definition;
this.context = context;
this.toolRegistry = toolRegistry;
this.promptRegistry = promptRegistry;
this.resourceRegistry = resourceRegistry;
this.onActivity = onActivity;
this.compressionService = new ChatCompressionService();
this.parentCallId = parentCallId;
@@ -491,7 +530,7 @@ export class LocalAgentExecutor<TOutput extends z.ZodTypeAny> {
const combinedSignal = AbortSignal.any([signal, deadlineTimer.signal]);
logAgentStart(
this.config,
this.context.config,
new AgentStartEvent(this.agentId, this.definition.name),
);
@@ -586,6 +625,15 @@ export class LocalAgentExecutor<TOutput extends z.ZodTypeAny> {
}
} finally {
this.config.userHintService.offUserHint(hintListener);
const globalMcpManager = this.context.config.getMcpClientManager();
if (globalMcpManager) {
globalMcpManager.removeRegistries({
toolRegistry: this.toolRegistry,
promptRegistry: this.promptRegistry,
resourceRegistry: this.resourceRegistry,
});
}
}
// === UNIFIED RECOVERY BLOCK ===
@@ -698,7 +746,7 @@ export class LocalAgentExecutor<TOutput extends z.ZodTypeAny> {
} finally {
deadlineTimer.abort();
logAgentFinish(
this.config,
this.context.config,
new AgentFinishEvent(
this.agentId,
this.definition.name,
@@ -1118,10 +1166,12 @@ export class LocalAgentExecutor<TOutput extends z.ZodTypeAny> {
this.config,
toolRequests,
{
schedulerId: this.agentId,
schedulerId: promptId,
subagent: this.definition.name,
parentCallId: this.parentCallId,
toolRegistry: this.toolRegistry,
promptRegistry: this.promptRegistry,
resourceRegistry: this.resourceRegistry,
signal,
onWaitingForConfirmation,
},
+13
View File
@@ -570,6 +570,19 @@ export class AgentRegistry {
},
};
if (overrides.tools) {
merged.toolConfig = {
tools: overrides.tools,
};
}
if (overrides.mcpServers) {
merged.mcpServers = {
...definition.mcpServers,
...overrides.mcpServers,
};
}
return merged;
}
+6
View File
@@ -14,6 +14,7 @@ import { type z } from 'zod';
import type { ModelConfig } from '../services/modelConfigService.js';
import type { AnySchema } from 'ajv';
import type { A2AAuthConfig } from './auth-provider/types.js';
import type { MCPServerConfig } from '../config/config.js';
/**
* Describes the possible termination modes for an agent.
@@ -130,6 +131,11 @@ export interface LocalAgentDefinition<
// Optional configs
toolConfig?: ToolConfig;
/**
* Optional inline MCP servers for this agent.
*/
mcpServers?: Record<string, MCPServerConfig>;
/**
* An optional function to process the raw output from the agent's final tool
* call into a string format.
+3
View File
@@ -98,6 +98,7 @@ vi.mock('../tools/mcp-client-manager.js', () => ({
McpClientManager: vi.fn().mockImplementation(() => ({
startConfiguredMcpServers: vi.fn(),
getMcpInstructions: vi.fn().mockReturnValue('MCP Instructions'),
setMainRegistries: vi.fn(),
})),
}));
@@ -368,6 +369,7 @@ describe('Server Config (config.ts)', () => {
mcpStarted = true;
}),
getMcpInstructions: vi.fn(),
setMainRegistries: vi.fn(),
}) as Partial<McpClientManager> as McpClientManager,
);
@@ -401,6 +403,7 @@ describe('Server Config (config.ts)', () => {
mcpStarted = true;
}),
getMcpInstructions: vi.fn(),
setMainRegistries: vi.fn(),
}) as Partial<McpClientManager> as McpClientManager,
);
+19 -2
View File
@@ -239,6 +239,8 @@ export interface AgentOverride {
modelConfig?: ModelConfig;
runConfig?: AgentRunConfig;
enabled?: boolean;
tools?: string[];
mcpServers?: Record<string, MCPServerConfig>;
}
export interface AgentSettings {
@@ -520,6 +522,7 @@ export interface ConfigParameters {
question?: string;
coreTools?: string[];
mainAgentTools?: string[];
/** @deprecated Use Policy Engine instead */
allowedTools?: string[];
/** @deprecated Use Policy Engine instead */
@@ -675,6 +678,7 @@ export class Config implements McpContext, AgentLoopContext {
readonly enableConseca: boolean;
private readonly coreTools: string[] | undefined;
private readonly mainAgentTools: string[] | undefined;
/** @deprecated Use Policy Engine instead */
private readonly allowedTools: string[] | undefined;
/** @deprecated Use Policy Engine instead */
@@ -888,6 +892,7 @@ export class Config implements McpContext, AgentLoopContext {
this.question = params.question;
this.coreTools = params.coreTools;
this.mainAgentTools = params.mainAgentTools;
this.allowedTools = params.allowedTools;
this.excludeTools = params.excludeTools;
this.toolDiscoveryCommand = params.toolDiscoveryCommand;
@@ -1231,10 +1236,14 @@ export class Config implements McpContext, AgentLoopContext {
discoverToolsHandle?.end();
this.mcpClientManager = new McpClientManager(
this.clientVersion,
this._toolRegistry,
this,
this.eventEmitter,
);
this.mcpClientManager.setMainRegistries({
toolRegistry: this._toolRegistry,
promptRegistry: this.promptRegistry,
resourceRegistry: this.resourceRegistry,
});
// We do not await this promise so that the CLI can start up even if
// MCP servers are slow to connect.
this.mcpInitializationPromise = Promise.allSettled([
@@ -1887,6 +1896,10 @@ export class Config implements McpContext, AgentLoopContext {
return this.coreTools;
}
getMainAgentTools(): string[] | undefined {
return this.mainAgentTools;
}
getAllowedTools(): string[] | undefined {
return this.allowedTools;
}
@@ -2982,7 +2995,11 @@ export class Config implements McpContext, AgentLoopContext {
}
async createToolRegistry(): Promise<ToolRegistry> {
const registry = new ToolRegistry(this, this.messageBus);
const registry = new ToolRegistry(
this,
this.messageBus,
/* isMainRegistry= */ true,
);
// helper to create & register core tools that are enabled
const maybeRegister = (
@@ -14,9 +14,11 @@ import {
type MockedObject,
} from 'vitest';
import { McpClientManager } from './mcp-client-manager.js';
import { McpClient, MCPDiscoveryState } from './mcp-client.js';
import { McpClient, MCPDiscoveryState, MCPServerStatus } from './mcp-client.js';
import type { ToolRegistry } from './tool-registry.js';
import type { Config, GeminiCLIExtension } from '../config/config.js';
import type { PromptRegistry } from '../prompts/prompt-registry.js';
import type { ResourceRegistry } from '../resources/resource-registry.js';
vi.mock('./mcp-client.js', async () => {
const originalModule = await vi.importActual('./mcp-client.js');
@@ -34,21 +36,25 @@ describe('McpClientManager', () => {
beforeEach(() => {
mockedMcpClient = vi.mockObject({
connect: vi.fn(),
discover: vi.fn(),
discoverInto: vi.fn(),
disconnect: vi.fn(),
getStatus: vi.fn(),
getStatus: vi.fn().mockReturnValue(MCPServerStatus.DISCONNECTED),
getServerConfig: vi.fn(),
getServerName: vi.fn().mockReturnValue('test-server'),
} as unknown as McpClient);
vi.mocked(McpClient).mockReturnValue(mockedMcpClient);
mockConfig = vi.mockObject({
isTrustedFolder: vi.fn().mockReturnValue(true),
getMcpServers: vi.fn().mockReturnValue({}),
getPromptRegistry: () => {},
getResourceRegistry: () => {},
getPromptRegistry: vi.fn().mockReturnValue({ registerPrompt: vi.fn() }),
getResourceRegistry: vi
.fn()
.mockReturnValue({ setResourcesForServer: vi.fn() }),
getDebugMode: () => false,
getWorkspaceContext: () => {},
getWorkspaceContext: () => ({ getDirectories: () => [] }),
getAllowedMcpServers: vi.fn().mockReturnValue([]),
getBlockedMcpServers: vi.fn().mockReturnValue([]),
getExcludedMcpServers: vi.fn().mockReturnValue([]),
getMcpServerCommand: vi.fn().mockReturnValue(''),
getMcpEnablementCallbacks: vi.fn().mockReturnValue(undefined),
getGeminiClient: vi.fn().mockReturnValue({
@@ -56,21 +62,39 @@ describe('McpClientManager', () => {
}),
refreshMcpContext: vi.fn(),
} as unknown as Config);
toolRegistry = {} as ToolRegistry;
toolRegistry = vi.mockObject({
registerTool: vi.fn(),
unregisterTool: vi.fn(),
sortTools: vi.fn(),
getMessageBus: vi.fn().mockReturnValue({}),
removeMcpToolsByServer: vi.fn(),
getToolsByServer: vi.fn().mockReturnValue([]),
} as unknown as ToolRegistry);
});
afterEach(() => {
vi.restoreAllMocks();
});
const setupManager = (manager: McpClientManager) => {
manager.setMainRegistries({
toolRegistry,
promptRegistry:
mockConfig.getPromptRegistry() as unknown as PromptRegistry,
resourceRegistry:
mockConfig.getResourceRegistry() as unknown as ResourceRegistry,
});
return manager;
};
it('should discover tools from all configured', async () => {
mockConfig.getMcpServers.mockReturnValue({
'test-server': { command: 'node' },
});
const manager = new McpClientManager('0.0.1', toolRegistry, mockConfig);
const manager = setupManager(new McpClientManager('0.0.1', mockConfig));
await manager.startConfiguredMcpServers();
expect(mockedMcpClient.connect).toHaveBeenCalledOnce();
expect(mockedMcpClient.discover).toHaveBeenCalledOnce();
expect(mockedMcpClient.discoverInto).toHaveBeenCalledOnce();
expect(mockConfig.refreshMcpContext).toHaveBeenCalledOnce();
});
@@ -80,12 +104,12 @@ describe('McpClientManager', () => {
'server-2': { command: 'node' },
'server-3': { command: 'node' },
});
const manager = new McpClientManager('0.0.1', toolRegistry, mockConfig);
const manager = setupManager(new McpClientManager('0.0.1', mockConfig));
await manager.startConfiguredMcpServers();
// Each client should be connected/discovered
expect(mockedMcpClient.connect).toHaveBeenCalledTimes(3);
expect(mockedMcpClient.discover).toHaveBeenCalledTimes(3);
expect(mockedMcpClient.discoverInto).toHaveBeenCalledTimes(3);
// But context refresh should happen only once
expect(mockConfig.refreshMcpContext).toHaveBeenCalledOnce();
@@ -95,7 +119,7 @@ describe('McpClientManager', () => {
mockConfig.getMcpServers.mockReturnValue({
'test-server': { command: 'node' },
});
const manager = new McpClientManager('0.0.1', toolRegistry, mockConfig);
const manager = setupManager(new McpClientManager('0.0.1', mockConfig));
expect(manager.getDiscoveryState()).toBe(MCPDiscoveryState.NOT_STARTED);
const promise = manager.startConfiguredMcpServers();
expect(manager.getDiscoveryState()).toBe(MCPDiscoveryState.IN_PROGRESS);
@@ -112,7 +136,7 @@ describe('McpClientManager', () => {
isFileEnabled: vi.fn().mockResolvedValue(false),
});
const manager = new McpClientManager('0.0.1', toolRegistry, mockConfig);
const manager = setupManager(new McpClientManager('0.0.1', mockConfig));
const promise = manager.startConfiguredMcpServers();
expect(manager.getDiscoveryState()).toBe(MCPDiscoveryState.IN_PROGRESS);
await promise;
@@ -120,7 +144,7 @@ describe('McpClientManager', () => {
expect(manager.getDiscoveryState()).toBe(MCPDiscoveryState.COMPLETED);
expect(manager.getMcpServerCount()).toBe(0);
expect(mockedMcpClient.connect).not.toHaveBeenCalled();
expect(mockedMcpClient.discover).not.toHaveBeenCalled();
expect(mockedMcpClient.discoverInto).not.toHaveBeenCalled();
});
it('should mark discovery completed when all configured servers are blocked', async () => {
@@ -129,7 +153,7 @@ describe('McpClientManager', () => {
});
mockConfig.getBlockedMcpServers.mockReturnValue(['test-server']);
const manager = new McpClientManager('0.0.1', toolRegistry, mockConfig);
const manager = setupManager(new McpClientManager('0.0.1', mockConfig));
const promise = manager.startConfiguredMcpServers();
expect(manager.getDiscoveryState()).toBe(MCPDiscoveryState.IN_PROGRESS);
await promise;
@@ -137,7 +161,7 @@ describe('McpClientManager', () => {
expect(manager.getDiscoveryState()).toBe(MCPDiscoveryState.COMPLETED);
expect(manager.getMcpServerCount()).toBe(0);
expect(mockedMcpClient.connect).not.toHaveBeenCalled();
expect(mockedMcpClient.discover).not.toHaveBeenCalled();
expect(mockedMcpClient.discoverInto).not.toHaveBeenCalled();
});
it('should not discover tools if folder is not trusted', async () => {
@@ -145,10 +169,10 @@ describe('McpClientManager', () => {
'test-server': { command: 'node' },
});
mockConfig.isTrustedFolder.mockReturnValue(false);
const manager = new McpClientManager('0.0.1', toolRegistry, mockConfig);
const manager = setupManager(new McpClientManager('0.0.1', mockConfig));
await manager.startConfiguredMcpServers();
expect(mockedMcpClient.connect).not.toHaveBeenCalled();
expect(mockedMcpClient.discover).not.toHaveBeenCalled();
expect(mockedMcpClient.discoverInto).not.toHaveBeenCalled();
});
it('should not start blocked servers', async () => {
@@ -156,10 +180,10 @@ describe('McpClientManager', () => {
'test-server': { command: 'node' },
});
mockConfig.getBlockedMcpServers.mockReturnValue(['test-server']);
const manager = new McpClientManager('0.0.1', toolRegistry, mockConfig);
const manager = setupManager(new McpClientManager('0.0.1', mockConfig));
await manager.startConfiguredMcpServers();
expect(mockedMcpClient.connect).not.toHaveBeenCalled();
expect(mockedMcpClient.discover).not.toHaveBeenCalled();
expect(mockedMcpClient.discoverInto).not.toHaveBeenCalled();
});
it('should only start allowed servers if allow list is not empty', async () => {
@@ -168,14 +192,14 @@ describe('McpClientManager', () => {
'another-server': { command: 'node' },
});
mockConfig.getAllowedMcpServers.mockReturnValue(['another-server']);
const manager = new McpClientManager('0.0.1', toolRegistry, mockConfig);
const manager = setupManager(new McpClientManager('0.0.1', mockConfig));
await manager.startConfiguredMcpServers();
expect(mockedMcpClient.connect).toHaveBeenCalledOnce();
expect(mockedMcpClient.discover).toHaveBeenCalledOnce();
expect(mockedMcpClient.discoverInto).toHaveBeenCalledOnce();
});
it('should start servers from extensions', async () => {
const manager = new McpClientManager('0.0.1', toolRegistry, mockConfig);
const manager = setupManager(new McpClientManager('0.0.1', mockConfig));
await manager.startExtension({
name: 'test-extension',
mcpServers: {
@@ -188,11 +212,11 @@ describe('McpClientManager', () => {
id: '123',
});
expect(mockedMcpClient.connect).toHaveBeenCalledOnce();
expect(mockedMcpClient.discover).toHaveBeenCalledOnce();
expect(mockedMcpClient.discoverInto).toHaveBeenCalledOnce();
});
it('should not start servers from disabled extensions', async () => {
const manager = new McpClientManager('0.0.1', toolRegistry, mockConfig);
const manager = setupManager(new McpClientManager('0.0.1', mockConfig));
await manager.startExtension({
name: 'test-extension',
mcpServers: {
@@ -205,7 +229,7 @@ describe('McpClientManager', () => {
id: '123',
});
expect(mockedMcpClient.connect).not.toHaveBeenCalled();
expect(mockedMcpClient.discover).not.toHaveBeenCalled();
expect(mockedMcpClient.discoverInto).not.toHaveBeenCalled();
});
it('should add blocked servers to the blockedMcpServers list', async () => {
@@ -213,7 +237,7 @@ describe('McpClientManager', () => {
'test-server': { command: 'node' },
});
mockConfig.getBlockedMcpServers.mockReturnValue(['test-server']);
const manager = new McpClientManager('0.0.1', toolRegistry, mockConfig);
const manager = setupManager(new McpClientManager('0.0.1', mockConfig));
await manager.startConfiguredMcpServers();
expect(manager.getBlockedMcpServers()).toEqual([
{ name: 'test-server', extensionName: '' },
@@ -224,10 +248,10 @@ describe('McpClientManager', () => {
mockConfig.getMcpServers.mockReturnValue({
'test-server': { excludeTools: ['dangerous_tool'] },
});
const manager = new McpClientManager('0.0.1', toolRegistry, mockConfig);
const manager = setupManager(new McpClientManager('0.0.1', mockConfig));
await manager.startConfiguredMcpServers();
expect(mockedMcpClient.connect).not.toHaveBeenCalled();
expect(mockedMcpClient.discover).not.toHaveBeenCalled();
expect(mockedMcpClient.discoverInto).not.toHaveBeenCalled();
// But it should still be tracked in allServerConfigs
expect(manager.getMcpServers()).toHaveProperty('test-server');
@@ -240,16 +264,16 @@ describe('McpClientManager', () => {
'test-server': serverConfig,
});
mockedMcpClient.getServerConfig.mockReturnValue(serverConfig);
const manager = new McpClientManager('0.0.1', toolRegistry, mockConfig);
const manager = setupManager(new McpClientManager('0.0.1', mockConfig));
await manager.startConfiguredMcpServers();
expect(mockedMcpClient.connect).toHaveBeenCalledTimes(1);
expect(mockedMcpClient.discover).toHaveBeenCalledTimes(1);
expect(mockedMcpClient.discoverInto).toHaveBeenCalledTimes(1);
await manager.restart();
expect(mockedMcpClient.disconnect).toHaveBeenCalledTimes(1);
expect(mockedMcpClient.connect).toHaveBeenCalledTimes(2);
expect(mockedMcpClient.discover).toHaveBeenCalledTimes(2);
expect(mockedMcpClient.discoverInto).toHaveBeenCalledTimes(2);
});
});
@@ -260,21 +284,21 @@ describe('McpClientManager', () => {
'test-server': serverConfig,
});
mockedMcpClient.getServerConfig.mockReturnValue(serverConfig);
const manager = new McpClientManager('0.0.1', toolRegistry, mockConfig);
const manager = setupManager(new McpClientManager('0.0.1', mockConfig));
await manager.startConfiguredMcpServers();
expect(mockedMcpClient.connect).toHaveBeenCalledTimes(1);
expect(mockedMcpClient.discover).toHaveBeenCalledTimes(1);
expect(mockedMcpClient.discoverInto).toHaveBeenCalledTimes(1);
await manager.restartServer('test-server');
expect(mockedMcpClient.disconnect).toHaveBeenCalledTimes(1);
expect(mockedMcpClient.connect).toHaveBeenCalledTimes(2);
expect(mockedMcpClient.discover).toHaveBeenCalledTimes(2);
expect(mockedMcpClient.discoverInto).toHaveBeenCalledTimes(2);
});
it('should throw an error if the server does not exist', async () => {
const manager = new McpClientManager('0.0.1', toolRegistry, mockConfig);
const manager = setupManager(new McpClientManager('0.0.1', mockConfig));
await expect(manager.restartServer('non-existent')).rejects.toThrow(
'No MCP server registered with the name "non-existent"',
);
@@ -296,7 +320,7 @@ describe('McpClientManager', () => {
});
mockedMcpClient.getServerConfig.mockReturnValue(originalConfig);
const manager = new McpClientManager('0.0.1', toolRegistry, mockConfig);
const manager = setupManager(new McpClientManager('0.0.1', mockConfig));
await manager.startConfiguredMcpServers();
// First call should use the original config
@@ -321,9 +345,10 @@ describe('McpClientManager', () => {
(name, config) =>
({
connect: vi.fn(),
discover: vi.fn(),
discoverInto: vi.fn(),
disconnect: vi.fn(),
getServerConfig: vi.fn().mockReturnValue(config),
getServerName: vi.fn().mockReturnValue(name),
getInstructions: vi
.fn()
.mockReturnValue(
@@ -333,12 +358,7 @@ describe('McpClientManager', () => {
),
}) as unknown as McpClient,
);
const manager = new McpClientManager(
'0.0.1',
{} as ToolRegistry,
mockConfig,
);
const manager = new McpClientManager('0.0.1', mockConfig);
mockConfig.getMcpServers.mockReturnValue({
'server-with-instructions': { command: 'node' },
@@ -373,11 +393,7 @@ describe('McpClientManager', () => {
'test-server': { command: 'node' },
});
const manager = new McpClientManager(
'0.0.1',
{} as ToolRegistry,
mockConfig,
);
const manager = new McpClientManager('0.0.1', mockConfig);
await expect(manager.startConfiguredMcpServers()).resolves.not.toThrow();
});
@@ -396,11 +412,8 @@ describe('McpClientManager', () => {
'test-server': { command: 'node' },
});
const manager = new McpClientManager(
'0.0.1',
{} as ToolRegistry,
mockConfig,
);
const manager = new McpClientManager('0.0.1', mockConfig);
await manager.startConfiguredMcpServers();
await expect(manager.restartServer('test-server')).resolves.not.toThrow();
@@ -409,7 +422,7 @@ describe('McpClientManager', () => {
describe('Extension handling', () => {
it('should remove mcp servers from allServerConfigs when stopExtension is called', async () => {
const manager = new McpClientManager('0.0.1', toolRegistry, mockConfig);
const manager = setupManager(new McpClientManager('0.0.1', mockConfig));
const mcpServers = {
'test-server': { command: 'node', args: ['server.js'] },
};
@@ -431,7 +444,7 @@ describe('McpClientManager', () => {
});
it('should merge extension configuration with an existing user-configured server', async () => {
const manager = new McpClientManager('0.0.1', toolRegistry, mockConfig);
const manager = setupManager(new McpClientManager('0.0.1', mockConfig));
const userConfig = { command: 'node', args: ['user-server.js'] };
mockConfig.getMcpServers.mockReturnValue({
@@ -468,7 +481,7 @@ describe('McpClientManager', () => {
});
it('should securely merge tool lists and env variables regardless of load order', async () => {
const manager = new McpClientManager('0.0.1', toolRegistry, mockConfig);
const manager = setupManager(new McpClientManager('0.0.1', mockConfig));
const userConfig = {
excludeTools: ['user-tool'],
@@ -523,7 +536,7 @@ describe('McpClientManager', () => {
// Reset for Case 2
vi.mocked(McpClient).mockClear();
const manager2 = new McpClientManager('0.0.1', toolRegistry, mockConfig);
const manager2 = setupManager(new McpClientManager('0.0.1', mockConfig));
// Case 2: User config loads first, then Extension loads
// This call will skip discovery because userConfig has no connection details
@@ -551,7 +564,7 @@ describe('McpClientManager', () => {
});
it('should result in empty includeTools if intersection is empty', async () => {
const manager = new McpClientManager('0.0.1', toolRegistry, mockConfig);
const manager = setupManager(new McpClientManager('0.0.1', mockConfig));
const userConfig = { includeTools: ['user-tool'] };
const extConfig = {
command: 'node',
@@ -567,7 +580,7 @@ describe('McpClientManager', () => {
});
it('should respect a single allowlist if only one is provided', async () => {
const manager = new McpClientManager('0.0.1', toolRegistry, mockConfig);
const manager = setupManager(new McpClientManager('0.0.1', mockConfig));
const userConfig = { includeTools: ['user-tool'] };
const extConfig = { command: 'node', args: ['ext.js'] };
@@ -579,7 +592,7 @@ describe('McpClientManager', () => {
});
it('should allow partial overrides of connection properties', async () => {
const manager = new McpClientManager('0.0.1', toolRegistry, mockConfig);
const manager = setupManager(new McpClientManager('0.0.1', mockConfig));
const extConfig = { command: 'node', args: ['ext.js'], timeout: 1000 };
const userOverride = { args: ['overridden.js'] };
@@ -599,7 +612,7 @@ describe('McpClientManager', () => {
});
it('should prevent one extension from hijacking another extension server name', async () => {
const manager = new McpClientManager('0.0.1', toolRegistry, mockConfig);
const manager = setupManager(new McpClientManager('0.0.1', mockConfig));
const extension1: GeminiCLIExtension = {
name: 'extension-1',
@@ -641,7 +654,7 @@ describe('McpClientManager', () => {
it('should remove servers from blockedMcpServers when stopExtension is called', async () => {
mockConfig.getBlockedMcpServers.mockReturnValue(['blocked-server']);
const manager = new McpClientManager('0.0.1', toolRegistry, mockConfig);
const manager = setupManager(new McpClientManager('0.0.1', mockConfig));
const mcpServers = {
'blocked-server': { command: 'node', args: ['server.js'] },
};
@@ -679,7 +692,7 @@ describe('McpClientManager', () => {
});
it('should emit hint instead of full error when user has not interacted with MCP', () => {
const manager = new McpClientManager('0.0.1', toolRegistry, mockConfig);
const manager = setupManager(new McpClientManager('0.0.1', mockConfig));
manager.emitDiagnostic(
'error',
'Something went wrong',
@@ -698,7 +711,7 @@ describe('McpClientManager', () => {
});
it('should emit full error when user has interacted with MCP', () => {
const manager = new McpClientManager('0.0.1', toolRegistry, mockConfig);
const manager = setupManager(new McpClientManager('0.0.1', mockConfig));
manager.setUserInteractedWithMcp();
manager.emitDiagnostic(
'error',
@@ -714,7 +727,7 @@ describe('McpClientManager', () => {
});
it('should still deduplicate diagnostic messages after user interaction', () => {
const manager = new McpClientManager('0.0.1', toolRegistry, mockConfig);
const manager = setupManager(new McpClientManager('0.0.1', mockConfig));
manager.setUserInteractedWithMcp();
manager.emitDiagnostic('error', 'Same error');
@@ -724,7 +737,7 @@ describe('McpClientManager', () => {
});
it('should only show hint once per session', () => {
const manager = new McpClientManager('0.0.1', toolRegistry, mockConfig);
const manager = setupManager(new McpClientManager('0.0.1', mockConfig));
manager.emitDiagnostic('error', 'Error 1');
manager.emitDiagnostic('error', 'Error 2');
@@ -737,7 +750,7 @@ describe('McpClientManager', () => {
});
it('should capture last error for a server even when silenced', () => {
const manager = new McpClientManager('0.0.1', toolRegistry, mockConfig);
const manager = setupManager(new McpClientManager('0.0.1', mockConfig));
manager.emitDiagnostic(
'error',
@@ -752,7 +765,7 @@ describe('McpClientManager', () => {
});
it('should show previously deduplicated errors after interaction clears state', () => {
const manager = new McpClientManager('0.0.1', toolRegistry, mockConfig);
const manager = setupManager(new McpClientManager('0.0.1', mockConfig));
manager.emitDiagnostic('error', 'Same error');
expect(coreEventsMock.emitFeedback).toHaveBeenCalledTimes(1); // The hint
+122 -41
View File
@@ -13,6 +13,7 @@ import type { ToolRegistry } from './tool-registry.js';
import {
McpClient,
MCPDiscoveryState,
MCPServerStatus,
populateMcpServerCommand,
} from './mcp-client.js';
import { getErrorMessage, isAuthenticationError } from '../utils/errors.js';
@@ -20,6 +21,11 @@ import type { EventEmitter } from 'node:events';
import { coreEvents } from '../utils/events.js';
import { debugLogger } from '../utils/debugLogger.js';
import { createHash } from 'node:crypto';
import { stableStringify } from '../policy/stable-stringify.js';
import type { PromptRegistry } from '../prompts/prompt-registry.js';
import type { ResourceRegistry } from '../resources/resource-registry.js';
/**
* Manages the lifecycle of multiple MCP clients, including local child processes.
* This class is responsible for starting, stopping, and discovering tools from
@@ -30,7 +36,6 @@ export class McpClientManager {
// Track all configured servers (including disabled ones) for UI display
private allServerConfigs: Map<string, MCPServerConfig> = new Map();
private readonly clientVersion: string;
private readonly toolRegistry: ToolRegistry;
private readonly cliConfig: Config;
// If we have ongoing MCP client discovery, this completes once that is done.
private discoveryPromise: Promise<void> | undefined;
@@ -42,6 +47,10 @@ export class McpClientManager {
extensionName: string;
}> = [];
private mainToolRegistry: ToolRegistry | undefined;
private mainPromptRegistry: PromptRegistry | undefined;
private mainResourceRegistry: ResourceRegistry | undefined;
/**
* Track whether the user has explicitly interacted with MCP in this session
* (e.g. by running an /mcp command).
@@ -66,16 +75,24 @@ export class McpClientManager {
constructor(
clientVersion: string,
toolRegistry: ToolRegistry,
cliConfig: Config,
eventEmitter?: EventEmitter,
) {
this.clientVersion = clientVersion;
this.toolRegistry = toolRegistry;
this.cliConfig = cliConfig;
this.eventEmitter = eventEmitter;
}
setMainRegistries(registries: {
toolRegistry: ToolRegistry;
promptRegistry: PromptRegistry;
resourceRegistry: ResourceRegistry;
}) {
this.mainToolRegistry = registries.toolRegistry;
this.mainPromptRegistry = registries.promptRegistry;
this.mainResourceRegistry = registries.resourceRegistry;
}
setUserInteractedWithMcp() {
this.userInteractedWithMcp = true;
}
@@ -147,6 +164,16 @@ export class McpClientManager {
return this.clients.get(serverName);
}
removeRegistries(registries: {
toolRegistry: ToolRegistry;
promptRegistry: PromptRegistry;
resourceRegistry: ResourceRegistry;
}): void {
for (const client of this.clients.values()) {
client.removeRegistries(registries);
}
}
/**
* For all the MCP servers associated with this extension:
*
@@ -236,16 +263,17 @@ export class McpClientManager {
return false;
}
private async disconnectClient(name: string, skipRefresh = false) {
const existing = this.clients.get(name);
private async disconnectClient(clientKey: string, skipRefresh = false) {
const existing = this.clients.get(clientKey);
if (existing) {
const serverName = existing.getServerName();
try {
this.clients.delete(name);
this.clients.delete(clientKey);
this.eventEmitter?.emit('mcp-client-update', this.clients);
await existing.disconnect();
} catch (error) {
debugLogger.warn(
`Error stopping client '${name}': ${getErrorMessage(error)}`,
`Error stopping client '${serverName}': ${getErrorMessage(error)}`,
);
} finally {
if (!skipRefresh) {
@@ -257,6 +285,16 @@ export class McpClientManager {
}
}
private getClientKey(name: string, config: MCPServerConfig): string {
const { extension, ...rest } = config;
const keyData = {
name,
config: rest,
extensionId: extension?.id,
};
return createHash('sha256').update(stableStringify(keyData)).digest('hex');
}
/**
* Merges two MCP configurations. The second configuration (override)
* takes precedence for scalar properties, but array properties are
@@ -305,6 +343,11 @@ export class McpClientManager {
async maybeDiscoverMcpServer(
name: string,
config: MCPServerConfig,
registries?: {
toolRegistry: ToolRegistry;
promptRegistry: PromptRegistry;
resourceRegistry: ResourceRegistry;
},
): Promise<void> {
const existingConfig = this.allServerConfigs.get(name);
if (
@@ -337,11 +380,27 @@ export class McpClientManager {
// Always track server config for UI display
this.allServerConfigs.set(name, finalConfig);
// Capture the existing client synchronously here before any asynchronous
// operations. This ensures that if multiple discovery turns happen
// concurrently, this turn only replaces/disconnects the client that was
// present when this specific configuration update request began.
const existing = this.clients.get(name);
const clientKey = this.getClientKey(name, finalConfig);
// If no registries are provided (main agent) and a server with this name already exists
// but with a different configuration, handle potential conflicts.
if (!registries) {
const existingSameName = Array.from(this.clients.values()).find(
(c) => c.getServerName() === name,
);
if (existingSameName) {
const existingConfigFromClient = existingSameName.getServerConfig();
const existingKey = this.getClientKey(name, existingConfigFromClient);
if (existingKey !== clientKey) {
// This is a configuration update (hot-reload).
// We should stop the old client before starting the new one.
await this.disconnectClient(existingKey, true);
}
}
}
const existing = this.clients.get(clientKey);
// If no connection details are provided, we can't discover this server.
// This often happens when a user provides only overrides (like excludeTools)
@@ -363,7 +422,7 @@ export class McpClientManager {
// User-disabled servers: disconnect if running, don't start
if (await this.isDisabledByUser(name)) {
if (existing) {
await this.disconnectClient(name);
await this.disconnectClient(clientKey);
}
return;
}
@@ -374,34 +433,48 @@ export class McpClientManager {
return;
}
const currentDiscoveryPromise = new Promise<void>((resolve, reject) => {
(async () => {
const currentDiscoveryPromise = new Promise<void>((resolve) => {
void (async () => {
try {
if (existing) {
this.clients.delete(name);
await existing.disconnect();
let client = existing;
if (!client) {
client = new McpClient(
name,
finalConfig,
this.cliConfig.getWorkspaceContext(),
this.cliConfig,
this.cliConfig.getDebugMode(),
this.clientVersion,
async () => {
debugLogger.log(
`🔔 Refreshing context for server '${name}'...`,
);
await this.scheduleMcpContextRefresh();
},
);
this.clients.set(clientKey, client);
this.eventEmitter?.emit('mcp-client-update', this.clients);
}
const client = new McpClient(
name,
finalConfig,
this.toolRegistry,
this.cliConfig.getPromptRegistry(),
this.cliConfig.getResourceRegistry(),
this.cliConfig.getWorkspaceContext(),
this.cliConfig,
this.cliConfig.getDebugMode(),
this.clientVersion,
async () => {
debugLogger.log(`🔔 Refreshing context for server '${name}'...`);
await this.scheduleMcpContextRefresh();
},
);
this.clients.set(name, client);
this.eventEmitter?.emit('mcp-client-update', this.clients);
const targetRegistries =
registries ??
(this.mainToolRegistry &&
this.mainPromptRegistry &&
this.mainResourceRegistry
? {
toolRegistry: this.mainToolRegistry,
promptRegistry: this.mainPromptRegistry,
resourceRegistry: this.mainResourceRegistry,
}
: undefined);
try {
await client.connect();
await client.discover(this.cliConfig);
if (client.getStatus() === MCPServerStatus.DISCONNECTED) {
await client.connect();
}
if (targetRegistries) {
await client.discoverInto(this.cliConfig, targetRegistries);
}
this.eventEmitter?.emit('mcp-client-update', this.clients);
} catch (error) {
this.eventEmitter?.emit('mcp-client-update', this.clients);
@@ -421,13 +494,13 @@ export class McpClientManager {
const errorMessage = getErrorMessage(error);
this.emitDiagnostic(
'error',
`Error initializing MCP server '${name}': ${errorMessage}`,
`Fatal error ensuring MCP server '${name}' is connected: ${errorMessage}`,
error,
);
} finally {
resolve();
}
})().catch(reject);
})();
});
if (this.discoveryPromise) {
@@ -510,6 +583,11 @@ export class McpClientManager {
* Restarts all MCP servers (including newly enabled ones).
*/
async restart(): Promise<void> {
const disconnectionPromises = Array.from(this.clients.keys()).map((key) =>
this.disconnectClient(key, true),
);
await Promise.all(disconnectionPromises);
await Promise.all(
Array.from(this.allServerConfigs.entries()).map(
async ([name, config]) => {
@@ -534,6 +612,8 @@ export class McpClientManager {
if (!config) {
throw new Error(`No MCP server registered with the name "${name}"`);
}
const clientKey = this.getClientKey(name, config);
await this.disconnectClient(clientKey, true);
await this.maybeDiscoverMcpServer(name, config);
await this.scheduleMcpContextRefresh();
}
@@ -578,11 +658,12 @@ export class McpClientManager {
getMcpInstructions(): string {
const instructions: string[] = [];
for (const [name, client] of this.clients) {
for (const client of this.clients.values()) {
const serverName = client.getServerName();
const clientInstructions = client.getInstructions();
if (clientInstructions) {
instructions.push(
`The following are instructions provided by the tool server '${name}':\n---[start of server instructions]---\n${clientInstructions}\n---[end of server instructions]---`,
`The following are instructions provided by the tool server '${serverName}':\n---[start of server instructions]---\n${clientInstructions}\n---[end of server instructions]---`,
);
}
}
+177 -134
View File
@@ -4,6 +4,7 @@
* SPDX-License-Identifier: Apache-2.0
*/
/* eslint-disable @typescript-eslint/no-explicit-any */
import * as ClientLib from '@modelcontextprotocol/sdk/client/index.js';
import { SSEClientTransport } from '@modelcontextprotocol/sdk/client/sse.js';
import * as SdkClientStdioLib from '@modelcontextprotocol/sdk/client/stdio.js';
@@ -160,16 +161,17 @@ describe('mcp-client', () => {
{
command: 'test-command',
},
mockedToolRegistry,
promptRegistry,
resourceRegistry,
workspaceContext,
MOCK_CONTEXT,
false,
'0.0.1',
);
await client.connect();
await client.discover(MOCK_CONTEXT);
await client.discoverInto(MOCK_CONTEXT, {
toolRegistry: mockedToolRegistry,
promptRegistry,
resourceRegistry,
});
expect(mockedClient.listTools).toHaveBeenCalledWith(
{},
expect.objectContaining({ timeout: 600000, progressReporter: client }),
@@ -244,16 +246,17 @@ describe('mcp-client', () => {
{
command: 'test-command',
},
mockedToolRegistry,
promptRegistry,
resourceRegistry,
workspaceContext,
MOCK_CONTEXT,
false,
'0.0.1',
);
await client.connect();
await client.discover(MOCK_CONTEXT);
await client.discoverInto(MOCK_CONTEXT, {
toolRegistry: mockedToolRegistry,
promptRegistry,
resourceRegistry,
});
expect(mockedToolRegistry.registerTool).toHaveBeenCalledTimes(2);
expect(consoleWarnSpy).not.toHaveBeenCalled();
consoleWarnSpy.mockRestore();
@@ -296,16 +299,19 @@ describe('mcp-client', () => {
{
command: 'test-command',
},
mockedToolRegistry,
promptRegistry,
resourceRegistry,
workspaceContext,
MOCK_CONTEXT,
false,
'0.0.1',
);
await client.connect();
await expect(client.discover(MOCK_CONTEXT)).rejects.toThrow('Test error');
await expect(
client.discoverInto(MOCK_CONTEXT, {
toolRegistry: mockedToolRegistry,
promptRegistry,
resourceRegistry,
}),
).rejects.toThrow('Test error');
expect(MOCK_CONTEXT.emitMcpDiagnostic).toHaveBeenCalledWith(
'error',
`Error discovering prompts from test-server: Test error`,
@@ -354,18 +360,19 @@ describe('mcp-client', () => {
{
command: 'test-command',
},
mockedToolRegistry,
promptRegistry,
resourceRegistry,
workspaceContext,
MOCK_CONTEXT,
false,
'0.0.1',
);
await client.connect();
await expect(client.discover(MOCK_CONTEXT)).rejects.toThrow(
'No prompts, tools, or resources found on the server.',
);
await expect(
client.discoverInto(MOCK_CONTEXT, {
toolRegistry: mockedToolRegistry,
promptRegistry,
resourceRegistry,
}),
).rejects.toThrow('No prompts, tools, or resources found on the server.');
});
it('should discover tools if server supports them', async () => {
@@ -417,16 +424,17 @@ describe('mcp-client', () => {
{
command: 'test-command',
},
mockedToolRegistry,
promptRegistry,
resourceRegistry,
workspaceContext,
MOCK_CONTEXT,
false,
'0.0.1',
);
await client.connect();
await client.discover(MOCK_CONTEXT);
await client.discoverInto(MOCK_CONTEXT, {
toolRegistry: mockedToolRegistry,
promptRegistry,
resourceRegistry,
});
expect(mockedToolRegistry.registerTool).toHaveBeenCalledOnce();
});
@@ -485,9 +493,6 @@ describe('mcp-client', () => {
const client = new McpClient(
'test-server',
{ command: 'test-command' },
mockedToolRegistry,
promptRegistry,
resourceRegistry,
workspaceContext,
MOCK_CONTEXT,
false,
@@ -495,7 +500,11 @@ describe('mcp-client', () => {
);
await client.connect();
await client.discover(mockConfig);
await client.discoverInto(mockConfig, {
toolRegistry: mockedToolRegistry,
promptRegistry,
resourceRegistry,
});
// Verify tool registration
expect(mockedToolRegistry.registerTool).toHaveBeenCalledOnce();
@@ -566,9 +575,6 @@ describe('mcp-client', () => {
const client = new McpClient(
'test-server',
{ command: 'test-command' },
mockedToolRegistry,
promptRegistry,
resourceRegistry,
workspaceContext,
MOCK_CONTEXT,
false,
@@ -576,7 +582,11 @@ describe('mcp-client', () => {
);
await client.connect();
await client.discover(mockConfig);
await client.discoverInto(mockConfig, {
toolRegistry: mockedToolRegistry,
promptRegistry,
resourceRegistry,
});
expect(mockedToolRegistry.registerTool).toHaveBeenCalledOnce();
expect(mockPolicyEngine.addRule).not.toHaveBeenCalled();
@@ -644,9 +654,6 @@ describe('mcp-client', () => {
const client = new McpClient(
'test-server',
{ command: 'test-command' },
mockedToolRegistry,
promptRegistry,
resourceRegistry,
workspaceContext,
MOCK_CONTEXT,
false,
@@ -654,7 +661,11 @@ describe('mcp-client', () => {
);
await client.connect();
await client.discover(mockConfig);
await client.discoverInto(mockConfig, {
toolRegistry: mockedToolRegistry,
promptRegistry,
resourceRegistry,
});
expect(mockedToolRegistry.registerTool).toHaveBeenCalledOnce();
@@ -733,16 +744,17 @@ describe('mcp-client', () => {
{
command: 'test-command',
},
mockedToolRegistry,
promptRegistry,
resourceRegistry,
workspaceContext,
MOCK_CONTEXT,
false,
'0.0.1',
);
await client.connect();
await client.discover(MOCK_CONTEXT);
await client.discoverInto(MOCK_CONTEXT, {
toolRegistry: mockedToolRegistry,
promptRegistry,
resourceRegistry,
});
expect(mockedToolRegistry.registerTool).toHaveBeenCalledOnce();
const registeredTool = vi.mocked(mockedToolRegistry.registerTool).mock
.calls[0][0];
@@ -818,16 +830,17 @@ describe('mcp-client', () => {
{
command: 'test-command',
},
mockedToolRegistry,
promptRegistry,
resourceRegistry,
workspaceContext,
MOCK_CONTEXT,
false,
'0.0.1',
);
await client.connect();
await client.discover(MOCK_CONTEXT);
await client.discoverInto(MOCK_CONTEXT, {
toolRegistry: mockedToolRegistry,
promptRegistry,
resourceRegistry,
});
expect(resourceRegistry.setResourcesForServer).toHaveBeenCalledWith(
'test-server',
[
@@ -907,16 +920,17 @@ describe('mcp-client', () => {
{
command: 'test-command',
},
mockedToolRegistry,
promptRegistry,
resourceRegistry,
workspaceContext,
MOCK_CONTEXT,
false,
'0.0.1',
);
await client.connect();
await client.discover(MOCK_CONTEXT);
await client.discoverInto(MOCK_CONTEXT, {
toolRegistry: mockedToolRegistry,
promptRegistry,
resourceRegistry,
});
expect(mockedClient.setNotificationHandler).toHaveBeenCalledTimes(2);
expect(resourceListHandler).toBeDefined();
@@ -996,16 +1010,17 @@ describe('mcp-client', () => {
{
command: 'test-command',
},
mockedToolRegistry,
promptRegistry,
resourceRegistry,
workspaceContext,
MOCK_CONTEXT,
false,
'0.0.1',
);
await client.connect();
await client.discover(MOCK_CONTEXT);
await client.discoverInto(MOCK_CONTEXT, {
toolRegistry: mockedToolRegistry,
promptRegistry,
resourceRegistry,
});
expect(mockedClient.setNotificationHandler).toHaveBeenCalledTimes(2);
expect(promptListHandler).toBeDefined();
@@ -1080,16 +1095,17 @@ describe('mcp-client', () => {
{
command: 'test-command',
},
mockedToolRegistry,
mockedPromptRegistry,
resourceRegistry,
workspaceContext,
MOCK_CONTEXT,
false,
'0.0.1',
);
await client.connect();
await client.discover(MOCK_CONTEXT);
await client.discoverInto(MOCK_CONTEXT, {
toolRegistry: mockedToolRegistry,
promptRegistry: mockedPromptRegistry,
resourceRegistry,
});
expect(mockedToolRegistry.registerTool).toHaveBeenCalledOnce();
expect(mockedPromptRegistry.registerPrompt).toHaveBeenCalledOnce();
@@ -1138,17 +1154,6 @@ describe('mcp-client', () => {
const client = new McpClient(
'test-server',
{ command: 'test-command' },
mockedToolRegistry,
{
getPromptsByServer: vi.fn().mockReturnValue([]),
registerPrompt: vi.fn(),
} as unknown as PromptRegistry,
{
getResourcesByServer: vi.fn().mockReturnValue([]),
registerResource: vi.fn(),
removeResourcesByServer: vi.fn(),
setResourcesForServer: vi.fn(),
} as unknown as ResourceRegistry,
workspaceContext,
MOCK_CONTEXT,
false,
@@ -1156,6 +1161,20 @@ describe('mcp-client', () => {
);
await client.connect();
// INJECTED REGISTRIES
(client as any).registeredRegistries?.add({
toolRegistry: mockedToolRegistry,
promptRegistry: {
getPromptsByServer: vi.fn().mockReturnValue([]),
registerPrompt: vi.fn(),
} as unknown as PromptRegistry,
resourceRegistry: {
getResourcesByServer: vi.fn().mockReturnValue([]),
registerResource: vi.fn(),
removeResourcesByServer: vi.fn(),
setResourcesForServer: vi.fn(),
} as unknown as ResourceRegistry,
});
expect(mockedClient.setNotificationHandler).toHaveBeenCalledWith(
ToolListChangedNotificationSchema,
@@ -1183,21 +1202,6 @@ describe('mcp-client', () => {
const client = new McpClient(
'test-server',
{ command: 'test-command' },
{
getToolsByServer: vi.fn().mockReturnValue([]),
registerTool: vi.fn(),
sortTools: vi.fn(),
} as unknown as ToolRegistry,
{
getPromptsByServer: vi.fn().mockReturnValue([]),
registerPrompt: vi.fn(),
} as unknown as PromptRegistry,
{
getResourcesByServer: vi.fn().mockReturnValue([]),
registerResource: vi.fn(),
removeResourcesByServer: vi.fn(),
setResourcesForServer: vi.fn(),
} as unknown as ResourceRegistry,
workspaceContext,
MOCK_CONTEXT,
false,
@@ -1205,6 +1209,24 @@ describe('mcp-client', () => {
);
await client.connect();
// INJECTED REGISTRIES
(client as any).registeredRegistries?.add({
toolRegistry: {
getToolsByServer: vi.fn().mockReturnValue([]),
registerTool: vi.fn(),
sortTools: vi.fn(),
} as unknown as ToolRegistry,
promptRegistry: {
getPromptsByServer: vi.fn().mockReturnValue([]),
registerPrompt: vi.fn(),
} as unknown as PromptRegistry,
resourceRegistry: {
getResourcesByServer: vi.fn().mockReturnValue([]),
registerResource: vi.fn(),
removeResourcesByServer: vi.fn(),
setResourcesForServer: vi.fn(),
} as unknown as ResourceRegistry,
});
// Should be called for ProgressNotificationSchema, even if no other capabilities
expect(mockedClient.setNotificationHandler).toHaveBeenCalled();
@@ -1234,21 +1256,6 @@ describe('mcp-client', () => {
const client = new McpClient(
'test-server',
{ command: 'test-command' },
{
getToolsByServer: vi.fn().mockReturnValue([]),
registerTool: vi.fn(),
sortTools: vi.fn(),
} as unknown as ToolRegistry,
{
getPromptsByServer: vi.fn().mockReturnValue([]),
registerPrompt: vi.fn(),
} as unknown as PromptRegistry,
{
getResourcesByServer: vi.fn().mockReturnValue([]),
registerResource: vi.fn(),
removeResourcesByServer: vi.fn(),
setResourcesForServer: vi.fn(),
} as unknown as ResourceRegistry,
workspaceContext,
MOCK_CONTEXT,
false,
@@ -1256,6 +1263,24 @@ describe('mcp-client', () => {
);
await client.connect();
// INJECTED REGISTRIES
(client as any).registeredRegistries?.add({
toolRegistry: {
getToolsByServer: vi.fn().mockReturnValue([]),
registerTool: vi.fn(),
sortTools: vi.fn(),
} as unknown as ToolRegistry,
promptRegistry: {
getPromptsByServer: vi.fn().mockReturnValue([]),
registerPrompt: vi.fn(),
} as unknown as PromptRegistry,
resourceRegistry: {
getResourcesByServer: vi.fn().mockReturnValue([]),
registerResource: vi.fn(),
removeResourcesByServer: vi.fn(),
setResourcesForServer: vi.fn(),
} as unknown as ResourceRegistry,
});
const toolUpdateCall =
mockedClient.setNotificationHandler.mock.calls.find(
@@ -1308,12 +1333,6 @@ describe('mcp-client', () => {
const client = new McpClient(
'test-server',
{ command: 'test-command' },
mockedToolRegistry,
{} as PromptRegistry,
{
removeMcpResourcesByServer: vi.fn(),
registerResource: vi.fn(),
} as unknown as ResourceRegistry,
workspaceContext,
MOCK_CONTEXT,
false,
@@ -1323,6 +1342,15 @@ describe('mcp-client', () => {
// 1. Connect (sets up listener)
await client.connect();
// INJECTED REGISTRIES
(client as any).registeredRegistries?.add({
toolRegistry: mockedToolRegistry,
promptRegistry: {} as PromptRegistry,
resourceRegistry: {
removeMcpResourcesByServer: vi.fn(),
registerResource: vi.fn(),
} as unknown as ResourceRegistry,
});
// 2. Extract the callback passed to setNotificationHandler for tools
const toolUpdateCall =
@@ -1388,9 +1416,6 @@ describe('mcp-client', () => {
const client = new McpClient(
'test-server',
{ command: 'test-command' },
mockedToolRegistry,
{} as PromptRegistry,
{} as ResourceRegistry,
workspaceContext,
MOCK_CONTEXT,
false,
@@ -1398,6 +1423,12 @@ describe('mcp-client', () => {
);
await client.connect();
// INJECTED REGISTRIES
(client as any).registeredRegistries?.add({
toolRegistry: mockedToolRegistry,
promptRegistry: {} as PromptRegistry,
resourceRegistry: {} as ResourceRegistry,
});
const toolUpdateCall =
mockedClient.setNotificationHandler.mock.calls.find(
@@ -1463,9 +1494,6 @@ describe('mcp-client', () => {
const clientA = new McpClient(
'server-A',
{ command: 'cmd-a' },
mockedToolRegistry,
{} as PromptRegistry,
{} as ResourceRegistry,
workspaceContext,
MOCK_CONTEXT,
false,
@@ -1476,9 +1504,6 @@ describe('mcp-client', () => {
const clientB = new McpClient(
'server-B',
{ command: 'cmd-b' },
mockedToolRegistry,
{} as PromptRegistry,
{} as ResourceRegistry,
workspaceContext,
MOCK_CONTEXT,
false,
@@ -1487,7 +1512,19 @@ describe('mcp-client', () => {
);
await clientA.connect();
// INJECTED REGISTRIES
(clientA as any).registeredRegistries?.add({
toolRegistry: mockedToolRegistry,
promptRegistry: {} as PromptRegistry,
resourceRegistry: {} as ResourceRegistry,
});
await clientB.connect();
// INJECTED REGISTRIES
(clientB as any).registeredRegistries?.add({
toolRegistry: mockedToolRegistry,
promptRegistry: {} as PromptRegistry,
resourceRegistry: {} as ResourceRegistry,
});
const toolUpdateCallA =
mockClientA.setNotificationHandler.mock.calls.find(
@@ -1572,18 +1609,6 @@ describe('mcp-client', () => {
'test-server',
// Set a very short timeout
{ command: 'test-command', timeout: 50 },
mockedToolRegistry,
{
getPromptsByServer: vi.fn().mockReturnValue([]),
registerPrompt: vi.fn(),
removePromptsByServer: vi.fn(),
} as unknown as PromptRegistry,
{
getResourcesByServer: vi.fn().mockReturnValue([]),
registerResource: vi.fn(),
removeResourcesByServer: vi.fn(),
setResourcesForServer: vi.fn(),
} as unknown as ResourceRegistry,
workspaceContext,
MOCK_CONTEXT,
false,
@@ -1591,6 +1616,21 @@ describe('mcp-client', () => {
);
await client.connect();
// INJECTED REGISTRIES
(client as any).registeredRegistries?.add({
toolRegistry: mockedToolRegistry,
promptRegistry: {
getPromptsByServer: vi.fn().mockReturnValue([]),
registerPrompt: vi.fn(),
removePromptsByServer: vi.fn(),
} as unknown as PromptRegistry,
resourceRegistry: {
getResourcesByServer: vi.fn().mockReturnValue([]),
registerResource: vi.fn(),
removeResourcesByServer: vi.fn(),
setResourcesForServer: vi.fn(),
} as unknown as ResourceRegistry,
});
const toolUpdateCall =
mockedClient.setNotificationHandler.mock.calls.find(
@@ -1648,18 +1688,6 @@ describe('mcp-client', () => {
const client = new McpClient(
'test-server',
{ command: 'test-command' },
mockedToolRegistry,
{
getPromptsByServer: vi.fn().mockReturnValue([]),
registerPrompt: vi.fn(),
removePromptsByServer: vi.fn(),
} as unknown as PromptRegistry,
{
getResourcesByServer: vi.fn().mockReturnValue([]),
registerResource: vi.fn(),
removeResourcesByServer: vi.fn(),
setResourcesForServer: vi.fn(),
} as unknown as ResourceRegistry,
workspaceContext,
MOCK_CONTEXT,
false,
@@ -1668,6 +1696,21 @@ describe('mcp-client', () => {
);
await client.connect();
// INJECTED REGISTRIES
(client as any).registeredRegistries?.add({
toolRegistry: mockedToolRegistry,
promptRegistry: {
getPromptsByServer: vi.fn().mockReturnValue([]),
registerPrompt: vi.fn(),
removePromptsByServer: vi.fn(),
} as unknown as PromptRegistry,
resourceRegistry: {
getResourcesByServer: vi.fn().mockReturnValue([]),
registerResource: vi.fn(),
removeResourcesByServer: vi.fn(),
setResourcesForServer: vi.fn(),
} as unknown as ResourceRegistry,
});
const toolUpdateCall =
mockedClient.setNotificationHandler.mock.calls.find(
+144 -95
View File
@@ -130,6 +130,12 @@ export interface McpProgressReporter {
unregisterProgressToken(token: string | number): void;
}
export interface RegistrySet {
toolRegistry: ToolRegistry;
promptRegistry: PromptRegistry;
resourceRegistry: ResourceRegistry;
}
/**
* A client for a single MCP server.
*
@@ -147,6 +153,8 @@ export class McpClient implements McpProgressReporter {
private isRefreshingPrompts: boolean = false;
private pendingPromptRefresh: boolean = false;
private readonly registeredRegistries = new Set<RegistrySet>();
/**
* Map of progress tokens to tool call IDs.
* This allows us to route progress notifications to the correct tool call.
@@ -156,9 +164,6 @@ export class McpClient implements McpProgressReporter {
constructor(
private readonly serverName: string,
private readonly serverConfig: MCPServerConfig,
private readonly toolRegistry: ToolRegistry,
private readonly promptRegistry: PromptRegistry,
private readonly resourceRegistry: ResourceRegistry,
private readonly workspaceContext: WorkspaceContext,
private readonly cliConfig: McpContext,
private readonly debugMode: boolean,
@@ -166,6 +171,10 @@ export class McpClient implements McpProgressReporter {
private readonly onContextUpdated?: (signal?: AbortSignal) => Promise<void>,
) {}
getServerName(): string {
return this.serverName;
}
/**
* Connects to the MCP server.
*/
@@ -210,27 +219,34 @@ export class McpClient implements McpProgressReporter {
}
/**
* Discovers tools and prompts from the MCP server.
* Discovers tools and prompts from the MCP server into the specified registries.
*/
async discover(cliConfig: McpContext): Promise<void> {
async discoverInto(
cliConfig: McpContext,
registries: RegistrySet,
): Promise<void> {
this.assertConnected();
this.registeredRegistries.add(registries);
const prompts = await this.fetchPrompts();
const tools = await this.discoverTools(cliConfig);
const tools = await this.discoverTools(
cliConfig,
registries.toolRegistry.getMessageBus(),
);
const resources = await this.discoverResources();
this.updateResourceRegistry(resources);
this.updateResourceRegistry(resources, registries.resourceRegistry);
if (prompts.length === 0 && tools.length === 0 && resources.length === 0) {
throw new Error('No prompts, tools, or resources found on the server.');
}
for (const prompt of prompts) {
this.promptRegistry.registerPrompt(prompt);
registries.promptRegistry.registerPrompt(prompt);
}
for (const tool of tools) {
this.toolRegistry.registerTool(tool);
registries.toolRegistry.registerTool(tool);
}
this.toolRegistry.sortTools();
registries.toolRegistry.sortTools();
// Validate MCP tool names in policy rules against discovered tools
try {
@@ -250,6 +266,14 @@ export class McpClient implements McpProgressReporter {
}
}
/**
* Unregisters registries so this client will no longer update them when it receives
* list_changed notifications from the server.
*/
removeRegistries(registries: RegistrySet): void {
this.registeredRegistries.delete(registries);
}
/**
* Disconnects from the MCP server.
*/
@@ -257,9 +281,11 @@ export class McpClient implements McpProgressReporter {
if (this.status !== MCPServerStatus.CONNECTED) {
return;
}
this.toolRegistry.removeMcpToolsByServer(this.serverName);
this.promptRegistry.removePromptsByServer(this.serverName);
this.resourceRegistry.removeResourcesByServer(this.serverName);
for (const registries of this.registeredRegistries) {
registries.toolRegistry.removeMcpToolsByServer(this.serverName);
registries.promptRegistry.removePromptsByServer(this.serverName);
registries.resourceRegistry.removeResourcesByServer(this.serverName);
}
this.updateStatus(MCPServerStatus.DISCONNECTING);
const client = this.client;
this.client = undefined;
@@ -294,6 +320,7 @@ export class McpClient implements McpProgressReporter {
private async discoverTools(
cliConfig: McpContext,
messageBus: MessageBus,
options?: { timeout?: number; signal?: AbortSignal },
): Promise<DiscoveredMCPTool[]> {
this.assertConnected();
@@ -302,7 +329,7 @@ export class McpClient implements McpProgressReporter {
this.serverConfig,
this.client!,
cliConfig,
this.toolRegistry.messageBus,
messageBus,
{
...(options ?? {
timeout: this.serverConfig.timeout ?? MCP_DEFAULT_TIMEOUT_MSEC,
@@ -329,8 +356,11 @@ export class McpClient implements McpProgressReporter {
return discoverResources(this.serverName, this.client!, this.cliConfig);
}
private updateResourceRegistry(resources: Resource[]): void {
this.resourceRegistry.setResourcesForServer(this.serverName, resources);
private updateResourceRegistry(
resources: Resource[],
resourceRegistry: ResourceRegistry,
): void {
resourceRegistry.setResourcesForServer(this.serverName, resources);
}
async readResource(
@@ -482,23 +512,32 @@ export class McpClient implements McpProgressReporter {
try {
newResources = await this.discoverResources();
// Verification Retry: If no resources are found or resources didn't change,
// wait briefly and try one more time. Some servers notify before they're fully ready.
const currentResources =
this.resourceRegistry.getResourcesByServer(this.serverName) || [];
const resourceMatch =
newResources.length === currentResources.length &&
newResources.every((nr: Resource) =>
currentResources.some((cr: MCPResource) => cr.uri === nr.uri),
);
for (const registries of this.registeredRegistries) {
// Verification Retry: If no resources are found or resources didn't change,
// wait briefly and try one more time. Some servers notify before they're fully ready.
const currentResources =
registries.resourceRegistry.getResourcesByServer(
this.serverName,
) || [];
const resourceMatch =
newResources.length === currentResources.length &&
newResources.every((nr: Resource) =>
currentResources.some((cr: MCPResource) => cr.uri === nr.uri),
);
if (resourceMatch && !this.pendingResourceRefresh) {
debugLogger.log(
`No resource changes detected for '${this.serverName}'. Retrying once in 500ms...`,
if (resourceMatch && !this.pendingResourceRefresh) {
debugLogger.log(
`No resource changes detected for '${this.serverName}'. Retrying once in 500ms...`,
);
const retryDelay = 500;
await new Promise((resolve) => setTimeout(resolve, retryDelay));
newResources = await this.discoverResources();
}
this.updateResourceRegistry(
newResources,
registries.resourceRegistry,
);
const retryDelay = 500;
await new Promise((resolve) => setTimeout(resolve, retryDelay));
newResources = await this.discoverResources();
}
} catch (err) {
debugLogger.error(
@@ -508,8 +547,6 @@ export class McpClient implements McpProgressReporter {
break;
}
this.updateResourceRegistry(newResources);
if (this.onContextUpdated) {
await this.onContextUpdated(abortController.signal);
}
@@ -575,30 +612,33 @@ export class McpClient implements McpProgressReporter {
signal: abortController.signal,
});
// Verification Retry: If no prompts are found or prompts didn't change,
// wait briefly and try one more time. Some servers notify before they're fully ready.
const currentPrompts =
this.promptRegistry.getPromptsByServer(this.serverName) || [];
const promptsMatch =
newPrompts.length === currentPrompts.length &&
newPrompts.every((np) =>
currentPrompts.some((cp) => cp.name === np.name),
);
for (const registries of this.registeredRegistries) {
// Verification Retry: If no prompts are found or prompts didn't change,
// wait briefly and try one more time. Some servers notify before they're fully ready.
const currentPrompts =
registries.promptRegistry.getPromptsByServer(this.serverName) ||
[];
const promptsMatch =
newPrompts.length === currentPrompts.length &&
newPrompts.every((np) =>
currentPrompts.some((cp) => cp.name === np.name),
);
if (promptsMatch && !this.pendingPromptRefresh) {
debugLogger.log(
`No prompt changes detected for '${this.serverName}'. Retrying once in 500ms...`,
);
const retryDelay = 500;
await new Promise((resolve) => setTimeout(resolve, retryDelay));
newPrompts = await this.fetchPrompts({
signal: abortController.signal,
});
}
if (promptsMatch && !this.pendingPromptRefresh) {
debugLogger.log(
`No prompt changes detected for '${this.serverName}'. Retrying once in 500ms...`,
);
const retryDelay = 500;
await new Promise((resolve) => setTimeout(resolve, retryDelay));
newPrompts = await this.fetchPrompts({
signal: abortController.signal,
});
}
this.promptRegistry.removePromptsByServer(this.serverName);
for (const prompt of newPrompts) {
this.promptRegistry.registerPrompt(prompt);
registries.promptRegistry.removePromptsByServer(this.serverName);
for (const prompt of newPrompts) {
registries.promptRegistry.registerPrompt(prompt);
}
}
} catch (err) {
debugLogger.error(
@@ -666,42 +706,58 @@ export class McpClient implements McpProgressReporter {
const abortController = new AbortController();
const timeoutId = setTimeout(() => abortController.abort(), timeoutMs);
let newTools;
try {
newTools = await this.discoverTools(this.cliConfig, {
signal: abortController.signal,
});
debugLogger.log(
`Refresh for '${this.serverName}' discovered ${newTools.length} tools.`,
);
// Verification Retry (Option 3): If no tools are found or tools didn't change,
// wait briefly and try one more time. Some servers notify before they're fully ready.
const currentTools =
this.toolRegistry.getToolsByServer(this.serverName) || [];
const toolNamesMatch =
newTools.length === currentTools.length &&
newTools.every((nt) =>
currentTools.some(
(ct) =>
ct.name === nt.name ||
(ct instanceof DiscoveredMCPTool &&
ct.serverToolName === nt.serverToolName),
),
for (const registries of this.registeredRegistries) {
let newTools = await this.discoverTools(
this.cliConfig,
registries.toolRegistry.getMessageBus(),
{
signal: abortController.signal,
},
);
debugLogger.log(
`Refresh for '${this.serverName}' discovered ${newTools.length} tools.`,
);
if (toolNamesMatch && !this.pendingToolRefresh) {
debugLogger.log(
`No tool changes detected for '${this.serverName}'. Retrying once in 500ms...`,
);
const retryDelay = 500;
await new Promise((resolve) => setTimeout(resolve, retryDelay));
newTools = await this.discoverTools(this.cliConfig, {
signal: abortController.signal,
});
debugLogger.log(
`Retry refresh for '${this.serverName}' discovered ${newTools.length} tools.`,
);
// Verification Retry (Option 3): If no tools are found or tools didn't change,
// wait briefly and try one more time. Some servers notify before they're fully ready.
const currentTools =
registries.toolRegistry.getToolsByServer(this.serverName) || [];
const toolNamesMatch =
newTools.length === currentTools.length &&
newTools.every((nt) =>
currentTools.some(
(ct) =>
ct.name === nt.name ||
(ct instanceof DiscoveredMCPTool &&
ct.serverToolName === nt.serverToolName),
),
);
if (toolNamesMatch && !this.pendingToolRefresh) {
debugLogger.log(
`No tool changes detected for '${this.serverName}'. Retrying once in 500ms...`,
);
const retryDelay = 500;
await new Promise((resolve) => setTimeout(resolve, retryDelay));
newTools = await this.discoverTools(
this.cliConfig,
registries.toolRegistry.getMessageBus(),
{
signal: abortController.signal,
},
);
debugLogger.log(
`Retry refresh for '${this.serverName}' discovered ${newTools.length} tools.`,
);
}
registries.toolRegistry.removeMcpToolsByServer(this.serverName);
for (const tool of newTools) {
registries.toolRegistry.registerTool(tool);
}
registries.toolRegistry.sortTools();
}
} catch (err) {
debugLogger.error(
@@ -711,13 +767,6 @@ export class McpClient implements McpProgressReporter {
break;
}
this.toolRegistry.removeMcpToolsByServer(this.serverName);
for (const tool of newTools) {
this.toolRegistry.registerTool(tool);
}
this.toolRegistry.sortTools();
if (this.onContextUpdated) {
await this.onContextUpdated(abortController.signal);
}
@@ -284,6 +284,26 @@ describe('ToolRegistry', () => {
});
});
describe('removeMcpToolsByServer', () => {
it('should remove all tools from a specific server', () => {
const serverName = 'test-server';
const mcpTool1 = createMCPTool(serverName, 'tool1', 'desc1');
const mcpTool2 = createMCPTool(serverName, 'tool2', 'desc2');
const otherTool = createMCPTool('other-server', 'tool3', 'desc3');
toolRegistry.registerTool(mcpTool1);
toolRegistry.registerTool(mcpTool2);
toolRegistry.registerTool(otherTool);
expect(toolRegistry.getToolsByServer(serverName)).toHaveLength(2);
toolRegistry.removeMcpToolsByServer(serverName);
expect(toolRegistry.getToolsByServer(serverName)).toHaveLength(0);
expect(toolRegistry.getToolsByServer('other-server')).toHaveLength(1);
});
});
describe('excluded tools', () => {
const simpleTool = new MockTool({
name: 'tool-a',
+21 -1
View File
@@ -223,10 +223,16 @@ export class ToolRegistry {
private allKnownTools: Map<string, AnyDeclarativeTool> = new Map();
private config: Config;
readonly messageBus: MessageBus;
private isMainRegistry: boolean;
constructor(config: Config, messageBus: MessageBus) {
constructor(
config: Config,
messageBus: MessageBus,
isMainRegistry: boolean = false,
) {
this.config = config;
this.messageBus = messageBus;
this.isMainRegistry = isMainRegistry;
}
getMessageBus(): MessageBus {
@@ -599,6 +605,10 @@ export class ToolRegistry {
const declarations: FunctionDeclaration[] = [];
const seenNames = new Set<string>();
const mainAgentTools = this.isMainRegistry
? this.config.getMainAgentTools()
: undefined;
this.getActiveTools().forEach((tool) => {
const toolName =
tool instanceof DiscoveredMCPTool
@@ -608,6 +618,16 @@ export class ToolRegistry {
if (seenNames.has(toolName)) {
return;
}
if (
mainAgentTools &&
!mainAgentTools.includes(toolName) &&
!mainAgentTools.includes(tool.constructor.name) &&
!mainAgentTools.some((t) => t.startsWith(`${tool.constructor.name}(`))
) {
return;
}
seenNames.add(toolName);
let schema = tool.getSchema(modelId);
+19
View File
@@ -427,6 +427,25 @@ export abstract class DeclarativeTool<
readonly extensionId?: string,
) {}
clone(messageBus?: MessageBus): this {
// Note: we cannot use structuredClone() here because it does not preserve
// prototype chains or handle non-serializable properties (like functions).
// eslint-disable-next-line @typescript-eslint/no-unsafe-type-assertion
const cloned = Object.assign(
// eslint-disable-next-line no-restricted-syntax
Object.create(Object.getPrototypeOf(this)),
this,
) as this;
if (messageBus) {
Object.defineProperty(cloned, 'messageBus', {
value: messageBus,
writable: false,
configurable: true,
});
}
return cloned;
}
get isReadOnly(): boolean {
return READ_ONLY_KINDS.includes(this.kind);
}