diff --git a/agent/harness/todo/todo.go b/agent/harness/todo/todo.go index f2db516e..51afe946 100644 --- a/agent/harness/todo/todo.go +++ b/agent/harness/todo/todo.go @@ -13,8 +13,10 @@ package todo import ( "context" "fmt" + "runtime" "strings" "sync" + "weak" "github.com/microsoft/agent-framework-go/agent" "github.com/microsoft/agent-framework-go/message" @@ -87,7 +89,7 @@ type Provider struct { instructions string suppressTodoMessage bool todoListMessageBuilder func([]Item) string - sessionLocks sync.Map // map[sessionKey]*sync.Mutex + sessionLocks sync.Map // map[weak.Pointer[agent.Session]]*sync.Mutex nullSessionLock sync.Mutex } @@ -166,20 +168,41 @@ func (p *Provider) saveState(opts []agent.Option, s *state) { session.Set(stateKey, *s) } -// getSessionLock returns a per-session mutex. If no session is available, -// a shared fallback lock is returned. This matches the .NET pattern of -// per-session SemaphoreSlim via ConditionalWeakTable. +// getSessionLock returns a per-session mutex that guards the todo state against +// concurrent tool invocations, which toolautocall runs on separate goroutines +// when AllowConcurrentInvocations is enabled. +// +// The registry is keyed by session object identity via a weak pointer, not by +// Session.ServiceID(): a service ID may be empty (causing unrelated sessions to +// collide on a shared lock), shared across distinct sessions, or mutated during +// a session's lifetime — any of which would break the guarantee that a given +// session always maps to the same lock. Identity keying also matches the .NET +// provider, which keys its per-session lock by object identity via +// ConditionalWeakTable. When no session is available, a shared fallback lock is +// returned so state access is still serialized. +// +// Weak keys do not keep sessions alive and do not remove map entries on their +// own, so a runtime cleanup deletes the entry once the session is collected, +// keeping the registry from growing unbounded. func (p *Provider) getSessionLock(opts []agent.Option) *sync.Mutex { session, ok := agent.GetOption(opts, agent.WithSession) - if !ok { + if !ok || session == nil { return &p.nullSessionLock } - key := session.ServiceID() - if key == "" { - key = "_default" + key := weak.Make(session) + if existing, ok := p.sessionLocks.Load(key); ok { + return existing.(*sync.Mutex) + } + actual, loaded := p.sessionLocks.LoadOrStore(key, &sync.Mutex{}) + if !loaded { + // First registration for this session: drop the entry when the session + // is garbage collected. The cleanup captures only the weak key, never + // the session itself, or it would keep the session alive. + runtime.AddCleanup(session, func(k weak.Pointer[agent.Session]) { + p.sessionLocks.Delete(k) + }, key) } - v, _ := p.sessionLocks.LoadOrStore(key, &sync.Mutex{}) - return v.(*sync.Mutex) + return actual.(*sync.Mutex) } func (p *Provider) provide(ctx context.Context, invoking agent.InvokingContext) ([]*message.Message, []agent.Option, error) { diff --git a/agent/harness/todo/todo_test.go b/agent/harness/todo/todo_test.go index f21983e0..5209486a 100644 --- a/agent/harness/todo/todo_test.go +++ b/agent/harness/todo/todo_test.go @@ -7,6 +7,7 @@ import ( "fmt" "slices" "strings" + "sync" "testing" "github.com/microsoft/agent-framework-go/agent" @@ -675,3 +676,47 @@ func TestCompleteTodos_EmptyReasonIsAccepted(t *testing.T) { }) } } + +// Concurrent tool invocations and public reads on a shared session must be +// serialized by the per-session lock rather than race on the session's todo +// state. Run under -race. +func TestTodo_ConcurrentSessionAccess_NoDataRace(t *testing.T) { + p := todo.New(nil) + opts := sessionOpts() + + _, outOpts, err := invokeProvider(p, context.Background(), newMessages("hi"), opts...) + if err != nil { + t.Fatal(err) + } + var addTool tool.FuncTool + for _, tt := range collectTools(outOpts) { + if tt.Name() == "todos_add" { + addTool, _ = tt.(tool.FuncTool) + } + } + if addTool == nil { + t.Fatal("todos_add tool not found") + } + + const n = 50 + var wg sync.WaitGroup + wg.Add(n * 2) + errs := make([]error, n*2) + for i := 0; i < n; i++ { + go func(idx int) { + defer wg.Done() + _, errs[idx] = addTool.Call(context.Background(), fmt.Sprintf(`{"Arg0":[{"title":"item-%d"}]}`, idx)) + }(i * 2) + go func(idx int) { + defer wg.Done() + _ = p.GetAllItems(opts...) + _ = p.GetRemainingItems(opts...) + }(i*2 + 1) + } + wg.Wait() + for _, e := range errs { + if e != nil { + t.Fatalf("concurrent todos_add failed: %v", e) + } + } +}