fix(session): make InMemorySessionStore.messageID atomic (#3592)

pkg/session/store.go

@@ -11,6 +11,7 @@ import (
 	"slices"
 	"strconv"
 	"strings"
+	"sync/atomic"
 	"time"
 
 	"github.com/docker/docker-agent/pkg/chat"
@@ -134,7 +135,7 @@ type Store interface {
 
 type InMemorySessionStore struct {
 	sessions  *concurrent.Map[string, *Session]
-	messageID int64 // simple counter for message IDs
+	messageID atomic.Int64 // counter for message IDs, incremented via Add(1)
 }
 
 func NewInMemorySessionStore() Store {
@@ -280,10 +281,10 @@ func (s *InMemorySessionStore) AddMessage(_ context.Context, sessionID string, m
 	// concurrently by another goroutine (snapshotItems → cloneMessage),
 	// and writing msg.ID directly races with those reads.
 	stored := cloneMessage(msg)
-	s.messageID++
-	stored.ID = s.messageID
+	id := s.messageID.Add(1)
+	stored.ID = id
 	session.AddMessage(stored)
-	return s.messageID, nil
+	return id, nil
 }
 
 // UpdateMessage updates an existing message by its ID.

pkg/session/store_messageid_race_test.go

@@ -0,0 +1,60 @@
+package session
+
+import (
+	"sync"
+	"testing"
+
+	"github.com/docker/docker-agent/pkg/chat"
+)
+
+// TestInMemoryStoreAddMessageConcurrentUniqueIDs pins the fix for
+// InMemorySessionStore.messageID: it used to be a plain int64 incremented
+// with `s.messageID++`, which loses increments under concurrent AddMessage
+// calls (HTTP handlers racing the runtime's PersistenceObserver) and can
+// hand out duplicate IDs. A duplicate ID makes UpdateMessage (which looks
+// up messages by ID) edit the wrong message. The counter must be an
+// atomic.Int64 incremented via Add(1) so every concurrent AddMessage call
+// observes a unique, monotonically increasing ID.
+func TestInMemoryStoreAddMessageConcurrentUniqueIDs(t *testing.T) {
+	t.Parallel()
+
+	ctx := t.Context()
+	store := NewInMemorySessionStore()
+	sess := New()
+	if err := store.AddSession(ctx, sess); err != nil {
+		t.Fatalf("AddSession: %v", err)
+	}
+
+	const n = 200
+	ids := make([]int64, n)
+	var wg sync.WaitGroup
+	for i := range n {
+		wg.Go(func() {
+			id, err := store.AddMessage(ctx, sess.ID, &Message{Message: chat.Message{Role: chat.MessageRoleUser, Content: "u"}})
+			if err != nil {
+				t.Errorf("AddMessage: %v", err)
+				return
+			}
+			ids[i] = id
+		})
+	}
+	wg.Wait()
+
+	seen := make(map[int64]bool, n)
+	for _, id := range ids {
+		if id == 0 {
+			t.Fatalf("got zero message ID, AddMessage call must have failed")
+		}
+		if seen[id] {
+			t.Fatalf("duplicate message ID %d assigned to two concurrent AddMessage calls", id)
+		}
+		seen[id] = true
+	}
+	if len(seen) != n {
+		t.Fatalf("expected %d unique message IDs, got %d", n, len(seen))
+	}
+
+	if got := sess.MessageCount(); got != n {
+		t.Errorf("expected %d messages recorded on the session, got %d", n, got)
+	}
+}