From d4b7d358c5b102489bda35eb6a9707b9905d09cf Mon Sep 17 00:00:00 2001 From: Akhilesh Kumar Date: Wed, 11 Mar 2026 16:58:30 +0000 Subject: [PATCH] feat(core): support inline MCP server definitions in subagent markdown --- packages/core/src/agents/agentLoader.test.ts | 54 +++++++++ packages/core/src/agents/agentLoader.ts | 60 ++++++++++ .../core/src/agents/local-executor.test.ts | 107 ++++++++++++++++-- packages/core/src/agents/local-executor.ts | 21 ++++ packages/core/src/agents/registry.ts | 13 +++ packages/core/src/agents/types.ts | 6 + packages/core/src/config/config.ts | 1 + 7 files changed, 250 insertions(+), 12 deletions(-) diff --git a/packages/core/src/agents/agentLoader.test.ts b/packages/core/src/agents/agentLoader.test.ts index a526382553..ea7ef0b2c3 100644 --- a/packages/core/src/agents/agentLoader.test.ts +++ b/packages/core/src/agents/agentLoader.test.ts @@ -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, diff --git a/packages/core/src/agents/agentLoader.ts b/packages/core/src/agents/agentLoader.ts index 12337c6248..8716c5604b 100644 --- a/packages/core/src/agents/agentLoader.ts +++ b/packages/core/src/agents/agentLoader.ts @@ -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; + cwd?: string; + url?: string; + http_url?: string; + headers?: Record; + 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; system_prompt: string; model?: string; temperature?: number; @@ -99,6 +118,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'), @@ -112,6 +148,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(), @@ -475,6 +512,28 @@ export function markdownToAgentDefinition( // If a model is specified, use it. Otherwise, inherit const modelName = markdown.model || 'inherit'; + const mcpServers: Record = {}; + 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, @@ -500,6 +559,7 @@ export function markdownToAgentDefinition( tools: markdown.tools, } : undefined, + mcpServers: Object.keys(mcpServers).length > 0 ? mcpServers : undefined, inputConfig, metadata, }; diff --git a/packages/core/src/agents/local-executor.test.ts b/packages/core/src/agents/local-executor.test.ts index f8758cd935..c2949b10b6 100644 --- a/packages/core/src/agents/local-executor.test.ts +++ b/packages/core/src/agents/local-executor.test.ts @@ -13,6 +13,37 @@ 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'; @@ -69,18 +100,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; @@ -2466,4 +2485,68 @@ describe('LocalAgentExecutor', () => { expect(mockSetHistory).toHaveBeenCalledWith(compressedHistory); }); }); + + 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, + }; + + await LocalAgentExecutor.create(definition, mockConfig); + + expect(mockMaybeDiscoverMcpServer).toHaveBeenCalledWith( + 'test-server', + mcpServers['test-server'], + ); + }); + + it('should stop McpClientManager when agent execution finishes', async () => { + const { MCPServerConfig } = await import('../config/config.js'); + const mcpServers = { + 'test-server': new MCPServerConfig('node', ['server.js']), + }; + + const definition = { + ...createTestDefinition(), + mcpServers, + }; + + const executor = await LocalAgentExecutor.create(definition, mockConfig); + + mockSendMessageStream.mockResolvedValueOnce({ + async *[Symbol.asyncIterator] () { + yield { + type: 'chunk', + value: { + candidates: [ + { + content: { + parts: [ + { + functionCall: { + name: TASK_COMPLETE_TOOL_NAME, + args: { result: 'Done' }, + id: 't1', + }, + }, + ], + }, + }, + ], + }, + }; + }, + }); + + await executor.run({ goal: 'test' }, signal); + + expect(mockStopMcp).toHaveBeenCalled(); + }); + }); }); diff --git a/packages/core/src/agents/local-executor.ts b/packages/core/src/agents/local-executor.ts index 9d9cebdd45..0eed4f2888 100644 --- a/packages/core/src/agents/local-executor.ts +++ b/packages/core/src/agents/local-executor.ts @@ -18,6 +18,7 @@ import { } from '@google/genai'; import { ToolRegistry } from '../tools/tool-registry.js'; import { DiscoveredMCPTool } from '../tools/mcp-tool.js'; +import { McpClientManager } from '../tools/mcp-client-manager.js'; import { CompressionStatus } from '../core/turn.js'; import { type ToolCallRequestInfo } from '../scheduler/types.js'; import { type Message } from '../confirmation-bus/types.js'; @@ -94,6 +95,7 @@ export class LocalAgentExecutor { private readonly agentId: string; private readonly toolRegistry: ToolRegistry; private readonly context: AgentLoopContext; + private readonly mcpClientManager?: McpClientManager; private readonly onActivity?: ActivityCallback; private readonly compressionService: ChatCompressionService; private readonly parentCallId?: string; @@ -141,6 +143,19 @@ export class LocalAgentExecutor { context.config, subagentMessageBus, ); + let mcpClientManager: McpClientManager | undefined; + if (definition.mcpServers) { + mcpClientManager = new McpClientManager( + await getVersion(), + agentToolRegistry, + context.config, + ); + + for (const [name, config] of Object.entries(definition.mcpServers)) { + await mcpClientManager.maybeDiscoverMcpServer(name, config); + } + } + const parentToolRegistry = context.toolRegistry; const allAgentNames = new Set( context.config.getAgentRegistry().getAllAgentNames(), @@ -210,6 +225,7 @@ export class LocalAgentExecutor { parentPromptId, parentCallId, onActivity, + mcpClientManager, ); } @@ -226,6 +242,7 @@ export class LocalAgentExecutor { parentPromptId: string | undefined, parentCallId: string | undefined, onActivity?: ActivityCallback, + mcpClientManager?: McpClientManager, ) { this.definition = definition; this.context = context; @@ -233,6 +250,7 @@ export class LocalAgentExecutor { this.onActivity = onActivity; this.compressionService = new ChatCompressionService(); this.parentCallId = parentCallId; + this.mcpClientManager = mcpClientManager; const randomIdPart = Math.random().toString(36).slice(2, 8); // parentPromptId will be undefined if this agent is invoked directly @@ -679,6 +697,9 @@ export class LocalAgentExecutor { throw error; // Re-throw other errors or external aborts. } finally { deadlineTimer.abort(); + if (this.mcpClientManager) { + await this.mcpClientManager.stop(); + } logAgentFinish( this.config, new AgentFinishEvent( diff --git a/packages/core/src/agents/registry.ts b/packages/core/src/agents/registry.ts index c4b08eba22..aff9e07c9f 100644 --- a/packages/core/src/agents/registry.ts +++ b/packages/core/src/agents/registry.ts @@ -537,6 +537,19 @@ export class AgentRegistry { ); } + if (overrides.tools) { + merged.toolConfig = { + tools: overrides.tools, + }; + } + + if (overrides.mcpServers) { + merged.mcpServers = { + ...definition.mcpServers, + ...overrides.mcpServers, + }; + } + return merged; } diff --git a/packages/core/src/agents/types.ts b/packages/core/src/agents/types.ts index ceac0909df..162d1095d1 100644 --- a/packages/core/src/agents/types.ts +++ b/packages/core/src/agents/types.ts @@ -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; + /** * An optional function to process the raw output from the agent's final tool * call into a string format. diff --git a/packages/core/src/config/config.ts b/packages/core/src/config/config.ts index e8029c495e..901eb269d8 100644 --- a/packages/core/src/config/config.ts +++ b/packages/core/src/config/config.ts @@ -236,6 +236,7 @@ export interface AgentOverride { runConfig?: AgentRunConfig; enabled?: boolean; tools?: string[]; + mcpServers?: Record; } export interface AgentSettings {