Merge pull request #2429 from dgageot/fix2

fix: add mutex to protect lastSelectedID in rule-based router

Author: dgageot

pkg/model/provider/rulebased/client.go

@@ -11,6 +11,7 @@ import (
 	"errors"
 	"fmt"
 	"log/slog"
+	"sync"
 
 	"github.com/blevesearch/bleve/v2"
 	"github.com/blevesearch/bleve/v2/mapping"
@@ -45,6 +46,7 @@ type Client struct {
 	routes         []Provider
 	fallback       Provider
 	index          bleve.Index
+	mu             sync.RWMutex
 	lastSelectedID string // ID of the provider selected by the most recent call
 }
 
@@ -165,10 +167,13 @@ func (c *Client) CreateChatCompletionStream(
 		return nil, errors.New("no provider available for routing")
 	}
 
-	c.lastSelectedID = provider.ID()
+	selectedID := provider.ID()
+	c.mu.Lock()
+	c.lastSelectedID = selectedID
+	c.mu.Unlock()
 	slog.Debug("Rule-based router selected model",
 		"router", c.ID(),
-		"selected_model", c.lastSelectedID,
+		"selected_model", selectedID,
 		"message_count", len(messages),
 	)
 
@@ -179,6 +184,8 @@ func (c *Client) CreateChatCompletionStream(
 // recent CreateChatCompletionStream call. This allows callers to display
 // the YAML-configured sub-model name for rule-based routing.
 func (c *Client) LastSelectedModelID() string {
+	c.mu.RLock()
+	defer c.mu.RUnlock()
 	return c.lastSelectedID
 }
 

pkg/model/provider/rulebased/client_race_test.go

@@ -0,0 +1,45 @@
+package rulebased
+
+import (
+	"sync"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+
+	"github.com/docker/docker-agent/pkg/chat"
+	"github.com/docker/docker-agent/pkg/config/latest"
+)
+
+func TestLastSelectedModelID_Concurrent(t *testing.T) {
+	t.Parallel()
+
+	cfg := &latest.ModelConfig{
+		Provider: "openai",
+		Model:    "gpt-4o",
+		Routing: []latest.RoutingRule{
+			{
+				Model:    "anthropic/claude-3-haiku",
+				Examples: []string{"hello", "hi there"},
+			},
+		},
+	}
+
+	client, err := NewClient(t.Context(), cfg, nil, nil, mockProviderFactory)
+	require.NoError(t, err)
+	defer client.Close()
+
+	var wg sync.WaitGroup
+	for range 100 {
+		wg.Add(2)
+		go func() {
+			defer wg.Done()
+			messages := []chat.Message{{Role: chat.MessageRoleUser, Content: "hello"}}
+			_, _ = client.CreateChatCompletionStream(t.Context(), messages, nil)
+		}()
+		go func() {
+			defer wg.Done()
+			_ = client.LastSelectedModelID()
+		}()
+	}
+	wg.Wait()
+}