diff --git a/pkg/session/store.go b/pkg/session/store.go index 8b9d70f0e..6dba7f5a8 100644 --- a/pkg/session/store.go +++ b/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. diff --git a/pkg/session/store_messageid_race_test.go b/pkg/session/store_messageid_race_test.go new file mode 100644 index 000000000..b7796baf6 --- /dev/null +++ b/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) + } +}