From f461c93cb507b303bf47d21158b58efda59624c2 Mon Sep 17 00:00:00 2001
From: PratikDhanave <i.pratikdhanave@gmail.com>
Date: Fri, 17 Jul 2026 10:39:42 +0530
Subject: [PATCH] Synchronize the in-memory checkpoint manager

inMemoryManager.Store (and its per-session caches) had no synchronization,
but the manager is keyed by session ID and NewInMemoryManager carries no
single-use caveat -- so sharing one manager across concurrent workflow runs
(distinct sessions) is a natural pattern. Concurrent Commit calls then race
on the map and can crash with 'fatal: concurrent map writes'.

Guard Store with a mutex in Commit, Lookup, and RetrieveIndex. The critical
sections are pure in-memory map/cache operations (no I/O), so holding the
lock across each method is cheap.

Adds TestInMemoryManager_ConcurrentSessions_NoRace: 64 goroutines commit and
read distinct sessions concurrently; it reports data races under -race before
this change and passes after.
---
 workflow/checkpoint/manager.go                | 11 ++++++
 .../checkpoint/manager_race_internal_test.go  | 38 +++++++++++++++++++
 2 files changed, 49 insertions(+)
 create mode 100644 workflow/checkpoint/manager_race_internal_test.go

diff --git a/workflow/checkpoint/manager.go b/workflow/checkpoint/manager.go
index b43bcc6d..dd866c9d 100644
--- a/workflow/checkpoint/manager.go
+++ b/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
diff --git a/workflow/checkpoint/manager_race_internal_test.go b/workflow/checkpoint/manager_race_internal_test.go
new file mode 100644
index 00000000..859294e8
--- /dev/null
+++ b/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()
+}