diff --git a/pkg/agent/context.go b/pkg/agent/context.go index 7bbf6b16e..92663a32d 100644 --- a/pkg/agent/context.go +++ b/pkg/agent/context.go @@ -18,10 +18,11 @@ import ( ) type ContextBuilder struct { - workspace string - skillsLoader *skills.SkillsLoader - memory *MemoryStore - toolDiscovery bool + workspace string + skillsLoader *skills.SkillsLoader + memory *MemoryStore + toolDiscoveryBM25 bool + toolDiscoveryRegex bool // Cache for system prompt to avoid rebuilding on every call. // This fixes issue #607: repeated reprocessing of the entire context. @@ -42,8 +43,9 @@ type ContextBuilder struct { skillFilesAtCache map[string]time.Time } -func (cb *ContextBuilder) WithToolDiscovery(enabled bool) *ContextBuilder { - cb.toolDiscovery = enabled +func (cb *ContextBuilder) WithToolDiscovery(useBM25, useRegex bool) *ContextBuilder { + cb.toolDiscoveryBM25 = useBM25 + cb.toolDiscoveryRegex = useRegex return cb } @@ -104,10 +106,22 @@ Your workspace is at: %s } func (cb *ContextBuilder) getDiscoveryRule() string { - if cb.toolDiscovery { - return `5. **Tool Discovery** - Your visible tools are limited to save memory, but a vast hidden library exists. If you lack the right tool for a task, BEFORE giving up, you MUST search using the "tool_search_tool_bm25" or "tool_search_tool_regex" tool. Do not refuse a request unless the search returns nothing. Found tools will temporarily unlock for your next turn.` + if !cb.toolDiscoveryBM25 && !cb.toolDiscoveryRegex { + return "" } - return "" + + var toolNames []string + if cb.toolDiscoveryBM25 { + toolNames = append(toolNames, `"tool_search_tool_bm25"`) + } + if cb.toolDiscoveryRegex { + toolNames = append(toolNames, `"tool_search_tool_regex"`) + } + + return fmt.Sprintf( + `5. **Tool Discovery** - Your visible tools are limited to save memory, but a vast hidden library exists. If you lack the right tool for a task, BEFORE giving up, you MUST search using the %s tool. Do not refuse a request unless the search returns nothing. Found tools will temporarily unlock for your next turn.`, + strings.Join(toolNames, " or "), + ) } func (cb *ContextBuilder) BuildSystemPrompt() string { diff --git a/pkg/agent/instance.go b/pkg/agent/instance.go index 193cc648b..cefc87d66 100644 --- a/pkg/agent/instance.go +++ b/pkg/agent/instance.go @@ -96,7 +96,10 @@ func NewAgentInstance( sessionsDir := filepath.Join(workspace, "sessions") sessionsManager := session.NewSessionManager(sessionsDir) - contextBuilder := NewContextBuilder(workspace).WithToolDiscovery(cfg.Tools.MCP.ToolConfig.Discovery.Enabled) + contextBuilder := NewContextBuilder(workspace).WithToolDiscovery( + cfg.Tools.MCP.Discovery.Enabled && cfg.Tools.MCP.Discovery.UseBM25, + cfg.Tools.MCP.Discovery.Enabled && cfg.Tools.MCP.Discovery.UseRegex, + ) agentID := routing.DefaultAgentID agentName := "" diff --git a/pkg/agent/loop.go b/pkg/agent/loop.go index 5103085b5..0881ce93c 100644 --- a/pkg/agent/loop.go +++ b/pkg/agent/loop.go @@ -11,7 +11,6 @@ import ( "encoding/json" "errors" "fmt" - "log" "path/filepath" "regexp" "strings" @@ -285,7 +284,7 @@ func (al *AgentLoop) Run(ctx context.Context) error { mcpTool := tools.NewMCPTool(mcpManager, serverName, tool) - if al.cfg.Tools.MCP.ToolConfig.Enabled { + if al.cfg.Tools.MCP.Discovery.Enabled { agent.Tools.RegisterHidden(mcpTool) } else { agent.Tools.Register(mcpTool) @@ -315,10 +314,10 @@ func (al *AgentLoop) Run(ctx context.Context) error { useBM25 := al.cfg.Tools.MCP.Discovery.UseBM25 useRegex := al.cfg.Tools.MCP.Discovery.UseRegex - // Fail fast: If discovery is enabled but no tool is turned on, break the app + // Fail fast: If discovery is enabled but no search method is turned on if !useBM25 && !useRegex { - log.Fatalf( - "Critical error: tool discovery is enabled but neither 'use_bm25' nor 'use_regex' is set to true in the configuration.", + return fmt.Errorf( + "tool discovery is enabled but neither 'use_bm25' nor 'use_regex' is set to true in the configuration", ) } @@ -328,6 +327,9 @@ func (al *AgentLoop) Run(ctx context.Context) error { } maxSearchResults := al.cfg.Tools.MCP.Discovery.MaxSearchResults + if maxSearchResults <= 0 { + maxSearchResults = 5 // Default value + } for _, agentID := range agentIDs { agent, ok := al.registry.GetAgent(agentID) @@ -931,9 +933,6 @@ func (al *AgentLoop) runLLMIteration( "max": agent.MaxIterations, }) - // Scale down the TTL of the discovered tools with each new round of the LLM - agent.Tools.TickTTL() - // Build tool definitions providerToolDefs := agent.Tools.ToProviderDefs() @@ -1299,6 +1298,11 @@ func (al *AgentLoop) runLLMIteration( // Save tool result message to session agent.Sessions.AddFullMessage(opts.SessionKey, toolResultMsg) } + + // Tick down TTL of discovered tools after processing tool results. + // Only reached when tool calls were made (the loop continues); + // the break on no-tool-call responses skips this. + agent.Tools.TickTTL() } return finalContent, iteration, nil diff --git a/pkg/config/config.go b/pkg/config/config.go index 7fbf56518..15fcfff64 100644 --- a/pkg/config/config.go +++ b/pkg/config/config.go @@ -568,8 +568,7 @@ type ToolDiscoveryConfig struct { } type ToolConfig struct { - Enabled bool `json:"enabled" env:"ENABLED"` - Discovery ToolDiscoveryConfig `json:"discovery"` + Enabled bool `json:"enabled" env:"ENABLED"` } type BraveConfig struct { @@ -720,7 +719,8 @@ type MCPServerConfig struct { // MCPConfig defines configuration for all MCP servers type MCPConfig struct { - ToolConfig `envPrefix:"PICOCLAW_TOOLS_MCP_"` + ToolConfig ` envPrefix:"PICOCLAW_TOOLS_MCP_"` + Discovery ToolDiscoveryConfig ` json:"discovery"` // Servers is a map of server name to server configuration Servers map[string]MCPServerConfig `json:"servers,omitempty"` } diff --git a/pkg/config/defaults.go b/pkg/config/defaults.go index 6049e0273..55508d886 100644 --- a/pkg/config/defaults.go +++ b/pkg/config/defaults.go @@ -410,13 +410,13 @@ func DefaultConfig() *Config { MCP: MCPConfig{ ToolConfig: ToolConfig{ Enabled: false, - Discovery: ToolDiscoveryConfig{ - Enabled: false, - TTL: 5, - MaxSearchResults: 5, - UseBM25: true, - UseRegex: false, - }, + }, + Discovery: ToolDiscoveryConfig{ + Enabled: false, + TTL: 5, + MaxSearchResults: 5, + UseBM25: true, + UseRegex: false, }, Servers: map[string]MCPServerConfig{}, }, diff --git a/pkg/tools/registry.go b/pkg/tools/registry.go index 7dbc65e80..a1107039d 100644 --- a/pkg/tools/registry.go +++ b/pkg/tools/registry.go @@ -5,6 +5,7 @@ import ( "fmt" "sort" "sync" + "sync/atomic" "time" "github.com/sipeed/picoclaw/pkg/logger" @@ -18,8 +19,9 @@ type ToolEntry struct { } type ToolRegistry struct { - tools map[string]*ToolEntry - mu sync.RWMutex + tools map[string]*ToolEntry + mu sync.RWMutex + version atomic.Uint64 // incremented on Register/RegisterHidden for cache invalidation } func NewToolRegistry() *ToolRegistry { @@ -41,6 +43,7 @@ func (r *ToolRegistry) Register(tool Tool) { IsCore: true, TTL: 0, // Core tools do not use TTL } + r.version.Add(1) } // RegisterHidden saves hidden tools (visible only via TTL) @@ -48,20 +51,28 @@ func (r *ToolRegistry) RegisterHidden(tool Tool) { r.mu.Lock() defer r.mu.Unlock() name := tool.Name() + if _, exists := r.tools[name]; exists { + logger.WarnCF("tools", "Hidden tool registration overwrites existing tool", + map[string]any{"name": name}) + } r.tools[name] = &ToolEntry{ Tool: tool, IsCore: false, TTL: 0, } + r.version.Add(1) } -// PromoteTool imposta il TTL solo se il tool NON รจ un core tool -func (r *ToolRegistry) PromoteTool(name string, ttl int) { +// PromoteTools atomically sets the TTL for multiple non-core tools. +// This prevents a concurrent TickTTL from decrementing between promotions. +func (r *ToolRegistry) PromoteTools(names []string, ttl int) { r.mu.Lock() defer r.mu.Unlock() - if entry, exists := r.tools[name]; exists { - if !entry.IsCore { - entry.TTL = ttl + for _, name := range names { + if entry, exists := r.tools[name]; exists { + if !entry.IsCore { + entry.TTL = ttl + } } } } @@ -84,6 +95,10 @@ func (r *ToolRegistry) Get(name string) (Tool, bool) { if !ok { return nil, false } + // Hidden tools with expired TTL are not callable. + if !entry.IsCore && entry.TTL <= 0 { + return nil, false + } return entry.Tool, true } diff --git a/pkg/tools/search_tool.go b/pkg/tools/search_tool.go index 89f982d6e..c3d661c85 100644 --- a/pkg/tools/search_tool.go +++ b/pkg/tools/search_tool.go @@ -6,10 +6,15 @@ import ( "fmt" "regexp" "strings" + "sync" "github.com/sipeed/picoclaw/pkg/utils" ) +const ( + MaxRegexPatternLength = 200 +) + type RegexSearchTool struct { registry *ToolRegistry ttl int @@ -49,6 +54,11 @@ func (t *RegexSearchTool) Execute(ctx context.Context, args map[string]any) *Too return ErrorResult("Missing or invalid 'pattern' argument. Must be a non-empty string.") } + if len(pattern) > MaxRegexPatternLength { + // Limit on length to avoid catastrophic patterns + return ErrorResult(fmt.Sprintf("Pattern too long: max %d characters allowed", MaxRegexPatternLength)) + } + res, err := t.registry.SearchRegex(pattern, t.maxSearchResults) if err != nil { return ErrorResult(fmt.Sprintf("Invalid regex pattern syntax: %v. Please fix your regex and try again.", err)) @@ -61,6 +71,11 @@ type BM25SearchTool struct { registry *ToolRegistry ttl int maxSearchResults int + + // Cache: rebuilt only when the registry version changes. + cacheMu sync.Mutex + cachedEngine *bm25CachedEngine + cacheVersion uint64 } func NewBM25SearchTool(r *ToolRegistry, ttl int, maxSearchResults int) *BM25SearchTool { @@ -96,17 +111,40 @@ func (t *BM25SearchTool) Execute(ctx context.Context, args map[string]any) *Tool return ErrorResult("Missing or invalid 'query' argument. Must be a non-empty string.") } - return formatDiscoveryResponse(t.registry, t.registry.SearchBM25(query, t.maxSearchResults), t.ttl) + cached := t.getOrBuildEngine() + if cached == nil { + return SilentResult("No tools found matching the query.") + } + + ranked := cached.engine.Search(query, t.maxSearchResults) + if len(ranked) == 0 { + return SilentResult("No tools found matching the query.") + } + + results := make([]ToolSearchResult, len(ranked)) + for i, r := range ranked { + results[i] = ToolSearchResult{ + Name: r.Document.Name, + Description: r.Document.Description, + } + } + + return formatDiscoveryResponse(t.registry, results, t.ttl) } // ToolSearchResult represents the result returned to the LLM. +// Parameters are omitted from the JSON response to save context tokens; +// the LLM will see full schemas via ToProviderDefs after promotion. type ToolSearchResult struct { - Name string `json:"name"` - Description string `json:"description"` - Parameters map[string]any `json:"parameters"` + Name string `json:"name"` + Description string `json:"description"` } func (r *ToolRegistry) SearchRegex(pattern string, maxSearchResults int) ([]ToolSearchResult, error) { + if maxSearchResults <= 0 { + return nil, nil + } + regex, err := regexp.Compile("(?i)" + pattern) if err != nil { return nil, fmt.Errorf("failed to compile regex pattern %q: %w", pattern, err) @@ -117,7 +155,9 @@ func (r *ToolRegistry) SearchRegex(pattern string, maxSearchResults int) ([]Tool var results []ToolSearchResult - for name, entry := range r.tools { + // Iterate in sorted order for deterministic results across calls. + for _, name := range r.sortedToolNames() { + entry := r.tools[name] // Search only among the hidden tools (Core tools are already visible) if !entry.IsCore { // Directly call interface methods! No reflection/unmarshalling needed. @@ -127,7 +167,6 @@ func (r *ToolRegistry) SearchRegex(pattern string, maxSearchResults int) ([]Tool results = append(results, ToolSearchResult{ Name: name, Description: desc, - Parameters: entry.Tool.Parameters(), }) if len(results) >= maxSearchResults { break // Stop searching once we hit the max! Saves CPU. @@ -144,9 +183,11 @@ func formatDiscoveryResponse(registry *ToolRegistry, results []ToolSearchResult, return SilentResult("No tools found matching the query.") } - for _, r := range results { - registry.PromoteTool(r.Name, ttl) + names := make([]string, len(results)) + for i, r := range results { + names[i] = r.Name } + registry.PromoteTools(names, ttl) b, err := json.Marshal(results) if err != nil { @@ -162,20 +203,64 @@ func formatDiscoveryResponse(registry *ToolRegistry, results []ToolSearchResult, return SilentResult(msg) } -// Lightweight internal type +// Lightweight internal type used as corpus document for BM25. type searchDoc struct { Name string Description string - Tool Tool // Hold the interface reference +} + +// bm25CachedEngine wraps a BM25Engine with its corpus snapshot. +type bm25CachedEngine struct { + engine *utils.BM25Engine[searchDoc] +} + +// getOrBuildEngine returns a cached BM25 engine, rebuilding it only when +// the registry version has changed (new tools registered). +func (t *BM25SearchTool) getOrBuildEngine() *bm25CachedEngine { + currentVersion := t.registry.version.Load() + + t.cacheMu.Lock() + defer t.cacheMu.Unlock() + + if t.cachedEngine != nil && t.cacheVersion == currentVersion { + return t.cachedEngine + } + + t.registry.mu.RLock() + snapshot := make([]searchDoc, 0, len(t.registry.tools)) + for name, entry := range t.registry.tools { + if !entry.IsCore { + snapshot = append(snapshot, searchDoc{ + Name: name, + Description: entry.Tool.Description(), + }) + } + } + t.registry.mu.RUnlock() + + if len(snapshot) == 0 { + t.cachedEngine = nil + t.cacheVersion = currentVersion + return nil + } + + engine := utils.NewBM25Engine( + snapshot, + func(doc searchDoc) string { + return doc.Name + " " + doc.Description + }, + ) + + cached := &bm25CachedEngine{engine: engine} + t.cachedEngine = cached + t.cacheVersion = currentVersion + return cached } // SearchBM25 ranks hidden tools against query using BM25 via utils.BM25Engine. // The corpus snapshot is built under the registry read-lock, then released // before scoring so the lock is not held during CPU-intensive work. func (r *ToolRegistry) SearchBM25(query string, maxSearchResults int) []ToolSearchResult { - // We copy only the lightweight searchDoc values (name, description, - // Tool interface reference). This keeps the lock window short and avoids - // holding it during BM25 indexing and scoring. r.mu.RLock() snapshot := make([]searchDoc, 0, len(r.tools)) for name, entry := range r.tools { @@ -183,7 +268,6 @@ func (r *ToolRegistry) SearchBM25(query string, maxSearchResults int) []ToolSear snapshot = append(snapshot, searchDoc{ Name: name, Description: entry.Tool.Description(), - Tool: entry.Tool, }) } } @@ -193,7 +277,6 @@ func (r *ToolRegistry) SearchBM25(query string, maxSearchResults int) []ToolSear return nil } - // Delegate scoring to the generic BM25 engine engine := utils.NewBM25Engine( snapshot, func(doc searchDoc) string { @@ -211,7 +294,6 @@ func (r *ToolRegistry) SearchBM25(query string, maxSearchResults int) []ToolSear out[i] = ToolSearchResult{ Name: r.Document.Name, Description: r.Document.Description, - Parameters: r.Document.Tool.Parameters(), } } return out diff --git a/pkg/tools/search_tools_test.go b/pkg/tools/search_tools_test.go index c1e2efbfc..471e7931a 100644 --- a/pkg/tools/search_tools_test.go +++ b/pkg/tools/search_tools_test.go @@ -138,6 +138,74 @@ func TestBM25SearchTool_Execute(t *testing.T) { }) } +func TestRegexSearchTool_PatternTooLong(t *testing.T) { + reg := setupPopulatedRegistry() + tool := NewRegexSearchTool(reg, 5, 10) + ctx := context.Background() + + longPattern := strings.Repeat("a", MaxRegexPatternLength+1) + res := tool.Execute(ctx, map[string]any{"pattern": longPattern}) + if !res.IsError || !strings.Contains(res.ForLLM, "Pattern too long") { + t.Errorf("Expected pattern too long error, got: %v", res.ForLLM) + } +} + +func TestSearchRegex_ZeroMaxResults(t *testing.T) { + reg := setupPopulatedRegistry() + + res, err := reg.SearchRegex("mcp", 0) + if err != nil { + t.Fatalf("SearchRegex failed: %v", err) + } + if len(res) != 0 { + t.Errorf("Expected 0 results with maxSearchResults=0, got %d", len(res)) + } +} + +func TestSearchBM25_ZeroMaxResults(t *testing.T) { + reg := setupPopulatedRegistry() + + res := reg.SearchBM25("read file", 0) + if len(res) != 0 { + t.Errorf("Expected 0 results with maxSearchResults=0, got %d", len(res)) + } +} + +func TestSearchRegex_DeterministicOrder(t *testing.T) { + reg := NewToolRegistry() + for i := 0; i < 20; i++ { + reg.RegisterHidden(&mockSearchableTool{ + name: fmt.Sprintf("tool_%02d", i), + desc: "searchable tool", + }) + } + + // Run the same search multiple times and verify order is stable + var firstRun []string + for attempt := 0; attempt < 10; attempt++ { + res, err := reg.SearchRegex("searchable", 20) + if err != nil { + t.Fatalf("SearchRegex failed: %v", err) + } + + names := make([]string, len(res)) + for i, r := range res { + names[i] = r.Name + } + + if attempt == 0 { + firstRun = names + } else { + for i, name := range names { + if name != firstRun[i] { + t.Fatalf("Non-deterministic order at attempt %d, index %d: got %q, want %q", + attempt, i, name, firstRun[i]) + } + } + } + } +} + func TestToolRegistry_SearchLimitsAndCoreFiltering(t *testing.T) { reg := NewToolRegistry()