From dc59b590cf3ddb2f4f6dadda53afbe401db4f43e Mon Sep 17 00:00:00 2001 From: PratikDhanave Date: Thu, 16 Jul 2026 10:04:33 +0530 Subject: [PATCH 1/2] Key todo provider session locks by identity, not service ID The todo context provider keyed its per-session lock registry by Session.ServiceID(), which breaks the synchronization guarantee the lock is meant to provide: - a service ID may be empty, collapsing unrelated sessions onto one shared "_default" lock; - distinct sessions can share a service ID and unexpectedly share a lock; - the service ID is mutable, so one session can resolve to different locks over its lifetime; - registry entries were never removed and accumulated for the provider's lifetime. Key the registry by session object identity via weak.Pointer[agent.Session] instead, with a runtime.AddCleanup to drop entries once a session is collected. This mirrors the agentmode fix in #491 and matches the .NET provider, which keys its per-session lock by object identity via ConditionalWeakTable. A nil/absent session uses the shared fallback lock. Adds internal tests covering: same session -> same lock; distinct sessions with empty service IDs -> distinct locks; distinct sessions sharing a service ID -> distinct locks; nil session -> fallback lock; stale entries drained after GC. --- .../todo/getsessionlock_internal_test.go | 108 ++++++++++++++++++ agent/harness/todo/todo.go | 43 +++++-- 2 files changed, 141 insertions(+), 10 deletions(-) create mode 100644 agent/harness/todo/getsessionlock_internal_test.go diff --git a/agent/harness/todo/getsessionlock_internal_test.go b/agent/harness/todo/getsessionlock_internal_test.go new file mode 100644 index 00000000..a37e0fd0 --- /dev/null +++ b/agent/harness/todo/getsessionlock_internal_test.go @@ -0,0 +1,108 @@ +// Copyright (c) Microsoft. All rights reserved. + +package todo + +import ( + "runtime" + "sync" + "testing" + "time" + + "github.com/microsoft/agent-framework-go/agent" + "github.com/microsoft/agent-framework-go/internal/agenttest" +) + +func sessionOption(s *agent.Session) []agent.Option { + return []agent.Option{agent.WithSession(s)} +} + +func countLocks(m *sync.Map) int { + n := 0 + m.Range(func(_, _ any) bool { + n++ + return true + }) + return n +} + +// The same session must always resolve to the same lock; otherwise concurrent +// tool invocations on that session would not be mutually excluded. +func TestGetSessionLock_SameSessionSameLock(t *testing.T) { + p := New(nil) + s := agenttest.CreateSession() + + first := p.getSessionLock(sessionOption(s)) + second := p.getSessionLock(sessionOption(s)) + if first != second { + t.Fatalf("same session resolved to different locks: %p vs %p", first, second) + } +} + +// Distinct sessions with empty service IDs must not share a lock. Keying by +// service ID would collapse them onto a single "_default" entry. +func TestGetSessionLock_EmptyServiceIDsGetDistinctLocks(t *testing.T) { + p := New(nil) + s1 := agenttest.CreateSession() + s2 := agenttest.CreateSession() + s1.SetServiceID("") + s2.SetServiceID("") + + if l1, l2 := p.getSessionLock(sessionOption(s1)), p.getSessionLock(sessionOption(s2)); l1 == l2 { + t.Fatalf("distinct sessions with empty service IDs shared a lock: %p", l1) + } +} + +// Distinct sessions that happen to share a service ID must not share a lock. +// Keying by service ID would incorrectly serialize unrelated sessions. +func TestGetSessionLock_SameServiceIDGetsDistinctLocks(t *testing.T) { + p := New(nil) + s1 := agenttest.CreateSession() + s2 := agenttest.CreateSession() + s1.SetServiceID("shared-id") + s2.SetServiceID("shared-id") + + if l1, l2 := p.getSessionLock(sessionOption(s1)), p.getSessionLock(sessionOption(s2)); l1 == l2 { + t.Fatalf("distinct sessions sharing a service ID shared a lock: %p", l1) + } +} + +// A missing or nil session must resolve to the shared fallback lock. +func TestGetSessionLock_NilSessionUsesFallback(t *testing.T) { + p := New(nil) + want := &p.nullSessionLock + + if got := p.getSessionLock(nil); got != want { + t.Errorf("no session: got %p, want fallback %p", got, want) + } + if got := p.getSessionLock(sessionOption(nil)); got != want { + t.Errorf("nil session: got %p, want fallback %p", got, want) + } +} + +// Registry entries for collected sessions must not accumulate indefinitely: the +// weak-key cleanup should drop them once the sessions are garbage collected. +func TestGetSessionLock_StaleEntriesDoNotAccumulate(t *testing.T) { + p := New(nil) + + const n = 200 + for i := 0; i < n; i++ { + s := agenttest.CreateSession() + _ = p.getSessionLock(sessionOption(s)) + // s is unreachable after this iteration, so its entry becomes eligible + // for cleanup on the next GC. + } + + // runtime.AddCleanup runs cleanups asynchronously after a GC, so poll with a + // bounded deadline rather than assuming a single GC drains everything. + var remaining int + for i := 0; i < 100; i++ { + runtime.GC() + if remaining = countLocks(&p.sessionLocks); remaining == 0 { + break + } + time.Sleep(time.Millisecond) + } + if remaining != 0 { + t.Fatalf("expected stale lock entries to be cleaned up, %d of %d remain", remaining, n) + } +} 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) { From 2d37e28d03f265be62c0a9fdaf05ad2708a1fb1c Mon Sep 17 00:00:00 2001 From: PratikDhanave Date: Fri, 17 Jul 2026 23:36:02 +0530 Subject: [PATCH 2/2] Black-box the todo session-lock coverage in todo_test.go Address review feedback (prefer black-box). Replace the white-box getSessionLock identity tests with a black-box concurrency test in todo_test.go that exercises the per-session lock via concurrent tool invocations and public reads on a shared session, mirroring the resolution merged for the agentmode provider (#491). --- .../todo/getsessionlock_internal_test.go | 108 ------------------ agent/harness/todo/todo_test.go | 45 ++++++++ 2 files changed, 45 insertions(+), 108 deletions(-) delete mode 100644 agent/harness/todo/getsessionlock_internal_test.go diff --git a/agent/harness/todo/getsessionlock_internal_test.go b/agent/harness/todo/getsessionlock_internal_test.go deleted file mode 100644 index a37e0fd0..00000000 --- a/agent/harness/todo/getsessionlock_internal_test.go +++ /dev/null @@ -1,108 +0,0 @@ -// Copyright (c) Microsoft. All rights reserved. - -package todo - -import ( - "runtime" - "sync" - "testing" - "time" - - "github.com/microsoft/agent-framework-go/agent" - "github.com/microsoft/agent-framework-go/internal/agenttest" -) - -func sessionOption(s *agent.Session) []agent.Option { - return []agent.Option{agent.WithSession(s)} -} - -func countLocks(m *sync.Map) int { - n := 0 - m.Range(func(_, _ any) bool { - n++ - return true - }) - return n -} - -// The same session must always resolve to the same lock; otherwise concurrent -// tool invocations on that session would not be mutually excluded. -func TestGetSessionLock_SameSessionSameLock(t *testing.T) { - p := New(nil) - s := agenttest.CreateSession() - - first := p.getSessionLock(sessionOption(s)) - second := p.getSessionLock(sessionOption(s)) - if first != second { - t.Fatalf("same session resolved to different locks: %p vs %p", first, second) - } -} - -// Distinct sessions with empty service IDs must not share a lock. Keying by -// service ID would collapse them onto a single "_default" entry. -func TestGetSessionLock_EmptyServiceIDsGetDistinctLocks(t *testing.T) { - p := New(nil) - s1 := agenttest.CreateSession() - s2 := agenttest.CreateSession() - s1.SetServiceID("") - s2.SetServiceID("") - - if l1, l2 := p.getSessionLock(sessionOption(s1)), p.getSessionLock(sessionOption(s2)); l1 == l2 { - t.Fatalf("distinct sessions with empty service IDs shared a lock: %p", l1) - } -} - -// Distinct sessions that happen to share a service ID must not share a lock. -// Keying by service ID would incorrectly serialize unrelated sessions. -func TestGetSessionLock_SameServiceIDGetsDistinctLocks(t *testing.T) { - p := New(nil) - s1 := agenttest.CreateSession() - s2 := agenttest.CreateSession() - s1.SetServiceID("shared-id") - s2.SetServiceID("shared-id") - - if l1, l2 := p.getSessionLock(sessionOption(s1)), p.getSessionLock(sessionOption(s2)); l1 == l2 { - t.Fatalf("distinct sessions sharing a service ID shared a lock: %p", l1) - } -} - -// A missing or nil session must resolve to the shared fallback lock. -func TestGetSessionLock_NilSessionUsesFallback(t *testing.T) { - p := New(nil) - want := &p.nullSessionLock - - if got := p.getSessionLock(nil); got != want { - t.Errorf("no session: got %p, want fallback %p", got, want) - } - if got := p.getSessionLock(sessionOption(nil)); got != want { - t.Errorf("nil session: got %p, want fallback %p", got, want) - } -} - -// Registry entries for collected sessions must not accumulate indefinitely: the -// weak-key cleanup should drop them once the sessions are garbage collected. -func TestGetSessionLock_StaleEntriesDoNotAccumulate(t *testing.T) { - p := New(nil) - - const n = 200 - for i := 0; i < n; i++ { - s := agenttest.CreateSession() - _ = p.getSessionLock(sessionOption(s)) - // s is unreachable after this iteration, so its entry becomes eligible - // for cleanup on the next GC. - } - - // runtime.AddCleanup runs cleanups asynchronously after a GC, so poll with a - // bounded deadline rather than assuming a single GC drains everything. - var remaining int - for i := 0; i < 100; i++ { - runtime.GC() - if remaining = countLocks(&p.sessionLocks); remaining == 0 { - break - } - time.Sleep(time.Millisecond) - } - if remaining != 0 { - t.Fatalf("expected stale lock entries to be cleaned up, %d of %d remain", remaining, n) - } -} 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) + } + } +}