mirror of
https://github.com/google-gemini/gemini-cli.git
synced 2026-08-02 13:11:03 -07:00
Compare commits
21 Commits
| Author | SHA1 | Date | |
|---|---|---|---|
| 9ec4add2a2 | |||
| cdf848d5c7 | |||
| e76d684210 | |||
| e8d4b73d4d | |||
| da3d65599a | |||
| ef5627eece | |||
| b49b6907e8 | |||
| 6a40da8d97 | |||
| a8ef876296 | |||
| 54a9bce2b7 | |||
| 7586efcf49 | |||
| ee425228fe | |||
| 3bf0a5579a | |||
| 7a08a4fbd5 | |||
| 164abf9d9b | |||
| b132791cd2 | |||
| 5a020e7720 | |||
| eb5d22848c | |||
| 2e6c81e7ad | |||
| d4b7d358c5 | |||
| c68a2cb933 |
@@ -1157,11 +1157,6 @@ their corresponding top-level category object in your `settings.json` file.
|
||||
- **Default:** `true`
|
||||
- **Requires restart:** Yes
|
||||
|
||||
- **`experimental.plannerSubagent`** (boolean):
|
||||
- **Description:** Use the new planner subagent for plan mode.
|
||||
- **Default:** `false`
|
||||
- **Requires restart:** Yes
|
||||
|
||||
- **`experimental.enableAgents`** (boolean):
|
||||
- **Description:** Enable local and remote subagents. Warning: Experimental
|
||||
feature, uses YOLO mode for subagents
|
||||
|
||||
+17
-1
@@ -51,6 +51,7 @@ export default tseslint.config(
|
||||
'evals/**',
|
||||
'packages/test-utils/**',
|
||||
'.gemini/skills/**',
|
||||
'**/*.d.ts',
|
||||
],
|
||||
},
|
||||
eslint.configs.recommended,
|
||||
@@ -206,11 +207,26 @@ export default tseslint.config(
|
||||
{
|
||||
// Rules that only apply to product code
|
||||
files: ['packages/*/src/**/*.{ts,tsx}'],
|
||||
ignores: ['**/*.test.ts', '**/*.test.tsx'],
|
||||
ignores: ['**/*.test.ts', '**/*.test.tsx', 'packages/*/src/test-utils/**'],
|
||||
rules: {
|
||||
'@typescript-eslint/no-unsafe-type-assertion': 'error',
|
||||
'@typescript-eslint/no-unsafe-assignment': 'error',
|
||||
'@typescript-eslint/no-unsafe-return': 'error',
|
||||
'no-restricted-syntax': [
|
||||
'error',
|
||||
...commonRestrictedSyntaxRules,
|
||||
{
|
||||
selector:
|
||||
'CallExpression[callee.object.name="Object"][callee.property.name="create"]',
|
||||
message:
|
||||
'Avoid using Object.create() in product code. Use object spread {...obj}, explicit class instantiation, structuredClone(), or copy constructors instead.',
|
||||
},
|
||||
{
|
||||
selector: 'Identifier[name="Reflect"]',
|
||||
message:
|
||||
'Avoid using Reflect namespace in product code. Do not use reflection to make copies. Instead, use explicit object copying or cloning (structuredClone() for values, new instance/clone function for classes).',
|
||||
},
|
||||
],
|
||||
},
|
||||
},
|
||||
{
|
||||
|
||||
@@ -1,96 +0,0 @@
|
||||
/**
|
||||
* @license
|
||||
* Copyright 2026 Google LLC
|
||||
* SPDX-License-Identifier: Apache-2.0
|
||||
*/
|
||||
|
||||
import { describe, it, expect, beforeEach, afterEach } from 'vitest';
|
||||
import { TestRig } from './test-helper.js';
|
||||
import fs from 'node:fs';
|
||||
import path from 'node:path';
|
||||
|
||||
describe('Planner Subagent E2E', () => {
|
||||
let rig: TestRig;
|
||||
|
||||
beforeEach(() => {
|
||||
rig = new TestRig();
|
||||
});
|
||||
|
||||
afterEach(async () => await rig.cleanup());
|
||||
|
||||
// TODO(joshualitt): Re-enable after planner subagent is working end-to-end.
|
||||
it.skip('should enter plan mode, create a plan, and exit plan mode', async () => {
|
||||
// We'll use a real run to see if it works.
|
||||
// We need to enable the planner agent and the required tools.
|
||||
rig.setup('planner-subagent-e2e', {
|
||||
settings: {
|
||||
experimental: {
|
||||
plan: true,
|
||||
agents: true,
|
||||
},
|
||||
tools: {
|
||||
core: [
|
||||
'enter_plan_mode',
|
||||
'exit_plan_mode',
|
||||
'write_file',
|
||||
'read_file',
|
||||
'list_directory',
|
||||
],
|
||||
},
|
||||
// Add policy to allow enter_plan_mode without confirmation in yolo mode
|
||||
// Actually TestRig.run uses --approval-mode=yolo by default if not specified.
|
||||
// But enter_plan_mode and exit_plan_mode might still ask if not explicitly allowed.
|
||||
policy: {
|
||||
rules: [
|
||||
{ toolName: 'enter_plan_mode', decision: 'ALLOW' },
|
||||
{ toolName: 'exit_plan_mode', decision: 'ALLOW' },
|
||||
{ toolName: 'write_file', decision: 'ALLOW' },
|
||||
],
|
||||
},
|
||||
},
|
||||
});
|
||||
|
||||
rig.mkdir('plans');
|
||||
rig.createFile('hello.ts', 'console.log("hello");');
|
||||
|
||||
// Prompt the agent to create a plan for refactoring hello.ts
|
||||
// We expect:
|
||||
// 1. Main agent calls enter_plan_mode
|
||||
// 2. Planner subagent is launched
|
||||
// 3. Planner subagent reads hello.ts
|
||||
// 4. Planner subagent writes plans/refactor.md
|
||||
// 5. Planner subagent calls exit_plan_mode
|
||||
// 6. Main agent resumes
|
||||
|
||||
await rig.run({
|
||||
stdin:
|
||||
'Please create a refactoring plan for hello.ts to use a function. Put the plan in plans/refactor.md and then approve it.',
|
||||
approvalMode: 'yolo',
|
||||
timeout: 60000, // Planning can take a while
|
||||
});
|
||||
|
||||
// 1. Verify enter_plan_mode was called
|
||||
const enterPlanCalled = await rig.waitForToolCall('enter_plan_mode');
|
||||
expect(enterPlanCalled).toBe(true);
|
||||
|
||||
// 2. Verify write_file was called for the plan
|
||||
const writeFileCalled = await rig.waitForToolCall(
|
||||
'write_file',
|
||||
20000,
|
||||
(args) => {
|
||||
return args.includes('plans/refactor.md');
|
||||
},
|
||||
);
|
||||
expect(writeFileCalled).toBe(true);
|
||||
|
||||
// 3. Verify exit_plan_mode was called
|
||||
const exitPlanCalled = await rig.waitForToolCall('exit_plan_mode');
|
||||
expect(exitPlanCalled).toBe(true);
|
||||
|
||||
// 4. Verify the plan file exists on disk
|
||||
const planPath = path.join(rig.testDir!, 'plans/refactor.md');
|
||||
expect(fs.existsSync(planPath)).toBe(true);
|
||||
const planContent = fs.readFileSync(planPath, 'utf-8');
|
||||
expect(planContent).toContain('hello.ts');
|
||||
});
|
||||
});
|
||||
Generated
+1
-26
@@ -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"
|
||||
},
|
||||
|
||||
@@ -802,7 +802,6 @@ export async function loadCliConfig(
|
||||
extensionLoader: extensionManager,
|
||||
extensionRegistryURI,
|
||||
enableExtensionReloading: settings.experimental?.extensionReloading,
|
||||
plannerSubagent: settings.experimental?.plannerSubagent,
|
||||
enableAgents: settings.experimental?.enableAgents,
|
||||
plan: settings.experimental?.plan,
|
||||
tracker: settings.experimental?.taskTracker,
|
||||
|
||||
@@ -1833,15 +1833,6 @@ const SETTINGS_SCHEMA = {
|
||||
},
|
||||
},
|
||||
},
|
||||
plannerSubagent: {
|
||||
type: 'boolean',
|
||||
label: 'Planner Subagent',
|
||||
category: 'Experimental',
|
||||
requiresRestart: true,
|
||||
default: false,
|
||||
description: 'Use the new planner subagent for plan mode.',
|
||||
showInDialog: false,
|
||||
},
|
||||
enableAgents: {
|
||||
type: 'boolean',
|
||||
label: 'Enable Agents',
|
||||
|
||||
@@ -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 │ │
|
||||
│ │ │ │
|
||||
│ └────────────────────────────────────────────────────────────────────────────────────────────┘ │
|
||||
│ │
|
||||
|
||||
@@ -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,16 +57,26 @@ 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.
|
||||
// eslint-disable-next-line @typescript-eslint/no-unsafe-assignment
|
||||
// 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,
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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,
|
||||
};
|
||||
|
||||
@@ -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();
|
||||
|
||||
@@ -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,
|
||||
@@ -26,7 +28,6 @@ import {
|
||||
} from '../tools/mcp-tool.js';
|
||||
import { CompressionStatus } from '../core/turn.js';
|
||||
import { type ToolCallRequestInfo } from '../scheduler/types.js';
|
||||
import { type Message } from '../confirmation-bus/types.js';
|
||||
import { ChatCompressionService } from '../services/chatCompressionService.js';
|
||||
import { getDirectoryContextString } from '../utils/environmentContext.js';
|
||||
import { promptIdContext } from '../utils/promptIdContext.js';
|
||||
@@ -99,19 +100,27 @@ 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;
|
||||
private readonly parentCallId?: string;
|
||||
private readonly completionToolName: string;
|
||||
private hasFailedCompressionAttempt = false;
|
||||
|
||||
private usesCustomCompletionTool(): boolean {
|
||||
return this.completionToolName !== TASK_COMPLETE_TOOL_NAME;
|
||||
}
|
||||
|
||||
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;
|
||||
}
|
||||
|
||||
/**
|
||||
@@ -133,25 +142,29 @@ export class LocalAgentExecutor<TOutput extends z.ZodTypeAny> {
|
||||
const parentMessageBus = context.messageBus;
|
||||
|
||||
// Create an override object to inject the subagent name into tool confirmation requests
|
||||
// eslint-disable-next-line @typescript-eslint/no-unsafe-type-assertion
|
||||
const subagentMessageBus = Object.create(
|
||||
parentMessageBus,
|
||||
) as typeof parentMessageBus;
|
||||
subagentMessageBus.publish = async (message: Message) => {
|
||||
if (message.type === 'tool-confirmation-request') {
|
||||
return parentMessageBus.publish({
|
||||
...message,
|
||||
subagent: definition.name,
|
||||
});
|
||||
}
|
||||
return parentMessageBus.publish(message);
|
||||
};
|
||||
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(),
|
||||
@@ -167,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) => {
|
||||
@@ -242,10 +257,12 @@ export class LocalAgentExecutor<TOutput extends z.ZodTypeAny> {
|
||||
return new LocalAgentExecutor(
|
||||
definition,
|
||||
context,
|
||||
agentToolRegistry,
|
||||
parentPromptId,
|
||||
parentCallId,
|
||||
agentToolRegistry,
|
||||
agentPromptRegistry,
|
||||
agentResourceRegistry,
|
||||
onActivity,
|
||||
parentCallId,
|
||||
);
|
||||
}
|
||||
|
||||
@@ -258,19 +275,21 @@ 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;
|
||||
this.completionToolName =
|
||||
definition.runConfig.completionToolName || TASK_COMPLETE_TOOL_NAME;
|
||||
|
||||
const randomIdPart = Math.random().toString(36).slice(2, 8);
|
||||
// parentPromptId will be undefined if this agent is invoked directly
|
||||
@@ -316,7 +335,7 @@ export class LocalAgentExecutor<TOutput extends z.ZodTypeAny> {
|
||||
// If the model stops calling tools without calling complete_task, it's an error.
|
||||
if (functionCalls.length === 0) {
|
||||
this.emitActivity('ERROR', {
|
||||
error: `Agent stopped calling tools but did not call '${this.completionToolName}' to finalize the session.`,
|
||||
error: `Agent stopped calling tools but did not call '${TASK_COMPLETE_TOOL_NAME}' to finalize the session.`,
|
||||
context: 'protocol_violation',
|
||||
});
|
||||
return {
|
||||
@@ -381,10 +400,7 @@ export class LocalAgentExecutor<TOutput extends z.ZodTypeAny> {
|
||||
default:
|
||||
throw new Error(`Unknown terminate reason: ${reason}`);
|
||||
}
|
||||
const explainInterrupted = this.usesCustomCompletionTool()
|
||||
? ''
|
||||
: ' and explain that your investigation was interrupted';
|
||||
return `${explanation} You have one final chance to complete the task with a short grace period. You MUST call \`${this.completionToolName}\` immediately with your best answer${explainInterrupted}. Do not call any other tools.`;
|
||||
return `${explanation} You have one final chance to complete the task with a short grace period. You MUST call \`${TASK_COMPLETE_TOOL_NAME}\` immediately with your best answer and explain that your investigation was interrupted. Do not call any other tools.`;
|
||||
}
|
||||
|
||||
/**
|
||||
@@ -514,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),
|
||||
);
|
||||
|
||||
@@ -609,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 ===
|
||||
@@ -651,7 +676,7 @@ export class LocalAgentExecutor<TOutput extends z.ZodTypeAny> {
|
||||
// The finalResult was already set by executeTurn, but we re-emit just in case.
|
||||
finalResult =
|
||||
finalResult ||
|
||||
`Agent stopped calling tools but did not call '${this.completionToolName}'.`;
|
||||
`Agent stopped calling tools but did not call '${TASK_COMPLETE_TOOL_NAME}'.`;
|
||||
this.emitActivity('ERROR', {
|
||||
error: finalResult,
|
||||
context: 'protocol_violation',
|
||||
@@ -721,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,
|
||||
@@ -910,54 +935,6 @@ export class LocalAgentExecutor<TOutput extends z.ZodTypeAny> {
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Validates and processes the final output of the agent.
|
||||
*/
|
||||
private validateAndProcessOutput(outputValue: unknown): {
|
||||
submittedOutput: string;
|
||||
success: boolean;
|
||||
error?: string;
|
||||
} {
|
||||
const { outputConfig, processOutput } = this.definition;
|
||||
|
||||
if (outputConfig) {
|
||||
const validationResult = outputConfig.schema.safeParse(outputValue);
|
||||
|
||||
if (!validationResult.success) {
|
||||
return {
|
||||
submittedOutput: '',
|
||||
success: false,
|
||||
error: `Output validation failed: ${JSON.stringify(validationResult.error.flatten())}`,
|
||||
};
|
||||
}
|
||||
|
||||
// eslint-disable-next-line @typescript-eslint/no-unsafe-assignment
|
||||
const validatedOutput = validationResult.data;
|
||||
if (processOutput) {
|
||||
return {
|
||||
submittedOutput: processOutput(validatedOutput),
|
||||
success: true,
|
||||
};
|
||||
} else {
|
||||
return {
|
||||
submittedOutput:
|
||||
typeof validatedOutput === 'string'
|
||||
? validatedOutput
|
||||
: JSON.stringify(validatedOutput, null, 2),
|
||||
success: true,
|
||||
};
|
||||
}
|
||||
}
|
||||
|
||||
return {
|
||||
submittedOutput:
|
||||
typeof outputValue === 'string'
|
||||
? outputValue
|
||||
: JSON.stringify(outputValue, null, 2),
|
||||
success: true,
|
||||
};
|
||||
}
|
||||
|
||||
/**
|
||||
* Executes function calls requested by the model and returns the results.
|
||||
*
|
||||
@@ -976,7 +953,7 @@ export class LocalAgentExecutor<TOutput extends z.ZodTypeAny> {
|
||||
}> {
|
||||
const allowedToolNames = new Set(this.toolRegistry.getAllToolNames());
|
||||
// Always allow the completion tool
|
||||
allowedToolNames.add(this.completionToolName);
|
||||
allowedToolNames.add(TASK_COMPLETE_TOOL_NAME);
|
||||
|
||||
let submittedOutput: string | null = null;
|
||||
let taskCompleted = false;
|
||||
@@ -1017,10 +994,7 @@ export class LocalAgentExecutor<TOutput extends z.ZodTypeAny> {
|
||||
callId,
|
||||
});
|
||||
|
||||
if (
|
||||
toolName === TASK_COMPLETE_TOOL_NAME &&
|
||||
!this.usesCustomCompletionTool()
|
||||
) {
|
||||
if (toolName === TASK_COMPLETE_TOOL_NAME) {
|
||||
if (taskCompleted) {
|
||||
const error =
|
||||
'Task already marked complete in this turn. Ignoring duplicate call.';
|
||||
@@ -1044,13 +1018,13 @@ export class LocalAgentExecutor<TOutput extends z.ZodTypeAny> {
|
||||
|
||||
if (outputConfig) {
|
||||
const outputName = outputConfig.outputName;
|
||||
const outputValue = args[outputName];
|
||||
if (outputValue !== undefined) {
|
||||
const result = this.validateAndProcessOutput(outputValue);
|
||||
if (args[outputName] !== undefined) {
|
||||
const outputValue = args[outputName];
|
||||
const validationResult = outputConfig.schema.safeParse(outputValue);
|
||||
|
||||
if (!result.success) {
|
||||
if (!validationResult.success) {
|
||||
taskCompleted = false; // Validation failed, revoke completion
|
||||
const error = result.error!;
|
||||
const error = `Output validation failed: ${JSON.stringify(validationResult.error.flatten())}`;
|
||||
syncResults.set(callId, {
|
||||
functionResponse: {
|
||||
name: TASK_COMPLETE_TOOL_NAME,
|
||||
@@ -1066,7 +1040,16 @@ export class LocalAgentExecutor<TOutput extends z.ZodTypeAny> {
|
||||
continue;
|
||||
}
|
||||
|
||||
submittedOutput = result.submittedOutput;
|
||||
// eslint-disable-next-line @typescript-eslint/no-unsafe-assignment
|
||||
const validatedOutput = validationResult.data;
|
||||
if (this.definition.processOutput) {
|
||||
submittedOutput = this.definition.processOutput(validatedOutput);
|
||||
} else {
|
||||
submittedOutput =
|
||||
typeof outputValue === 'string'
|
||||
? outputValue
|
||||
: JSON.stringify(outputValue, null, 2);
|
||||
}
|
||||
syncResults.set(callId, {
|
||||
functionResponse: {
|
||||
name: TASK_COMPLETE_TOOL_NAME,
|
||||
@@ -1183,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,
|
||||
},
|
||||
@@ -1201,34 +1186,6 @@ export class LocalAgentExecutor<TOutput extends z.ZodTypeAny> {
|
||||
id: call.request.callId,
|
||||
output: call.response.resultDisplay,
|
||||
});
|
||||
|
||||
// Check if this was a custom completion tool and it signaled task completion
|
||||
if (
|
||||
toolName === this.completionToolName &&
|
||||
call.response.isTaskCompletion
|
||||
) {
|
||||
debugLogger.log(
|
||||
`[LocalAgentExecutor] Custom completion tool '${toolName}' signaled task completion.`,
|
||||
);
|
||||
const outputValue = call.response.data;
|
||||
const result = this.validateAndProcessOutput(outputValue);
|
||||
|
||||
if (result.success) {
|
||||
taskCompleted = true;
|
||||
submittedOutput = result.submittedOutput;
|
||||
} else {
|
||||
// If validation fails, we still mark it as completed but use the raw output
|
||||
// and maybe log a warning.
|
||||
debugLogger.warn(
|
||||
`[LocalAgentExecutor] Custom completion tool '${toolName}' returned data that failed output validation: ${result.error}`,
|
||||
);
|
||||
taskCompleted = true;
|
||||
submittedOutput =
|
||||
typeof outputValue === 'string'
|
||||
? outputValue
|
||||
: JSON.stringify(outputValue, null, 2);
|
||||
}
|
||||
}
|
||||
} else if (call.status === 'error') {
|
||||
this.emitActivity('ERROR', {
|
||||
context: 'tool_call',
|
||||
@@ -1298,45 +1255,43 @@ export class LocalAgentExecutor<TOutput extends z.ZodTypeAny> {
|
||||
toolsList.push(...this.toolRegistry.getFunctionDeclarations());
|
||||
}
|
||||
|
||||
// Always inject completion tool if it's complete_task.
|
||||
// If it's a custom tool, it should be in toolConfig.
|
||||
if (!this.usesCustomCompletionTool()) {
|
||||
const completeTool: FunctionDeclaration = {
|
||||
name: TASK_COMPLETE_TOOL_NAME,
|
||||
description: outputConfig
|
||||
? 'Call this tool to submit your final answer and complete the task. This is the ONLY way to finish.'
|
||||
: 'Call this tool to submit your final findings and complete the task. This is the ONLY way to finish.',
|
||||
parameters: {
|
||||
type: Type.OBJECT,
|
||||
properties: {},
|
||||
required: [],
|
||||
},
|
||||
// Always inject complete_task.
|
||||
// Configure its schema based on whether output is expected.
|
||||
const completeTool: FunctionDeclaration = {
|
||||
name: TASK_COMPLETE_TOOL_NAME,
|
||||
description: outputConfig
|
||||
? 'Call this tool to submit your final answer and complete the task. This is the ONLY way to finish.'
|
||||
: 'Call this tool to submit your final findings and complete the task. This is the ONLY way to finish.',
|
||||
parameters: {
|
||||
type: Type.OBJECT,
|
||||
properties: {},
|
||||
required: [],
|
||||
},
|
||||
};
|
||||
|
||||
if (outputConfig) {
|
||||
const jsonSchema = zodToJsonSchema(outputConfig.schema);
|
||||
const {
|
||||
$schema: _$schema,
|
||||
definitions: _definitions,
|
||||
...schema
|
||||
} = jsonSchema;
|
||||
completeTool.parameters!.properties![outputConfig.outputName] =
|
||||
// eslint-disable-next-line @typescript-eslint/no-unsafe-type-assertion
|
||||
schema as Schema;
|
||||
completeTool.parameters!.required!.push(outputConfig.outputName);
|
||||
} else {
|
||||
completeTool.parameters!.properties!['result'] = {
|
||||
type: Type.STRING,
|
||||
description:
|
||||
'Your final results or findings to return to the orchestrator. ' +
|
||||
'Ensure this is comprehensive and follows any formatting requested in your instructions.',
|
||||
};
|
||||
|
||||
if (outputConfig) {
|
||||
const jsonSchema = zodToJsonSchema(outputConfig.schema);
|
||||
const {
|
||||
$schema: _$schema,
|
||||
definitions: _definitions,
|
||||
...schema
|
||||
} = jsonSchema;
|
||||
completeTool.parameters!.properties![outputConfig.outputName] =
|
||||
// eslint-disable-next-line @typescript-eslint/no-unsafe-type-assertion
|
||||
schema as Schema;
|
||||
completeTool.parameters!.required!.push(outputConfig.outputName);
|
||||
} else {
|
||||
completeTool.parameters!.properties!['result'] = {
|
||||
type: Type.STRING,
|
||||
description:
|
||||
'Your final results or findings to return to the orchestrator. ' +
|
||||
'Ensure this is comprehensive and follows any formatting requested in your instructions.',
|
||||
};
|
||||
completeTool.parameters!.required!.push('result');
|
||||
}
|
||||
|
||||
toolsList.push(completeTool);
|
||||
completeTool.parameters!.required!.push('result');
|
||||
}
|
||||
|
||||
toolsList.push(completeTool);
|
||||
|
||||
return toolsList;
|
||||
}
|
||||
|
||||
@@ -1361,12 +1316,7 @@ Important Rules:
|
||||
* Work systematically using available tools to complete your task.
|
||||
* Always use absolute paths for file operations. Construct them using the provided "Environment Context".`;
|
||||
|
||||
if (this.usesCustomCompletionTool()) {
|
||||
finalPrompt += `
|
||||
* When you have completed your task, you MUST call the \`${this.completionToolName}\` tool to signal completion.
|
||||
* Do not call any other tools in the same turn as \`${this.completionToolName}\`.
|
||||
* This is the ONLY way to complete your mission. If you stop calling tools without calling this, you have failed.`;
|
||||
} else if (this.definition.outputConfig) {
|
||||
if (this.definition.outputConfig) {
|
||||
finalPrompt += `
|
||||
* When you have completed your task, you MUST call the \`${TASK_COMPLETE_TOOL_NAME}\` tool with your structured output.
|
||||
* Do not call any other tools in the same turn as \`${TASK_COMPLETE_TOOL_NAME}\`.
|
||||
|
||||
@@ -1,15 +0,0 @@
|
||||
/**
|
||||
* @license
|
||||
* Copyright 2026 Google LLC
|
||||
* SPDX-License-Identifier: Apache-2.0
|
||||
*/
|
||||
|
||||
import { z } from 'zod';
|
||||
|
||||
export const PlannerAgentSchema = z.object({
|
||||
plan_path: z
|
||||
.string()
|
||||
.describe('The path to the finalized and approved plan.'),
|
||||
});
|
||||
|
||||
export type PlannerAgentOutput = z.infer<typeof PlannerAgentSchema>;
|
||||
@@ -1,86 +0,0 @@
|
||||
/**
|
||||
* @license
|
||||
* Copyright 2026 Google LLC
|
||||
* SPDX-License-Identifier: Apache-2.0
|
||||
*/
|
||||
|
||||
import type { Config } from '../config/config.js';
|
||||
import { getCoreSystemPrompt } from '../core/prompts.js';
|
||||
import type { LocalAgentDefinition } from './types.js';
|
||||
import { PlannerAgentSchema } from './planner-schema.js';
|
||||
import {
|
||||
GLOB_TOOL_NAME,
|
||||
GREP_TOOL_NAME,
|
||||
LS_TOOL_NAME,
|
||||
READ_FILE_TOOL_NAME,
|
||||
ASK_USER_TOOL_NAME,
|
||||
EXIT_PLAN_MODE_TOOL_NAME,
|
||||
WRITE_FILE_TOOL_NAME,
|
||||
WEB_SEARCH_TOOL_NAME,
|
||||
WEB_FETCH_TOOL_NAME,
|
||||
} from '../tools/tool-names.js';
|
||||
|
||||
/**
|
||||
* A specialized subagent for research and planning.
|
||||
* It operates with read-only tools (mostly) and is restricted to writing plans
|
||||
* until user approval is received.
|
||||
*/
|
||||
export const PlannerAgent = (
|
||||
config: Config,
|
||||
): LocalAgentDefinition<typeof PlannerAgentSchema> => ({
|
||||
kind: 'local',
|
||||
name: 'planner',
|
||||
displayName: 'Planner Agent',
|
||||
description:
|
||||
'A specialized subagent for research and planning. It explores the codebase, designs solutions, and drafts detailed implementation plans for user approval.',
|
||||
inputConfig: {
|
||||
inputSchema: {
|
||||
type: 'object',
|
||||
properties: {
|
||||
reason: {
|
||||
type: 'string',
|
||||
description: 'The reason for entering plan mode or the task to plan.',
|
||||
},
|
||||
},
|
||||
required: ['reason'],
|
||||
},
|
||||
},
|
||||
outputConfig: {
|
||||
outputName: 'plan_path',
|
||||
description: 'The path to the finalized and approved plan.',
|
||||
schema: PlannerAgentSchema,
|
||||
},
|
||||
processOutput: (output) => output.plan_path,
|
||||
modelConfig: {
|
||||
model: 'inherit',
|
||||
},
|
||||
runConfig: {
|
||||
maxTimeMinutes: 10,
|
||||
maxTurns: 20,
|
||||
completionToolName: EXIT_PLAN_MODE_TOOL_NAME,
|
||||
hasCustomEntryPoint: true,
|
||||
},
|
||||
toolConfig: {
|
||||
tools: [
|
||||
LS_TOOL_NAME,
|
||||
READ_FILE_TOOL_NAME,
|
||||
GLOB_TOOL_NAME,
|
||||
GREP_TOOL_NAME,
|
||||
ASK_USER_TOOL_NAME,
|
||||
EXIT_PLAN_MODE_TOOL_NAME,
|
||||
WRITE_FILE_TOOL_NAME,
|
||||
WEB_SEARCH_TOOL_NAME,
|
||||
WEB_FETCH_TOOL_NAME,
|
||||
],
|
||||
},
|
||||
promptConfig: {
|
||||
get systemPrompt() {
|
||||
return getCoreSystemPrompt(
|
||||
config,
|
||||
undefined, // userMemory
|
||||
false, // interactiveOverride
|
||||
);
|
||||
},
|
||||
query: 'Start planning for: ${reason}',
|
||||
},
|
||||
});
|
||||
@@ -274,7 +274,6 @@ describe('AgentRegistry', () => {
|
||||
codebase_investigator: { enabled: false },
|
||||
cli_help: { enabled: false },
|
||||
generalist: { enabled: false },
|
||||
planner: { enabled: false },
|
||||
},
|
||||
},
|
||||
});
|
||||
@@ -336,15 +335,6 @@ describe('AgentRegistry', () => {
|
||||
expect(registry.getDefinition('generalist')).toBeDefined();
|
||||
});
|
||||
|
||||
it('should NOT register planner agent if plannerSubagent is false', async () => {
|
||||
const config = makeMockedConfig({ plannerSubagent: false });
|
||||
const registry = new TestableAgentRegistry(config);
|
||||
|
||||
await registry.initialize();
|
||||
|
||||
expect(registry.getDefinition('planner')).toBeUndefined();
|
||||
});
|
||||
|
||||
it('should NOT register a non-experimental agent if enabled is false', async () => {
|
||||
// CLI help is NOT experimental, but we explicitly disable it via enabled: false
|
||||
const config = makeMockedConfig({
|
||||
|
||||
@@ -12,7 +12,6 @@ import { loadAgentsFromDirectory } from './agentLoader.js';
|
||||
import { CodebaseInvestigatorAgent } from './codebase-investigator.js';
|
||||
import { CliHelpAgent } from './cli-help-agent.js';
|
||||
import { GeneralistAgent } from './generalist-agent.js';
|
||||
import { PlannerAgent } from './planner.js';
|
||||
import { BrowserAgentDefinition } from './browser/browserAgentDefinition.js';
|
||||
import { A2AClientManager } from './a2a-client-manager.js';
|
||||
import { A2AAuthProviderFactory } from './auth-provider/factory.js';
|
||||
@@ -244,9 +243,6 @@ export class AgentRegistry {
|
||||
this.registerLocalAgent(CodebaseInvestigatorAgent(this.config));
|
||||
this.registerLocalAgent(CliHelpAgent(this.config));
|
||||
this.registerLocalAgent(GeneralistAgent(this.config));
|
||||
if (this.config.isPlannerSubagentEnabled()) {
|
||||
this.registerLocalAgent(PlannerAgent(this.config));
|
||||
}
|
||||
|
||||
// Register the browser agent if enabled in settings.
|
||||
// Tools are configured dynamically at invocation time via browserAgentFactory.
|
||||
@@ -524,22 +520,67 @@ export class AgentRegistry {
|
||||
return definition;
|
||||
}
|
||||
|
||||
// Use Object.create to preserve lazy getters on the definition object
|
||||
// eslint-disable-next-line @typescript-eslint/no-unsafe-assignment
|
||||
const merged: LocalAgentDefinition<TOutput> = Object.create(definition);
|
||||
// Preserve lazy getters on the definition object by wrapping in a new object with getters
|
||||
const merged: LocalAgentDefinition<TOutput> = {
|
||||
get kind() {
|
||||
return definition.kind;
|
||||
},
|
||||
get name() {
|
||||
return definition.name;
|
||||
},
|
||||
get displayName() {
|
||||
return definition.displayName;
|
||||
},
|
||||
get description() {
|
||||
return definition.description;
|
||||
},
|
||||
get experimental() {
|
||||
return definition.experimental;
|
||||
},
|
||||
get metadata() {
|
||||
return definition.metadata;
|
||||
},
|
||||
get inputConfig() {
|
||||
return definition.inputConfig;
|
||||
},
|
||||
get outputConfig() {
|
||||
return definition.outputConfig;
|
||||
},
|
||||
get promptConfig() {
|
||||
return definition.promptConfig;
|
||||
},
|
||||
get toolConfig() {
|
||||
return definition.toolConfig;
|
||||
},
|
||||
get processOutput() {
|
||||
return definition.processOutput;
|
||||
},
|
||||
get runConfig() {
|
||||
return overrides.runConfig
|
||||
? { ...definition.runConfig, ...overrides.runConfig }
|
||||
: definition.runConfig;
|
||||
},
|
||||
get modelConfig() {
|
||||
return overrides.modelConfig
|
||||
? ModelConfigService.merge(
|
||||
definition.modelConfig,
|
||||
overrides.modelConfig,
|
||||
)
|
||||
: definition.modelConfig;
|
||||
},
|
||||
};
|
||||
|
||||
if (overrides.runConfig) {
|
||||
merged.runConfig = {
|
||||
...definition.runConfig,
|
||||
...overrides.runConfig,
|
||||
if (overrides.tools) {
|
||||
merged.toolConfig = {
|
||||
tools: overrides.tools,
|
||||
};
|
||||
}
|
||||
|
||||
if (overrides.modelConfig) {
|
||||
merged.modelConfig = ModelConfigService.merge(
|
||||
definition.modelConfig,
|
||||
overrides.modelConfig,
|
||||
);
|
||||
if (overrides.mcpServers) {
|
||||
merged.mcpServers = {
|
||||
...definition.mcpServers,
|
||||
...overrides.mcpServers,
|
||||
};
|
||||
}
|
||||
|
||||
return merged;
|
||||
|
||||
@@ -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.
|
||||
@@ -231,14 +237,4 @@ export interface RunConfig {
|
||||
* If not specified, defaults to DEFAULT_MAX_TURNS (30).
|
||||
*/
|
||||
maxTurns?: number;
|
||||
/**
|
||||
* The name of the tool that signals task completion.
|
||||
* Defaults to 'complete_task'.
|
||||
*/
|
||||
completionToolName?: string;
|
||||
/**
|
||||
* Whether or not this subagent has a custom entry point.
|
||||
* Defaults to `false`, and wraps subagents in a type-safe entry-point.
|
||||
*/
|
||||
hasCustomEntryPoint?: boolean;
|
||||
}
|
||||
|
||||
@@ -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,
|
||||
);
|
||||
|
||||
|
||||
@@ -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 */
|
||||
@@ -618,7 +621,6 @@ export interface ConfigParameters {
|
||||
disabledHooks?: string[];
|
||||
projectHooks?: { [K in HookEventName]?: HookDefinition[] };
|
||||
enableAgents?: boolean;
|
||||
plannerSubagent?: boolean;
|
||||
enableEventDrivenScheduler?: boolean;
|
||||
skillsSupport?: boolean;
|
||||
disabledSkills?: string[];
|
||||
@@ -676,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 */
|
||||
@@ -838,7 +841,6 @@ export class Config implements McpContext, AgentLoopContext {
|
||||
overageStrategy: OverageStrategy;
|
||||
};
|
||||
|
||||
private readonly plannerSubagent: boolean;
|
||||
private readonly enableAgents: boolean;
|
||||
private agents: AgentSettings;
|
||||
private readonly enableEventDrivenScheduler: boolean;
|
||||
@@ -890,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;
|
||||
@@ -950,7 +953,6 @@ export class Config implements McpContext, AgentLoopContext {
|
||||
this.model = params.model;
|
||||
this.disableLoopDetection = params.disableLoopDetection ?? false;
|
||||
this._activeModel = params.model;
|
||||
this.plannerSubagent = params.plannerSubagent ?? false;
|
||||
this.enableAgents = params.enableAgents ?? false;
|
||||
this.agents = params.agents ?? {};
|
||||
this.disableLLMCorrection = params.disableLLMCorrection ?? true;
|
||||
@@ -1234,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([
|
||||
@@ -1890,6 +1896,10 @@ export class Config implements McpContext, AgentLoopContext {
|
||||
return this.coreTools;
|
||||
}
|
||||
|
||||
getMainAgentTools(): string[] | undefined {
|
||||
return this.mainAgentTools;
|
||||
}
|
||||
|
||||
getAllowedTools(): string[] | undefined {
|
||||
return this.allowedTools;
|
||||
}
|
||||
@@ -2484,10 +2494,6 @@ export class Config implements McpContext, AgentLoopContext {
|
||||
this.approvedPlanPath = path;
|
||||
}
|
||||
|
||||
isPlannerSubagentEnabled(): boolean {
|
||||
return this.plannerSubagent && this.isPlanEnabled();
|
||||
}
|
||||
|
||||
isAgentsEnabled(): boolean {
|
||||
return this.enableAgents;
|
||||
}
|
||||
@@ -2989,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 = (
|
||||
@@ -3135,13 +3145,7 @@ export class Config implements McpContext, AgentLoopContext {
|
||||
|
||||
for (const definition of definitions) {
|
||||
try {
|
||||
if (
|
||||
definition.kind === 'local' &&
|
||||
definition.runConfig.hasCustomEntryPoint === true
|
||||
) {
|
||||
continue;
|
||||
}
|
||||
const tool = new SubagentTool(definition, this, this._messageBus);
|
||||
const tool = new SubagentTool(definition, this, this.messageBus);
|
||||
registry.registerTool(tool);
|
||||
} catch (e: unknown) {
|
||||
debugLogger.warn(
|
||||
|
||||
@@ -262,4 +262,90 @@ describe('MessageBus', () => {
|
||||
);
|
||||
});
|
||||
});
|
||||
|
||||
describe('derive', () => {
|
||||
it('should receive responses from parent bus on derived bus', async () => {
|
||||
vi.spyOn(policyEngine, 'check').mockResolvedValue({
|
||||
decision: PolicyDecision.ASK_USER,
|
||||
});
|
||||
|
||||
const subagentName = 'test-subagent';
|
||||
const subagentBus = messageBus.derive(subagentName);
|
||||
|
||||
const request: Omit<ToolConfirmationRequest, 'correlationId'> = {
|
||||
type: MessageBusType.TOOL_CONFIRMATION_REQUEST,
|
||||
toolCall: { name: 'test-tool', args: {} },
|
||||
};
|
||||
|
||||
const requestPromise = subagentBus.request<
|
||||
ToolConfirmationRequest,
|
||||
ToolConfirmationResponse
|
||||
>(request, MessageBusType.TOOL_CONFIRMATION_RESPONSE, 2000);
|
||||
|
||||
// Wait for request on root bus and respond
|
||||
await new Promise<void>((resolve) => {
|
||||
messageBus.subscribe<ToolConfirmationRequest>(
|
||||
MessageBusType.TOOL_CONFIRMATION_REQUEST,
|
||||
(msg) => {
|
||||
if (msg.subagent === subagentName) {
|
||||
void messageBus.publish({
|
||||
type: MessageBusType.TOOL_CONFIRMATION_RESPONSE,
|
||||
correlationId: msg.correlationId,
|
||||
confirmed: true,
|
||||
});
|
||||
resolve();
|
||||
}
|
||||
},
|
||||
);
|
||||
});
|
||||
|
||||
await expect(requestPromise).resolves.toEqual(
|
||||
expect.objectContaining({
|
||||
type: MessageBusType.TOOL_CONFIRMATION_RESPONSE,
|
||||
confirmed: true,
|
||||
}),
|
||||
);
|
||||
});
|
||||
|
||||
it('should correctly chain subagent names for nested subagents', async () => {
|
||||
vi.spyOn(policyEngine, 'check').mockResolvedValue({
|
||||
decision: PolicyDecision.ASK_USER,
|
||||
});
|
||||
|
||||
const subagentBus1 = messageBus.derive('agent1');
|
||||
const subagentBus2 = subagentBus1.derive('agent2');
|
||||
|
||||
const request: Omit<ToolConfirmationRequest, 'correlationId'> = {
|
||||
type: MessageBusType.TOOL_CONFIRMATION_REQUEST,
|
||||
toolCall: { name: 'test-tool', args: {} },
|
||||
};
|
||||
|
||||
const requestPromise = subagentBus2.request<
|
||||
ToolConfirmationRequest,
|
||||
ToolConfirmationResponse
|
||||
>(request, MessageBusType.TOOL_CONFIRMATION_RESPONSE, 2000);
|
||||
|
||||
await new Promise<void>((resolve) => {
|
||||
messageBus.subscribe<ToolConfirmationRequest>(
|
||||
MessageBusType.TOOL_CONFIRMATION_REQUEST,
|
||||
(msg) => {
|
||||
if (msg.subagent === 'agent1/agent2') {
|
||||
void messageBus.publish({
|
||||
type: MessageBusType.TOOL_CONFIRMATION_RESPONSE,
|
||||
correlationId: msg.correlationId,
|
||||
confirmed: true,
|
||||
});
|
||||
resolve();
|
||||
}
|
||||
},
|
||||
);
|
||||
});
|
||||
|
||||
await expect(requestPromise).resolves.toEqual(
|
||||
expect.objectContaining({
|
||||
confirmed: true,
|
||||
}),
|
||||
);
|
||||
});
|
||||
});
|
||||
});
|
||||
|
||||
@@ -40,6 +40,37 @@ export class MessageBus extends EventEmitter {
|
||||
this.emit(message.type, message);
|
||||
}
|
||||
|
||||
/**
|
||||
* Derives a child message bus scoped to a specific subagent.
|
||||
*/
|
||||
derive(subagentName: string): MessageBus {
|
||||
const bus = new MessageBus(this.policyEngine, this.debug);
|
||||
|
||||
bus.publish = async (message: Message) => {
|
||||
if (message.type === MessageBusType.TOOL_CONFIRMATION_REQUEST) {
|
||||
return this.publish({
|
||||
...message,
|
||||
subagent: message.subagent
|
||||
? `${subagentName}/${message.subagent}`
|
||||
: subagentName,
|
||||
});
|
||||
}
|
||||
return this.publish(message);
|
||||
};
|
||||
|
||||
// Delegate subscription methods to the parent bus
|
||||
bus.subscribe = this.subscribe.bind(this);
|
||||
bus.unsubscribe = this.unsubscribe.bind(this);
|
||||
bus.on = this.on.bind(this);
|
||||
bus.off = this.off.bind(this);
|
||||
bus.emit = this.emit.bind(this);
|
||||
bus.once = this.once.bind(this);
|
||||
bus.removeListener = this.removeListener.bind(this);
|
||||
bus.listenerCount = this.listenerCount.bind(this);
|
||||
|
||||
return bus;
|
||||
}
|
||||
|
||||
async publish(message: Message): Promise<void> {
|
||||
if (this.debug) {
|
||||
debugLogger.debug(`[MESSAGE_BUS] publish: ${safeJsonStringify(message)}`);
|
||||
|
||||
@@ -364,7 +364,6 @@ export class ToolExecutor {
|
||||
errorType: undefined,
|
||||
outputFile,
|
||||
contentLength: typeof content === 'string' ? content.length : undefined,
|
||||
isTaskCompletion: toolResult.isTaskCompletion,
|
||||
data: toolResult.data,
|
||||
};
|
||||
|
||||
|
||||
@@ -57,7 +57,6 @@ export interface ToolCallResponseInfo {
|
||||
errorType: ToolErrorType | undefined;
|
||||
outputFile?: string | undefined;
|
||||
contentLength?: number;
|
||||
isTaskCompletion?: boolean;
|
||||
/**
|
||||
* Optional data payload for passing structured information back to the caller.
|
||||
*/
|
||||
|
||||
@@ -10,40 +10,25 @@ import { createMockMessageBus } from '../test-utils/mock-message-bus.js';
|
||||
import type { Config } from '../config/config.js';
|
||||
import type { MessageBus } from '../confirmation-bus/message-bus.js';
|
||||
import { ToolConfirmationOutcome } from './tools.js';
|
||||
import { SubagentToolWrapper } from '../agents/subagent-tool-wrapper.js';
|
||||
import type { LocalAgentDefinition } from '../agents/types.js';
|
||||
|
||||
vi.mock('../agents/subagent-tool-wrapper.js');
|
||||
import { ApprovalMode } from '../policy/types.js';
|
||||
|
||||
describe('EnterPlanModeTool', () => {
|
||||
let tool: EnterPlanModeTool;
|
||||
let mockMessageBus: ReturnType<typeof createMockMessageBus>;
|
||||
let mockConfig: Config;
|
||||
let mockPlannerDefinition: LocalAgentDefinition;
|
||||
let mockConfig: Partial<Config>;
|
||||
|
||||
beforeEach(() => {
|
||||
mockMessageBus = createMockMessageBus();
|
||||
vi.mocked(mockMessageBus.publish).mockResolvedValue(undefined);
|
||||
|
||||
mockPlannerDefinition = {
|
||||
kind: 'local',
|
||||
name: 'planner',
|
||||
description: 'Mock Planner',
|
||||
inputConfig: { inputSchema: {} },
|
||||
} as LocalAgentDefinition;
|
||||
|
||||
mockConfig = {
|
||||
setApprovalMode: vi.fn(),
|
||||
isPlannerSubagentEnabled: vi.fn().mockReturnValue(true),
|
||||
getAgentRegistry: vi.fn().mockReturnValue({
|
||||
getDefinition: vi.fn().mockReturnValue(mockPlannerDefinition),
|
||||
}),
|
||||
storage: {
|
||||
getPlansDir: vi.fn().mockReturnValue('/mock/plans/dir'),
|
||||
} as unknown as Config['storage'],
|
||||
} as unknown as Config;
|
||||
};
|
||||
tool = new EnterPlanModeTool(
|
||||
mockConfig,
|
||||
mockConfig as Config,
|
||||
mockMessageBus as unknown as MessageBus,
|
||||
);
|
||||
});
|
||||
@@ -56,6 +41,7 @@ describe('EnterPlanModeTool', () => {
|
||||
it('should return info confirmation details when policy says ASK_USER', async () => {
|
||||
const invocation = tool.build({});
|
||||
|
||||
// Mock getMessageBusDecision to return ASK_USER
|
||||
vi.spyOn(
|
||||
invocation as unknown as {
|
||||
getMessageBusDecision: () => Promise<string>;
|
||||
@@ -78,41 +64,73 @@ describe('EnterPlanModeTool', () => {
|
||||
);
|
||||
}
|
||||
});
|
||||
|
||||
it('should return false when policy decision is ALLOW', async () => {
|
||||
const invocation = tool.build({});
|
||||
|
||||
// Mock getMessageBusDecision to return ALLOW
|
||||
vi.spyOn(
|
||||
invocation as unknown as {
|
||||
getMessageBusDecision: () => Promise<string>;
|
||||
},
|
||||
'getMessageBusDecision',
|
||||
).mockResolvedValue('ALLOW');
|
||||
|
||||
const result = await invocation.shouldConfirmExecute(
|
||||
new AbortController().signal,
|
||||
);
|
||||
|
||||
expect(result).toBe(false);
|
||||
});
|
||||
|
||||
it('should throw error when policy decision is DENY', async () => {
|
||||
const invocation = tool.build({});
|
||||
|
||||
// Mock getMessageBusDecision to return DENY
|
||||
vi.spyOn(
|
||||
invocation as unknown as {
|
||||
getMessageBusDecision: () => Promise<string>;
|
||||
},
|
||||
'getMessageBusDecision',
|
||||
).mockResolvedValue('DENY');
|
||||
|
||||
await expect(
|
||||
invocation.shouldConfirmExecute(new AbortController().signal),
|
||||
).rejects.toThrow(/denied by policy/);
|
||||
});
|
||||
});
|
||||
|
||||
describe('execute', () => {
|
||||
it('should delegate to planner subagent', async () => {
|
||||
const invocation = tool.build({ reason: 'test reason' });
|
||||
|
||||
const mockSubInvocation = {
|
||||
execute: vi.fn().mockResolvedValue({
|
||||
llmContent: 'Plan created.',
|
||||
returnDisplay: 'Plan Display',
|
||||
}),
|
||||
};
|
||||
|
||||
vi.mocked(SubagentToolWrapper).prototype.build = vi
|
||||
.fn()
|
||||
.mockReturnValue(mockSubInvocation);
|
||||
it('should set approval mode to PLAN and return message', async () => {
|
||||
const invocation = tool.build({});
|
||||
|
||||
const result = await invocation.execute(new AbortController().signal);
|
||||
|
||||
expect(mockConfig.getAgentRegistry().getDefinition).toHaveBeenCalledWith(
|
||||
'planner',
|
||||
expect(mockConfig.setApprovalMode).toHaveBeenCalledWith(
|
||||
ApprovalMode.PLAN,
|
||||
);
|
||||
expect(SubagentToolWrapper).toHaveBeenCalledWith(
|
||||
mockPlannerDefinition,
|
||||
mockConfig,
|
||||
mockMessageBus,
|
||||
expect(result.llmContent).toContain('Switching to Plan mode');
|
||||
expect(result.returnDisplay).toBe('Switching to Plan mode');
|
||||
});
|
||||
|
||||
it('should include optional reason in output display but not in llmContent', async () => {
|
||||
const reason = 'Design new database schema';
|
||||
const invocation = tool.build({ reason });
|
||||
|
||||
const result = await invocation.execute(new AbortController().signal);
|
||||
|
||||
expect(mockConfig.setApprovalMode).toHaveBeenCalledWith(
|
||||
ApprovalMode.PLAN,
|
||||
);
|
||||
expect(mockSubInvocation.execute).toHaveBeenCalled();
|
||||
expect(result.llmContent).toBe('Plan created.');
|
||||
expect(result.returnDisplay).toBe('Plan Display');
|
||||
expect(result.llmContent).toBe('Switching to Plan mode.');
|
||||
expect(result.llmContent).not.toContain(reason);
|
||||
expect(result.returnDisplay).toContain(reason);
|
||||
});
|
||||
|
||||
it('should not enter plan mode if cancelled', async () => {
|
||||
const invocation = tool.build({});
|
||||
|
||||
// Simulate getting confirmation details
|
||||
vi.spyOn(
|
||||
invocation as unknown as {
|
||||
getMessageBusDecision: () => Promise<string>;
|
||||
@@ -123,14 +141,30 @@ describe('EnterPlanModeTool', () => {
|
||||
const details = await invocation.shouldConfirmExecute(
|
||||
new AbortController().signal,
|
||||
);
|
||||
expect(details).not.toBe(false);
|
||||
|
||||
if (details) {
|
||||
// Simulate user cancelling
|
||||
await details.onConfirm(ToolConfirmationOutcome.Cancel);
|
||||
}
|
||||
|
||||
const result = await invocation.execute(new AbortController().signal);
|
||||
|
||||
expect(mockConfig.setApprovalMode).not.toHaveBeenCalled();
|
||||
expect(result.returnDisplay).toBe('Cancelled');
|
||||
expect(result.llmContent).toContain('User cancelled');
|
||||
});
|
||||
});
|
||||
|
||||
describe('validateToolParams', () => {
|
||||
it('should allow empty params', () => {
|
||||
const result = tool.validateToolParams({});
|
||||
expect(result).toBeNull();
|
||||
});
|
||||
|
||||
it('should allow reason param', () => {
|
||||
const result = tool.validateToolParams({ reason: 'test' });
|
||||
expect(result).toBeNull();
|
||||
});
|
||||
});
|
||||
});
|
||||
|
||||
@@ -11,7 +11,6 @@ import {
|
||||
Kind,
|
||||
type ToolInfoConfirmationDetails,
|
||||
ToolConfirmationOutcome,
|
||||
type ToolLiveOutput,
|
||||
} from './tools.js';
|
||||
import type { MessageBus } from '../confirmation-bus/message-bus.js';
|
||||
import type { Config } from '../config/config.js';
|
||||
@@ -19,7 +18,6 @@ import { ENTER_PLAN_MODE_TOOL_NAME } from './tool-names.js';
|
||||
import { ApprovalMode } from '../policy/types.js';
|
||||
import { ENTER_PLAN_MODE_DEFINITION } from './definitions/coreTools.js';
|
||||
import { resolveToolDeclaration } from './definitions/resolver.js';
|
||||
import { SubagentToolWrapper } from '../agents/subagent-tool-wrapper.js';
|
||||
|
||||
export interface EnterPlanModeParams {
|
||||
reason?: string;
|
||||
@@ -42,8 +40,6 @@ export class EnterPlanModeTool extends BaseDeclarativeTool<
|
||||
Kind.Plan,
|
||||
ENTER_PLAN_MODE_DEFINITION.base.parametersJsonSchema,
|
||||
messageBus,
|
||||
/* isOutputMarkdown */ true,
|
||||
/* canUpdateOutput */ true,
|
||||
);
|
||||
}
|
||||
|
||||
@@ -116,10 +112,7 @@ export class EnterPlanModeInvocation extends BaseToolInvocation<
|
||||
};
|
||||
}
|
||||
|
||||
async execute(
|
||||
signal: AbortSignal,
|
||||
updateOutput?: (output: ToolLiveOutput) => void,
|
||||
): Promise<ToolResult> {
|
||||
async execute(_signal: AbortSignal): Promise<ToolResult> {
|
||||
if (this.confirmationOutcome === ToolConfirmationOutcome.Cancel) {
|
||||
return {
|
||||
llmContent: 'User cancelled entering Plan Mode.',
|
||||
@@ -129,31 +122,11 @@ export class EnterPlanModeInvocation extends BaseToolInvocation<
|
||||
|
||||
this.config.setApprovalMode(ApprovalMode.PLAN);
|
||||
|
||||
if (!this.config.isPlannerSubagentEnabled()) {
|
||||
return {
|
||||
llmContent: 'Switching to Plan mode.',
|
||||
returnDisplay: this.params.reason
|
||||
? `Switching to Plan mode: ${this.params.reason}`
|
||||
: 'Switching to Plan mode',
|
||||
};
|
||||
}
|
||||
|
||||
const plannerDefinition = this.config
|
||||
.getAgentRegistry()
|
||||
.getDefinition('planner');
|
||||
if (!plannerDefinition) {
|
||||
throw new Error('Planner agent not found.');
|
||||
}
|
||||
|
||||
const wrapper = new SubagentToolWrapper(
|
||||
plannerDefinition,
|
||||
this.config,
|
||||
this.messageBus,
|
||||
);
|
||||
const subInvocation = wrapper.build({
|
||||
reason: this.params.reason || 'Requested by main agent',
|
||||
});
|
||||
|
||||
return subInvocation.execute(signal, updateOutput);
|
||||
return {
|
||||
llmContent: 'Switching to Plan mode.',
|
||||
returnDisplay: this.params.reason
|
||||
? `Switching to Plan mode: ${this.params.reason}`
|
||||
: 'Switching to Plan mode',
|
||||
};
|
||||
}
|
||||
}
|
||||
|
||||
@@ -44,7 +44,6 @@ describe('ExitPlanModeTool', () => {
|
||||
getTargetDir: vi.fn().mockReturnValue(tempRootDir),
|
||||
setApprovalMode: vi.fn(),
|
||||
setApprovedPlanPath: vi.fn(),
|
||||
isPlannerSubagentEnabled: vi.fn().mockReturnValue(true),
|
||||
storage: {
|
||||
getPlansDir: vi.fn().mockReturnValue(mockPlansDir),
|
||||
} as unknown as Config['storage'],
|
||||
@@ -201,10 +200,6 @@ describe('ExitPlanModeTool', () => {
|
||||
const expectedPath = path.join(mockPlansDir, 'test.md');
|
||||
|
||||
expect(result).toEqual({
|
||||
data: {
|
||||
plan_path: expectedPath,
|
||||
},
|
||||
isTaskCompletion: true,
|
||||
llmContent: `Plan approved. Switching to Default mode (edits will require confirmation).
|
||||
|
||||
The approved implementation plan is stored at: ${expectedPath}
|
||||
@@ -233,10 +228,6 @@ Read and follow the plan strictly during implementation.`,
|
||||
const expectedPath = path.join(mockPlansDir, 'test.md');
|
||||
|
||||
expect(result).toEqual({
|
||||
data: {
|
||||
plan_path: expectedPath,
|
||||
},
|
||||
isTaskCompletion: true,
|
||||
llmContent: `Plan approved. Switching to Auto-Edit mode (edits will be applied automatically).
|
||||
|
||||
The approved implementation plan is stored at: ${expectedPath}
|
||||
|
||||
@@ -26,9 +26,10 @@ import { PlanExecutionEvent } from '../telemetry/types.js';
|
||||
import { getExitPlanModeDefinition } from './definitions/coreTools.js';
|
||||
import { resolveToolDeclaration } from './definitions/resolver.js';
|
||||
import { getPlanModeExitMessage } from '../utils/approvalModeUtils.js';
|
||||
import { type PlannerAgentOutput } from '../agents/planner-schema.js';
|
||||
|
||||
export type ExitPlanModeParams = PlannerAgentOutput;
|
||||
export interface ExitPlanModeParams {
|
||||
plan_path: string;
|
||||
}
|
||||
|
||||
export class ExitPlanModeTool extends BaseDeclarativeTool<
|
||||
ExitPlanModeParams,
|
||||
@@ -224,20 +225,14 @@ export class ExitPlanModeInvocation extends BaseToolInvocation<
|
||||
logPlanExecution(this.config, new PlanExecutionEvent(newMode));
|
||||
|
||||
const exitMessage = getPlanModeExitMessage(newMode);
|
||||
const result: ToolResult = {
|
||||
|
||||
return {
|
||||
llmContent: `${exitMessage}
|
||||
|
||||
The approved implementation plan is stored at: ${resolvedPlanPath}
|
||||
Read and follow the plan strictly during implementation.`,
|
||||
returnDisplay: `Plan approved: ${resolvedPlanPath}`,
|
||||
};
|
||||
|
||||
if (this.config.isPlannerSubagentEnabled()) {
|
||||
result.isTaskCompletion = true;
|
||||
result.data = { plan_path: resolvedPlanPath };
|
||||
}
|
||||
|
||||
return result;
|
||||
} else {
|
||||
const feedback = payload?.feedback?.trim();
|
||||
if (feedback) {
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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]---`,
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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(
|
||||
|
||||
@@ -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',
|
||||
|
||||
@@ -223,16 +223,31 @@ 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 {
|
||||
return this.messageBus;
|
||||
}
|
||||
|
||||
/**
|
||||
* Creates a shallow clone of the registry and its current known tools.
|
||||
*/
|
||||
clone(): ToolRegistry {
|
||||
const clone = new ToolRegistry(this.config, this.messageBus);
|
||||
clone.allKnownTools = new Map(this.allKnownTools);
|
||||
return clone;
|
||||
}
|
||||
|
||||
/**
|
||||
* Registers a tool definition.
|
||||
*
|
||||
@@ -590,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
|
||||
@@ -599,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);
|
||||
|
||||
@@ -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);
|
||||
}
|
||||
@@ -680,12 +699,6 @@ export interface ToolResult {
|
||||
type?: ToolErrorType; // An optional machine-readable error type (e.g., 'FILE_NOT_FOUND').
|
||||
};
|
||||
|
||||
/**
|
||||
* If true, this tool result signals an agent indicating they have completed a
|
||||
* task.
|
||||
*/
|
||||
isTaskCompletion?: boolean;
|
||||
|
||||
/**
|
||||
* Optional data payload for passing structured information back to the caller.
|
||||
*/
|
||||
|
||||
@@ -77,43 +77,55 @@ export function patchStdio(): () => void {
|
||||
};
|
||||
}
|
||||
|
||||
/**
|
||||
* Type guard to check if a property key exists on an object.
|
||||
*/
|
||||
function isKey<T extends object>(
|
||||
key: string | symbol | number,
|
||||
obj: T,
|
||||
): key is keyof T {
|
||||
return key in obj;
|
||||
}
|
||||
|
||||
/**
|
||||
* Creates proxies for process.stdout and process.stderr that use the real write methods
|
||||
* (writeToStdout and writeToStderr) bypassing any monkey patching.
|
||||
* This is used to write to the real output even when stdio is patched.
|
||||
*/
|
||||
export function createWorkingStdio() {
|
||||
const inkStdout = new Proxy(process.stdout, {
|
||||
get(target, prop, receiver) {
|
||||
const stdoutHandler: ProxyHandler<typeof process.stdout> = {
|
||||
get(target, prop) {
|
||||
if (prop === 'write') {
|
||||
return writeToStdout;
|
||||
}
|
||||
// eslint-disable-next-line @typescript-eslint/no-unsafe-assignment
|
||||
const value = Reflect.get(target, prop, receiver);
|
||||
if (typeof value === 'function') {
|
||||
// eslint-disable-next-line @typescript-eslint/no-unsafe-return
|
||||
return value.bind(target);
|
||||
if (isKey(prop, target)) {
|
||||
const value = target[prop];
|
||||
if (typeof value === 'function') {
|
||||
return value.bind(target);
|
||||
}
|
||||
return value;
|
||||
}
|
||||
// eslint-disable-next-line @typescript-eslint/no-unsafe-return
|
||||
return value;
|
||||
return undefined;
|
||||
},
|
||||
});
|
||||
};
|
||||
const inkStdout = new Proxy(process.stdout, stdoutHandler);
|
||||
|
||||
const inkStderr = new Proxy(process.stderr, {
|
||||
get(target, prop, receiver) {
|
||||
const stderrHandler: ProxyHandler<typeof process.stderr> = {
|
||||
get(target, prop) {
|
||||
if (prop === 'write') {
|
||||
return writeToStderr;
|
||||
}
|
||||
// eslint-disable-next-line @typescript-eslint/no-unsafe-assignment
|
||||
const value = Reflect.get(target, prop, receiver);
|
||||
if (typeof value === 'function') {
|
||||
// eslint-disable-next-line @typescript-eslint/no-unsafe-return
|
||||
return value.bind(target);
|
||||
if (isKey(prop, target)) {
|
||||
const value = target[prop];
|
||||
if (typeof value === 'function') {
|
||||
return value.bind(target);
|
||||
}
|
||||
return value;
|
||||
}
|
||||
// eslint-disable-next-line @typescript-eslint/no-unsafe-return
|
||||
return value;
|
||||
return undefined;
|
||||
},
|
||||
});
|
||||
};
|
||||
const inkStderr = new Proxy(process.stderr, stderrHandler);
|
||||
|
||||
return { stdout: inkStdout, stderr: inkStderr };
|
||||
}
|
||||
|
||||
@@ -243,10 +243,10 @@ export class GeminiCliSession {
|
||||
|
||||
const loopContext: AgentLoopContext = this.config;
|
||||
const originalRegistry = loopContext.toolRegistry;
|
||||
// eslint-disable-next-line @typescript-eslint/no-unsafe-assignment
|
||||
const scopedRegistry: ToolRegistry = Object.create(originalRegistry);
|
||||
const scopedRegistry: ToolRegistry = originalRegistry.clone();
|
||||
const originalGetTool = scopedRegistry.getTool.bind(scopedRegistry);
|
||||
scopedRegistry.getTool = (name: string) => {
|
||||
const tool = originalRegistry.getTool(name);
|
||||
const tool = originalGetTool(name);
|
||||
if (tool instanceof SdkTool) {
|
||||
return tool.bindContext(context);
|
||||
}
|
||||
|
||||
@@ -1968,13 +1968,6 @@
|
||||
},
|
||||
"additionalProperties": false
|
||||
},
|
||||
"plannerSubagent": {
|
||||
"title": "Planner Subagent",
|
||||
"description": "Use the new planner subagent for plan mode.",
|
||||
"markdownDescription": "Use the new planner subagent for plan mode.\n\n- Category: `Experimental`\n- Requires restart: `yes`\n- Default: `false`",
|
||||
"default": false,
|
||||
"type": "boolean"
|
||||
},
|
||||
"enableAgents": {
|
||||
"title": "Enable Agents",
|
||||
"description": "Enable local and remote subagents. Warning: Experimental feature, uses YOLO mode for subagents",
|
||||
|
||||
Reference in New Issue
Block a user