test(runtime): a server fallback turn is billed per attempt

The lifecycle test now prices each model differently: the handoff turn
costs the declined Opus attempt plus the Sonnet one, and the sticky turn
the Sonnet attempt only.

Signed-off-by: Arnaud Héritier <aheritier@apache.org>
This commit is contained in:
Arnaud Héritier authored and Arnaud Héritier committed 2026-10-26 19:07:28 +01:00
1 parent ae3548f817
commit 4b00a02896
3 files changed
+27 -12

No files matched your search

@@ -28,12 +28,13 @@ import (
// registry, against server, with a 40k-token context window.
func claude5Runtime(t *testing.T, server *anthropictest.Server, providerOpts map[string]any, store session.Store) *LocalRuntime {
t.Helper()
return anthropicRuntime(t, server, "claude-sonnet-5-5", providerOpts, store)
return anthropicRuntime(t, server, "claude-sonnet-5-5", providerOpts, store,
mockModelStoreWithCostAndLimit{limit: 40_000, cost: modelsdev.Cost{Input: 1, Output: 2}})
}
// anthropicRuntime runs a real Anthropic client for model, built by the
// provider registry, against server, with a 40k-token context window.
func anthropicRuntime(t *testing.T, server *anthropictest.Server, modelName string, providerOpts map[string]any, store session.Store, agentOpts ...agent.Opt) *LocalRuntime {
// provider registry, against server, priced and sized by models.
func anthropicRuntime(t *testing.T, server *anthropictest.Server, modelName string, providerOpts map[string]any, store session.Store, models ModelStore, agentOpts ...agent.Opt) *LocalRuntime {
t.Helper()
target, err := url.Parse(server.URL())
require.NoError(t, err)
@@ -52,7 +53,7 @@ func anthropicRuntime(t *testing.T, server *anthropictest.Server, modelName stri
root := agent.New("root", "You are terse.", append([]agent.Opt{agent.WithModel(model)}, agentOpts...)...)
rt, err := NewLocalRuntime(t.Context(), team.New(team.WithAgents(root)),
WithSessionStore(store),
WithModelStore(mockModelStoreWithCostAndLimit{limit: 40_000, cost: modelsdev.Cost{Input: 1, Output: 2}}))
WithModelStore(models))
require.NoError(t, err)
t.Cleanup(func() { assert.NoError(t, rt.Close()) })
return rt
+21 -7
View File
@@ -13,6 +13,7 @@ import (
"github.com/docker/docker-agent/pkg/agent"
"github.com/docker/docker-agent/pkg/chat"
"github.com/docker/docker-agent/pkg/model/provider/anthropic/anthropictest"
"github.com/docker/docker-agent/pkg/modelsdev"
"github.com/docker/docker-agent/pkg/session"
"github.com/docker/docker-agent/pkg/tools"
)
@@ -67,7 +68,11 @@ func TestServerFallbackLifecycle(t *testing.T) {
}})
opts := map[string]any{"fallbacks": []any{fallback}}
dbPath := filepath.Join(t.TempDir(), "sessions.db")
rt := anthropicRuntime(t, server, primary, opts, openSessionStore(t, dbPath), record)
models := catalogStore{costs: map[string]modelsdev.Cost{
"anthropic/" + primary: {Input: 10, Output: 50},
"anthropic/" + fallback: {Input: 3, Output: 15},
}}
rt := anthropicRuntime(t, server, primary, opts, openSessionStore(t, dbPath), models, record)
sess := session.New(session.WithUserMessage("u1"), session.WithToolsApproved(true))
requireNoErrorEvents(t, drainRunStream(t, rt, sess))
@@ -78,16 +83,24 @@ func TestServerFallbackLifecycle(t *testing.T) {
loaded, err := openSessionStore(t, dbPath).GetSession(t.Context(), sess.ID)
require.NoError(t, err)
var models []string
var served []string
var costs []float64
for _, item := range loaded.Messages {
if item.Message != nil && item.Message.Message.Role == chat.MessageRoleAssistant {
models = append(models, item.Message.Message.Model)
served = append(served, item.Message.Message.Model)
costs = append(costs, item.Message.Message.Cost)
}
}
assert.Equal(t, []string{"anthropic/" + fallback, "anthropic/" + fallback}, models, "the marker and sticky routing both name the serving model")
assert.Equal(t, []string{"anthropic/" + fallback, "anthropic/" + fallback}, served, "the marker and sticky routing both name the serving model")
// Each scripted attempt reports 10 input and 7 output tokens. The handoff
// turn bills the declined Opus attempt and the Sonnet one; the sticky turn
// bills Sonnet only.
require.Len(t, costs, 2)
assert.InDelta(t, (10*10+7*50+10*3+7*15)/1e6, costs[0], 1e-12)
assert.InDelta(t, (10*3+7*15)/1e6, costs[1], 1e-12)
loaded.AddMessage(session.UserMessage("u3"))
reloaded := anthropicRuntime(t, server, primary, opts, openSessionStore(t, dbPath), record)
reloaded := anthropicRuntime(t, server, primary, opts, openSessionStore(t, dbPath), models, record)
requireNoErrorEvents(t, drainRunStream(t, reloaded, loaded))
requests := server.Requests()
@@ -127,7 +140,8 @@ func TestSanitizedTurnReplayAfterReload(t *testing.T) {
opts["fallbacks"] = []any{"claude-sonnet-4-6"}
}
dbPath := filepath.Join(t.TempDir(), "sessions.db")
rt := anthropicRuntime(t, server, "claude-opus-4-8", opts, openSessionStore(t, dbPath), record)
models := mockModelStoreWithCostAndLimit{limit: 40_000, cost: modelsdev.Cost{Input: 1, Output: 2}}
rt := anthropicRuntime(t, server, "claude-opus-4-8", opts, openSessionStore(t, dbPath), models, record)
sess := session.New(session.WithUserMessage("u1"), session.WithToolsApproved(true))
requireNoErrorEvents(t, drainRunStream(t, rt, sess))
require.EqualValues(t, 1, calls.Load())
@@ -152,7 +166,7 @@ func TestSanitizedTurnReplayAfterReload(t *testing.T) {
require.NoError(t, err)
loaded.AddMessage(session.UserMessage("u2"))
reloaded := anthropicRuntime(t, server, "claude-opus-4-8", opts, openSessionStore(t, dbPath), record)
reloaded := anthropicRuntime(t, server, "claude-opus-4-8", opts, openSessionStore(t, dbPath), models, record)
requireNoErrorEvents(t, drainRunStream(t, reloaded, loaded))
storedAfter, err := json.Marshal(stored)
require.NoError(t, err)
@@ -35,7 +35,7 @@ func TestClaude5EarlyStickyAttribution(t *testing.T) {
server := anthropictest.NewServer(t, replies...)
dbPath := filepath.Join(t.TempDir(), "sessions.db")
store := openSessionStore(t, dbPath)
rt := anthropicRuntime(t, server, tc.model, tc.opts, store)
rt := anthropicRuntime(t, server, tc.model, tc.opts, store, mockModelStoreWithCostAndLimit{limit: 40_000})
sess := session.New(session.WithUserMessage("q"), session.WithTitle("Sticky attribution"))
if tc.history {
requireNoErrorEvents(t, drainRunStream(t, rt, sess))