improvements and optimizations

pkg/agent/context.go

@@ -21,7 +21,8 @@ type ContextBuilder struct {
 	workspace          string
 	skillsLoader       *skills.SkillsLoader
 	memory             *MemoryStore
-	toolDiscovery bool
+	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 {
+func (cb *ContextBuilder) WithToolDiscovery(useBM25, useRegex bool) *ContextBuilder {
-	cb.toolDiscovery = enabled
+	cb.toolDiscoveryBM25 = useBM25
+	cb.toolDiscoveryRegex = useRegex
 	return cb
 }
 
@@ -104,12 +106,24 @@ Your workspace is at: %s
 }
 
 func (cb *ContextBuilder) getDiscoveryRule() string {
-	if cb.toolDiscovery {
+	if !cb.toolDiscoveryBM25 && !cb.toolDiscoveryRegex {
-		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.`
-	}
 		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 {
 	parts := []string{}
 

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 := ""

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(
+					return fmt.Errorf(
-						"Critical error: tool discovery is enabled but neither 'use_bm25' nor 'use_regex' is set to true in the configuration.",
+						"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

pkg/config/config.go

@@ -569,7 +569,6 @@ type ToolDiscoveryConfig struct {
 
 type ToolConfig struct {
 	Enabled bool `json:"enabled" env:"ENABLED"`
-	Discovery ToolDiscoveryConfig `json:"discovery"`
 }
 
 type BraveConfig struct {
@@ -721,6 +720,7 @@ type MCPServerConfig struct {
 // MCPConfig defines configuration for all MCP servers
 type MCPConfig struct {
 	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"`
 }

pkg/config/defaults.go

@@ -410,6 +410,7 @@ func DefaultConfig() *Config {
 			MCP: MCPConfig{
 				ToolConfig: ToolConfig{
 					Enabled: false,
+				},
 				Discovery: ToolDiscoveryConfig{
 					Enabled:          false,
 					TTL:              5,
@@ -417,7 +418,6 @@ func DefaultConfig() *Config {
 					UseBM25:          true,
 					UseRegex:         false,
 				},
-				},
 				Servers: map[string]MCPServerConfig{},
 			},
 			AppendFile: ToolConfig{

pkg/tools/registry.go

@@ -5,6 +5,7 @@ import (
 	"fmt"
 	"sort"
 	"sync"
+	"sync/atomic"
 	"time"
 
 	"github.com/sipeed/picoclaw/pkg/logger"
@@ -20,6 +21,7 @@ type ToolEntry struct {
 type ToolRegistry struct {
 	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,23 +51,31 @@ 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
+// PromoteTools atomically sets the TTL for multiple non-core tools.
-func (r *ToolRegistry) PromoteTool(name string, ttl int) {
+// This prevents a concurrent TickTTL from decrementing between promotions.
+func (r *ToolRegistry) PromoteTools(names []string, ttl int) {
 	r.mu.Lock()
 	defer r.mu.Unlock()
+	for _, name := range names {
 		if entry, exists := r.tools[name]; exists {
 			if !entry.IsCore {
 				entry.TTL = ttl
 			}
 		}
 	}
+}
 
 // TickTTL decreases TTL only for non-core tools
 func (r *ToolRegistry) TickTTL() {
@@ -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
 }
 

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"`
 }
 
 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 {
+	names := make([]string, len(results))
-		registry.PromoteTool(r.Name, ttl)
+	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

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()