diff --git a/pkg/runtime/claude5_compaction_lifecycle_test.go b/pkg/runtime/claude5_compaction_lifecycle_test.go index 0b65e3adf..3a750e2fc 100644 --- a/pkg/runtime/claude5_compaction_lifecycle_test.go +++ b/pkg/runtime/claude5_compaction_lifecycle_test.go @@ -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 diff --git a/pkg/runtime/claude5_server_fallback_test.go b/pkg/runtime/claude5_server_fallback_test.go index 2c33c23b4..69c5fdacc 100644 --- a/pkg/runtime/claude5_server_fallback_test.go +++ b/pkg/runtime/claude5_server_fallback_test.go @@ -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) diff --git a/pkg/runtime/claude5_sticky_attribution_test.go b/pkg/runtime/claude5_sticky_attribution_test.go index 8bfc5c3bc..d90461625 100644 --- a/pkg/runtime/claude5_sticky_attribution_test.go +++ b/pkg/runtime/claude5_sticky_attribution_test.go @@ -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))