Synchronize the in-memory checkpoint manager

workflow/checkpoint/manager.go

@@ -7,6 +7,7 @@ import (
 	"encoding/json"
 	"fmt"
 	"slices"
+	"sync"
 
 	"github.com/microsoft/agent-framework-go/workflow"
 	"github.com/microsoft/agent-framework-go/workflow/internal/checkpoint"
@@ -33,6 +34,10 @@ func NewJSONManager(store Store[json.RawMessage]) Manager {
 }
 
 type inMemoryManager struct {
+	// mu guards Store (and the per-session caches within it) so a single
+	// manager can be shared across concurrent workflow runs — which its
+	// session-keyed design implies — without racing on the map or the caches.
+	mu    sync.Mutex
 	Store map[string]*checkpoint.SessionCache[*checkpoint.Checkpoint]
 }
 
@@ -61,6 +66,8 @@ func (s *inMemoryManager) Commit(_ context.Context, sessionID string, checkpoint
 		return workflow.CheckpointInfo{}, fmt.Errorf("checkpoint: parent sessionID %q does not match sessionID %q", checkpoint.Parent.SessionID, sessionID)
 	}
 
+	s.mu.Lock()
+	defer s.mu.Unlock()
 	store := s.sessionStore(sessionID)
 	return store.Add(sessionID, checkpoint), nil
 }
@@ -73,6 +80,8 @@ func (s *inMemoryManager) Lookup(_ context.Context, sessionID string, checkpoint
 		return nil, fmt.Errorf("checkpoint: checkpoint sessionID %q does not match sessionID %q", checkpointInfo.SessionID, sessionID)
 	}
 
+	s.mu.Lock()
+	defer s.mu.Unlock()
 	store := s.sessionStore(sessionID)
 	v, ok := store.Get(checkpointInfo)
 	if !ok {
@@ -89,6 +98,8 @@ func (s *inMemoryManager) RetrieveIndex(_ context.Context, sessionID string, wit
 		return nil, fmt.Errorf("checkpoint: parent sessionID %q does not match sessionID %q", withParent.SessionID, sessionID)
 	}
 
+	s.mu.Lock()
+	defer s.mu.Unlock()
 	store := s.sessionStore(sessionID)
 	if withParent == nil {
 		return slices.Clone(store.CheckpointIndex), nil

workflow/checkpoint/manager_race_internal_test.go

@@ -0,0 +1,38 @@
+// Copyright (c) Microsoft. All rights reserved.
+
+package checkpoint
+
+import (
+	"context"
+	"fmt"
+	"sync"
+	"testing"
+
+	icheckpoint "github.com/microsoft/agent-framework-go/workflow/internal/checkpoint"
+)
+
+// The in-memory manager is keyed by session ID, so a single manager can be
+// shared across concurrent workflow runs (distinct sessions). Its Store map and
+// per-session caches must be synchronized: without a lock, concurrent Commit
+// calls race on the map and can crash with "fatal: concurrent map writes".
+func TestInMemoryManager_ConcurrentSessions_NoRace(t *testing.T) {
+	mgr := &inMemoryManager{}
+	ctx := context.Background()
+
+	const n = 64
+	var wg sync.WaitGroup
+	wg.Add(n)
+	for i := range n {
+		go func(i int) {
+			defer wg.Done()
+			sid := fmt.Sprintf("session-%d", i)
+			if _, err := mgr.Commit(ctx, sid, &icheckpoint.Checkpoint{}); err != nil {
+				t.Errorf("Commit(%s): %v", sid, err)
+			}
+			if _, err := mgr.RetrieveIndex(ctx, sid, nil); err != nil {
+				t.Errorf("RetrieveIndex(%s): %v", sid, err)
+			}
+		}(i)
+	}
+	wg.Wait()
+}