From d6e5ccb12693c2e13306f139b5560bbc52b57ff7 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Arnaud=20H=C3=A9ritier?= Date: Tue, 13 Oct 2026 13:39:53 +0000 Subject: [PATCH 1/6] fix: account routed turns using the serving model MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: Arnaud Héritier --- docs/configuration/routing/index.md | 7 + pkg/model/provider/contracts/contracts.go | 12 + pkg/model/provider/provider.go | 6 + pkg/model/provider/rulebased/client.go | 44 +++- .../provider/rulebased/client_race_test.go | 23 +- pkg/model/provider/rulebased/client_test.go | 4 +- pkg/model/provider/rulebased/serving_test.go | 133 +++++++++++ pkg/runtime/fallback.go | 19 +- pkg/runtime/loop.go | 20 +- pkg/runtime/routing_pricing_test.go | 215 ++++++++++++++++++ pkg/runtime/streaming.go | 1 + 11 files changed, 457 insertions(+), 27 deletions(-) create mode 100644 pkg/model/provider/rulebased/serving_test.go create mode 100644 pkg/runtime/routing_pricing_test.go diff --git a/docs/configuration/routing/index.md b/docs/configuration/routing/index.md index f6dca57e1a..b96cf4054f 100644 --- a/docs/configuration/routing/index.md +++ b/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: diff --git a/pkg/model/provider/contracts/contracts.go b/pkg/model/provider/contracts/contracts.go index 40b26a5662..2a71beecf1 100644 --- a/pkg/model/provider/contracts/contracts.go +++ b/pkg/model/provider/contracts/contracts.go @@ -128,6 +128,18 @@ 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. +type ServingStream interface { + chat.MessageStream + ServingModel() ServingModel +} + // EmbeddingProvider is a provider that supports embeddings. type EmbeddingProvider interface { Provider diff --git a/pkg/model/provider/provider.go b/pkg/model/provider/provider.go index 233722319e..a9b47ed9e8 100644 --- a/pkg/model/provider/provider.go +++ b/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 diff --git a/pkg/model/provider/rulebased/client.go b/pkg/model/provider/rulebased/client.go index ee152ad685..bb70f9aebd 100644 --- a/pkg/model/provider/rulebased/client.go +++ b/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 +} diff --git a/pkg/model/provider/rulebased/client_race_test.go b/pkg/model/provider/rulebased/client_race_test.go index 88b656f6d0..b2e1f8534a 100644 --- a/pkg/model/provider/rulebased/client_race_test.go +++ b/pkg/model/provider/rulebased/client_race_test.go @@ -1,6 +1,7 @@ package rulebased import ( + "context" "sync" "testing" @@ -8,6 +9,10 @@ import ( "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() diff --git a/pkg/model/provider/rulebased/client_test.go b/pkg/model/provider/rulebased/client_test.go index 05e944e4f9..92d39777c8 100644 --- a/pkg/model/provider/rulebased/client_test.go +++ b/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()) } diff --git a/pkg/model/provider/rulebased/serving_test.go b/pkg/model/provider/rulebased/serving_test.go new file mode 100644 index 0000000000..ac67e22c60 --- /dev/null +++ b/pkg/model/provider/rulebased/serving_test.go @@ -0,0 +1,133 @@ +package rulebased + +import ( + "context" + "errors" + "io" + "testing" + + "github.com/stretchr/testify/assert" + "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" + "github.com/docker/docker-agent/pkg/modelsdev" + "github.com/docker/docker-agent/pkg/tools" +) + +type servingProvider struct { + base.Config + + stream chat.MessageStream + err error +} + +func (p *servingProvider) CreateChatCompletionStream(context.Context, []chat.Message, []tools.Tool) (chat.MessageStream, error) { + return p.stream, p.err +} + +type servingTestStream struct { + closed bool +} + +func (*servingTestStream) Recv() (chat.MessageStreamResponse, error) { + return chat.MessageStreamResponse{}, io.EOF +} + +func (s *servingTestStream) Close() { s.closed = true } + +func TestServingModel(t *testing.T) { + t.Parallel() + + for _, tc := range []struct { + name string + matched bool + leafCost *latest.CostConfig + leafURL string + inner bool + nilStream bool + err error + }{ + {name: "matched leaf", matched: true}, + {name: "no match"}, + {name: "no match preserves leaf overrides", leafCost: &latest.CostConfig{Input: 2}, leafURL: "https://leaf.example/v1"}, + {name: "inner serving stream", matched: true, inner: true}, + {name: "nil stream", matched: true, nilStream: true}, + {name: "error", matched: true, err: errors.New("failed")}, + } { + t.Run(tc.name, func(t *testing.T) { + t.Parallel() + routerConfig := latest.ModelConfig{ + Provider: "openai", Model: "router", DisplayModel: "display-router", + BaseURL: "https://router.example/v1", Cost: &latest.CostConfig{Input: 99}, + Routing: []latest.RoutingRule{{Model: "leaf", Examples: []string{"poetry"}}}, + } + underlying := &servingTestStream{} + var stream chat.MessageStream = underlying + inner := contracts.ServingModel{ID: modelsdev.NewID("test", "inner"), Config: base.Config{BaseURL: "https://inner.example"}} + if tc.inner { + stream = &routedStream{MessageStream: stream, serving: inner} + } + if tc.nilStream { + stream = nil + } + leaf := &servingProvider{Config: base.Config{ + ModelConfig: latest.ModelConfig{Provider: "anthropic", Model: "leaf", Cost: tc.leafCost, BaseURL: tc.leafURL, ProviderOpts: map[string]any{"service_tier": "default"}}, + BaseURL: "https://api.openai.com/v1", + }, stream: stream, err: tc.err} + client, err := NewClient(t.Context(), &routerConfig, nil, nil, + func(context.Context, string, map[string]latest.ModelConfig, environment.Provider, ...options.Opt) (Provider, error) { + return leaf, nil + }) + require.NoError(t, err) + prompt := "unrelated" + if tc.matched { + prompt = "poetry" + } + got, err := client.CreateChatCompletionStream(t.Context(), []chat.Message{{Role: chat.MessageRoleUser, Content: prompt}}, nil) + if tc.err != nil { + require.ErrorIs(t, err, tc.err) + assert.Same(t, stream, got) + return + } + require.NoError(t, err) + if tc.nilStream { + assert.Nil(t, got) + return + } + serving := got.(contracts.ServingStream).ServingModel() + switch { + case tc.inner: + assert.Equal(t, inner, serving) + case tc.matched: + assert.Equal(t, leaf.ID(), serving.ID) + assert.Equal(t, leaf.BaseConfig(), serving.Config) + default: + assert.Equal(t, client.ID(), serving.ID) + assert.Equal(t, leaf.BaseConfig().BaseURL, serving.Config.BaseURL) + assert.Empty(t, serving.Config.ModelConfig.Routing) + assert.Equal(t, leaf.ModelConfig.ProviderOpts, serving.Config.ModelConfig.ProviderOpts) + if tc.leafCost != nil { + assert.Equal(t, tc.leafCost, serving.Config.ModelConfig.Cost) + } else { + assert.Equal(t, routerConfig.Cost, serving.Config.ModelConfig.Cost) + } + if tc.leafURL != "" { + assert.Equal(t, tc.leafURL, serving.Config.ModelConfig.BaseURL) + } else { + assert.Equal(t, routerConfig.BaseURL, serving.Config.ModelConfig.BaseURL) + } + } + _, err = got.Recv() + require.ErrorIs(t, err, io.EOF) + got.Close() + assert.True(t, underlying.closed) + assert.Equal(t, tc.leafCost, leaf.ModelConfig.Cost, "metadata does not mutate leaf config") + assert.Equal(t, tc.leafURL, leaf.ModelConfig.BaseURL) + }) + } +} diff --git a/pkg/runtime/fallback.go b/pkg/runtime/fallback.go index 2430ca3468..ad3f7753d0 100644 --- a/pkg/runtime/fallback.go +++ b/pkg/runtime/fallback.go @@ -369,18 +369,13 @@ func (e *fallbackExecutor) execute( slog.DebugContext(ctx, "Processing stream", "agent", a.Name(), "model", modelEntry.provider.ID().String()) - // If the provider is a rule-based router, notify the sidebar - // of the selected sub-model's YAML-configured name. - if rp, ok := modelEntry.provider.(interface { - LastSelectedModelID() modelsdev.ID - }); ok { - if selected := rp.LastSelectedModelID(); !selected.IsZero() { - events.Emit(AgentInfo(a.Name(), selected.String(), a.Description(), a.WelcomeMessage())) - } + if serving := servingModelFrom(stream); serving != nil && serving.ID != modelEntry.provider.ID() { + events.Emit(AgentInfo(a.Name(), serving.ID.String(), a.Description(), a.WelcomeMessage())) } idleTimeout := streamIdleTimeoutFor(modelEntry.provider) res, err := handleStream(streamCtx, streamCancel, stream, a, attemptTools, sess, events, idleTimeout) + res.Serving = servingModelFrom(stream) streamCancel(nil) // always release the child context if err != nil { lastErr = err @@ -419,6 +414,7 @@ func (e *fallbackExecutor) execute( stream, err = modelEntry.provider.CreateChatCompletionStream(streamCtx, attemptMessages, attemptTools) if err == nil { res, err = handleStream(streamCtx, streamCancel, stream, a, attemptTools, sess, events, idleTimeout) + res.Serving = servingModelFrom(stream) } streamCancel(nil) if err == nil { @@ -559,3 +555,10 @@ func (e *fallbackExecutor) handleModelError( "error", err) return retryDecisionContinue } + +func servingModelFrom(stream chat.MessageStream) *provider.ServingModel { + if stream, ok := stream.(provider.ServingStream); ok { + return new(stream.ServingModel()) + } + return nil +} diff --git a/pkg/runtime/loop.go b/pkg/runtime/loop.go index 4d9a39ca8a..3a36585e0b 100644 --- a/pkg/runtime/loop.go +++ b/pkg/runtime/loop.go @@ -937,18 +937,24 @@ func (r *LocalRuntime) runTurn( if usedModel != nil { agentTools = toolsForProvider(ctx, usedModel, agentTools) - if usedModel.ID() != modelID { - slog.InfoContext(ctx, "Used fallback model", "agent", a.Name(), "primary", modelID.String(), "used", usedModel.ID().String()) - modelID = usedModel.ID() + servingID, servingConfig := usedModel.ID(), usedModel.BaseConfig() + if res.Serving != nil { + servingID, servingConfig = res.Serving.ID, res.Serving.Config + } + if servingID != modelID { + if res.Serving == nil { + slog.InfoContext(ctx, "Used fallback model", "agent", a.Name(), "primary", modelID.String(), "used", servingID.String()) + events.Emit(AgentInfo(a.Name(), servingID.String(), a.Description(), a.WelcomeMessage())) + } + modelID = servingID m, err = r.modelsStore.GetModel(ctx, modelID) if err != nil { slog.DebugContext(ctx, "Failed to get fallback model definition", "model_id", modelID.String(), "error", err) m = nil } - events.Emit(AgentInfo(a.Name(), modelID.String(), a.Description(), a.WelcomeMessage())) } - // Fallbacks may share an ID but have different endpoints or cost overrides. - m = applyModelCost(m, modelID, res.Usage, usedModel.BaseConfig()) + // Serving endpoints may share an ID but have different rates or cost overrides. + m = applyModelCost(m, modelID, res.Usage, servingConfig) } // A successful model call resets the overflow compaction counter. @@ -1266,7 +1272,7 @@ func (r *LocalRuntime) Run(ctx context.Context, sess *session.Session) ([]sessio func applyModelCost(m *modelsdev.Model, id modelsdev.ID, usage *chat.Usage, config base.Config) *modelsdev.Model { cfg := config.ModelConfig - // Routers cannot identify the serving endpoint here; custom endpoints need their own rates. + // Unreported routing endpoints and custom endpoints cannot use vendor tier multipliers. isCustomEndpoint := func(endpoint string) bool { return endpoint != "" && strings.TrimRight(endpoint, "/") != "https://api.openai.com/v1" } diff --git a/pkg/runtime/routing_pricing_test.go b/pkg/runtime/routing_pricing_test.go new file mode 100644 index 0000000000..c5472dd486 --- /dev/null +++ b/pkg/runtime/routing_pricing_test.go @@ -0,0 +1,215 @@ +package runtime + +import ( + "context" + "errors" + "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/hooks" + "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/session" + "github.com/docker/docker-agent/pkg/team" +) + +type routingPricingProvider struct { + *countingProvider + + config base.Config +} + +func (p *routingPricingProvider) BaseConfig() base.Config { return p.config } + +func TestRoutedPricing(t *testing.T) { + t.Parallel() + + for _, tc := range []struct { + name string + routerID string + leafID string + routerCost *latest.CostConfig + leafCost *latest.CostConfig + routerURL string + leafURL string + resolvedLeafURL string + tier string + noMatch bool + fallback bool + want float64 + unpriced bool + }{ + {name: "catalog leaf", want: 0.43}, + {name: "unknown router", routerID: "local/router", want: 0.43}, + {name: "router override does not leak", routerCost: &latest.CostConfig{Input: 99, Output: 99}, want: 0.43}, + {name: "leaf override and cache classes", leafCost: &latest.CostConfig{Input: 1, Output: 2, CacheRead: 0.1, CacheWrite: 1.25}, want: 0.0995}, + {name: "free leaf", leafCost: &latest.CostConfig{}, want: 0}, + {name: "same ID response tier", routerID: "openai/gpt-6-astra", leafID: "openai/gpt-6-astra", tier: "ultrafast", want: 2.58}, + {name: "response downgrade", leafID: "openai/gpt-6-astra", tier: "default", want: 0.43}, + {name: "configured leaf endpoint", leafID: "openai/gpt-6-astra", leafURL: "https://leaf.example/v1", tier: "ultrafast", want: 0.43}, + {name: "resolved leaf endpoint", leafID: "openai/gpt-6-astra", resolvedLeafURL: "https://leaf.example/v1", tier: "ultrafast", want: 0.43}, + {name: "router endpoint does not suppress leaf tier", leafID: "openai/gpt-6-astra", routerURL: "https://router.example/v1", tier: "ultrafast", want: 2.58}, + {name: "unknown custom leaf", leafID: "local/unknown", leafURL: "https://leaf.example/v1", unpriced: true}, + {name: "unknown custom leaf override", leafID: "local/unknown", leafURL: "https://leaf.example/v1", leafCost: &latest.CostConfig{Input: 1, Output: 2, CacheRead: 0.1, CacheWrite: 1.25}, want: 0.0995}, + {name: "no match", noMatch: true, want: 0.072}, + {name: "no match override", noMatch: true, routerCost: &latest.CostConfig{Input: 1, Output: 2, CacheRead: 0.1, CacheWrite: 1.25}, want: 0.0995}, + {name: "no match tier", noMatch: true, routerID: "openai/gpt-6-astra", tier: "ultrafast", want: 2.58}, + {name: "no match declared endpoint", noMatch: true, routerID: "openai/gpt-6-astra", routerURL: "https://router.example/v1", tier: "ultrafast", want: 0.43}, + {name: "no match resolved endpoint", noMatch: true, routerID: "openai/gpt-6-astra", resolvedLeafURL: "https://leaf.example/v1", tier: "ultrafast", want: 0.43}, + {name: "failed route outer fallback and cooldown", fallback: true, tier: "ultrafast", want: 2.58}, + } { + t.Run(tc.name, func(t *testing.T) { + t.Parallel() + store := modelsdev.NewDatabaseStore(&modelsdev.Database{Providers: map[string]modelsdev.Provider{ + "openai": {Models: map[string]modelsdev.Model{ + "router": {Name: "Router", Cost: &modelsdev.Cost{Input: 1, Output: 1, CacheRead: 0.1, CacheWrite: 0.5}}, + "gpt-6-astra": {Name: "Astra", Cost: &modelsdev.Cost{Input: 4, Output: 8, CacheRead: 0.5, CacheWrite: 6}}, + }}, + "anthropic": {Models: map[string]modelsdev.Model{ + "leaf": {Name: "Leaf", Cost: &modelsdev.Cost{Input: 4, Output: 8, CacheRead: 0.5, CacheWrite: 6}}, + }}, + }}) + routerID := modelsdev.ParseIDOrZero(tc.routerID) + if routerID.IsZero() { + routerID = modelsdev.NewID("openai", "router") + } + leafID := modelsdev.ParseIDOrZero(tc.leafID) + if leafID.IsZero() { + leafID = modelsdev.NewID("anthropic", "leaf") + } + resolvedURL := tc.resolvedLeafURL + if resolvedURL == "" { + resolvedURL = "https://api.openai.com/v1" + } + leaf := &routingPricingProvider{countingProvider: &countingProvider{id: leafID.String()}, config: base.Config{ + ModelConfig: latest.ModelConfig{Provider: leafID.Provider, Model: leafID.Model, Cost: tc.leafCost, BaseURL: tc.leafURL, ProviderOpts: map[string]any{"service_tier": "fast"}}, BaseURL: resolvedURL, + }} + noMatch := &routingPricingProvider{countingProvider: &countingProvider{id: routerID.String()}, config: base.Config{BaseURL: resolvedURL}} + fallback := &routingPricingProvider{countingProvider: &countingProvider{id: "openai/gpt-6-astra"}, config: base.Config{BaseURL: "https://api.openai.com/v1"}} + router, err := rulebased.NewClient(t.Context(), &latest.ModelConfig{ + Provider: routerID.Provider, Model: routerID.Model, Cost: tc.routerCost, BaseURL: tc.routerURL, + Routing: []latest.RoutingRule{{Model: "leaf", Examples: []string{"poetry"}}}, + }, nil, nil, func(_ context.Context, spec string, _ map[string]latest.ModelConfig, _ environment.Provider, _ ...options.Opt) (rulebased.Provider, error) { + if spec == "leaf" { + return leaf, nil + } + return noMatch, nil + }) + require.NoError(t, err) + opts := []agent.Opt{agent.WithModel(router), agent.WithHooks(&latest.HooksConfig{AfterLLMCall: []latest.HookDefinition{{Type: "builtin", Command: "capture-routing"}}})} + if tc.fallback { + leaf.failCount, leaf.err = 100, errors.New("401 unauthorized") + opts = append(opts, agent.WithFallbackModel(fallback), agent.WithFallbackRetries(-1)) + } + root := agent.New("root", "test", opts...) + rec := &recordingTelemetry{} + rt, err := NewLocalRuntime(t.Context(), team.New(team.WithAgents(root)), WithTelemetry(rec), WithModelStore(store), WithSessionCompaction(false), WithBudget(&latest.BudgetConfig{MaxCost: 100})) + require.NoError(t, err) + var captured *hooks.Input + require.NoError(t, rt.hooksRegistry.RegisterBuiltin("capture-routing", func(_ context.Context, in *hooks.Input, _ []string) (*hooks.Output, error) { + snapshot := *in + captured = &snapshot + return nil, nil + })) + prompt, wantID, wantName := "poetry", leafID.String(), "Leaf" + switch leafID.Provider { + case "openai": + wantName = "Astra" + case "local": + wantName = "unknown" + if tc.leafCost != nil { + wantName = leafID.Model + } + } + if tc.noMatch { + prompt, wantID, wantName = "unrelated", routerID.String(), "Router" + if routerID.Model == "gpt-6-astra" { + wantName = "Astra" + } + } else if tc.fallback { + wantID, wantName = fallback.id, "Astra" + } + sess := session.New(session.WithTitle("Routing pricing test"), session.WithUserMessage(prompt)) + for turn := range 2 { + for _, p := range []*routingPricingProvider{leaf, noMatch, fallback} { + stream := newStreamBuilder().AddContent("ok").AddStopWithUsage(0, 0).Build() + stream.responses[len(stream.responses)-1].Usage = &chat.Usage{InputTokens: 50_000, CachedInputTokens: 20_000, CacheWriteTokens: 30_000, OutputTokens: 5_000, ServiceTier: tc.tier} + p.stream = stream + } + captured = nil + var message *chat.Message + var usage *MessageUsage + var budget *BudgetStatus + for ev := range rt.RunStream(t.Context(), sess) { + switch ev := ev.(type) { + case *ErrorEvent: + t.Errorf("unexpected runtime error: %s", ev.Error) + case *MessageAddedEvent: + if ev.Message != nil && ev.Message.Message.Role == chat.MessageRoleAssistant { + message = &ev.Message.Message + } + case *TokenUsageEvent: + if ev.Usage != nil && ev.Usage.LastMessage != nil { + usage = ev.Usage.LastMessage + } + case *BudgetUsageEvent: + for _, b := range ev.Budgets { + if b.Name == runBudgetName { + budget = &b + } + } + } + } + require.NotNil(t, captured) + assert.Equal(t, wantID, captured.ModelID) + if tc.unpriced { + assert.Nil(t, captured.Cost) + } else { + require.NotNil(t, captured.Cost) + assert.InDelta(t, tc.want, *captured.Cost, 1e-9) + } + require.NotNil(t, message) + assert.Equal(t, wantID, message.Model) + assert.InDelta(t, tc.want, message.Cost, 1e-9) + require.NotNil(t, usage) + assert.Equal(t, wantID, usage.Model) + assert.InDelta(t, tc.want, usage.Cost, 1e-9) + require.NotNil(t, budget) + assert.Equal(t, tc.unpriced, budget.Unpriced) + assert.InDelta(t, float64(turn+1)*tc.want, budget.Cost, 1e-9) + assert.InDelta(t, float64(turn+1)*tc.want, sess.OwnCost(), 1e-9) + records := rec.snapshot().tokenUsages + require.Len(t, records, turn+1) + assert.Equal(t, wantName, records[turn].Model) + assert.InDelta(t, tc.want, records[turn].Cost, 1e-9) + assert.Equal(t, int64(100_000), records[turn].InputTokens) + assert.Equal(t, int64(5_000), records[turn].OutputTokens) + } + 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: + assert.Equal(t, 2, noMatch.callCount) + assert.Zero(t, leaf.callCount) + default: + assert.Equal(t, 2, leaf.callCount, "successful router is not cooled down") + assert.Zero(t, noMatch.callCount) + } + }) + } +} + +func TestServingModelFromPlainStream(t *testing.T) { + t.Parallel() + assert.Nil(t, servingModelFrom(newStreamBuilder().Build())) +} diff --git a/pkg/runtime/streaming.go b/pkg/runtime/streaming.go index 8652320f9e..1b69e27e50 100644 --- a/pkg/runtime/streaming.go +++ b/pkg/runtime/streaming.go @@ -75,6 +75,7 @@ type streamResult struct { Stopped bool FinishReason chat.FinishReason Usage *chat.Usage + Serving *provider.ServingModel } // handleStream reads a chat.MessageStream to completion, emitting streaming From d63d780de095487fee7e8aea6e1d7ad91636ce74 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Arnaud=20H=C3=A9ritier?= Date: Tue, 13 Oct 2026 13:57:03 +0000 Subject: [PATCH 2/6] fix: preserve serving metadata and fallback sidebar updates MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: Arnaud Héritier --- pkg/model/provider/contracts/contracts.go | 1 + pkg/model/provider/instrument.go | 17 +++- pkg/model/provider/serving_test.go | 98 +++++++++++++++++++++++ pkg/runtime/fallback.go | 8 +- pkg/runtime/loop.go | 2 +- pkg/runtime/routing_pricing_test.go | 63 ++++++++++++++- 6 files changed, 184 insertions(+), 5 deletions(-) create mode 100644 pkg/model/provider/serving_test.go diff --git a/pkg/model/provider/contracts/contracts.go b/pkg/model/provider/contracts/contracts.go index 2a71beecf1..cacdd3e04e 100644 --- a/pkg/model/provider/contracts/contracts.go +++ b/pkg/model/provider/contracts/contracts.go @@ -135,6 +135,7 @@ type ServingModel struct { } // 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 diff --git a/pkg/model/provider/instrument.go b/pkg/model/provider/instrument.go index 210b96e250..7d8712a2e5 100644 --- a/pkg/model/provider/instrument.go +++ b/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 +} diff --git a/pkg/model/provider/serving_test.go b/pkg/model/provider/serving_test.go new file mode 100644 index 0000000000..6a91e077dd --- /dev/null +++ b/pkg/model/provider/serving_test.go @@ -0,0 +1,98 @@ +package provider + +import ( + "context" + "io" + "testing" + + "github.com/stretchr/testify/assert" + "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/options" + "github.com/docker/docker-agent/pkg/modelsdev" + "github.com/docker/docker-agent/pkg/tools" +) + +type registryServingProvider struct { + base.Config + + stream chat.MessageStream +} + +func (p *registryServingProvider) CreateChatCompletionStream(context.Context, []chat.Message, []tools.Tool) (chat.MessageStream, error) { + return p.stream, nil +} + +func TestServingModelThroughRegistry(t *testing.T) { + t.Parallel() + + for _, tc := range []struct { + name string + routed, noMatch, silent bool + }{ + {name: "direct"}, + {name: "routed", routed: true}, + {name: "routed silent leaf", routed: true, silent: true}, + {name: "no match reported leaf", routed: true, noMatch: true}, + {name: "no match silent leaf", routed: true, noMatch: true, silent: true}, + } { + t.Run(tc.name, func(t *testing.T) { + t.Parallel() + underlying := &instrumentTestStream{} + serving := ServingModel{ID: modelsdev.NewID("test", "actual"), Config: base.Config{ + ModelConfig: latest.ModelConfig{Provider: "test", Model: "actual", Cost: &latest.CostConfig{Input: 1, CacheRead: 0.1}}, + BaseURL: "https://actual.example/v1", + }} + registry := NewRegistry(map[string]Factory{"test": func(_ context.Context, cfg *latest.ModelConfig, _ environment.Provider, _ ...options.Opt) (Provider, error) { + var stream chat.MessageStream = underlying + if !tc.silent { + stream = &tracedServingStream{MessageStream: stream, serving: serving} + } + return ®istryServingProvider{Config: base.Config{ModelConfig: *cfg, BaseURL: "https://resolved.example/v1"}, stream: stream}, nil + }}) + cfg := &latest.ModelConfig{Provider: "test", Model: "child"} + if tc.routed { + cfg.Model = "router" + cfg.Routing = []latest.RoutingRule{{Model: "test/child", Examples: []string{"poetry"}}} + } + if tc.noMatch { + cfg.DisplayModel = "display-router" + cfg.BaseURL = "https://declared.example/v1" + cfg.Cost = &latest.CostConfig{Input: 9} + } + p, err := registry.New(t.Context(), cfg, nil) + require.NoError(t, err) + prompt := "poetry" + if tc.noMatch { + prompt = "unrelated" + } + stream, err := p.CreateChatCompletionStream(t.Context(), []chat.Message{{Role: chat.MessageRoleUser, Content: prompt}}, nil) + require.NoError(t, err) + reporter, ok := stream.(ServingStream) + require.True(t, ok) + if tc.silent { + got := reporter.ServingModel() + assert.Equal(t, "https://resolved.example/v1", got.Config.BaseURL) + if tc.noMatch { + assert.Equal(t, p.ID(), got.ID) + assert.Equal(t, cfg.Cost, got.Config.ModelConfig.Cost) + assert.Equal(t, cfg.BaseURL, got.Config.ModelConfig.BaseURL) + } else { + assert.Equal(t, modelsdev.NewID("test", "child"), got.ID) + assert.Nil(t, got.Config.ModelConfig.Cost) + } + } else { + assert.Equal(t, serving, reporter.ServingModel()) + } + _, err = stream.Recv() + require.ErrorIs(t, err, io.EOF) + stream.Close() + assert.Equal(t, 1, underlying.received) + assert.Equal(t, 1, underlying.closed) + }) + } +} diff --git a/pkg/runtime/fallback.go b/pkg/runtime/fallback.go index ad3f7753d0..c70daa50ec 100644 --- a/pkg/runtime/fallback.go +++ b/pkg/runtime/fallback.go @@ -369,9 +369,12 @@ func (e *fallbackExecutor) execute( slog.DebugContext(ctx, "Processing stream", "agent", a.Name(), "model", modelEntry.provider.ID().String()) - if serving := servingModelFrom(stream); serving != nil && serving.ID != modelEntry.provider.ID() { - events.Emit(AgentInfo(a.Name(), serving.ID.String(), a.Description(), a.WelcomeMessage())) + announceServing := func(stream chat.MessageStream) { + if serving := servingModelFrom(stream); serving != nil && serving.ID != modelEntry.provider.ID() { + events.Emit(AgentInfo(a.Name(), serving.ID.String(), a.Description(), a.WelcomeMessage())) + } } + announceServing(stream) idleTimeout := streamIdleTimeoutFor(modelEntry.provider) res, err := handleStream(streamCtx, streamCancel, stream, a, attemptTools, sess, events, idleTimeout) @@ -413,6 +416,7 @@ func (e *fallbackExecutor) execute( fbSpan.IncrementAttempt() stream, err = modelEntry.provider.CreateChatCompletionStream(streamCtx, attemptMessages, attemptTools) if err == nil { + announceServing(stream) res, err = handleStream(streamCtx, streamCancel, stream, a, attemptTools, sess, events, idleTimeout) res.Serving = servingModelFrom(stream) } diff --git a/pkg/runtime/loop.go b/pkg/runtime/loop.go index 3a36585e0b..42ff97892c 100644 --- a/pkg/runtime/loop.go +++ b/pkg/runtime/loop.go @@ -942,7 +942,7 @@ func (r *LocalRuntime) runTurn( servingID, servingConfig = res.Serving.ID, res.Serving.Config } if servingID != modelID { - if res.Serving == nil { + if servingID == usedModel.ID() { slog.InfoContext(ctx, "Used fallback model", "agent", a.Name(), "primary", modelID.String(), "used", servingID.String()) events.Emit(AgentInfo(a.Name(), servingID.String(), a.Description(), a.WelcomeMessage())) } diff --git a/pkg/runtime/routing_pricing_test.go b/pkg/runtime/routing_pricing_test.go index c5472dd486..aad12b7ae1 100644 --- a/pkg/runtime/routing_pricing_test.go +++ b/pkg/runtime/routing_pricing_test.go @@ -4,6 +4,7 @@ import ( "context" "errors" "testing" + "time" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" @@ -13,6 +14,7 @@ import ( "github.com/docker/docker-agent/pkg/config/latest" "github.com/docker/docker-agent/pkg/environment" "github.com/docker/docker-agent/pkg/hooks" + "github.com/docker/docker-agent/pkg/model/provider" "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" @@ -44,6 +46,7 @@ func TestRoutedPricing(t *testing.T) { tier string noMatch bool fallback bool + outerRouter bool want float64 unpriced bool }{ @@ -65,6 +68,8 @@ func TestRoutedPricing(t *testing.T) { {name: "no match declared endpoint", noMatch: true, routerID: "openai/gpt-6-astra", routerURL: "https://router.example/v1", tier: "ultrafast", want: 0.43}, {name: "no match resolved endpoint", noMatch: true, routerID: "openai/gpt-6-astra", resolvedLeafURL: "https://leaf.example/v1", tier: "ultrafast", want: 0.43}, {name: "failed route outer fallback and cooldown", fallback: true, tier: "ultrafast", want: 2.58}, + {name: "outer fallback router no match and cooldown", outerRouter: true, noMatch: true, want: 0.072}, + {name: "outer fallback router matched and cooldown", outerRouter: true, want: 0.43}, } { t.Run(tc.name, func(t *testing.T) { t.Parallel() @@ -104,11 +109,19 @@ func TestRoutedPricing(t *testing.T) { return noMatch, nil }) require.NoError(t, err) - opts := []agent.Opt{agent.WithModel(router), agent.WithHooks(&latest.HooksConfig{AfterLLMCall: []latest.HookDefinition{{Type: "builtin", Command: "capture-routing"}}})} + model := rulebased.Provider(router) + primary := &countingProvider{id: "test/failed", failCount: 100, err: errors.New("401 unauthorized")} + if tc.outerRouter { + model = primary + } + opts := []agent.Opt{agent.WithModel(model), agent.WithHooks(&latest.HooksConfig{AfterLLMCall: []latest.HookDefinition{{Type: "builtin", Command: "capture-routing"}}})} if tc.fallback { leaf.failCount, leaf.err = 100, errors.New("401 unauthorized") opts = append(opts, agent.WithFallbackModel(fallback), agent.WithFallbackRetries(-1)) } + if tc.outerRouter { + opts = append(opts, agent.WithFallbackModel(router), agent.WithFallbackRetries(-1)) + } root := agent.New("root", "test", opts...) rec := &recordingTelemetry{} rt, err := NewLocalRuntime(t.Context(), team.New(team.WithAgents(root)), WithTelemetry(rec), WithModelStore(store), WithSessionCompaction(false), WithBudget(&latest.BudgetConfig{MaxCost: 100})) @@ -145,11 +158,14 @@ func TestRoutedPricing(t *testing.T) { p.stream = stream } captured = nil + var sidebarModels []string var message *chat.Message var usage *MessageUsage var budget *BudgetStatus for ev := range rt.RunStream(t.Context(), sess) { switch ev := ev.(type) { + case *AgentInfoEvent: + sidebarModels = append(sidebarModels, ev.Model) case *ErrorEvent: t.Errorf("unexpected runtime error: %s", ev.Error) case *MessageAddedEvent: @@ -170,6 +186,11 @@ func TestRoutedPricing(t *testing.T) { } require.NotNil(t, captured) assert.Equal(t, wantID, captured.ModelID) + require.NotEmpty(t, sidebarModels) + assert.Equal(t, wantID, sidebarModels[len(sidebarModels)-1]) + if wantID != model.ID().String() { + assert.Equal(t, []string{model.ID().String(), wantID}, sidebarModels) + } if tc.unpriced { assert.Nil(t, captured.Cost) } else { @@ -193,6 +214,9 @@ func TestRoutedPricing(t *testing.T) { assert.Equal(t, int64(100_000), records[turn].InputTokens) assert.Equal(t, int64(5_000), records[turn].OutputTokens) } + if tc.outerRouter { + assert.Equal(t, 1, primary.callCount, "cooldown skips failed primary before successful fallback router") + } switch { case tc.fallback: assert.Equal(t, leafID, router.LastSelectedModelID(), "failed route selection must not price the successful outer fallback") @@ -213,3 +237,40 @@ func TestServingModelFromPlainStream(t *testing.T) { t.Parallel() assert.Nil(t, servingModelFrom(newStreamBuilder().Build())) } + +type pricingServingStream struct { + chat.MessageStream + + serving provider.ServingModel +} + +func (s *pricingServingStream) ServingModel() provider.ServingModel { return s.serving } + +func TestRoutedIdleRetryServingModel(t *testing.T) { + t.Parallel() + + first := provider.ServingModel{ID: modelsdev.NewID("test", "first")} + second := provider.ServingModel{ID: modelsdev.NewID("test", "second"), Config: base.Config{BaseURL: "https://second.example/v1"}} + recovered := newStreamBuilder().AddContent("recovered").AddStopWithUsage(1, 1).Build() + model := &idleThenResultProvider{id: "test/router", streams: []chat.MessageStream{ + &pricingServingStream{MessageStream: &idleErrorStream{}, serving: first}, + &pricingServingStream{MessageStream: recovered, serving: second}, + }} + a := agent.New("delegate", "test", agent.WithModel(model), agent.WithFallbackRetries(-1)) + executor := newFallbackExecutor() + executor.cooldowns = newCooldownManager(time.Now) + sink := &collectSink{} + result, used, err := executor.execute(t.Context(), a, model, nil, nil, session.New(), sink, + &idleStreamRetryAllowance{remaining: true, parentSessionID: "parent"}, func() error { return nil }) + require.NoError(t, err) + assert.Equal(t, model, used) + require.NotNil(t, result.Serving) + assert.Equal(t, second, *result.Serving) + var announced []string + for _, event := range sink.events { + if event, ok := event.(*AgentInfoEvent); ok { + announced = append(announced, event.Model) + } + } + assert.Equal(t, []string{first.ID.String(), second.ID.String()}, announced) +} From bd2bb64df196f767cee6ea8673cf384b7fb439a6 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Arnaud=20H=C3=A9ritier?= Date: Tue, 13 Oct 2026 14:13:01 +0000 Subject: [PATCH 3/6] fix: reconcile serving sidebar identity across retries MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: Arnaud Héritier --- pkg/runtime/fallback.go | 18 +++-- pkg/runtime/loop.go | 1 - pkg/runtime/routing_pricing_test.go | 108 +++++++++++++++++++++++----- 3 files changed, 103 insertions(+), 24 deletions(-) diff --git a/pkg/runtime/fallback.go b/pkg/runtime/fallback.go index c70daa50ec..eed608e1a0 100644 --- a/pkg/runtime/fallback.go +++ b/pkg/runtime/fallback.go @@ -288,10 +288,23 @@ func (e *fallbackExecutor) execute( var lastErr error primaryFailedWithNonRetryable := false hasFallbacks := len(fallbackModels) > 0 + // The loop announces the primary at turn start; track changes across the whole chain. + announcedID := primaryModel.ID() for chainIdx := startIndex; chainIdx < len(modelChain); chainIdx++ { modelEntry := modelChain[chainIdx] + announceServing := func(stream chat.MessageStream) { + id := modelEntry.provider.ID() + if serving := servingModelFrom(stream); serving != nil { + id = serving.ID + } + if id != announcedID { + events.Emit(AgentInfo(a.Name(), id.String(), a.Description(), a.WelcomeMessage())) + announcedID = id + } + } + // Each model in the chain gets (1 + retries) attempts for retryable errors. // Non-retryable errors (429 with fallbacks, 4xx) skip immediately to the next model. // 429 without fallbacks is retried directly on the same model. @@ -369,11 +382,6 @@ func (e *fallbackExecutor) execute( slog.DebugContext(ctx, "Processing stream", "agent", a.Name(), "model", modelEntry.provider.ID().String()) - announceServing := func(stream chat.MessageStream) { - if serving := servingModelFrom(stream); serving != nil && serving.ID != modelEntry.provider.ID() { - events.Emit(AgentInfo(a.Name(), serving.ID.String(), a.Description(), a.WelcomeMessage())) - } - } announceServing(stream) idleTimeout := streamIdleTimeoutFor(modelEntry.provider) diff --git a/pkg/runtime/loop.go b/pkg/runtime/loop.go index 42ff97892c..4e41a7d71b 100644 --- a/pkg/runtime/loop.go +++ b/pkg/runtime/loop.go @@ -944,7 +944,6 @@ func (r *LocalRuntime) runTurn( if servingID != modelID { if servingID == usedModel.ID() { slog.InfoContext(ctx, "Used fallback model", "agent", a.Name(), "primary", modelID.String(), "used", servingID.String()) - events.Emit(AgentInfo(a.Name(), servingID.String(), a.Description(), a.WelcomeMessage())) } modelID = servingID m, err = r.modelsStore.GetModel(ctx, modelID) diff --git a/pkg/runtime/routing_pricing_test.go b/pkg/runtime/routing_pricing_test.go index aad12b7ae1..65b21b72a7 100644 --- a/pkg/runtime/routing_pricing_test.go +++ b/pkg/runtime/routing_pricing_test.go @@ -4,6 +4,7 @@ import ( "context" "errors" "testing" + "testing/synctest" "time" "github.com/stretchr/testify/assert" @@ -246,31 +247,102 @@ type pricingServingStream struct { func (s *pricingServingStream) ServingModel() provider.ServingModel { return s.serving } +type pricingErrorStream struct{ err error } + +func (s *pricingErrorStream) Recv() (chat.MessageStreamResponse, error) { + return chat.MessageStreamResponse{}, s.err +} +func (*pricingErrorStream) Close() {} + func TestRoutedIdleRetryServingModel(t *testing.T) { t.Parallel() - first := provider.ServingModel{ID: modelsdev.NewID("test", "first")} - second := provider.ServingModel{ID: modelsdev.NewID("test", "second"), Config: base.Config{BaseURL: "https://second.example/v1"}} - recovered := newStreamBuilder().AddContent("recovered").AddStopWithUsage(1, 1).Build() - model := &idleThenResultProvider{id: "test/router", streams: []chat.MessageStream{ - &pricingServingStream{MessageStream: &idleErrorStream{}, serving: first}, - &pricingServingStream{MessageStream: recovered, serving: second}, + for _, tc := range []struct { + name, secondModel string + plain, ordinary bool + }{ + {name: "different leaf", secondModel: "second"}, + {name: "dispatch model", secondModel: "router"}, + {name: "plain stream", secondModel: "router", plain: true}, + {name: "ordinary retry dispatch model", secondModel: "router", ordinary: true}, + {name: "ordinary retry plain stream", secondModel: "router", ordinary: true, plain: true}, + } { + t.Run(tc.name, func(t *testing.T) { + t.Parallel() + synctest.Test(t, func(t *testing.T) { + first := provider.ServingModel{ID: modelsdev.NewID("test", "first")} + second := provider.ServingModel{ID: modelsdev.NewID("test", tc.secondModel), Config: base.Config{BaseURL: "https://second.example/v1"}} + var recovered chat.MessageStream = newStreamBuilder().AddContent("recovered").AddStopWithUsage(1, 1).Build() + if !tc.plain { + recovered = &pricingServingStream{MessageStream: recovered, serving: second} + } + var failed chat.MessageStream = &idleErrorStream{} + retries := -1 + allowance := &idleStreamRetryAllowance{remaining: true, parentSessionID: "parent"} + if tc.ordinary { + failed = &pricingErrorStream{err: errors.New("500 internal server error")} + retries = 1 + allowance = nil + } + model := &idleThenResultProvider{id: "test/router", streams: []chat.MessageStream{ + &pricingServingStream{MessageStream: failed, serving: first}, recovered, + }} + a := agent.New("delegate", "test", agent.WithModel(model), agent.WithFallbackRetries(retries)) + executor := newFallbackExecutor() + executor.cooldowns = newCooldownManager(time.Now) + sink := &collectSink{} + result, used, err := executor.execute(t.Context(), a, model, nil, nil, session.New(), sink, + allowance, func() error { return nil }) + require.NoError(t, err) + assert.Equal(t, model, used) + if tc.plain { + assert.Nil(t, result.Serving) + } else { + require.NotNil(t, result.Serving) + assert.Equal(t, second, *result.Serving) + } + var announced []string + for _, event := range sink.events { + if event, ok := event.(*AgentInfoEvent); ok { + announced = append(announced, event.Model) + } + } + assert.Equal(t, []string{first.ID.String(), second.ID.String()}, announced) + }) + }) + } +} + +func TestRoutedSameIDFallbackServingModel(t *testing.T) { + t.Parallel() + + first := provider.ServingModel{ID: modelsdev.NewID("test", "leaf")} + primary := &idleThenResultProvider{id: "test/router", streams: []chat.MessageStream{ + &pricingServingStream{MessageStream: &pricingErrorStream{err: errors.New("401 unauthorized")}, serving: first}, }} - a := agent.New("delegate", "test", agent.WithModel(model), agent.WithFallbackRetries(-1)) - executor := newFallbackExecutor() - executor.cooldowns = newCooldownManager(time.Now) - sink := &collectSink{} - result, used, err := executor.execute(t.Context(), a, model, nil, nil, session.New(), sink, - &idleStreamRetryAllowance{remaining: true, parentSessionID: "parent"}, func() error { return nil }) + fallback := &routingPricingProvider{countingProvider: &countingProvider{ + id: "test/router", stream: newStreamBuilder().AddContent("ok").AddStopWithUsage(50_000, 5_000).Build(), + }, config: base.Config{ModelConfig: latest.ModelConfig{Cost: &latest.CostConfig{Input: 2, Output: 3}}}} + a := agent.New("root", "test", agent.WithModel(primary), agent.WithFallbackModel(fallback), agent.WithFallbackRetries(-1)) + rt, err := NewLocalRuntime(t.Context(), team.New(team.WithAgents(a)), WithModelStore(modelsdev.NewDatabaseStore(&modelsdev.Database{})), WithSessionCompaction(false)) require.NoError(t, err) - assert.Equal(t, model, used) - require.NotNil(t, result.Serving) - assert.Equal(t, second, *result.Serving) + sess := session.New(session.WithTitle("Same ID fallback"), session.WithUserMessage("hi")) var announced []string - for _, event := range sink.events { - if event, ok := event.(*AgentInfoEvent); ok { + var assistant *chat.Message + for event := range rt.RunStream(t.Context(), sess) { + switch event := event.(type) { + case *ErrorEvent: + t.Errorf("runtime error: %s", event.Error) + case *AgentInfoEvent: announced = append(announced, event.Model) + case *MessageAddedEvent: + if event.Message != nil && event.Message.Message.Role == chat.MessageRoleAssistant { + assistant = &event.Message.Message + } } } - assert.Equal(t, []string{first.ID.String(), second.ID.String()}, announced) + assert.Equal(t, []string{"test/router", "test/leaf", "test/router"}, announced) + require.NotNil(t, assistant) + assert.Equal(t, "test/router", assistant.Model) + assert.InDelta(t, 0.115, assistant.Cost, 1e-9) } From 990978008ecb3d777870e90480ab1c16c3def0d6 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Arnaud=20H=C3=A9ritier?= Date: Tue, 13 Oct 2026 14:22:19 +0000 Subject: [PATCH 4/6] fix: track fallback notifications in sidebar reconciliation MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: Arnaud Héritier --- pkg/runtime/fallback.go | 3 ++- pkg/runtime/routing_pricing_test.go | 14 ++++++++++++-- 2 files changed, 14 insertions(+), 3 deletions(-) diff --git a/pkg/runtime/fallback.go b/pkg/runtime/fallback.go index eed608e1a0..dee3d1e0fb 100644 --- a/pkg/runtime/fallback.go +++ b/pkg/runtime/fallback.go @@ -288,7 +288,7 @@ func (e *fallbackExecutor) execute( var lastErr error primaryFailedWithNonRetryable := false hasFallbacks := len(fallbackModels) > 0 - // The loop announces the primary at turn start; track changes across the whole chain. + // Track the sidebar model set by AgentInfo and ModelFallback, starting with the primary. announcedID := primaryModel.ID() for chainIdx := startIndex; chainIdx < len(modelChain); chainIdx++ { @@ -349,6 +349,7 @@ func (e *fallbackExecutor) execute( attempt+1, maxAttempts, )) + announcedID = modelEntry.provider.ID() } slog.DebugContext(ctx, "Creating chat completion stream", diff --git a/pkg/runtime/routing_pricing_test.go b/pkg/runtime/routing_pricing_test.go index 65b21b72a7..2bcd7b2ad8 100644 --- a/pkg/runtime/routing_pricing_test.go +++ b/pkg/runtime/routing_pricing_test.go @@ -71,10 +71,14 @@ func TestRoutedPricing(t *testing.T) { {name: "failed route outer fallback and cooldown", fallback: true, tier: "ultrafast", want: 2.58}, {name: "outer fallback router no match and cooldown", outerRouter: true, noMatch: true, want: 0.072}, {name: "outer fallback router matched and cooldown", outerRouter: true, want: 0.43}, + {name: "fallback router serves primary ID", outerRouter: true, leafID: "test/failed", want: 0.43}, } { t.Run(tc.name, func(t *testing.T) { t.Parallel() store := modelsdev.NewDatabaseStore(&modelsdev.Database{Providers: map[string]modelsdev.Provider{ + "test": {Models: map[string]modelsdev.Model{ + "failed": {Name: "Leaf", Cost: &modelsdev.Cost{Input: 4, Output: 8, CacheRead: 0.5, CacheWrite: 6}}, + }}, "openai": {Models: map[string]modelsdev.Model{ "router": {Name: "Router", Cost: &modelsdev.Cost{Input: 1, Output: 1, CacheRead: 0.1, CacheWrite: 0.5}}, "gpt-6-astra": {Name: "Astra", Cost: &modelsdev.Cost{Input: 4, Output: 8, CacheRead: 0.5, CacheWrite: 6}}, @@ -159,6 +163,7 @@ func TestRoutedPricing(t *testing.T) { p.stream = stream } captured = nil + var sidebarModel string var sidebarModels []string var message *chat.Message var usage *MessageUsage @@ -167,6 +172,9 @@ func TestRoutedPricing(t *testing.T) { switch ev := ev.(type) { case *AgentInfoEvent: sidebarModels = append(sidebarModels, ev.Model) + sidebarModel = ev.Model + case *ModelFallbackEvent: + sidebarModel = ev.FallbackModel case *ErrorEvent: t.Errorf("unexpected runtime error: %s", ev.Error) case *MessageAddedEvent: @@ -188,8 +196,8 @@ func TestRoutedPricing(t *testing.T) { require.NotNil(t, captured) assert.Equal(t, wantID, captured.ModelID) require.NotEmpty(t, sidebarModels) - assert.Equal(t, wantID, sidebarModels[len(sidebarModels)-1]) - if wantID != model.ID().String() { + assert.Equal(t, wantID, sidebarModel, "both AgentInfo and ModelFallback update the sidebar") + if wantID != model.ID().String() && (!tc.fallback && !tc.outerRouter || turn > 0) { assert.Equal(t, []string{model.ID().String(), wantID}, sidebarModels) } if tc.unpriced { @@ -335,6 +343,8 @@ func TestRoutedSameIDFallbackServingModel(t *testing.T) { t.Errorf("runtime error: %s", event.Error) case *AgentInfoEvent: announced = append(announced, event.Model) + case *ModelFallbackEvent: + announced = append(announced, event.FallbackModel) case *MessageAddedEvent: if event.Message != nil && event.Message.Message.Role == chat.MessageRoleAssistant { assistant = &event.Message.Message From d7420c36746a2b31ecab1a5f4f88c7e79af9260c Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Arnaud=20H=C3=A9ritier?= Date: Tue, 13 Oct 2026 14:34:57 +0000 Subject: [PATCH 5/6] fix: notify all UI consumers of fallback serving models MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: Arnaud Héritier --- pkg/leantui/serving_model_test.go | 93 +++++++++++++++++++++++++++++ pkg/runtime/fallback.go | 9 ++- pkg/runtime/routing_pricing_test.go | 3 +- 3 files changed, 101 insertions(+), 4 deletions(-) create mode 100644 pkg/leantui/serving_model_test.go diff --git a/pkg/leantui/serving_model_test.go b/pkg/leantui/serving_model_test.go new file mode 100644 index 0000000000..f0678b9317 --- /dev/null +++ b/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) + } + }) + } +} diff --git a/pkg/runtime/fallback.go b/pkg/runtime/fallback.go index dee3d1e0fb..2bd3700f9c 100644 --- a/pkg/runtime/fallback.go +++ b/pkg/runtime/fallback.go @@ -288,8 +288,9 @@ func (e *fallbackExecutor) execute( var lastErr error primaryFailedWithNonRetryable := false hasFallbacks := len(fallbackModels) > 0 - // Track the sidebar model set by AgentInfo and ModelFallback, starting with the primary. + // The loop announces the primary before dispatch. announcedID := primaryModel.ID() + needsAnnouncement := false for chainIdx := startIndex; chainIdx < len(modelChain); chainIdx++ { modelEntry := modelChain[chainIdx] @@ -299,9 +300,10 @@ func (e *fallbackExecutor) execute( if serving := servingModelFrom(stream); serving != nil { id = serving.ID } - if id != announcedID { + if needsAnnouncement || id != announcedID { events.Emit(AgentInfo(a.Name(), id.String(), a.Description(), a.WelcomeMessage())) announcedID = id + needsAnnouncement = false } } @@ -349,7 +351,8 @@ func (e *fallbackExecutor) execute( attempt+1, maxAttempts, )) - announcedID = modelEntry.provider.ID() + // Not every consumer updates its model on ModelFallback. + needsAnnouncement = true } slog.DebugContext(ctx, "Creating chat completion stream", diff --git a/pkg/runtime/routing_pricing_test.go b/pkg/runtime/routing_pricing_test.go index 2bcd7b2ad8..5b674de942 100644 --- a/pkg/runtime/routing_pricing_test.go +++ b/pkg/runtime/routing_pricing_test.go @@ -196,6 +196,7 @@ func TestRoutedPricing(t *testing.T) { require.NotNil(t, captured) assert.Equal(t, wantID, captured.ModelID) require.NotEmpty(t, sidebarModels) + assert.Equal(t, wantID, sidebarModels[len(sidebarModels)-1], "AgentInfo-only consumers also receive the serving identity") assert.Equal(t, wantID, sidebarModel, "both AgentInfo and ModelFallback update the sidebar") if wantID != model.ID().String() && (!tc.fallback && !tc.outerRouter || turn > 0) { assert.Equal(t, []string{model.ID().String(), wantID}, sidebarModels) @@ -351,7 +352,7 @@ func TestRoutedSameIDFallbackServingModel(t *testing.T) { } } } - assert.Equal(t, []string{"test/router", "test/leaf", "test/router"}, announced) + assert.Equal(t, []string{"test/router", "test/leaf", "test/router", "test/router"}, announced) require.NotNil(t, assistant) assert.Equal(t, "test/router", assistant.Model) assert.InDelta(t, 0.115, assistant.Cost, 1e-9) From 0404f5c291d1ed02cb14a863a52b4b1ae2781528 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Arnaud=20H=C3=A9ritier?= Date: Wed, 14 Oct 2026 09:16:19 +0000 Subject: [PATCH 6/6] fix: attribute fallback spans to the serving model MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: Arnaud Héritier --- pkg/runtime/fallback.go | 11 +++- pkg/runtime/fallback_span_test.go | 91 +++++++++++++++++++++++++++++++ 2 files changed, 100 insertions(+), 2 deletions(-) create mode 100644 pkg/runtime/fallback_span_test.go diff --git a/pkg/runtime/fallback.go b/pkg/runtime/fallback.go index 2bd3700f9c..66ba8d6a83 100644 --- a/pkg/runtime/fallback.go +++ b/pkg/runtime/fallback.go @@ -435,7 +435,7 @@ func (e *fallbackExecutor) execute( streamCancel(nil) if err == nil { e.recordSuccess(a, modelEntry, primaryFailedWithNonRetryable) - fbSpan.SetFinalModel(modelEntry.provider.ID().Model) + fbSpan.SetFinalModel(finalModelFor(res.Serving, modelEntry.provider)) fbSpan.SetOutcome(genai.FallbackOutcomeSuccess) return res, modelEntry.provider, nil } @@ -453,7 +453,7 @@ func (e *fallbackExecutor) execute( } e.recordSuccess(a, modelEntry, primaryFailedWithNonRetryable) - fbSpan.SetFinalModel(modelEntry.provider.ID().Model) + fbSpan.SetFinalModel(finalModelFor(res.Serving, modelEntry.provider)) fbSpan.SetOutcome(genai.FallbackOutcomeSuccess) return res, modelEntry.provider, nil } @@ -572,6 +572,13 @@ func (e *fallbackExecutor) handleModelError( return retryDecisionContinue } +func finalModelFor(serving *provider.ServingModel, dispatched provider.Provider) string { + if serving != nil && serving.ID.Model != "" { + return serving.ID.Model + } + return dispatched.ID().Model +} + func servingModelFrom(stream chat.MessageStream) *provider.ServingModel { if stream, ok := stream.(provider.ServingStream); ok { return new(stream.ServingModel()) diff --git a/pkg/runtime/fallback_span_test.go b/pkg/runtime/fallback_span_test.go new file mode 100644 index 0000000000..ffc7fc5ec8 --- /dev/null +++ b/pkg/runtime/fallback_span_test.go @@ -0,0 +1,91 @@ +package runtime + +import ( + "context" + "testing" + "time" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "go.opentelemetry.io/otel" + "go.opentelemetry.io/otel/attribute" + sdktrace "go.opentelemetry.io/otel/sdk/trace" + "go.opentelemetry.io/otel/sdk/trace/tracetest" + + "github.com/docker/docker-agent/pkg/agent" + "github.com/docker/docker-agent/pkg/chat" + "github.com/docker/docker-agent/pkg/model/provider" + "github.com/docker/docker-agent/pkg/modelsdev" + "github.com/docker/docker-agent/pkg/session" + "github.com/docker/docker-agent/pkg/telemetry/genai" +) + +// Serial: the recording tracer replaces global OTel state. +func TestFallbackSpanServingModel(t *testing.T) { + exporter := tracetest.NewInMemoryExporter() + tp := sdktrace.NewTracerProvider(sdktrace.WithSyncer(exporter)) + previous := otel.GetTracerProvider() + otel.SetTracerProvider(tp) + t.Cleanup(func() { + otel.SetTracerProvider(previous) + require.NoError(t, tp.Shutdown(context.WithoutCancel(t.Context()))) + }) + + for _, tc := range []struct { + name, servingModel, wantModel string + plain, idleRetry bool + }{ + {name: "routed leaf", servingModel: "leaf", wantModel: "leaf"}, + {name: "plain stream", plain: true, wantModel: "router"}, + {name: "dispatch identity", servingModel: "router", wantModel: "router"}, + {name: "idle retry routed leaf", idleRetry: true, servingModel: "second", wantModel: "second"}, + {name: "idle retry plain stream", idleRetry: true, plain: true, wantModel: "router"}, + {name: "empty serving model", wantModel: "router"}, + } { + t.Run(tc.name, func(t *testing.T) { + exporter.Reset() + var stream chat.MessageStream = newStreamBuilder().AddContent("ok").AddStopWithUsage(1, 1).Build() + if !tc.plain { + stream = &pricingServingStream{MessageStream: stream, serving: provider.ServingModel{ID: modelsdev.NewID("test", tc.servingModel)}} + } + model := &idleThenResultProvider{id: "test/router", streams: []chat.MessageStream{stream}} + var allowance *idleStreamRetryAllowance + wantAttempts := 1 + if tc.idleRetry { + model.streams = append([]chat.MessageStream{&pricingServingStream{ + MessageStream: &idleErrorStream{}, serving: provider.ServingModel{ID: modelsdev.NewID("test", "failed")}, + }}, model.streams...) + allowance = &idleStreamRetryAllowance{remaining: true, parentSessionID: "parent"} + wantAttempts = 2 + } + a := agent.New("delegate", "test", agent.WithModel(model), agent.WithFallbackRetries(-1)) + executor := newFallbackExecutor() + executor.cooldowns = newCooldownManager(time.Now) + _, used, err := executor.execute(t.Context(), a, model, nil, nil, session.New(), &collectSink{}, + allowance, func() error { return nil }) + require.NoError(t, err) + assert.Same(t, model, used) + assert.Equal(t, wantAttempts, model.calls) + assert.Nil(t, executor.cooldowns.Get(a.Name())) + + var fallbackSpans tracetest.SpanStubs + for _, span := range exporter.GetSpans() { + if span.Name == "runtime.fallback" { + fallbackSpans = append(fallbackSpans, span) + } + } + require.Len(t, fallbackSpans, 1) + attrs := attribute.NewSet(fallbackSpans[0].Attributes...) + for key, want := range map[string]attribute.Value{ + genai.AttrFallbackPrimaryModel: attribute.StringValue("router"), + genai.AttrFallbackFinalModel: attribute.StringValue(tc.wantModel), + genai.AttrFallbackAttempts: attribute.IntValue(wantAttempts), + genai.AttrFallbackOutcome: attribute.StringValue(genai.FallbackOutcomeSuccess), + } { + got, ok := attrs.Value(attribute.Key(key)) + assert.True(t, ok, "missing span attribute %s", key) + assert.Equal(t, want, got, "span attribute %s", key) + } + }) + } +}