Merge pull request #2428 from dgageot/fix1

fix: add mutex lock to AddMessageUsageRecord to prevent data race

Author: dgageot

pkg/session/session.go

@@ -472,6 +472,8 @@ func (s *Session) AddMessageUsageRecord(agentName, model string, cost float64, u
 	if usage == nil {
 		return
 	}
+	s.mu.Lock()
+	defer s.mu.Unlock()
 	s.MessageUsageHistory = append(s.MessageUsageHistory, MessageUsageRecord{
 		AgentName: agentName,
 		Model:     model,

pkg/session/session_race_test.go

@@ -0,0 +1,22 @@
+package session
+
+import (
+	"sync"
+	"testing"
+
+	"github.com/docker/docker-agent/pkg/chat"
+)
+
+func TestAddMessageUsageRecordConcurrent(t *testing.T) {
+	s := New()
+	var wg sync.WaitGroup
+	for range 100 {
+		wg.Go(func() {
+			s.AddMessageUsageRecord("agent", "model", 0.1, &chat.Usage{InputTokens: 10, OutputTokens: 5})
+		})
+	}
+	wg.Wait()
+	if got := len(s.MessageUsageHistory); got != 100 {
+		t.Errorf("expected 100 records, got %d", got)
+	}
+}