fix(tools): consolidate discoveredMu into mu in ToolRegistry

DrainDiscovered previously held discoveredMu while acquiring mu.RLock,
creating a potential lock-order inversion with MarkDiscovered (which held
discoveredMu alone). Consolidate discoveredTools under the existing mu
RWMutex: MarkDiscovered and DrainDiscovered use mu.Lock, eliminating the
separate mutex and the two-lock sequence entirely.
This commit is contained in:
ZanzyTHEbar 2026-02-21 01:06:34 +00:00
parent 4f740dd024
commit 8223bc37a9

View file

@ -19,7 +19,6 @@ type ToolRegistry struct {
// in the current session. PrepareStep drains this set to dynamically // in the current session. PrepareStep drains this set to dynamically
// promote discovered tools to native callables. // promote discovered tools to native callables.
discoveredTools map[string]bool discoveredTools map[string]bool
discoveredMu sync.Mutex
} }
func NewToolRegistry() *ToolRegistry { func NewToolRegistry() *ToolRegistry {
@ -207,8 +206,8 @@ func (r *ToolRegistry) GetSummaries() []string {
// MarkDiscovered records that a tool was returned by tool_search. // MarkDiscovered records that a tool was returned by tool_search.
// Thread-safe; called from tool_search's Execute path. // Thread-safe; called from tool_search's Execute path.
func (r *ToolRegistry) MarkDiscovered(names ...string) { func (r *ToolRegistry) MarkDiscovered(names ...string) {
r.discoveredMu.Lock() r.mu.Lock()
defer r.discoveredMu.Unlock() defer r.mu.Unlock()
for _, name := range names { for _, name := range names {
if name == "tool_search" || name == "tool_call" { if name == "tool_search" || name == "tool_call" {
continue continue
@ -224,16 +223,14 @@ func (r *ToolRegistry) MarkDiscovered(names ...string) {
// since the last drain. PrepareStep calls this to promote discovered tools // since the last drain. PrepareStep calls this to promote discovered tools
// to native callables for the next inference step. // to native callables for the next inference step.
func (r *ToolRegistry) DrainDiscovered() []Tool { func (r *ToolRegistry) DrainDiscovered() []Tool {
r.discoveredMu.Lock() r.mu.Lock()
defer r.mu.Unlock()
names := make([]string, 0, len(r.discoveredTools)) names := make([]string, 0, len(r.discoveredTools))
for name := range r.discoveredTools { for name := range r.discoveredTools {
names = append(names, name) names = append(names, name)
} }
r.discoveredTools = make(map[string]bool) r.discoveredTools = make(map[string]bool)
r.discoveredMu.Unlock()
r.mu.RLock()
defer r.mu.RUnlock()
promoted := make([]Tool, 0, len(names)) promoted := make([]Tool, 0, len(names))
for _, name := range names { for _, name := range names {
if tool, ok := r.tools[name]; ok { if tool, ok := r.tools[name]; ok {