Skip to content

Commit 66b6c6b

Browse files
committed
refactor(core): update MCP tool filtering implementation
- Moved tool filtering logic to agent/runner layer - Removed server-side filtering and context coupling - Updated test suite to reflect new behavior
1 parent 2ddaea8 commit 66b6c6b

File tree

2 files changed

+103
-10
lines changed

2 files changed

+103
-10
lines changed

packages/agents-core/src/mcp.ts

Lines changed: 59 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -14,7 +14,11 @@ import {
1414
JsonObjectSchemaStrict,
1515
UnknownContext,
1616
} from './types';
17-
import type { ToolFilterCallable, ToolFilterStatic } from './mcpUtil';
17+
import type {
18+
ToolFilterCallable,
19+
ToolFilterStatic,
20+
ToolFilterContext,
21+
} from './mcpUtil';
1822
import type { RunContext } from './runContext';
1923
import type { Agent } from './agent';
2024

@@ -30,6 +34,7 @@ export const DEFAULT_STREAMABLE_HTTP_MCP_CLIENT_LOGGER_NAME =
3034
*/
3135
export interface MCPServer {
3236
cacheToolsList: boolean;
37+
toolFilter?: ToolFilterCallable | ToolFilterStatic;
3338
connect(): Promise<void>;
3439
readonly name: string;
3540
close(): Promise<void>;
@@ -47,7 +52,7 @@ export interface MCPServer {
4752
export abstract class BaseMCPServerStdio implements MCPServer {
4853
public cacheToolsList: boolean;
4954
protected _cachedTools: any[] | undefined = undefined;
50-
protected toolFilter?: ToolFilterCallable | ToolFilterStatic;
55+
public toolFilter?: ToolFilterCallable | ToolFilterStatic;
5156

5257
protected logger: Logger;
5358
constructor(options: MCPServerStdioOptions) {
@@ -85,7 +90,7 @@ export abstract class BaseMCPServerStdio implements MCPServer {
8590
export abstract class BaseMCPServerStreamableHttp implements MCPServer {
8691
public cacheToolsList: boolean;
8792
protected _cachedTools: any[] | undefined = undefined;
88-
protected toolFilter?: ToolFilterCallable | ToolFilterStatic;
93+
public toolFilter?: ToolFilterCallable | ToolFilterStatic;
8994

9095
protected logger: Logger;
9196
constructor(options: MCPServerStreamableHttpOptions) {
@@ -158,13 +163,13 @@ export class MCPServerStdio extends BaseMCPServerStdio {
158163
return this.underlying.close();
159164
}
160165
async listTools(
161-
runContext?: RunContext<any>,
162-
agent?: Agent<any, any>,
166+
_runContext?: RunContext<any>,
167+
_agent?: Agent<any, any>,
163168
): Promise<MCPTool[]> {
164169
if (this.cacheToolsList && this._cachedTools) {
165170
return this._cachedTools;
166171
}
167-
const tools = await this.underlying.listTools(runContext, agent);
172+
const tools = await this.underlying.listTools();
168173
if (this.cacheToolsList) {
169174
this._cachedTools = tools;
170175
}
@@ -197,13 +202,13 @@ export class MCPServerStreamableHttp extends BaseMCPServerStreamableHttp {
197202
return this.underlying.close();
198203
}
199204
async listTools(
200-
runContext?: RunContext<any>,
201-
agent?: Agent<any, any>,
205+
_runContext?: RunContext<any>,
206+
_agent?: Agent<any, any>,
202207
): Promise<MCPTool[]> {
203208
if (this.cacheToolsList && this._cachedTools) {
204209
return this._cachedTools;
205210
}
206-
const tools = await this.underlying.listTools(runContext, agent);
211+
const tools = await this.underlying.listTools();
207212
if (this.cacheToolsList) {
208213
this._cachedTools = tools;
209214
}
@@ -279,7 +284,16 @@ async function getFunctionToolsFromServer<TContext = UnknownContext>(
279284
}
280285
return withMCPListToolsSpan(
281286
async (span) => {
282-
const mcpTools = await server.listTools(runContext, agent);
287+
let mcpTools = await server.listTools(runContext, agent);
288+
if (server.toolFilter) {
289+
mcpTools = await filterMcpTools(
290+
mcpTools,
291+
server.toolFilter as ToolFilterCallable<TContext> | ToolFilterStatic,
292+
runContext,
293+
agent,
294+
server.name,
295+
);
296+
}
283297
span.spanData.result = mcpTools.map((t) => t.name);
284298
const tools: FunctionTool<TContext, any, string>[] = mcpTools.map((t) =>
285299
mcpToFunctionTool(t, server, convertSchemasToStrict),
@@ -383,6 +397,41 @@ function ensureStrictJsonSchema(
383397
return out;
384398
}
385399

400+
async function filterMcpTools<TContext = UnknownContext>(
401+
tools: MCPTool[],
402+
filter: ToolFilterCallable<TContext> | ToolFilterStatic,
403+
runContext: RunContext<TContext> | undefined,
404+
agent: Agent<TContext, any> | undefined,
405+
serverName: string,
406+
): Promise<MCPTool[]> {
407+
if (typeof filter === 'function') {
408+
if (!runContext || !agent) {
409+
return tools;
410+
}
411+
const ctx = {
412+
runContext,
413+
agent,
414+
serverName,
415+
} as ToolFilterContext<TContext>;
416+
const result: MCPTool[] = [];
417+
for (const tool of tools) {
418+
if (await filter(ctx, tool)) {
419+
result.push(tool);
420+
}
421+
}
422+
return result;
423+
}
424+
return tools.filter((t) => {
425+
if (filter.allowedToolNames && !filter.allowedToolNames.includes(t.name)) {
426+
return false;
427+
}
428+
if (filter.blockedToolNames && filter.blockedToolNames.includes(t.name)) {
429+
return false;
430+
}
431+
return true;
432+
});
433+
}
434+
386435
/**
387436
* Abstract base class for MCP servers that use a ClientSession for communication.
388437
* Handles session management, tool listing, tool calling, and cleanup.

packages/agents-core/test/mcpToolFilter.test.ts

Lines changed: 44 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@ import { describe, it, expect } from 'vitest';
22
import { withTrace } from '../src/tracing';
33
import { NodeMCPServerStdio } from '../src/shims/mcp-server/node';
44
import { createStaticToolFilter } from '../src/mcpUtil';
5+
import { getAllMcpTools } from '../src/mcp';
56
import { Agent } from '../src/agent';
67
import { RunContext } from '../src/runContext';
78

@@ -198,4 +199,47 @@ describe('MCP tool filtering', () => {
198199
expect(result.map((t) => t.name)).toEqual(['x']);
199200
});
200201
});
202+
203+
it('applies filter in getAllMcpTools', async () => {
204+
await withTrace('test', async () => {
205+
const tools = [
206+
{
207+
name: 'allow',
208+
description: '',
209+
inputSchema: {
210+
type: 'object',
211+
properties: {},
212+
required: [],
213+
additionalProperties: false,
214+
},
215+
},
216+
{
217+
name: 'block',
218+
description: '',
219+
inputSchema: {
220+
type: 'object',
221+
properties: {},
222+
required: [],
223+
additionalProperties: false,
224+
},
225+
},
226+
];
227+
const server = new StubServer(
228+
'filter',
229+
tools,
230+
createStaticToolFilter({ allowed: ['allow'] }),
231+
);
232+
const agent = new Agent({
233+
name: 'agent',
234+
instructions: '',
235+
model: '',
236+
modelSettings: {},
237+
tools: [],
238+
mcpServers: [server],
239+
});
240+
const runContext = new RunContext();
241+
const result = await getAllMcpTools([server], false, runContext, agent);
242+
expect(result.map((t) => t.name)).toEqual(['allow']);
243+
});
244+
});
201245
});

0 commit comments

Comments
 (0)