opencode/picoclaw
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

improvements and optimizations

Browse source
This commit is contained in:
afjcjsbx 2026-03-08 16:03:38 +01:00
parent 097df699eb
commit 725bec8247
8 changed files with 237 additions and 51 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

26
pkg/agent/context.go
View file

@ -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 {
cb.toolDiscovery = enabled
func (cb *ContextBuilder) WithToolDiscovery(useBM25, useRegex bool) *ContextBuilder {
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 {
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 ""
}
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{}

5
pkg/agent/instance.go
View file

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

20
pkg/agent/loop.go
View file

@ -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

2
pkg/config/config.go
View file

@ -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"`
}

2
pkg/config/defaults.go
View file

@ -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{

19
pkg/tools/registry.go
View file

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

110
pkg/tools/search_tool.go
View file

@ -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 {
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

68
pkg/tools/search_tools_test.go
View file

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