Fix data race in agentmode provider session state

agent/harness/agentmode/agentmode.go

@@ -12,7 +12,10 @@ package agentmode
 import (
 	"context"
 	"fmt"
+	"runtime"
 	"strings"
+	"sync"
+	"weak"
 
 	"github.com/microsoft/agent-framework-go/agent"
 	"github.com/microsoft/agent-framework-go/message"
@@ -161,6 +164,44 @@ type Provider struct {
 	defaultMode  string
 	instructions string
 	validModes   map[string]struct{}
+
+	sessionLocks    sync.Map // map[weak.Pointer[agent.Session]]*sync.Mutex
+	nullSessionLock sync.Mutex
+}
+
+// getSessionLock returns a per-session mutex that guards the session 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), 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. 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 || session == nil {
+		return &p.nullSessionLock
+	}
+	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: arrange to drop the entry when
+		// the session is garbage collected. The cleanup must not capture the
+		// session itself (only the weak key), or it would keep it alive.
+		runtime.AddCleanup(session, func(k weak.Pointer[agent.Session]) {
+			p.sessionLocks.Delete(k)
+		}, key)
+	}
+	return actual.(*sync.Mutex)
 }
 
 func (p *Provider) Invoking(ctx context.Context, invoking agent.InvokingContext) ([]*message.Message, []agent.Option, error) {
@@ -193,22 +234,15 @@ func (p *Provider) saveState(opts []agent.Option, s *state) {
 
 func (p *Provider) provide(ctx context.Context, invoking agent.InvokingContext) ([]*message.Message, []agent.Option, error) {
 	opts := invoking.Options
+
+	var outMessages []*message.Message
+
+	mu := p.getSessionLock(opts)
+	mu.Lock()
 	st := p.loadState(opts)
 	// Persist the initial state so SetMode can read it.
 	p.saveState(opts, st)
 
-	tools := p.createTools(opts, st)
-	var outOpts []agent.Option
-	for _, t := range tools {
-		outOpts = append(outOpts, agent.WithTool(t))
-	}
-
-	// Add instructions with mode info.
-	instructionText := p.buildInstructions(st.CurrentMode)
-	outOpts = append(outOpts, agent.WithInstructions(instructionText))
-
-	var outMessages []*message.Message
-
 	// If the mode was changed externally (e.g. via SetMode), inject a notification
 	// so the agent clearly sees the change in conversation context.
 	if st.PreviousMode != "" {
@@ -219,6 +253,18 @@ func (p *Provider) provide(ctx context.Context, invoking agent.InvokingContext)
 		st.PreviousMode = ""
 		p.saveState(opts, st)
 	}
+	currentMode := st.CurrentMode
+	mu.Unlock()
+
+	tools := p.createTools(opts)
+	var outOpts []agent.Option
+	for _, t := range tools {
+		outOpts = append(outOpts, agent.WithTool(t))
+	}
+
+	// Add instructions with mode info.
+	instructionText := p.buildInstructions(currentMode)
+	outOpts = append(outOpts, agent.WithInstructions(instructionText))
 
 	return outMessages, outOpts, nil
 }
@@ -235,7 +281,7 @@ func (p *Provider) buildInstructions(currentMode string) string {
 	return result
 }
 
-func (p *Provider) createTools(opts []agent.Option, st *state) []tool.FuncTool {
+func (p *Provider) createTools(opts []agent.Option) []tool.FuncTool {
 	modeNames := make([]string, len(p.modes))
 	for i, m := range p.modes {
 		modeNames[i] = m.Name
@@ -251,6 +297,10 @@ func (p *Provider) createTools(opts []agent.Option, st *state) []tool.FuncTool {
 			if _, ok := p.validModes[mode]; !ok {
 				return "", fmt.Errorf("invalid mode: %q. Supported modes: \"%s\"", mode, modeNamesDisplay)
 			}
+			mu := p.getSessionLock(opts)
+			mu.Lock()
+			defer mu.Unlock()
+			st := p.loadState(opts)
 			st.CurrentMode = mode
 			p.saveState(opts, st)
 			return fmt.Sprintf("Mode changed to %q.", mode), nil
@@ -263,6 +313,10 @@ func (p *Provider) createTools(opts []agent.Option, st *state) []tool.FuncTool {
 			Description: "Get the agent's current operating mode.",
 		},
 		func(ctx context.Context, _ struct{}) (string, error) {
+			mu := p.getSessionLock(opts)
+			mu.Lock()
+			defer mu.Unlock()
+			st := p.loadState(opts)
 			return st.CurrentMode, nil
 		},
 	)
@@ -273,6 +327,9 @@ func (p *Provider) createTools(opts []agent.Option, st *state) []tool.FuncTool {
 // GetMode returns the current operating mode from the session.
 // If no state has been persisted yet, it returns the configured default mode.
 func (p *Provider) GetMode(opts ...agent.Option) string {
+	mu := p.getSessionLock(opts)
+	mu.Lock()
+	defer mu.Unlock()
 	session, ok := agent.GetOption(opts, agent.WithSession)
 	if !ok {
 		return p.defaultMode
@@ -291,6 +348,9 @@ func (p *Provider) SetMode(mode string, opts ...agent.Option) error {
 	if _, ok := p.validModes[mode]; !ok {
 		return fmt.Errorf("agentmode: invalid mode %q", mode)
 	}
+	mu := p.getSessionLock(opts)
+	mu.Lock()
+	defer mu.Unlock()
 	session, ok := agent.GetOption(opts, agent.WithSession)
 	if !ok {
 		return fmt.Errorf("agentmode: no session available")

agent/harness/agentmode/agentmode_test.go

@@ -6,6 +6,7 @@ import (
 	"context"
 	"slices"
 	"strings"
+	"sync"
 	"testing"
 
 	"github.com/microsoft/agent-framework-go/agent"
@@ -48,6 +49,67 @@ func collectInstructions(opts []agent.Option) string {
 	return sb.String()
 }
 
+// TestConcurrentToolInvocations_NoDataRace mirrors toolautocall with
+// AllowConcurrentInvocations enabled, which invokes multiple tool calls from a
+// single model response on separate goroutines that share one session. The
+// mode_set/mode_get tools must therefore be safe for concurrent use. Run with
+// -race to surface unsynchronized session/state access.
+func TestConcurrentToolInvocations_NoDataRace(t *testing.T) {
+	p := agentmode.New(agentmode.Config{})
+	opts := sessionOpts()
+
+	_, outOpts, err := invokeProvider(p, context.Background(), newMessages("hi"), opts...)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	var setTool, getTool tool.FuncTool
+	for _, tt := range collectTools(outOpts) {
+		ft, ok := tt.(tool.FuncTool)
+		if !ok {
+			continue
+		}
+		switch ft.Name() {
+		case "mode_set":
+			setTool = ft
+		case "mode_get":
+			getTool = ft
+		}
+	}
+	if setTool == nil || getTool == nil {
+		t.Fatalf("expected mode_set and mode_get tools, got set=%v get=%v", setTool, getTool)
+	}
+
+	const n = 50
+	var wg sync.WaitGroup
+	wg.Add(n * 2)
+	// Each goroutine records its error in a distinct slot so the slice itself
+	// is not a source of races; assert them all after the goroutines finish so
+	// a functional regression (e.g. argument-decode failure) fails the test
+	// deterministically, independent of the race detector.
+	errs := make([]error, n*2)
+	for i := 0; i < n; i++ {
+		mode := "plan"
+		if i%2 == 0 {
+			mode = "execute"
+		}
+		go func(idx int, mode string) {
+			defer wg.Done()
+			_, errs[idx] = setTool.Call(context.Background(), `{"Arg0":"`+mode+`"}`)
+		}(i*2, mode)
+		go func(idx int) {
+			defer wg.Done()
+			_, errs[idx] = getTool.Call(context.Background(), "")
+		}(i*2 + 1)
+	}
+	wg.Wait()
+	for _, err := range errs {
+		if err != nil {
+			t.Fatalf("concurrent tool call failed: %v", err)
+		}
+	}
+}
+
 // 1. ProvideAIContextAsync_ReturnsToolsAndInstructions
 func TestProvide_ReturnsToolsAndInstructions(t *testing.T) {
 	p := agentmode.New(agentmode.Config{})