chore: separate the LoadSession function in the MCP loader (#966)

enhance: add the ability to specify allow tools for MCP servers

Signed-off-by: Donnie Adams <[email protected]>

Author: thedadams

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)

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",