From 52dd1425636aabb36b5345923499f50bcdea6e1f Mon Sep 17 00:00:00 2001 From: Susana Ferreira Date: Mon, 10 Aug 2026 16:33:48 +0100 Subject: [PATCH] fix: only write AI model prices that changed (#27923) Previously, the AI Gateway price seeder rewrote every row of `ai_model_prices` on each server start, because `ON CONFLICT` fires on a key conflict rather than on a value difference. `updated_at` therefore recorded when the server last restarted rather than when a price last changed. Guard the `DO UPDATE` branch so a conflicting row is only rewritten when one of its four prices differs. The comparison uses `IS DISTINCT FROM` rather than `<>` because the price columns are nullable, and `<>` yields NULL when either side is NULL, which would skip the update and leave a stale price in place. Related to https://linear.app/codercom/issue/AIGOV-567/experimental-cli-command-to-set-prices-for-unpriced-ai-models > [!NOTE] > Initially generated by Claude Opus 5, modified and reviewed by @ssncferreira (cherry picked from commit 5efa7abe7d8f0cbf0c853a234c972357fce1a752) --- coderd/aibridge/prices/prices_test.go | 107 +++++++++++++++++++++- coderd/database/querier.go | 3 + coderd/database/queries.sql.go | 14 +++ coderd/database/queries/aicostcontrol.sql | 16 +++- 4 files changed, 137 insertions(+), 3 deletions(-) diff --git a/coderd/aibridge/prices/prices_test.go b/coderd/aibridge/prices/prices_test.go index 1ce642e2084..ae30b7e449e 100644 --- a/coderd/aibridge/prices/prices_test.go +++ b/coderd/aibridge/prices/prices_test.go @@ -1,6 +1,7 @@ package prices_test import ( + "fmt" "testing" "github.com/prometheus/client_golang/prometheus" @@ -58,6 +59,7 @@ func TestSeedFromBytes(t *testing.T) { require.Equal(t, int64(25_000_000), opus.OutputPrice.Int64) require.Equal(t, int64(500_000), opus.CacheReadPrice.Int64) require.Equal(t, int64(6_250_000), opus.CacheWritePrice.Int64) + require.Equal(t, opus.CreatedAt, opus.UpdatedAt) // Spot-check a row where the seed has a NULL price (OpenAI does not // publish a cache_write_price). The column should land as SQL NULL. @@ -90,11 +92,11 @@ func TestSeedFromBytes(t *testing.T) { }) require.NoError(t, err) - // Prices must be identical across runs and CreatedAt must be - // preserved (only updated_at moves on a no-op upsert). + // A re-seed that changes nothing must not touch the row at all. require.Equal(t, first.InputPrice, second.InputPrice) require.Equal(t, first.OutputPrice, second.OutputPrice) require.Equal(t, first.CreatedAt, second.CreatedAt) + require.Equal(t, first.UpdatedAt, second.UpdatedAt) }) t.Run("OverwritesExistingPrices", func(t *testing.T) { @@ -114,6 +116,10 @@ func TestSeedFromBytes(t *testing.T) { "cache_read_price": 3, "cache_write_price": 4 }]`))) + before, err := db.GetAIModelPriceByProviderModel(ctx, database.GetAIModelPriceByProviderModelParams{ + Provider: "openai", Model: "gpt-4o", + }) + require.NoError(t, err) require.NoError(t, prices.SeedFromBytes(ctx, db, []byte(testSeedJSON))) @@ -126,6 +132,8 @@ func TestSeedFromBytes(t *testing.T) { require.Equal(t, int64(1_250_000), got.CacheReadPrice.Int64) require.False(t, got.CacheWritePrice.Valid) require.Zero(t, got.CacheWritePrice.Int64) + require.Equal(t, before.CreatedAt, got.CreatedAt) + require.True(t, got.UpdatedAt.After(before.UpdatedAt)) }) t.Run("LeavesOrphanRowsUntouched", func(t *testing.T) { @@ -174,6 +182,101 @@ func TestSeedFromBytes(t *testing.T) { require.True(t, got.InputPrice.Valid) require.Equal(t, int64(2_500_000), got.InputPrice.Int64) }) + + // Every price column counts toward the comparison, and a NULL on either + // side counts as a difference. + t.Run("UpdatedAtTracksPriceChanges", func(t *testing.T) { + t.Parallel() + + key := database.GetAIModelPriceByProviderModelParams{Provider: "openai", Model: "gpt-4o"} + seed := func(priceFields string) []byte { + return fmt.Appendf(nil, `[{"provider": %q, "model": %q, %s}]`, key.Provider, key.Model, priceFields) + } + + tests := []struct { + name string + initial, updated string + }{ + { + name: "InputPriceChanged", + initial: `"input_price": 100, "output_price": 200, "cache_read_price": 300, "cache_write_price": 400`, + updated: `"input_price": 111, "output_price": 200, "cache_read_price": 300, "cache_write_price": 400`, + }, + { + name: "InputPriceSetFromNull", + initial: `"input_price": null, "output_price": 200, "cache_read_price": 300, "cache_write_price": 400`, + updated: `"input_price": 100, "output_price": 200, "cache_read_price": 300, "cache_write_price": 400`, + }, + { + name: "InputPriceClearedToNull", + initial: `"input_price": 100, "output_price": 200, "cache_read_price": 300, "cache_write_price": 400`, + updated: `"input_price": null, "output_price": 200, "cache_read_price": 300, "cache_write_price": 400`, + }, + { + name: "OutputPriceChanged", + initial: `"input_price": 100, "output_price": 200, "cache_read_price": 300, "cache_write_price": 400`, + updated: `"input_price": 100, "output_price": 222, "cache_read_price": 300, "cache_write_price": 400`, + }, + { + name: "OutputPriceSetFromNull", + initial: `"input_price": 100, "output_price": null, "cache_read_price": 300, "cache_write_price": 400`, + updated: `"input_price": 100, "output_price": 200, "cache_read_price": 300, "cache_write_price": 400`, + }, + { + name: "OutputPriceClearedToNull", + initial: `"input_price": 100, "output_price": 200, "cache_read_price": 300, "cache_write_price": 400`, + updated: `"input_price": 100, "output_price": null, "cache_read_price": 300, "cache_write_price": 400`, + }, + { + name: "CacheReadPriceChanged", + initial: `"input_price": 100, "output_price": 200, "cache_read_price": 300, "cache_write_price": 400`, + updated: `"input_price": 100, "output_price": 200, "cache_read_price": 333, "cache_write_price": 400`, + }, + { + name: "CacheReadPriceSetFromNull", + initial: `"input_price": 100, "output_price": 200, "cache_read_price": null, "cache_write_price": 400`, + updated: `"input_price": 100, "output_price": 200, "cache_read_price": 300, "cache_write_price": 400`, + }, + { + name: "CacheReadPriceClearedToNull", + initial: `"input_price": 100, "output_price": 200, "cache_read_price": 300, "cache_write_price": 400`, + updated: `"input_price": 100, "output_price": 200, "cache_read_price": null, "cache_write_price": 400`, + }, + { + name: "CacheWritePriceChanged", + initial: `"input_price": 100, "output_price": 200, "cache_read_price": 300, "cache_write_price": 400`, + updated: `"input_price": 100, "output_price": 200, "cache_read_price": 300, "cache_write_price": 444`, + }, + { + name: "CacheWritePriceSetFromNull", + initial: `"input_price": 100, "output_price": 200, "cache_read_price": 300, "cache_write_price": null`, + updated: `"input_price": 100, "output_price": 200, "cache_read_price": 300, "cache_write_price": 400`, + }, + { + name: "CacheWritePriceClearedToNull", + initial: `"input_price": 100, "output_price": 200, "cache_read_price": 300, "cache_write_price": 400`, + updated: `"input_price": 100, "output_price": 200, "cache_read_price": 300, "cache_write_price": null`, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + ctx := testutil.Context(t, testutil.WaitShort) + db, _ := dbtestutil.NewDB(t) + + require.NoError(t, prices.SeedFromBytes(ctx, db, seed(tt.initial))) + before, err := db.GetAIModelPriceByProviderModel(ctx, key) + require.NoError(t, err) + + require.NoError(t, prices.SeedFromBytes(ctx, db, seed(tt.updated))) + after, err := db.GetAIModelPriceByProviderModel(ctx, key) + require.NoError(t, err) + + require.True(t, after.UpdatedAt.After(before.UpdatedAt), "updated_at should advance when a price changes") + }) + } + }) } // TestSeed exercises the real embedded prices.json so we catch a corrupted, diff --git a/coderd/database/querier.go b/coderd/database/querier.go index bcb6b722fd2..527168ba4a4 100644 --- a/coderd/database/querier.go +++ b/coderd/database/querier.go @@ -1635,6 +1635,9 @@ type sqlcQuerier interface { // Upsert a batch of (provider, model) rows from a JSON array. Each element // must have provider, model, and the four price fields; null prices are // written as SQL NULL. + // A conflicting row is only rewritten when a price differs, so updated_at + // records when a price last changed. Prices are nullable and a NULL on + // either side counts as a difference. UpsertAIModelPrices(ctx context.Context, seed json.RawMessage) error // Returns true if a new rows was inserted, false otherwise. UpsertAISeatState(ctx context.Context, arg UpsertAISeatStateParams) (bool, error) diff --git a/coderd/database/queries.sql.go b/coderd/database/queries.sql.go index 72f3812d4a3..07ccd2b813a 100644 --- a/coderd/database/queries.sql.go +++ b/coderd/database/queries.sql.go @@ -3406,11 +3406,25 @@ ON CONFLICT (provider, model) DO UPDATE SET cache_read_price = EXCLUDED.cache_read_price, cache_write_price = EXCLUDED.cache_write_price, updated_at = NOW() +WHERE ( + ai_model_prices.input_price, + ai_model_prices.output_price, + ai_model_prices.cache_read_price, + ai_model_prices.cache_write_price +) IS DISTINCT FROM ( + EXCLUDED.input_price, + EXCLUDED.output_price, + EXCLUDED.cache_read_price, + EXCLUDED.cache_write_price +) ` // Upsert a batch of (provider, model) rows from a JSON array. Each element // must have provider, model, and the four price fields; null prices are // written as SQL NULL. +// A conflicting row is only rewritten when a price differs, so updated_at +// records when a price last changed. Prices are nullable and a NULL on +// either side counts as a difference. func (q *sqlQuerier) UpsertAIModelPrices(ctx context.Context, seed json.RawMessage) error { _, err := q.db.ExecContext(ctx, upsertAIModelPrices, seed) return err diff --git a/coderd/database/queries/aicostcontrol.sql b/coderd/database/queries/aicostcontrol.sql index 43c5970fe2d..5552e8c7af4 100644 --- a/coderd/database/queries/aicostcontrol.sql +++ b/coderd/database/queries/aicostcontrol.sql @@ -2,6 +2,9 @@ -- Upsert a batch of (provider, model) rows from a JSON array. Each element -- must have provider, model, and the four price fields; null prices are -- written as SQL NULL. +-- A conflicting row is only rewritten when a price differs, so updated_at +-- records when a price last changed. Prices are nullable and a NULL on +-- either side counts as a difference. INSERT INTO ai_model_prices ( provider, model, input_price, output_price, cache_read_price, cache_write_price ) @@ -18,7 +21,18 @@ ON CONFLICT (provider, model) DO UPDATE SET output_price = EXCLUDED.output_price, cache_read_price = EXCLUDED.cache_read_price, cache_write_price = EXCLUDED.cache_write_price, - updated_at = NOW(); + updated_at = NOW() +WHERE ( + ai_model_prices.input_price, + ai_model_prices.output_price, + ai_model_prices.cache_read_price, + ai_model_prices.cache_write_price +) IS DISTINCT FROM ( + EXCLUDED.input_price, + EXCLUDED.output_price, + EXCLUDED.cache_read_price, + EXCLUDED.cache_write_price +); -- name: GetAIModelPriceByProviderModel :one SELECT *