fix(agent): suppress MCP discovery when no servers are selectable
This commit is contained in:
parent
ffa184d183
commit
c6a09a35e2
5 changed files with 183 additions and 2 deletions
|
|
@ -250,6 +250,9 @@ func (al *AgentLoop) ensureMCPInitialized(ctx context.Context) error {
|
||||||
if !ok {
|
if !ok {
|
||||||
continue
|
continue
|
||||||
}
|
}
|
||||||
|
if !agentHasDiscoverableMCPServers(al.cfg, agent.MCPServerAllowlist) {
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
|
||||||
if useRegex {
|
if useRegex {
|
||||||
agent.Tools.Register(tools.NewRegexSearchTool(agent.Tools, ttl, maxSearchResults))
|
agent.Tools.Register(tools.NewRegexSearchTool(agent.Tools, ttl, maxSearchResults))
|
||||||
|
|
@ -334,6 +337,21 @@ func filterMCPConfigServers(
|
||||||
return filtered
|
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
|
// serverIsDeferred reports whether an MCP server's tools should be registered
|
||||||
// as hidden (deferred/discovery mode).
|
// as hidden (deferred/discovery mode).
|
||||||
//
|
//
|
||||||
|
|
|
||||||
|
|
@ -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) {
|
func TestEnsureMCPInitialized_LoadFailureSetsInitErr(t *testing.T) {
|
||||||
al, cfg, _, _, cleanup := newTestAgentLoop(t)
|
al, cfg, _, _, cleanup := newTestAgentLoop(t)
|
||||||
defer cleanup()
|
defer cleanup()
|
||||||
|
|
|
||||||
|
|
@ -128,7 +128,7 @@ func NewAgentInstance(
|
||||||
sessionsDir := filepath.Join(workspace, "sessions")
|
sessionsDir := filepath.Join(workspace, "sessions")
|
||||||
sessions := initSessionStore(sessionsDir)
|
sessions := initSessionStore(sessionsDir)
|
||||||
|
|
||||||
mcpDiscoveryActive := cfg.Tools.MCP.Enabled && cfg.Tools.MCP.Discovery.Enabled
|
mcpDiscoveryActive := agentHasDiscoverableMCPServers(cfg, agentMCPServerAllowlist)
|
||||||
contextBuilder := NewContextBuilder(workspace).
|
contextBuilder := NewContextBuilder(workspace).
|
||||||
WithToolDiscovery(
|
WithToolDiscovery(
|
||||||
mcpDiscoveryActive && cfg.Tools.MCP.Discovery.UseBM25,
|
mcpDiscoveryActive && cfg.Tools.MCP.Discovery.UseBM25,
|
||||||
|
|
|
||||||
|
|
@ -10,6 +10,7 @@ import (
|
||||||
"github.com/sipeed/picoclaw/pkg/config"
|
"github.com/sipeed/picoclaw/pkg/config"
|
||||||
"github.com/sipeed/picoclaw/pkg/media"
|
"github.com/sipeed/picoclaw/pkg/media"
|
||||||
"github.com/sipeed/picoclaw/pkg/providers"
|
"github.com/sipeed/picoclaw/pkg/providers"
|
||||||
|
"github.com/sipeed/picoclaw/pkg/tools"
|
||||||
)
|
)
|
||||||
|
|
||||||
func TestNewAgentInstance_UsesDefaultsTemperatureAndMaxTokens(t *testing.T) {
|
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) {
|
func TestNewAgentInstance_InvalidFrontmatterFailsClosedForToolsAndMCPServers(t *testing.T) {
|
||||||
workspace := setupWorkspace(t, map[string]string{
|
workspace := setupWorkspace(t, map[string]string{
|
||||||
"AGENT.md": `---
|
"AGENT.md": `---
|
||||||
|
|
|
||||||
|
|
@ -164,7 +164,7 @@ func resolveAgentMCPServerAllowlist(definition AgentContextDefinition) map[strin
|
||||||
if frontmatterParseFailed(definition) {
|
if frontmatterParseFailed(definition) {
|
||||||
return map[string]struct{}{}
|
return map[string]struct{}{}
|
||||||
}
|
}
|
||||||
if definition.Agent == nil || definition.Agent.Frontmatter.MCPServers == nil {
|
if definition.Agent == nil || !frontmatterDeclaresField(definition, "mcpServers") {
|
||||||
return nil
|
return nil
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|
@ -180,6 +180,14 @@ func resolveAgentMCPServerAllowlist(definition AgentContextDefinition) map[strin
|
||||||
return allowlist
|
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 {
|
func frontmatterParseFailed(definition AgentContextDefinition) bool {
|
||||||
if definition.Agent == nil {
|
if definition.Agent == nil {
|
||||||
return false
|
return false
|
||||||
|
|
|
||||||
Loading…
Reference in a new issue