From c6a09a35e23ded5bcba79b5af3445f6e04f8b0fe Mon Sep 17 00:00:00 2001 From: afjcjsbx Date: Fri, 8 May 2026 13:48:47 +0200 Subject: [PATCH] fix(agent): suppress MCP discovery when no servers are selectable --- pkg/agent/agent_mcp.go | 18 ++++++++ pkg/agent/agent_mcp_test.go | 65 +++++++++++++++++++++++++++ pkg/agent/instance.go | 2 +- pkg/agent/instance_test.go | 90 +++++++++++++++++++++++++++++++++++++ pkg/agent/tool_allowlist.go | 10 ++++- 5 files changed, 183 insertions(+), 2 deletions(-) diff --git a/pkg/agent/agent_mcp.go b/pkg/agent/agent_mcp.go index 3d569b2bd..e8cdf81c8 100644 --- a/pkg/agent/agent_mcp.go +++ b/pkg/agent/agent_mcp.go @@ -250,6 +250,9 @@ func (al *AgentLoop) ensureMCPInitialized(ctx context.Context) error { if !ok { continue } + if !agentHasDiscoverableMCPServers(al.cfg, agent.MCPServerAllowlist) { + continue + } if useRegex { agent.Tools.Register(tools.NewRegexSearchTool(agent.Tools, ttl, maxSearchResults)) @@ -334,6 +337,21 @@ func filterMCPConfigServers( return filtered } +func agentHasDiscoverableMCPServers(cfg *config.Config, allowed map[string]struct{}) bool { + if cfg == nil || !cfg.Tools.MCP.Enabled || !cfg.Tools.MCP.Discovery.Enabled { + return false + } + + filtered := filterMCPConfigServers(cfg.Tools.MCP, allowed) + for _, serverCfg := range filtered.Servers { + if serverCfg.Enabled && serverIsDeferred(cfg.Tools.MCP.Discovery.Enabled, serverCfg) { + return true + } + } + + return false +} + // serverIsDeferred reports whether an MCP server's tools should be registered // as hidden (deferred/discovery mode). // diff --git a/pkg/agent/agent_mcp_test.go b/pkg/agent/agent_mcp_test.go index 7c8a4cd28..f85861146 100644 --- a/pkg/agent/agent_mcp_test.go +++ b/pkg/agent/agent_mcp_test.go @@ -204,6 +204,71 @@ func TestFilterMCPConfigServersCaseInsensitivePreservesOriginalKeys(t *testing.T } } +func TestAgentHasDiscoverableMCPServers(t *testing.T) { + deferredFalse := false + cfg := &config.Config{ + Tools: config.ToolsConfig{ + MCP: config.MCPConfig{ + ToolConfig: config.ToolConfig{Enabled: true}, + Discovery: config.ToolDiscoveryConfig{ + Enabled: true, + UseBM25: true, + UseRegex: false, + }, + Servers: map[string]config.MCPServerConfig{ + "github": {Enabled: true}, + "filesystem": {Enabled: true, Deferred: &deferredFalse}, + }, + }, + }, + } + + tests := []struct { + name string + allowed map[string]struct{} + want bool + }{ + { + name: "nil allowlist includes discoverable enabled server", + want: true, + }, + { + name: "empty allowlist denies all servers", + allowed: map[string]struct{}{}, + want: false, + }, + { + name: "selected server discoverable", + allowed: map[string]struct{}{ + "github": {}, + }, + want: true, + }, + { + name: "selected server opted out of discovery", + allowed: map[string]struct{}{ + "filesystem": {}, + }, + want: false, + }, + { + name: "unknown allowlist server matches nothing", + allowed: map[string]struct{}{ + "slack": {}, + }, + want: false, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + if got := agentHasDiscoverableMCPServers(cfg, tt.allowed); got != tt.want { + t.Fatalf("agentHasDiscoverableMCPServers() = %v, want %v", got, tt.want) + } + }) + } +} + func TestEnsureMCPInitialized_LoadFailureSetsInitErr(t *testing.T) { al, cfg, _, _, cleanup := newTestAgentLoop(t) defer cleanup() diff --git a/pkg/agent/instance.go b/pkg/agent/instance.go index 63aac150b..4ed713035 100644 --- a/pkg/agent/instance.go +++ b/pkg/agent/instance.go @@ -128,7 +128,7 @@ func NewAgentInstance( sessionsDir := filepath.Join(workspace, "sessions") sessions := initSessionStore(sessionsDir) - mcpDiscoveryActive := cfg.Tools.MCP.Enabled && cfg.Tools.MCP.Discovery.Enabled + mcpDiscoveryActive := agentHasDiscoverableMCPServers(cfg, agentMCPServerAllowlist) contextBuilder := NewContextBuilder(workspace). WithToolDiscovery( mcpDiscoveryActive && cfg.Tools.MCP.Discovery.UseBM25, diff --git a/pkg/agent/instance_test.go b/pkg/agent/instance_test.go index 97a5dde67..76e1b7f2d 100644 --- a/pkg/agent/instance_test.go +++ b/pkg/agent/instance_test.go @@ -10,6 +10,7 @@ import ( "github.com/sipeed/picoclaw/pkg/config" "github.com/sipeed/picoclaw/pkg/media" "github.com/sipeed/picoclaw/pkg/providers" + "github.com/sipeed/picoclaw/pkg/tools" ) func TestNewAgentInstance_UsesDefaultsTemperatureAndMaxTokens(t *testing.T) { @@ -717,6 +718,95 @@ model: claude-frontmatter } } +func TestNewAgentInstance_SuppressesToolDiscoveryPromptWhenNoMCPServersSelected(t *testing.T) { + workspace := setupWorkspace(t, map[string]string{ + "AGENT.md": `--- +mcpServers: [] +--- +# Agent +`, + }) + defer cleanupWorkspace(t, workspace) + + cfg := &config.Config{ + Agents: config.AgentsConfig{ + Defaults: config.AgentDefaults{ + Workspace: workspace, + ModelName: "default-model", + }, + }, + Tools: config.ToolsConfig{ + MCP: config.MCPConfig{ + ToolConfig: config.ToolConfig{Enabled: true}, + Discovery: config.ToolDiscoveryConfig{ + Enabled: true, + UseBM25: true, + UseRegex: false, + }, + Servers: map[string]config.MCPServerConfig{ + "github": {Enabled: true}, + }, + }, + }, + } + + agent := NewAgentInstance(&config.AgentConfig{ + ID: "research", + Workspace: workspace, + }, &cfg.Agents.Defaults, cfg, &mockProvider{}) + + if agent.AllowsMCPServer("github") { + t.Fatal("expected empty mcpServers allowlist to deny all servers") + } + messages := agent.ContextBuilder.BuildMessagesFromPrompt(PromptBuildRequest{CurrentMessage: "hello"}) + if prompt := messages[0].Content; strings.Contains(prompt, tools.BM25SearchToolName) { + t.Fatalf("expected no tool discovery prompt when no MCP servers are selected, got %q", prompt) + } +} + +func TestNewAgentInstance_IncludesToolDiscoveryPromptWhenDiscoverableMCPServerSelected(t *testing.T) { + workspace := setupWorkspace(t, map[string]string{ + "AGENT.md": `--- +mcpServers: [github] +--- +# Agent +`, + }) + defer cleanupWorkspace(t, workspace) + + cfg := &config.Config{ + Agents: config.AgentsConfig{ + Defaults: config.AgentDefaults{ + Workspace: workspace, + ModelName: "default-model", + }, + }, + Tools: config.ToolsConfig{ + MCP: config.MCPConfig{ + ToolConfig: config.ToolConfig{Enabled: true}, + Discovery: config.ToolDiscoveryConfig{ + Enabled: true, + UseBM25: true, + UseRegex: false, + }, + Servers: map[string]config.MCPServerConfig{ + "github": {Enabled: true}, + }, + }, + }, + } + + agent := NewAgentInstance(&config.AgentConfig{ + ID: "research", + Workspace: workspace, + }, &cfg.Agents.Defaults, cfg, &mockProvider{}) + + messages := agent.ContextBuilder.BuildMessagesFromPrompt(PromptBuildRequest{CurrentMessage: "hello"}) + if prompt := messages[0].Content; !strings.Contains(prompt, tools.BM25SearchToolName) { + t.Fatalf("expected tool discovery prompt when a discoverable MCP server is selected, got %q", prompt) + } +} + func TestNewAgentInstance_InvalidFrontmatterFailsClosedForToolsAndMCPServers(t *testing.T) { workspace := setupWorkspace(t, map[string]string{ "AGENT.md": `--- diff --git a/pkg/agent/tool_allowlist.go b/pkg/agent/tool_allowlist.go index 7a020c82c..ad7394c7d 100644 --- a/pkg/agent/tool_allowlist.go +++ b/pkg/agent/tool_allowlist.go @@ -164,7 +164,7 @@ func resolveAgentMCPServerAllowlist(definition AgentContextDefinition) map[strin if frontmatterParseFailed(definition) { return map[string]struct{}{} } - if definition.Agent == nil || definition.Agent.Frontmatter.MCPServers == nil { + if definition.Agent == nil || !frontmatterDeclaresField(definition, "mcpServers") { return nil } @@ -180,6 +180,14 @@ func resolveAgentMCPServerAllowlist(definition AgentContextDefinition) map[strin return allowlist } +func frontmatterDeclaresField(definition AgentContextDefinition, field string) bool { + if definition.Agent == nil || definition.Agent.Frontmatter.Fields == nil { + return false + } + _, ok := definition.Agent.Frontmatter.Fields[field] + return ok +} + func frontmatterParseFailed(definition AgentContextDefinition) bool { if definition.Agent == nil { return false