From 27bd816b1c832fdfae662c64842cfdb183e02d14 Mon Sep 17 00:00:00 2001 From: afjcjsbx Date: Thu, 7 May 2026 13:49:23 +0200 Subject: [PATCH] fix(agent): validate AGENT tool declarations from registry --- pkg/agent/agent_init.go | 2 + pkg/agent/instance.go | 4 +- pkg/agent/tool_allowlist.go | 102 +++++++++++++------------------ pkg/agent/tool_allowlist_test.go | 55 ++++++++++++++--- 4 files changed, 92 insertions(+), 71 deletions(-) diff --git a/pkg/agent/agent_init.go b/pkg/agent/agent_init.go index e95fbe7f8..8420cd101 100644 --- a/pkg/agent/agent_init.go +++ b/pkg/agent/agent_init.go @@ -352,5 +352,7 @@ func registerSharedTools( }) agent.Tools.Register(delegateTool) } + + warnOnUnknownAgentToolDeclarations(agentID, agent.Workspace, agent.Definition, agent.Tools) } } diff --git a/pkg/agent/instance.go b/pkg/agent/instance.go index 6d629ac57..ac2955334 100644 --- a/pkg/agent/instance.go +++ b/pkg/agent/instance.go @@ -38,6 +38,7 @@ type AgentInstance struct { Sessions session.SessionStore ContextBuilder *ContextBuilder Tools *tools.ToolRegistry + Definition AgentContextDefinition Subagents *config.SubagentsConfig SkillsFilter []string MCPServerAllowlist map[string]struct{} @@ -149,7 +150,7 @@ func NewAgentInstance( subagents = agentCfg.Subagents skillsFilter = resolveAgentSkillsFilter(agentCfg, definition) } - warnOnUnknownAgentDeclarations(agentID, workspace, cfg, definition) + warnOnUnknownAgentMCPServerDeclarations(agentID, workspace, cfg, definition) maxIter := defaults.MaxToolIterations if maxIter == 0 { @@ -256,6 +257,7 @@ func NewAgentInstance( Sessions: sessions, ContextBuilder: contextBuilder, Tools: toolsRegistry, + Definition: definition, Subagents: subagents, SkillsFilter: skillsFilter, MCPServerAllowlist: agentMCPServerAllowlist, diff --git a/pkg/agent/tool_allowlist.go b/pkg/agent/tool_allowlist.go index b220f1903..87ae2ee4c 100644 --- a/pkg/agent/tool_allowlist.go +++ b/pkg/agent/tool_allowlist.go @@ -6,11 +6,31 @@ import ( "github.com/sipeed/picoclaw/pkg/config" "github.com/sipeed/picoclaw/pkg/logger" + "github.com/sipeed/picoclaw/pkg/tools" ) const dynamicMCPToolPrefix = "mcp_" -func warnOnUnknownAgentDeclarations( +func warnOnUnknownAgentToolDeclarations( + agentID, workspace string, + definition AgentContextDefinition, + registry *tools.ToolRegistry, +) { + if registry == nil || frontmatterParseFailed(definition) { + return + } + + if unknownTools := unknownAgentToolNames(registry, definition); len(unknownTools) > 0 { + logger.WarnCF("agent", "AGENT.md declares unregistered tool names", + map[string]any{ + "agent_id": agentID, + "workspace": workspace, + "tools": unknownTools, + }) + } +} + +func warnOnUnknownAgentMCPServerDeclarations( agentID, workspace string, cfg *config.Config, definition AgentContextDefinition, @@ -19,15 +39,6 @@ func warnOnUnknownAgentDeclarations( return } - if unknownTools := unknownAgentToolNames(cfg, definition); len(unknownTools) > 0 { - logger.WarnCF("agent", "AGENT.md declares unknown tool names", - map[string]any{ - "agent_id": agentID, - "workspace": workspace, - "tools": unknownTools, - }) - } - if unknownServers := unknownAgentMCPServerNames(cfg, definition); len(unknownServers) > 0 { logger.WarnCF("agent", "AGENT.md declares unknown MCP server names", map[string]any{ @@ -38,12 +49,15 @@ func warnOnUnknownAgentDeclarations( } } -func unknownAgentToolNames(cfg *config.Config, definition AgentContextDefinition) []string { +func unknownAgentToolNames( + registry *tools.ToolRegistry, + definition AgentContextDefinition, +) []string { if definition.Agent == nil || definition.Agent.Frontmatter.Tools == nil { return nil } - known := knownRuntimeToolNames(cfg) + known := registeredRuntimeToolNames(registry) unknown := make(map[string]struct{}) for _, raw := range definition.Agent.Frontmatter.Tools { name := strings.ToLower(strings.TrimSpace(raw)) @@ -59,6 +73,21 @@ func unknownAgentToolNames(cfg *config.Config, definition AgentContextDefinition return sortedKeys(unknown) } +func registeredRuntimeToolNames(registry *tools.ToolRegistry) map[string]struct{} { + known := make(map[string]struct{}) + if registry == nil { + return known + } + for _, raw := range registry.List() { + name := strings.ToLower(strings.TrimSpace(raw)) + if name == "" { + continue + } + known[name] = struct{}{} + } + return known +} + func unknownAgentMCPServerNames(cfg *config.Config, definition AgentContextDefinition) []string { if cfg == nil || definition.Agent == nil || definition.Agent.Frontmatter.MCPServers == nil { return nil @@ -79,55 +108,6 @@ func unknownAgentMCPServerNames(cfg *config.Config, definition AgentContextDefin return sortedKeys(unknown) } -func knownRuntimeToolNames(cfg *config.Config) map[string]struct{} { - known := make(map[string]struct{}) - if cfg == nil { - return known - } - - addKnownToolIfEnabled(known, cfg.Tools.IsToolEnabled("read_file"), "read_file") - addKnownToolIfEnabled(known, cfg.Tools.IsToolEnabled("write_file"), "write_file") - addKnownToolIfEnabled(known, cfg.Tools.IsToolEnabled("list_dir"), "list_dir") - addKnownToolIfEnabled(known, cfg.Tools.IsToolEnabled("exec"), "exec") - addKnownToolIfEnabled(known, cfg.Tools.IsToolEnabled("edit_file"), "edit_file") - addKnownToolIfEnabled(known, cfg.Tools.IsToolEnabled("append_file"), "append_file") - addKnownToolIfEnabled(known, cfg.Tools.IsToolEnabled("cron"), "cron") - addKnownToolIfEnabled(known, cfg.Tools.IsToolEnabled("web"), "web_search") - addKnownToolIfEnabled(known, cfg.Tools.IsToolEnabled("web_fetch"), "web_fetch") - addKnownToolIfEnabled(known, cfg.Tools.IsToolEnabled("i2c"), "i2c") - addKnownToolIfEnabled(known, cfg.Tools.IsToolEnabled("spi"), "spi") - addKnownToolIfEnabled(known, cfg.Tools.IsToolEnabled("message"), "message") - addKnownToolIfEnabled(known, cfg.Tools.IsToolEnabled("send_file"), "send_file") - addKnownToolIfEnabled( - known, - cfg.Tools.IsToolEnabled("skills") && cfg.Tools.IsToolEnabled("find_skills"), - "find_skills", - ) - addKnownToolIfEnabled( - known, - cfg.Tools.IsToolEnabled("skills") && cfg.Tools.IsToolEnabled("install_skill"), - "install_skill", - ) - if cfg.Tools.IsToolEnabled("subagent") { - addKnownToolIfEnabled(known, cfg.Tools.IsToolEnabled("spawn"), "spawn") - addKnownToolIfEnabled(known, cfg.Tools.IsToolEnabled("subagent"), "subagent") - addKnownToolIfEnabled(known, cfg.Tools.IsToolEnabled("spawn_status"), "spawn_status") - } - if cfg.Tools.IsToolEnabled("mcp") && cfg.Tools.MCP.Discovery.Enabled { - addKnownToolIfEnabled(known, cfg.Tools.MCP.Discovery.UseRegex, "tool_search_tool_regex") - addKnownToolIfEnabled(known, cfg.Tools.MCP.Discovery.UseBM25, "tool_search_tool_bm25") - } - - return known -} - -func addKnownToolIfEnabled(known map[string]struct{}, enabled bool, name string) { - if !enabled { - return - } - known[name] = struct{}{} -} - func sortedKeys(values map[string]struct{}) []string { if len(values) == 0 { return nil diff --git a/pkg/agent/tool_allowlist_test.go b/pkg/agent/tool_allowlist_test.go index 059ee9344..46bbac2bc 100644 --- a/pkg/agent/tool_allowlist_test.go +++ b/pkg/agent/tool_allowlist_test.go @@ -1,11 +1,32 @@ package agent import ( + "context" "testing" "github.com/sipeed/picoclaw/pkg/config" + agenttools "github.com/sipeed/picoclaw/pkg/tools" ) +type allowlistTestTool struct { + name string +} + +func (t *allowlistTestTool) Name() string { return t.name } + +func (t *allowlistTestTool) Description() string { return "test tool" } + +func (t *allowlistTestTool) Parameters() map[string]any { + return map[string]any{"type": "object"} +} + +func (t *allowlistTestTool) Execute( + _ context.Context, + _ map[string]any, +) *agenttools.ToolResult { + return agenttools.NewToolResult("ok") +} + func TestUnknownAgentToolNames(t *testing.T) { workspace := setupWorkspace(t, map[string]string{ "AGENT.md": `--- @@ -16,21 +37,37 @@ tools: [read_file, web_serach, mcp_github_search] }) defer cleanupWorkspace(t, workspace) - cfg := &config.Config{ - Tools: config.ToolsConfig{ - ReadFile: config.ReadFileToolConfig{Enabled: true}, - Web: config.WebToolsConfig{ - ToolConfig: config.ToolConfig{Enabled: true}, - }, - }, - } + registry := agenttools.NewToolRegistry() + registry.Register(&allowlistTestTool{name: "read_file"}) + registry.Register(&allowlistTestTool{name: "web_search"}) - unknown := unknownAgentToolNames(cfg, loadAgentDefinition(workspace)) + unknown := unknownAgentToolNames(registry, loadAgentDefinition(workspace)) if len(unknown) != 1 || unknown[0] != "web_serach" { t.Fatalf("unknownAgentToolNames() = %v, want [web_serach]", unknown) } } +func TestUnknownAgentToolNamesUsesRegisteredRuntimeTools(t *testing.T) { + workspace := setupWorkspace(t, map[string]string{ + "AGENT.md": `--- +tools: [serial, reaction, send_tts, load_image, delegate, made_up] +--- +# Agent +`, + }) + defer cleanupWorkspace(t, workspace) + + registry := agenttools.NewToolRegistry() + for _, name := range []string{"serial", "reaction", "send_tts", "load_image", "delegate"} { + registry.Register(&allowlistTestTool{name: name}) + } + + unknown := unknownAgentToolNames(registry, loadAgentDefinition(workspace)) + if len(unknown) != 1 || unknown[0] != "made_up" { + t.Fatalf("unknownAgentToolNames() = %v, want [made_up]", unknown) + } +} + func TestUnknownAgentMCPServerNames(t *testing.T) { workspace := setupWorkspace(t, map[string]string{ "AGENT.md": `---