diff --git a/src/config/schema.ts b/src/config/schema.ts index 81f7e6e..ff8368e 100644 --- a/src/config/schema.ts +++ b/src/config/schema.ts @@ -13,6 +13,8 @@ const stdioServerSchema = z.object({ command: z.string(), args: z.array(z.string()).optional(), env: z.record(z.string()).optional(), + includeTools: z.array(z.string()).optional(), + excludeTools: z.array(z.string()).optional(), }); /** @@ -22,6 +24,8 @@ const httpServerSchema = z.object({ type: z.literal("http"), url: z.string().url(), headers: z.record(z.string()).optional(), + includeTools: z.array(z.string()).optional(), + excludeTools: z.array(z.string()).optional(), }); /** @@ -31,6 +35,8 @@ const sseServerSchema = z.object({ type: z.literal("sse"), url: z.string().url(), headers: z.record(z.string()).optional(), + includeTools: z.array(z.string()).optional(), + excludeTools: z.array(z.string()).optional(), }); /** diff --git a/src/config/types.ts b/src/config/types.ts index 94d752e..7b8d21d 100644 --- a/src/config/types.ts +++ b/src/config/types.ts @@ -13,6 +13,19 @@ export type ServerType = "stdio" | "http" | "sse"; */ export interface BaseServerConfig { type: ServerType; + /** + * Optional list of tool names to include from this server. + * If specified, only these tools will be served. + * Tool names must be exact matches. + */ + includeTools?: string[]; + /** + * Optional list of tool names to exclude from this server. + * If specified, these tools will not be served. + * Tool names must be exact matches. + * Exclude filters are applied after include filters. + */ + excludeTools?: string[]; } /** diff --git a/src/gateway/aggregator.ts b/src/gateway/aggregator.ts index 35bdd6e..6954c81 100644 --- a/src/gateway/aggregator.ts +++ b/src/gateway/aggregator.ts @@ -6,6 +6,7 @@ import type { ToolscriptConfig } from "../config/types.ts"; import { McpClient } from "./mcp-client.ts"; import { getLogger } from "../utils/logger.ts"; import { UnauthorizedError } from "@modelcontextprotocol/sdk/client/auth.js"; +import { shouldIncludeTool, type ToolFilters } from "./tool-filter.ts"; const logger = getLogger("aggregator"); @@ -27,6 +28,7 @@ export interface ToolInfo { export class ServerAggregator implements AsyncDisposable { private clients: Map = new Map(); private tools: Map = new Map(); + private serverFilters: Map = new Map(); /** * Initialize and connect to all configured servers @@ -38,6 +40,18 @@ export class ServerAggregator implements AsyncDisposable { // Connect to all servers in parallel const connectionPromises = Object.entries(config.mcpServers).map( async ([name, serverConfig]) => { + // Store filter configuration + if (serverConfig.includeTools || serverConfig.excludeTools) { + this.serverFilters.set(name, { + includeTools: serverConfig.includeTools + ? new Set(serverConfig.includeTools) + : undefined, + excludeTools: serverConfig.excludeTools + ? new Set(serverConfig.excludeTools) + : undefined, + }); + } + try { const client = new McpClient(name, serverConfig); await client.connect(); @@ -77,12 +91,20 @@ export class ServerAggregator implements AsyncDisposable { continue; } + let loadedCount = 0; for (const tool of tools) { if (!tool || !tool.name) { logger.warn(`Skipping invalid tool from ${serverName}:`, tool); continue; } + // Apply includeTools and excludeTools filters + const filters = this.serverFilters.get(serverName); + if (filters && !shouldIncludeTool(tool.name, filters)) { + logger.debug(`Tool ${tool.name} from ${serverName} filtered out`); + continue; + } + const qualifiedName = `${serverName}__${tool.name}`; this.tools.set(qualifiedName, { serverName, @@ -92,8 +114,9 @@ export class ServerAggregator implements AsyncDisposable { inputSchema: tool.inputSchema, outputSchema: tool.outputSchema, }); + loadedCount++; } - logger.info(`Loaded ${tools.length} tools from server: ${serverName}`); + logger.info(`Loaded ${loadedCount}/${tools.length} tools from server: ${serverName}`); } catch (error) { logger.error(`Failed to list tools from server ${serverName}: ${error}`); if (error instanceof Error) { diff --git a/src/gateway/tool-filter.test.ts b/src/gateway/tool-filter.test.ts new file mode 100644 index 0000000..a30ad25 --- /dev/null +++ b/src/gateway/tool-filter.test.ts @@ -0,0 +1,127 @@ +/** + * Tests for tool filtering utilities. + */ + +import { assertEquals } from "@std/assert"; +import { shouldIncludeTool, type ToolFilters } from "./tool-filter.ts"; + +Deno.test("shouldIncludeTool should include all tools when no filters", () => { + const filters: ToolFilters = {}; + + assertEquals(shouldIncludeTool("any_tool", filters), true); + assertEquals(shouldIncludeTool("another_tool", filters), true); + assertEquals(shouldIncludeTool("delete_something", filters), true); +}); + +Deno.test("shouldIncludeTool should match exact tool names in includeTools", () => { + const filters: ToolFilters = { + includeTools: new Set(["create_issue", "list_repos"]), + }; + + assertEquals(shouldIncludeTool("create_issue", filters), true); + assertEquals(shouldIncludeTool("list_repos", filters), true); + assertEquals(shouldIncludeTool("delete_repo", filters), false); + assertEquals(shouldIncludeTool("get_issue", filters), false); +}); + +Deno.test("shouldIncludeTool should match multiple exact tool names in includeTools", () => { + const filters: ToolFilters = { + includeTools: new Set(["get_issue", "get_repo", "list_repos"]), + }; + + assertEquals(shouldIncludeTool("get_issue", filters), true); + assertEquals(shouldIncludeTool("get_repo", filters), true); + assertEquals(shouldIncludeTool("list_repos", filters), true); + assertEquals(shouldIncludeTool("create_issue", filters), false); + assertEquals(shouldIncludeTool("delete_repo", filters), false); +}); + +Deno.test("shouldIncludeTool should exclude exact tool names in excludeTools", () => { + const filters: ToolFilters = { + excludeTools: new Set(["delete_repo", "remove_file"]), + }; + + assertEquals(shouldIncludeTool("get_issue", filters), true); + assertEquals(shouldIncludeTool("create_issue", filters), true); + assertEquals(shouldIncludeTool("delete_repo", filters), false); + assertEquals(shouldIncludeTool("remove_file", filters), false); +}); + +Deno.test("shouldIncludeTool should exclude multiple exact tool names in excludeTools", () => { + const filters: ToolFilters = { + excludeTools: new Set(["delete_repo", "delete_issue", "remove_file"]), + }; + + assertEquals(shouldIncludeTool("get_issue", filters), true); + assertEquals(shouldIncludeTool("create_issue", filters), true); + assertEquals(shouldIncludeTool("delete_repo", filters), false); + assertEquals(shouldIncludeTool("delete_issue", filters), false); + assertEquals(shouldIncludeTool("remove_file", filters), false); +}); + +Deno.test("shouldIncludeTool should apply both include and exclude filters", () => { + const filters: ToolFilters = { + includeTools: new Set(["get_issue", "create_issue", "create_repo", "delete_issue"]), + excludeTools: new Set(["delete_issue"]), + }; + + // In include list, not excluded + assertEquals(shouldIncludeTool("get_issue", filters), true); + assertEquals(shouldIncludeTool("create_issue", filters), true); + assertEquals(shouldIncludeTool("create_repo", filters), true); + + // In include list but also excluded (exclude wins) + assertEquals(shouldIncludeTool("delete_issue", filters), false); + + // Not in include list + assertEquals(shouldIncludeTool("list_users", filters), false); + assertEquals(shouldIncludeTool("get_user", filters), false); +}); + +Deno.test("shouldIncludeTool should work with include filters", () => { + const filters: ToolFilters = { + includeTools: new Set(["get_user", "get_public_data"]), + excludeTools: new Set(["get_private_key"]), + }; + + assertEquals(shouldIncludeTool("get_user", filters), true); + assertEquals(shouldIncludeTool("get_public_data", filters), true); + assertEquals(shouldIncludeTool("get_private_key", filters), false); + assertEquals(shouldIncludeTool("create_user", filters), false); +}); + +Deno.test("shouldIncludeTool should work with exclude filters", () => { + const filters: ToolFilters = { + excludeTools: new Set(["delete_dangerous", "get_internal"]), + }; + + assertEquals(shouldIncludeTool("get_user", filters), true); + assertEquals(shouldIncludeTool("create_user", filters), true); + assertEquals(shouldIncludeTool("delete_dangerous", filters), false); + assertEquals(shouldIncludeTool("get_internal", filters), false); +}); + +Deno.test("shouldIncludeTool should handle empty includeTools Set", () => { + const filters: ToolFilters = { + includeTools: new Set([]), + }; + + // Empty includeTools means include all (same as not specifying it) + assertEquals(shouldIncludeTool("any_tool", filters), true); +}); + +Deno.test("shouldIncludeTool should handle empty excludeTools Set", () => { + const filters: ToolFilters = { + excludeTools: new Set([]), + }; + + // Empty excludeTools means exclude none (same as not specifying it) + assertEquals(shouldIncludeTool("any_tool", filters), true); +}); + +Deno.test("shouldIncludeTool should work with empty ToolFilters", () => { + const filters: ToolFilters = {}; + + assertEquals(shouldIncludeTool("any_tool", filters), true); + assertEquals(shouldIncludeTool("another_tool", filters), true); +}); diff --git a/src/gateway/tool-filter.ts b/src/gateway/tool-filter.ts new file mode 100644 index 0000000..898f1d3 --- /dev/null +++ b/src/gateway/tool-filter.ts @@ -0,0 +1,47 @@ +/** + * Tool filtering utilities for MCP server configurations. + */ + +/** + * Tool filter configuration + */ +export interface ToolFilters { + includeTools?: Set; + excludeTools?: Set; +} + +/** + * Check if a tool should be included based on includeTools and excludeTools filters. + * + * @param toolName - The name of the tool to check + * @param filters - The filter configuration + * @returns true if the tool should be included, false otherwise + * + * Filter logic: + * 1. If includeTools is specified, tool must be in the list + * 2. If excludeTools is specified, tool must not be in the list + * 3. Exclude filters are applied after include filters + * 4. If no filters are specified, all tools are included + */ +export function shouldIncludeTool( + toolName: string, + filters: ToolFilters, +): boolean { + const { includeTools, excludeTools } = filters; + + // If includeTools is specified, tool must be in the list + if (includeTools && includeTools.size > 0) { + if (!includeTools.has(toolName)) { + return false; + } + } + + // If excludeTools is specified, tool must not be in the list + if (excludeTools && excludeTools.size > 0) { + if (excludeTools.has(toolName)) { + return false; + } + } + + return true; +}