diff --git a/pkg/runtime/ontoolschanged_race_test.go b/pkg/runtime/ontoolschanged_race_test.go new file mode 100644 index 000000000..520f720ca --- /dev/null +++ b/pkg/runtime/ontoolschanged_race_test.go @@ -0,0 +1,44 @@ +package runtime + +import ( + "sync" + "sync/atomic" + "testing" + + "github.com/stretchr/testify/require" + + "github.com/docker/docker-agent/pkg/agent" + "github.com/docker/docker-agent/pkg/team" +) + +// TestOnToolsChangedConcurrentRegisterAndEmit pins the fix for +// onToolsChanged: it used to be a plain func(Event) field, written by +// OnToolsChanged and read by emitToolsChanged, with no lock — unlike its +// sibling onBackgroundEvent, which is guarded by backgroundEventMu because +// background tasks read it from their own goroutines. MCP change- +// notification goroutines call emitToolsChanged concurrently with the +// handler being (re)registered, so onToolsChanged needs the same guard. +func TestOnToolsChangedConcurrentRegisterAndEmit(t *testing.T) { + t.Parallel() + + prov := &mockProvider{id: "test/mock-model"} + root := agent.New("root", "test", agent.WithModel(prov)) + tm := team.New(team.WithAgents(root)) + + rt, err := NewLocalRuntime(t.Context(), tm, WithModelStore(mockModelStore{})) + require.NoError(t, err) + + var calls atomic.Int32 + var wg sync.WaitGroup + for range 100 { + wg.Go(func() { + rt.OnToolsChanged(func(Event) { + calls.Add(1) + }) + }) + wg.Go(func() { + rt.emitToolsChanged() + }) + } + wg.Wait() +} diff --git a/pkg/runtime/runtime.go b/pkg/runtime/runtime.go index fa6a5365f..18d3a4291 100644 --- a/pkg/runtime/runtime.go +++ b/pkg/runtime/runtime.go @@ -303,7 +303,11 @@ type LocalRuntime struct { recallMu sync.RWMutex recallHandler RecallHandler - // onToolsChanged is called when an MCP toolset reports a tool list change. + // onToolsChanged is called when an MCP toolset reports a tool list + // change. Protected by toolsChangedMu because MCP change-notification + // goroutines call emitToolsChanged concurrently, mirroring + // onBackgroundEvent/backgroundEventMu below. + toolsChangedMu sync.RWMutex onToolsChanged func(Event) // onBackgroundEvent is called for events surfaced from detached @@ -1428,7 +1432,9 @@ func (r *LocalRuntime) ResetStartupInfo() { // reports a tool list change outside of a RunStream. This allows the UI // to update the tool count immediately. func (r *LocalRuntime) OnToolsChanged(handler func(Event)) { + r.toolsChangedMu.Lock() r.onToolsChanged = handler + r.toolsChangedMu.Unlock() for _, name := range r.team.AgentNames() { a, err := r.team.Agent(name) @@ -1446,7 +1452,10 @@ func (r *LocalRuntime) OnToolsChanged(handler func(Event)) { // emitToolsChanged is the callback registered on MCP toolsets. It re-reads // the current agent's full tool list and pushes a ToolsetInfo event. func (r *LocalRuntime) emitToolsChanged() { - if r.onToolsChanged == nil { + r.toolsChangedMu.RLock() + handler := r.onToolsChanged + r.toolsChangedMu.RUnlock() + if handler == nil { return } ctx, cancel := context.WithTimeout(r.ctx(), toolsChangedTimeout) @@ -1456,7 +1465,7 @@ func (r *LocalRuntime) emitToolsChanged() { if err != nil { return } - r.onToolsChanged(ToolsetInfo(len(agentTools), false, r.currentAgentName())) + handler(ToolsetInfo(len(agentTools), false, r.currentAgentName())) } // OnBackgroundEvent registers a handler that receives events surfaced from