fix: price routed turns using the serving model

docs/configuration/routing/index.md

@@ -58,6 +58,13 @@ agents:
     instruction: You are a helpful assistant.
 ```
 
+## Cost and Attribution
+
+Routed turns are priced and attributed to the model that served them, using its
+catalogue rates or `cost:` table and the response's service tier when applicable.
+The router's `cost:` table does not override a matched target's pricing. When no
+rule matches, the router's base model and accounting overrides apply.
+
 ## Routing Rules
 
 Each routing rule has:

pkg/leantui/serving_model_test.go

@@ -0,0 +1,93 @@
+package leantui
+
+import (
+	"context"
+	"errors"
+	"io"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+
+	"github.com/docker/docker-agent/pkg/agent"
+	"github.com/docker/docker-agent/pkg/chat"
+	"github.com/docker/docker-agent/pkg/config/latest"
+	"github.com/docker/docker-agent/pkg/environment"
+	"github.com/docker/docker-agent/pkg/model/provider/base"
+	"github.com/docker/docker-agent/pkg/model/provider/options"
+	"github.com/docker/docker-agent/pkg/model/provider/rulebased"
+	"github.com/docker/docker-agent/pkg/modelsdev"
+	"github.com/docker/docker-agent/pkg/runtime"
+	"github.com/docker/docker-agent/pkg/session"
+	"github.com/docker/docker-agent/pkg/team"
+	"github.com/docker/docker-agent/pkg/tools"
+)
+
+type footerProvider struct {
+	base.Config
+
+	err error
+}
+
+func (p *footerProvider) CreateChatCompletionStream(context.Context, []chat.Message, []tools.Tool) (chat.MessageStream, error) {
+	if p.err != nil {
+		return nil, p.err
+	}
+	return &footerStream{}, nil
+}
+
+type footerStream struct{ done bool }
+
+func (s *footerStream) Recv() (chat.MessageStreamResponse, error) {
+	if s.done {
+		return chat.MessageStreamResponse{}, io.EOF
+	}
+	s.done = true
+	return chat.MessageStreamResponse{Choices: []chat.MessageStreamChoice{{Delta: chat.MessageDelta{Content: "ok"}, FinishReason: chat.FinishReasonStop}}, Usage: &chat.Usage{InputTokens: 1, OutputTokens: 1}}, nil
+}
+func (*footerStream) Close() {}
+
+func TestServingModelFallbackFooter(t *testing.T) {
+	t.Parallel()
+
+	for _, tc := range []struct {
+		name, prompt, want string
+		router             bool
+	}{
+		{name: "direct", prompt: "hi", want: "test/fallback"},
+		{name: "router no match", router: true, prompt: "hi", want: "test/fallback"},
+		{name: "router matched", router: true, prompt: "poetry", want: "test/primary"},
+	} {
+		t.Run(tc.name, func(t *testing.T) {
+			t.Parallel()
+			primary := &footerProvider{Config: base.Config{ModelConfig: latest.ModelConfig{Provider: "test", Model: "primary"}}, err: errors.New("401 unauthorized")}
+			fallback := rulebased.Provider(&footerProvider{Config: base.Config{ModelConfig: latest.ModelConfig{Provider: "test", Model: "fallback", Cost: &latest.CostConfig{Input: 1}}}})
+			if tc.router {
+				cfg := fallback.BaseConfig().ModelConfig
+				cfg.Routing = []latest.RoutingRule{{Model: "test/primary", Examples: []string{"poetry"}}}
+				var err error
+				fallback, err = rulebased.NewClient(t.Context(), &cfg, nil, nil, func(_ context.Context, spec string, _ map[string]latest.ModelConfig, _ environment.Provider, _ ...options.Opt) (rulebased.Provider, error) {
+					id := modelsdev.ParseIDOrZero(spec)
+					return &footerProvider{Config: base.Config{ModelConfig: latest.ModelConfig{Provider: id.Provider, Model: id.Model, Cost: &latest.CostConfig{Input: 1}}}}, nil
+				})
+				require.NoError(t, err)
+			}
+			root := agent.New("root", "test", agent.WithModel(primary), agent.WithFallbackModel(fallback), agent.WithFallbackRetries(-1))
+			rt, err := runtime.NewLocalRuntime(t.Context(), team.New(team.WithAgents(root)), runtime.WithModelStore(modelsdev.NewDatabaseStore(&modelsdev.Database{})), runtime.WithSessionCompaction(false))
+			require.NoError(t, err)
+			sess := session.New(session.WithTitle("Serving footer test"), session.WithUserMessage(tc.prompt))
+			m := bareModel(24)
+			for range 2 {
+				for event := range rt.RunStream(t.Context(), sess) {
+					switch event := event.(type) {
+					case *runtime.ErrorEvent:
+						t.Errorf("runtime error: %s", event.Error)
+					case *runtime.AgentInfoEvent, *runtime.ModelFallbackEvent:
+						m.handleEvent(t.Context(), event)
+					}
+				}
+				assert.Equal(t, tc.want, m.status.Model)
+			}
+		})
+	}
+}

pkg/model/provider/contracts/contracts.go

@@ -128,6 +128,19 @@ type Provider interface {
 	BaseConfig() Config
 }
 
+// ServingModel identifies the endpoint that produced a completion stream.
+type ServingModel struct {
+	ID     modelsdev.ID
+	Config Config
+}
+
+// ServingStream reports the serving endpoint when it differs from the dispatched provider.
+// The metadata is available at stream creation and remains stable after Close.
+type ServingStream interface {
+	chat.MessageStream
+	ServingModel() ServingModel
+}
+
 // EmbeddingProvider is a provider that supports embeddings.
 type EmbeddingProvider interface {
 	Provider

pkg/model/provider/instrument.go

@@ -133,7 +133,11 @@ func (t *tracedChat) CreateChatCompletionStream(ctx context.Context, messages []
 		span.End()
 		return nil, err
 	}
-	return genai.WrapStream(span, stream), nil
+	wrapped := genai.WrapStream(span, stream)
+	if serving, ok := stream.(ServingStream); ok {
+		return &tracedServingStream{MessageStream: wrapped, serving: serving.ServingModel()}, nil
+	}
+	return wrapped, nil
 }
 
 // embeddingRequestForConfig builds an EmbeddingRequest from the inner
@@ -283,3 +287,14 @@ type tracedBatchEmbedRerank struct {
 	tracedBatchEmbed
 	tracedRerankOps
 }
+
+// tracedServingStream preserves metadata while Recv and Close retain span bookkeeping.
+type tracedServingStream struct {
+	chat.MessageStream
+
+	serving ServingModel
+}
+
+func (s *tracedServingStream) ServingModel() ServingModel {
+	return s.serving
+}

pkg/model/provider/provider.go

@@ -20,6 +20,12 @@ import (
 // Provider is the common interface implemented by model providers.
 type Provider = contracts.Provider
 
+// ServingModel identifies the endpoint that produced a completion stream.
+type ServingModel = contracts.ServingModel
+
+// ServingStream reports the serving endpoint when it differs from the dispatched provider.
+type ServingStream = contracts.ServingStream
+
 // EmbeddingProvider is a provider that supports embeddings.
 type EmbeddingProvider = contracts.EmbeddingProvider
 

pkg/model/provider/rulebased/client.go

@@ -121,7 +121,7 @@ func (c *Client) CreateChatCompletionStream(
 	messages []chat.Message,
 	availableTools []tools.Tool,
 ) (chat.MessageStream, error) {
-	provider := c.selectProvider(messages)
+	provider, selfServe := c.selectProvider(messages)
 	if provider == nil {
 		return nil, errors.New("no provider available for routing")
 	}
@@ -134,12 +134,32 @@ func (c *Client) CreateChatCompletionStream(
 		"message_count", len(messages),
 	)
 
-	return provider.CreateChatCompletionStream(ctx, messages, availableTools)
+	stream, err := provider.CreateChatCompletionStream(ctx, messages, availableTools)
+	if err != nil || stream == nil {
+		return stream, err
+	}
+	serving := contracts.ServingModel{ID: selectedID, Config: provider.BaseConfig()}
+	if selfServe {
+		serving.ID = c.ID()
+		// The default provider is built from provider/model alone; preserve
+		// the router's accounting overrides without losing the resolved endpoint.
+		if serving.Config.ModelConfig.Cost == nil {
+			serving.Config.ModelConfig.Cost = c.ModelConfig.Cost
+		}
+		if serving.Config.ModelConfig.BaseURL == "" {
+			serving.Config.ModelConfig.BaseURL = c.ModelConfig.BaseURL
+		}
+	}
+	if inner, ok := stream.(contracts.ServingStream); ok {
+		serving = inner.ServingModel()
+	}
+	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.
+// 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
@@ -148,20 +168,20 @@ func (c *Client) LastSelectedModelID() modelsdev.ID {
 }
 
 // selectProvider finds the best matching provider for the messages.
-func (c *Client) selectProvider(messages []chat.Message) Provider {
+func (c *Client) selectProvider(messages []chat.Message) (Provider, bool) {
 	userMessage := getLastUserMessage(messages)
 	if userMessage == "" {
-		return c.defaultProvider()
+		return c.defaultProvider(), true
 	}
 
 	routeIdx, ok := c.matcher.bestRoute(userMessage)
 	if !ok || routeIdx >= len(c.routes) {
-		return c.defaultProvider()
+		return c.defaultProvider(), true
 	}
 
 	selected := c.routes[routeIdx]
 	slog.Debug("Route matched", "model", selected.ID().String())
-	return selected
+	return selected, false
 }
 
 func (c *Client) defaultProvider() Provider {
@@ -192,3 +212,13 @@ func (c *Client) BaseConfig() base.Config {
 func (c *Client) Close() error {
 	return nil
 }
+
+type routedStream struct {
+	chat.MessageStream
+
+	serving contracts.ServingModel
+}
+
+func (s *routedStream) ServingModel() contracts.ServingModel {
+	return s.serving
+}

pkg/model/provider/rulebased/client_race_test.go

@@ -1,13 +1,18 @@
 package rulebased
 
 import (
+	"context"
 	"sync"
 	"testing"
 
 	"github.com/stretchr/testify/require"
 
 	"github.com/docker/docker-agent/pkg/chat"
 	"github.com/docker/docker-agent/pkg/config/latest"
+	"github.com/docker/docker-agent/pkg/environment"
+	"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"
 )
 
 func TestLastSelectedModelID_Concurrent(t *testing.T) {
@@ -24,15 +29,27 @@ func TestLastSelectedModelID_Concurrent(t *testing.T) {
 		},
 	}
 
-	client, err := NewClient(t.Context(), cfg, nil, nil, mockProviderFactory)
+	client, err := NewClient(t.Context(), cfg, nil, nil,
+		func(ctx context.Context, spec string, models map[string]latest.ModelConfig, env environment.Provider, opts ...options.Opt) (Provider, error) {
+			p, err := mockProviderFactory(ctx, spec, models, env, opts...)
+			return &servingProvider{Config: base.Config{ModelConfig: latest.ModelConfig{Provider: p.ID().Provider, Model: p.ID().Model}}, stream: &servingTestStream{}}, err
+		})
 	require.NoError(t, err)
 	defer client.Close()
 
 	var wg sync.WaitGroup
 	for range 100 {
 		wg.Go(func() {
-			messages := []chat.Message{{Role: chat.MessageRoleUser, Content: "hello"}}
-			_, _ = client.CreateChatCompletionStream(t.Context(), messages, nil)
+			for _, prompt := range []string{"hello", "unrelated"} {
+				messages := []chat.Message{{Role: chat.MessageRoleUser, Content: prompt}}
+				stream, err := client.CreateChatCompletionStream(t.Context(), messages, nil)
+				require.NoError(t, err)
+				want := client.ID()
+				if prompt == "hello" {
+					want = client.routes[0].ID()
+				}
+				require.Equal(t, want, stream.(contracts.ServingStream).ServingModel().ID)
+			}
 		})
 		wg.Go(func() {
 			_ = client.LastSelectedModelID()

pkg/model/provider/rulebased/client_test.go

@@ -199,7 +199,7 @@ func TestClient_SelectProvider(t *testing.T) {
 			defer client.Close()
 
 			messages := []chat.Message{{Role: chat.MessageRoleUser, Content: tt.message}}
-			provider := client.selectProvider(messages)
+			provider, _ := client.selectProvider(messages)
 			require.NotNil(t, provider)
 			assert.Equal(t, tt.expectedModel, provider.ID().String())
 		})
@@ -311,7 +311,7 @@ func TestClient_DefaultProvider(t *testing.T) {
 	require.NoError(t, err)
 	defer client.Close()
 
-	provider := client.selectProvider(nil)
+	provider, _ := client.selectProvider(nil)
 	assert.Equal(t, "openai/gpt-4o", provider.ID().String())
 }