From 9ac5eaa765fe6e62c326b35d60eec2d521188e03 Mon Sep 17 00:00:00 2001 From: Susana Cardoso Ferreira Date: Wed, 19 Aug 2026 20:31:16 +0000 Subject: [PATCH] feat: support custom prices for models in the price book --- coderd/aibridge/prices/prices.go | 30 -------- coderd/aibridge/prices/prices_test.go | 53 ------------- docs/ai-coder/ai-gateway/cost-controls.md | 30 +++++--- enterprise/cli/exp_aimodelprices.go | 4 +- enterprise/cli/exp_aimodelprices_test.go | 28 +++++-- enterprise/coderd/aimodelprices.go | 17 +---- .../coderd/aimodelprices_internal_test.go | 8 +- enterprise/coderd/aimodelprices_test.go | 75 +++++++++++++++++++ 8 files changed, 121 insertions(+), 124 deletions(-) diff --git a/coderd/aibridge/prices/prices.go b/coderd/aibridge/prices/prices.go index eff2551ae04..0eed74a0926 100644 --- a/coderd/aibridge/prices/prices.go +++ b/coderd/aibridge/prices/prices.go @@ -6,7 +6,6 @@ import ( "context" _ "embed" "encoding/json" - "sync" "golang.org/x/xerrors" @@ -64,32 +63,3 @@ func parseSeed(data []byte) ([]seedRow, error) { } return rows, nil } - -// modelKey identifies a priced model. -type modelKey struct { - provider string - model string -} - -// defaultPricedModels indexes the embedded price book by provider and model. -// Built on first use, since a deployment that never sets a price never needs -// it. -var defaultPricedModels = sync.OnceValue(func() map[modelKey]struct{} { - rows, err := parseSeed(seedJSON) - if err != nil { - panic(xerrors.Errorf("parse embedded price seed: %w", err)) - } - index := make(map[modelKey]struct{}, len(rows)) - for _, row := range rows { - index[modelKey{provider: row.Provider, model: row.Model}] = struct{}{} - } - return index -}) - -// IsDefaultPriced reports whether the embedded price book already carries a -// price for the model. Coder owns those prices and re-applies them on every -// startup, so an operator price set for one would not survive a restart. -func IsDefaultPriced(provider, model string) bool { - _, ok := defaultPricedModels()[modelKey{provider: provider, model: model}] - return ok -} diff --git a/coderd/aibridge/prices/prices_test.go b/coderd/aibridge/prices/prices_test.go index d4d14cd5e6d..f97ee2e081d 100644 --- a/coderd/aibridge/prices/prices_test.go +++ b/coderd/aibridge/prices/prices_test.go @@ -355,56 +355,3 @@ func TestSeed(t *testing.T) { db, _ := dbtestutil.NewDB(t) require.NoError(t, prices.Seed(ctx, db)) } - -// TestIsDefaultPriced reads the real embedded price book, so it uses a model -// the generator injects rather than one that could drift out of upstream. -func TestIsDefaultPriced(t *testing.T) { - t.Parallel() - - tests := []struct { - name string - provider string - model string - want bool - }{ - { - name: "ModelInThePriceBook", - provider: "anthropic", - model: "claude-opus-5", - want: true, - }, - { - name: "ModelNotInThePriceBook", - provider: "anthropic", - model: "not-a-real-model", - want: false, - }, - { - // The book is keyed on both columns, so the same model under - // another provider is a different entry. - name: "SameModelUnderAnotherProvider", - provider: "openai", - model: "claude-opus-5", - want: false, - }, - { - name: "UnknownProvider", - provider: "unknown-provider", - model: "claude-opus-5", - want: false, - }, - { - name: "Empty", - provider: "", - model: "", - want: false, - }, - } - - for _, tt := range tests { - t.Run(tt.name, func(t *testing.T) { - t.Parallel() - require.Equal(t, tt.want, prices.IsDefaultPriced(tt.provider, tt.model)) - }) - } -} diff --git a/docs/ai-coder/ai-gateway/cost-controls.md b/docs/ai-coder/ai-gateway/cost-controls.md index 4af6414ca38..6eeddfc1b54 100644 --- a/docs/ai-coder/ai-gateway/cost-controls.md +++ b/docs/ai-coder/ai-gateway/cost-controls.md @@ -198,14 +198,21 @@ https://github.com/coder/coder/blob/release//coderd/aibridge/prices/dat Replace `` with your Coder minor version, for example `2.36`. +To use your own price for any of these models, see +[Set model prices](#set-model-prices). + > [!IMPORTANT] -> Approximate spend can differ from provider-reported amounts, and some usage might not count toward spend: +> Spend is an approximation. It can differ from what the provider bills, and +> some usage does not count toward it at all: > -> - Approximate spend excludes negotiated discounts, committed-use pricing, and -> provider-specific billing rules. -> - Requests to models that are missing from the price table record token usage -> but add nothing to a user's spend. A user who only calls unpriced models is -> effectively unlimited. +> - Prices default to the price book's list prices. A custom price brings spend +> closer to the rates a deployment actually pays, though billing rules that +> are not a per-token rate, such as committed-use discounts, cannot be +> represented. +> - A model with no price adds nothing to spend. Its token usage is still +> recorded, but it never counts toward a limit, so a user who calls only +> unpriced models is effectively unlimited. Setting a price for the model +> closes the gap. Monitor `coder_ai_gateway_cost_control_unpriced_token_usage_records_total`, labeled by `provider`, `provider_type`, and `model`, to detect unpriced usage. @@ -216,10 +223,10 @@ price for it yourself. ### Set model prices -Use the experimental `coder exp ai-model-prices` command to set prices for -models the price book does not cover. It requires AI Governance, which is -included with a Premium license, and the `ai_model_price:update` permission. -Run `coder exp ai-model-prices --help` for the full reference. +Use the experimental `coder exp ai-model-prices` command to set model prices +for your deployment. It requires AI Governance, which is included with a +Premium license, and the `ai_model_price:update` permission. Run +`coder exp ai-model-prices --help` for the full reference. List the prices this deployment holds, optionally narrowed to one provider or model: @@ -249,7 +256,8 @@ coder exp ai-model-prices update prices.json > > - Prices are not retroactive. Usage recorded before you set a price stays > unpriced, so past spend does not change. -> - You can only set prices for models the price book does not cover. +> - A price you set takes precedence over the price book and stays in effect +> across upgrades, so it does not pick up price book updates. > - This command is experimental and can change without notice. ## Monitor spend diff --git a/enterprise/cli/exp_aimodelprices.go b/enterprise/cli/exp_aimodelprices.go index c3b39734e87..4cc07690c98 100644 --- a/enterprise/cli/exp_aimodelprices.go +++ b/enterprise/cli/exp_aimodelprices.go @@ -32,8 +32,8 @@ func (r *RootCmd) aiModelPricesCommand() *serpent.Command { } } -const modelPricesUpdateDescriptionLong = `Sets prices for models that Coder's price book does not cover. Models the -price book covers cannot be changed. +const modelPricesUpdateDescriptionLong = `Sets model prices for this deployment. A price set here takes effect over +Coder's price book. The JSON document is an array of model prices, in the same shape as Coder's price book: diff --git a/enterprise/cli/exp_aimodelprices_test.go b/enterprise/cli/exp_aimodelprices_test.go index 258c34eaf57..85491e516dc 100644 --- a/enterprise/cli/exp_aimodelprices_test.go +++ b/enterprise/cli/exp_aimodelprices_test.go @@ -293,7 +293,7 @@ func TestAIModelPricesUpdate(t *testing.T) { require.Contains(t, apply(), "No changes to apply.") }) - t.Run("RejectsAModelInThePriceBook", func(t *testing.T) { + t.Run("OverridesAModelInThePriceBook", func(t *testing.T) { t.Parallel() // Given: a model Coder already prices. @@ -301,17 +301,29 @@ func TestAIModelPricesUpdate(t *testing.T) { inv, conf := newCLI(t, "exp", "ai-model-prices", "update", "--provider", "anthropic", "--model", "claude-opus-5", - "--input-price", "100", "--output-price", "null", - "--cache-read-price", "null", "--cache-write-price", "null", + "--input-price", "100", "--output-price", "200", + "--cache-read-price", "300", "--cache-write-price", "null", "--yes", ) clitest.SetupConfig(t, client, conf) //nolint:gocritic // requires owner - inv.Stdout = &bytes.Buffer{} - // When: it is priced. Then: the server rejects it. - err := inv.Run() - require.Error(t, err) - require.Contains(t, err.Error(), "price book") + var stdout bytes.Buffer + inv.Stdout = &stdout + + // When: it is priced through the CLI. + require.NoError(t, inv.Run()) + require.Contains(t, stdout.String(), "Updated prices for 1 model(s).") + + // Then: the override is what the deployment reports. + ctx := testutil.Context(t, testutil.WaitLong) + prices, err := codersdk.NewExperimentalClient(client).ListAIModelPrices(ctx, + codersdk.AIModelPricesFilter{Provider: "anthropic", Model: "claude-opus-5"}) + require.NoError(t, err) + require.Len(t, prices, 1) + require.Equal(t, int64(100), *prices[0].InputPrice) + require.Equal(t, int64(200), *prices[0].OutputPrice) + require.Equal(t, int64(300), *prices[0].CacheReadPrice) + require.Nil(t, prices[0].CacheWritePrice) }) } diff --git a/enterprise/coderd/aimodelprices.go b/enterprise/coderd/aimodelprices.go index 0499502f0a7..d3f98484a49 100644 --- a/enterprise/coderd/aimodelprices.go +++ b/enterprise/coderd/aimodelprices.go @@ -10,7 +10,6 @@ import ( "strings" "cdr.dev/slog/v3" - "github.com/coder/coder/v2/coderd/aibridge/prices" "github.com/coder/coder/v2/coderd/aibridge/prices/providers" "github.com/coder/coder/v2/coderd/database" "github.com/coder/coder/v2/coderd/database/db2sdk" @@ -159,9 +158,8 @@ type modelKey struct { } // validateAIModelPrices reports every problem with the requested prices: a -// supported provider, a model Coder's price book does not already cover, all -// four price keys, non-negative prices with at least one set, and no repeated -// model. +// supported provider, a model, all four price keys, non-negative prices with +// at least one set, and no repeated model. func validateAIModelPrices(requested []codersdk.AIModelPriceUpsert, raw []map[string]json.RawMessage) []codersdk.ValidationError { if len(requested) == 0 { return []codersdk.ValidationError{{ @@ -196,17 +194,6 @@ func validateAIModelPrices(requested []codersdk.AIModelPriceUpsert, raw []map[st Detail: "Model is required.", }) } - // The price book is re-applied on every server start, so a price set for - // a model it covers would not survive a restart. - // TODO(ssncferreira): drop this once custom pricing is supported - // (AIGOV-589). - if prices.IsDefaultPriced(price.Provider, price.Model) { - validations = append(validations, codersdk.ValidationError{ - Field: field, - Detail: fmt.Sprintf("%s/%s is priced by Coder's default price book. Overriding a default price is not supported.", price.Provider, price.Model), - }) - } - named := []struct { name string value *int64 diff --git a/enterprise/coderd/aimodelprices_internal_test.go b/enterprise/coderd/aimodelprices_internal_test.go index 75d09e5fb09..f90210ddb98 100644 --- a/enterprise/coderd/aimodelprices_internal_test.go +++ b/enterprise/coderd/aimodelprices_internal_test.go @@ -75,13 +75,11 @@ func TestValidateAIModelPrices(t *testing.T) { }, }, { - // The price book is re-applied on every restart, so this price - // would not survive one. + // A model the price book covers is accepted, and stored as a + // custom price alongside the book's. name: "ModelInPriceBook", body: `{"prices":[{"provider":"anthropic","model":"claude-opus-5",` + allPrices + `}]}`, - want: []codersdk.ValidationError{ - {Field: "prices[0]", Detail: "anthropic/claude-opus-5 is priced by Coder's default price book. Overriding a default price is not supported."}, - }, + want: nil, }, { name: "NegativePrice", diff --git a/enterprise/coderd/aimodelprices_test.go b/enterprise/coderd/aimodelprices_test.go index 0902d2a0601..94fa4561a97 100644 --- a/enterprise/coderd/aimodelprices_test.go +++ b/enterprise/coderd/aimodelprices_test.go @@ -259,6 +259,81 @@ func TestUpsertAIModelPrices(t *testing.T) { require.Equal(t, int64(200), *prices[0].InputPrice) }) + // anthropic/claude-opus-5 is covered by the seeded price book, so these + // write over a price Coder ships. + t.Run("OverridesAPriceBookModel", func(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + price codersdk.AIModelPriceUpsert + want codersdk.AIModelPrice + }{ + { + name: "ReplacesEveryPrice", + price: codersdk.AIModelPriceUpsert{ + Provider: "anthropic", Model: "claude-opus-5", + InputPrice: ptr.Ref(int64(100)), + OutputPrice: ptr.Ref(int64(200)), + CacheReadPrice: ptr.Ref(int64(300)), + CacheWritePrice: ptr.Ref(int64(400)), + }, + want: codersdk.AIModelPrice{ + InputPrice: ptr.Ref(int64(100)), + OutputPrice: ptr.Ref(int64(200)), + CacheReadPrice: ptr.Ref(int64(300)), + CacheWritePrice: ptr.Ref(int64(400)), + }, + }, + { + name: "StoresNullPricesAsNull", + price: newAIModelPrice("anthropic", "claude-opus-5", 100), + want: codersdk.AIModelPrice{ + InputPrice: ptr.Ref(int64(100)), + OutputPrice: nil, + CacheReadPrice: nil, + CacheWritePrice: nil, + }, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + + // Given: a deployment holding the book's price for the model. + ownerClient, _ := setupAIModelPricesTest(t) + exp := codersdk.NewExperimentalClient(ownerClient) + ctx := testutil.Context(t, testutil.WaitLong) + + //nolint:gocritic // Managing AI model prices is owner-only. + seeded, err := exp.ListAIModelPrices(ctx, codersdk.AIModelPricesFilter{ + Provider: "anthropic", Model: "claude-opus-5", + }) + require.NoError(t, err) + require.Len(t, seeded, 1) + require.NotEqual(t, int64(100), *seeded[0].InputPrice) + + // When: it is priced through the endpoint. + require.NoError(t, exp.UpsertAIModelPrices(ctx, codersdk.UpsertAIModelPricesRequest{ + Prices: []codersdk.AIModelPriceUpsert{tt.price}, + })) + + // Then: the request is the price in effect, and nothing carries + // over from the book. + prices, err := exp.ListAIModelPrices(ctx, codersdk.AIModelPricesFilter{ + Provider: "anthropic", Model: "claude-opus-5", + }) + require.NoError(t, err) + require.Len(t, prices, 1) + require.Equal(t, tt.want.InputPrice, prices[0].InputPrice) + require.Equal(t, tt.want.OutputPrice, prices[0].OutputPrice) + require.Equal(t, tt.want.CacheReadPrice, prices[0].CacheReadPrice) + require.Equal(t, tt.want.CacheWritePrice, prices[0].CacheWritePrice) + }) + } + }) + t.Run("StoresAModelNameWithASeparator", func(t *testing.T) { t.Parallel()