diff --git a/pkg/mcp/loader.go b/pkg/mcp/loader.go index 0eb713e5..72c0ebf4 100644 --- a/pkg/mcp/loader.go +++ b/pkg/mcp/loader.go @@ -56,6 +56,7 @@ type ServerConfig struct { BaseURL string `json:"baseURL,omitempty"` Headers []string `json:"headers"` Scope string `json:"scope"` + AllowedTools []string `json:"allowedTools"` } func (s *ServerConfig) GetBaseURL() string { @@ -99,18 +100,31 @@ func (l *Local) Load(ctx context.Context, tool types.Tool) (result []types.Tool, } for server := range maps.Keys(servers.MCPServers) { - session, err := l.loadSession(servers.MCPServers[server]) + tools, err := l.LoadTools(ctx, servers.MCPServers[server], tool.Name) if err != nil { return nil, fmt.Errorf("failed to load MCP session for server %s: %w", server, err) } - return l.sessionToTools(ctx, session, tool.Name) + return tools, nil } // This should never happen, but just in case return nil, fmt.Errorf("no MCP server configuration found in tool instructions: %s", configData) } +func (l *Local) LoadTools(ctx context.Context, server ServerConfig, toolName string) ([]types.Tool, error) { + allowedTools := server.AllowedTools + // Reset so we don't start a new MCP server, no reason to if one is already running and the allowed tools change. + server.AllowedTools = nil + + session, err := l.loadSession(server) + if err != nil { + return nil, err + } + + return l.sessionToTools(ctx, session, toolName, allowedTools) +} + func (l *Local) Close() error { if l == nil { return nil @@ -139,7 +153,9 @@ func (l *Local) Close() error { return errors.Join(errs...) } -func (l *Local) sessionToTools(ctx context.Context, session *Session, toolName string) ([]types.Tool, error) { +func (l *Local) sessionToTools(ctx context.Context, session *Session, toolName string, allowedTools []string) ([]types.Tool, error) { + allToolsAllowed := len(allowedTools) == 0 || slices.Contains(allowedTools, "*") + tools, err := session.Client.ListTools(ctx, mcp.ListToolsRequest{}) if err != nil { return nil, fmt.Errorf("failed to list tools: %w", err) @@ -149,6 +165,10 @@ func (l *Local) sessionToTools(ctx context.Context, session *Session, toolName s var toolNames []string for _, tool := range tools.Tools { + if !allToolsAllowed && !slices.Contains(allowedTools, tool.Name) { + continue + } + var schema openapi3.Schema schemaData, err := json.Marshal(tool.InputSchema) diff --git a/pkg/tests/runner2_test.go b/pkg/tests/runner2_test.go index c531c661..75253a7b 100644 --- a/pkg/tests/runner2_test.go +++ b/pkg/tests/runner2_test.go @@ -313,7 +313,7 @@ name: mcp ], "type": "object" }, - "instructions": "#!sys.mcp.invoke.append_insight 607ca64476abf0288ef49061557243e43735fd4de4bc5fdcd51d93049ffa023e", + "instructions": "#!sys.mcp.invoke.append_insight c358c2eb93fa9a98631cd9e4f324d7b59f56aee11c7ae32a00984ad5844dc32c", "id": "inline:append_insight", "localTools": { "append_insight": "inline:append_insight", @@ -346,7 +346,7 @@ name: mcp ], "type": "object" }, - "instructions": "#!sys.mcp.invoke.create_table 607ca64476abf0288ef49061557243e43735fd4de4bc5fdcd51d93049ffa023e", + "instructions": "#!sys.mcp.invoke.create_table c358c2eb93fa9a98631cd9e4f324d7b59f56aee11c7ae32a00984ad5844dc32c", "id": "inline:create_table", "localTools": { "append_insight": "inline:append_insight", @@ -379,7 +379,7 @@ name: mcp ], "type": "object" }, - "instructions": "#!sys.mcp.invoke.describe_table 607ca64476abf0288ef49061557243e43735fd4de4bc5fdcd51d93049ffa023e", + "instructions": "#!sys.mcp.invoke.describe_table c358c2eb93fa9a98631cd9e4f324d7b59f56aee11c7ae32a00984ad5844dc32c", "id": "inline:describe_table", "localTools": { "append_insight": "inline:append_insight", @@ -403,7 +403,7 @@ name: mcp "arguments": { "type": "object" }, - "instructions": "#!sys.mcp.invoke.list_tables 607ca64476abf0288ef49061557243e43735fd4de4bc5fdcd51d93049ffa023e", + "instructions": "#!sys.mcp.invoke.list_tables c358c2eb93fa9a98631cd9e4f324d7b59f56aee11c7ae32a00984ad5844dc32c", "id": "inline:list_tables", "localTools": { "append_insight": "inline:append_insight", @@ -505,7 +505,7 @@ name: mcp ], "type": "object" }, - "instructions": "#!sys.mcp.invoke.read_query 607ca64476abf0288ef49061557243e43735fd4de4bc5fdcd51d93049ffa023e", + "instructions": "#!sys.mcp.invoke.read_query c358c2eb93fa9a98631cd9e4f324d7b59f56aee11c7ae32a00984ad5844dc32c", "id": "inline:read_query", "localTools": { "append_insight": "inline:append_insight", @@ -538,7 +538,7 @@ name: mcp ], "type": "object" }, - "instructions": "#!sys.mcp.invoke.write_query 607ca64476abf0288ef49061557243e43735fd4de4bc5fdcd51d93049ffa023e", + "instructions": "#!sys.mcp.invoke.write_query c358c2eb93fa9a98631cd9e4f324d7b59f56aee11c7ae32a00984ad5844dc32c", "id": "inline:write_query", "localTools": { "append_insight": "inline:append_insight",