Skip to content
Closed
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
10 changes: 9 additions & 1 deletion pkg/tools/registry.go
Original file line number Diff line number Diff line change
Expand Up @@ -30,6 +30,10 @@ type mediaStoreAware interface {
SetMediaStore(store media.MediaStore)
}

type registryCloneAware interface {
CloneForRegistry(registry *ToolRegistry) Tool
}

func NewToolRegistry() *ToolRegistry {
return &ToolRegistry{
tools: make(map[string]*ToolEntry),
Expand Down Expand Up @@ -412,8 +416,12 @@ func (r *ToolRegistry) Clone() *ToolRegistry {
mediaStore: r.mediaStore,
}
for name, entry := range r.tools {
tool := entry.Tool
if aware, ok := entry.Tool.(registryCloneAware); ok {
tool = aware.CloneForRegistry(clone)
}
clone.tools[name] = &ToolEntry{
Tool: entry.Tool,
Tool: tool,
IsCore: entry.IsCore,
TTL: entry.TTL,
}
Expand Down
26 changes: 26 additions & 0 deletions pkg/tools/registry_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -515,6 +515,32 @@ func TestToolRegistry_Clone_PreservesTTLValue(t *testing.T) {
}
}

func TestToolRegistry_Clone_RebindsDiscoveryToolsToClone(t *testing.T) {
parent := NewToolRegistry()
parent.RegisterHidden(newMockTool("mcp_research", "deep research report tool"))
parent.Register(NewBM25SearchTool(parent, 3, 5))

clone := parent.Clone()

searchTool, ok := clone.Get("tool_search_tool_bm25")
if !ok {
t.Fatal("expected cloned registry to expose BM25 search tool")
}
result := searchTool.Execute(context.Background(), map[string]any{
"query": "deep research",
})
if result == nil || result.IsError {
t.Fatalf("search result error: %+v", result)
}

if _, ok := clone.Get("mcp_research"); !ok {
t.Fatal("expected search in clone to promote hidden tool in clone")
}
if _, ok := parent.Get("mcp_research"); ok {
t.Fatal("expected search in clone not to promote hidden tool in parent")
}
}

func TestToolRegistry_ConcurrentAccess(t *testing.T) {
r := NewToolRegistry()
var wg sync.WaitGroup
Expand Down
14 changes: 14 additions & 0 deletions pkg/tools/search_tool.go
Original file line number Diff line number Diff line change
Expand Up @@ -26,6 +26,13 @@ func NewRegexSearchTool(r *ToolRegistry, ttl int, maxSearchResults int) *RegexSe
return &RegexSearchTool{registry: r, ttl: ttl, maxSearchResults: maxSearchResults}
}

func (t *RegexSearchTool) CloneForRegistry(registry *ToolRegistry) Tool {
if t == nil {
return NewRegexSearchTool(registry, 0, 0)
}
return NewRegexSearchTool(registry, t.ttl, t.maxSearchResults)
}

func (t *RegexSearchTool) Name() string {
return "tool_search_tool_regex"
}
Expand Down Expand Up @@ -95,6 +102,13 @@ func NewBM25SearchTool(r *ToolRegistry, ttl int, maxSearchResults int) *BM25Sear
return &BM25SearchTool{registry: r, ttl: ttl, maxSearchResults: maxSearchResults}
}

func (t *BM25SearchTool) CloneForRegistry(registry *ToolRegistry) Tool {
if t == nil {
return NewBM25SearchTool(registry, 0, 0)
}
return NewBM25SearchTool(registry, t.ttl, t.maxSearchResults)
}

func (t *BM25SearchTool) Name() string {
return "tool_search_tool_bm25"
}
Expand Down