Merge remote-tracking branch 'seratch/vrtnis-feat/mcp-tool-filtering-js' into feat/mcp-tool-filtering-js

Author: vrtnis

examples/mcp/tool-filter-example.ts

@@ -24,8 +24,7 @@ async function main() {
     await withTrace('MCP Tool Filter Example', async () => {
       const agent = new Agent({
         name: 'MCP Assistant',
-        instructions:
-          'Use the filesystem tools to answer questions. The write_file tool is blocked via toolFilter.',
+        instructions: 'Use the filesystem tools to answer questions.',
         mcpServers: [mcpServer],
       });
 
@@ -36,7 +35,7 @@ async function main() {
       console.log('\nAttempting to write a file (should be blocked):');
       result = await run(
         agent,
-        'Create a file named test.txt with the text "hello"',
+        'Create a file named sample_files/test.txt with the text "hello"',
       );
       console.log(result.finalOutput);
     });

packages/agents-core/src/agent.ts

@@ -518,7 +518,7 @@ export class Agent<
     runContext: RunContext<TContext>,
   ): Promise<Tool<TContext>[]> {
     if (this.mcpServers.length > 0) {
-      return getAllMcpTools(this.mcpServers, false, runContext, this);
+      return getAllMcpTools(this.mcpServers, runContext, this, false);
     }
 
     return [];

packages/agents-core/src/mcp.ts

@@ -30,13 +30,11 @@ export const DEFAULT_STREAMABLE_HTTP_MCP_CLIENT_LOGGER_NAME =
  */
 export interface MCPServer {
   cacheToolsList: boolean;
+  toolFilter?: MCPToolFilterCallable | MCPToolFilterStatic;
   connect(): Promise<void>;
   readonly name: string;
   close(): Promise<void>;
-  listTools(
-    runContext?: RunContext<any>,
-    agent?: Agent<any, any>,
-  ): Promise<MCPTool[]>;
+  listTools(): Promise<MCPTool[]>;
   callTool(
     toolName: string,
     args: Record<string, unknown> | null,
@@ -46,7 +44,7 @@ export interface MCPServer {
 export abstract class BaseMCPServerStdio implements MCPServer {
   public cacheToolsList: boolean;
   protected _cachedTools: any[] | undefined = undefined;
-  protected toolFilter?: MCPToolFilterCallable | MCPToolFilterStatic;
+  public toolFilter?: MCPToolFilterCallable | MCPToolFilterStatic;
 
   protected logger: Logger;
   constructor(options: MCPServerStdioOptions) {
@@ -59,10 +57,7 @@ export abstract class BaseMCPServerStdio implements MCPServer {
   abstract get name(): string;
   abstract connect(): Promise<void>;
   abstract close(): Promise<void>;
-  abstract listTools(
-    runContext?: RunContext<any>,
-    agent?: Agent<any, any>,
-  ): Promise<any[]>;
+  abstract listTools(): Promise<any[]>;
   abstract callTool(
     _toolName: string,
     _args: Record<string, unknown> | null,
@@ -83,7 +78,7 @@ export abstract class BaseMCPServerStdio implements MCPServer {
 export abstract class BaseMCPServerStreamableHttp implements MCPServer {
   public cacheToolsList: boolean;
   protected _cachedTools: any[] | undefined = undefined;
-  protected toolFilter?: MCPToolFilterCallable | MCPToolFilterStatic;
+  public toolFilter?: MCPToolFilterCallable | MCPToolFilterStatic;
 
   protected logger: Logger;
   constructor(options: MCPServerStreamableHttpOptions) {
@@ -97,10 +92,7 @@ export abstract class BaseMCPServerStreamableHttp implements MCPServer {
   abstract get name(): string;
   abstract connect(): Promise<void>;
   abstract close(): Promise<void>;
-  abstract listTools(
-    runContext?: RunContext<any>,
-    agent?: Agent<any, any>,
-  ): Promise<any[]>;
+  abstract listTools(): Promise<any[]>;
   abstract callTool(
     _toolName: string,
     _args: Record<string, unknown> | null,
@@ -154,14 +146,11 @@ export class MCPServerStdio extends BaseMCPServerStdio {
   close(): Promise<void> {
     return this.underlying.close();
   }
-  async listTools(
-    runContext?: RunContext<any>,
-    agent?: Agent<any, any>,
-  ): Promise<MCPTool[]> {
+  async listTools(): Promise<MCPTool[]> {
     if (this.cacheToolsList && this._cachedTools) {
       return this._cachedTools;
     }
-    const tools = await this.underlying.listTools(runContext, agent);
+    const tools = await this.underlying.listTools();
     if (this.cacheToolsList) {
       this._cachedTools = tools;
     }
@@ -190,14 +179,11 @@ export class MCPServerStreamableHttp extends BaseMCPServerStreamableHttp {
   close(): Promise<void> {
     return this.underlying.close();
   }
-  async listTools(
-    runContext?: RunContext<any>,
-    agent?: Agent<any, any>,
-  ): Promise<MCPTool[]> {
+  async listTools(): Promise<MCPTool[]> {
     if (this.cacheToolsList && this._cachedTools) {
       return this._cachedTools;
     }
-    const tools = await this.underlying.listTools(runContext, agent);
+    const tools = await this.underlying.listTools();
     if (this.cacheToolsList) {
       this._cachedTools = tools;
     }
@@ -217,18 +203,18 @@ export class MCPServerStreamableHttp extends BaseMCPServerStreamableHttp {
  */
 export async function getAllMcpFunctionTools<TContext = UnknownContext>(
   mcpServers: MCPServer[],
+  runContext: RunContext<TContext>,
+  agent: Agent<any, any>,
   convertSchemasToStrict = false,
-  runContext?: RunContext<TContext>,
-  agent?: Agent<TContext, any>,
 ): Promise<Tool<TContext>[]> {
   const allTools: Tool<TContext>[] = [];
   const toolNames = new Set<string>();
   for (const server of mcpServers) {
     const serverTools = await getFunctionToolsFromServer(
       server,
-      convertSchemasToStrict,
       runContext,
       agent,
+      convertSchemasToStrict,
     );
     const serverToolNames = new Set(serverTools.map((t) => t.name));
     const intersection = [...serverToolNames].filter((n) => toolNames.has(n));
@@ -259,16 +245,62 @@ export function invalidateServerToolsCache(serverName: string) {
  */
 async function getFunctionToolsFromServer<TContext = UnknownContext>(
   server: MCPServer,
+  runContext: RunContext<TContext>,
+  agent: Agent<any, any>,
   convertSchemasToStrict: boolean,
-  runContext?: RunContext<TContext>,
-  agent?: Agent<TContext, any>,
 ): Promise<FunctionTool<TContext, any, unknown>[]> {
   if (server.cacheToolsList && _cachedTools[server.name]) {
     return _cachedTools[server.name];
   }
   return withMCPListToolsSpan(
     async (span) => {
-      const mcpTools = await server.listTools(runContext, agent);
+      const fetchedMcpTools = await server.listTools();
+      const mcpTools: MCPTool[] = [];
+      const context = {
+        runContext,
+        agent,
+        serverName: server.name,
+      };
+      for (const tool of fetchedMcpTools) {
+        const filter = server.toolFilter;
+        if (filter) {
+          if (filter && typeof filter === 'function') {
+            const filtered = await filter(context, tool);
+            if (!filtered) {
+              globalLogger.debug(
+                `MCP Tool (server: ${server.name}, tool: ${tool.name}) is blocked by the callable filter.`,
+              );
+              continue; // skip this tool
+            }
+          } else {
+            const allowedToolNames = filter.allowedToolNames ?? [];
+            const blockedToolNames = filter.blockedToolNames ?? [];
+            if (allowedToolNames.length > 0 || blockedToolNames.length > 0) {
+              const allowed =
+                allowedToolNames.length > 0
+                  ? allowedToolNames.includes(tool.name)
+                  : true;
+              const blocked =
+                blockedToolNames.length > 0
+                  ? blockedToolNames.includes(tool.name)
+                  : false;
+              if (!allowed || blocked) {
+                if (blocked) {
+                  globalLogger.debug(
+                    `MCP Tool (server: ${server.name}, tool: ${tool.name}) is blocked by the static filter.`,
+                  );
+                } else if (!allowed) {
+                  globalLogger.debug(
+                    `MCP Tool (server: ${server.name}, tool: ${tool.name}) is not allowed by the static filter.`,
+                  );
+                }
+                continue; // skip this tool
+              }
+            }
+          }
+        }
+        mcpTools.push(tool);
+      }
       span.spanData.result = mcpTools.map((t) => t.name);
       const tools: FunctionTool<TContext, any, string>[] = mcpTools.map((t) =>
         mcpToFunctionTool(t, server, convertSchemasToStrict),
@@ -287,15 +319,15 @@ async function getFunctionToolsFromServer<TContext = UnknownContext>(
  */
 export async function getAllMcpTools<TContext = UnknownContext>(
   mcpServers: MCPServer[],
+  runContext: RunContext<TContext>,
+  agent: Agent<TContext, any>,
   convertSchemasToStrict = false,
-  runContext?: RunContext<TContext>,
-  agent?: Agent<TContext, any>,
 ): Promise<Tool<TContext>[]> {
   return getAllMcpFunctionTools(
     mcpServers,
-    convertSchemasToStrict,
     runContext,
     agent,
+    convertSchemasToStrict,
   );
 }
 

packages/agents-core/src/mcpUtil.ts

@@ -17,7 +17,7 @@ export interface MCPToolFilterContext<TContext = UnknownContext> {
 export type MCPToolFilterCallable<TContext = UnknownContext> = (
   context: MCPToolFilterContext<TContext>,
   tool: MCPTool,
-) => boolean | Promise<boolean>;
+) => Promise<boolean>;
 
 /** Static tool filter configuration using allow and block lists. */
 export interface MCPToolFilterStatic {

packages/agents-core/src/runState.ts

@@ -557,8 +557,8 @@ export class RunState<TContext, TAgent extends Agent<any, any>> {
       ? await deserializeProcessedResponse(
           agentMap,
           state._currentAgent,
-          stateJson.lastProcessedResponse,
           state._context,
+          stateJson.lastProcessedResponse,
         )
       : undefined;
 
@@ -708,12 +708,12 @@ export function deserializeItem(
 async function deserializeProcessedResponse<TContext = UnknownContext>(
   agentMap: Map<string, Agent<any, any>>,
   currentAgent: Agent<TContext, any>,
+  context: RunContext<TContext>,
   serializedProcessedResponse: z.infer<
     typeof serializedProcessedResponseSchema
   >,
-  runContext: RunContext<TContext>,
 ): Promise<ProcessedResponse<TContext>> {
-  const allTools = await currentAgent.getAllTools(runContext);
+  const allTools = await currentAgent.getAllTools(context);
   const tools = new Map(
     allTools
       .filter((tool) => tool.type === 'function')

packages/agents-core/src/shims/mcp-server/browser.ts

@@ -6,8 +6,6 @@ import {
   MCPServerStreamableHttpOptions,
   MCPTool,
 } from '../../mcp';
-import type { RunContext } from '../../runContext';
-import type { Agent } from '../../agent';
 
 export class MCPServerStdio extends BaseMCPServerStdio {
   constructor(params: MCPServerStdioOptions) {
@@ -22,10 +20,7 @@ export class MCPServerStdio extends BaseMCPServerStdio {
   close(): Promise<void> {
     throw new Error('Method not implemented.');
   }
-  listTools(
-    _runContext?: RunContext<any>,
-    _agent?: Agent<any, any>,
-  ): Promise<MCPTool[]> {
+  listTools(): Promise<MCPTool[]> {
     throw new Error('Method not implemented.');
   }
   callTool(
@@ -49,10 +44,7 @@ export class MCPServerStreamableHttp extends BaseMCPServerStreamableHttp {
   close(): Promise<void> {
     throw new Error('Method not implemented.');
   }
-  listTools(
-    _runContext?: RunContext<any>,
-    _agent?: Agent<any, any>,
-  ): Promise<MCPTool[]> {
+  listTools(): Promise<MCPTool[]> {
     throw new Error('Method not implemented.');
   }
   callTool(

packages/agents-core/src/shims/mcp-server/node.ts

@@ -12,8 +12,6 @@ import {
   invalidateServerToolsCache,
 } from '../../mcp';
 import logger from '../../logger';
-import type { RunContext } from '../../runContext';
-import type { Agent } from '../../agent';
 
 export interface SessionMessage {
   message: any;
@@ -98,10 +96,7 @@ export class NodeMCPServerStdio extends BaseMCPServerStdio {
     this._cacheDirty = true;
   }
 
-  async listTools(
-    _runContext?: RunContext<any>,
-    _agent?: Agent<any, any>,
-  ): Promise<MCPTool[]> {
+  async listTools(): Promise<MCPTool[]> {
     const { ListToolsResultSchema } = await import(
       '@modelcontextprotocol/sdk/types.js'
     ).catch(failedToImport);
@@ -218,10 +213,7 @@ export class NodeMCPServerStreamableHttp extends BaseMCPServerStreamableHttp {
     this._cacheDirty = true;
   }
 
-  async listTools(
-    _runContext?: RunContext<any>,
-    _agent?: Agent<any, any>,
-  ): Promise<MCPTool[]> {
+  async listTools(): Promise<MCPTool[]> {
     const { ListToolsResultSchema } = await import(
       '@modelcontextprotocol/sdk/types.js'
     ).catch(failedToImport);

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

@@ -3,6 +3,8 @@ import { getAllMcpTools } from '../src/mcp';
 import { withTrace } from '../src/tracing';
 import { NodeMCPServerStdio } from '../src/shims/mcp-server/node';
 import type { CallToolResultContent } from '../src/mcp';
+import { RunContext } from '../src/runContext';
+import { Agent } from '../src/agent';
 
 class StubServer extends NodeMCPServerStdio {
   public toolList: any[];
@@ -48,15 +50,27 @@ describe('MCP tools cache invalidation', () => {
       ];
       const server = new StubServer('server', toolsA);
 
-      let tools = await getAllMcpTools([server]);
+      let tools = await getAllMcpTools(
+        [server],
+        new RunContext({}),
+        new Agent({ name: 'test' }),
+      );
       expect(tools.map((t) => t.name)).toEqual(['a']);
 
       server.toolList = toolsB;
-      tools = await getAllMcpTools([server]);
+      tools = await getAllMcpTools(
+        [server],
+        new RunContext({}),
+        new Agent({ name: 'test' }),
+      );
       expect(tools.map((t) => t.name)).toEqual(['a']);
 
       server.invalidateToolsCache();
-      tools = await getAllMcpTools([server]);
+      tools = await getAllMcpTools(
+        [server],
+        new RunContext({}),
+        new Agent({ name: 'test' }),
+      );
       expect(tools.map((t) => t.name)).toEqual(['b']);
     });
   });