diff --git a/pkg/model/provider/rulebased/client.go b/pkg/model/provider/rulebased/client.go index 421d6e5123..b62e0e0129 100644 --- a/pkg/model/provider/rulebased/client.go +++ b/pkg/model/provider/rulebased/client.go @@ -13,7 +13,6 @@ import ( "fmt" "log/slog" "slices" - "sync/atomic" "github.com/docker/docker-agent/pkg/chat" "github.com/docker/docker-agent/pkg/config/latest" @@ -21,7 +20,6 @@ import ( "github.com/docker/docker-agent/pkg/model/provider/base" "github.com/docker/docker-agent/pkg/model/provider/contracts" "github.com/docker/docker-agent/pkg/model/provider/options" - "github.com/docker/docker-agent/pkg/modelsdev" "github.com/docker/docker-agent/pkg/tools" ) @@ -39,10 +37,6 @@ type Client struct { routes []Provider fallback Provider matcher *matcher - // lastSelectedID holds the ID of the provider selected by the most recent - // call. A plain atomic pointer suffices because it is the only mutable - // field; everything else is set once at construction and read-only after. - lastSelectedID atomic.Pointer[modelsdev.ID] } // NewClient creates a new rule-based routing client. @@ -120,8 +114,9 @@ func filterOutMaxTokens(opts []options.Opt) []options.Opt { return filtered } -// CreateChatCompletionStream selects a provider based on input and delegates the call. -// The selected provider's ID is recorded in LastSelectedModelID. +// CreateChatCompletionStream selects a provider based on input and delegates +// the call. The returned stream reports the selected provider as its +// ServingModel. func (c *Client) CreateChatCompletionStream( ctx context.Context, messages []chat.Message, @@ -133,7 +128,6 @@ func (c *Client) CreateChatCompletionStream( } selectedID := provider.ID() - c.lastSelectedID.Store(&selectedID) slog.DebugContext(ctx, "Rule-based router selected model", "router", c.ID().String(), "selected_model", selectedID.String(), @@ -162,17 +156,6 @@ func (c *Client) CreateChatCompletionStream( return &routedStream{MessageStream: stream, serving: serving}, nil } -// LastSelectedModelID returns the ID of the provider selected by the most -// recent CreateChatCompletionStream call. This allows callers to display -// the YAML-configured sub-model name for rule-based routing. Accounting uses -// the request-bound ServingModel on the returned stream instead. -func (c *Client) LastSelectedModelID() modelsdev.ID { - if id := c.lastSelectedID.Load(); id != nil { - return *id - } - return modelsdev.ID{} -} - // selectProvider finds the best matching provider for the messages. func (c *Client) selectProvider(messages []chat.Message) (Provider, bool) { userMessage := getLastUserMessage(messages) diff --git a/pkg/model/provider/rulebased/client_race_test.go b/pkg/model/provider/rulebased/client_race_test.go index b2e1f8534a..2d12e1520a 100644 --- a/pkg/model/provider/rulebased/client_race_test.go +++ b/pkg/model/provider/rulebased/client_race_test.go @@ -15,7 +15,9 @@ import ( "github.com/docker/docker-agent/pkg/model/provider/options" ) -func TestLastSelectedModelID_Concurrent(t *testing.T) { +// One router serves every conversation of a team load: each stream must report +// its own route, never another request's. +func TestServingModel_ConcurrentRoutes(t *testing.T) { t.Parallel() cfg := &latest.ModelConfig{ @@ -51,9 +53,6 @@ func TestLastSelectedModelID_Concurrent(t *testing.T) { require.Equal(t, want, stream.(contracts.ServingStream).ServingModel().ID) } }) - wg.Go(func() { - _ = client.LastSelectedModelID() - }) } wg.Wait() } diff --git a/pkg/runtime/routing_pricing_test.go b/pkg/runtime/routing_pricing_test.go index d758fe7bc2..fc8c6baa04 100644 --- a/pkg/runtime/routing_pricing_test.go +++ b/pkg/runtime/routing_pricing_test.go @@ -259,7 +259,6 @@ func TestRoutedPricing(t *testing.T) { } switch { case tc.fallback: - assert.Equal(t, leafID, router.LastSelectedModelID(), "failed route selection must not price the successful outer fallback") assert.Equal(t, 1, leaf.callCount, "cooldown stays keyed to the failed outer router") assert.Equal(t, 2, fallback.callCount) case tc.noMatch: