From af6ad094ea527d6d91316b1d5985ef11cf7548f1 Mon Sep 17 00:00:00 2001
From: Susana Cardoso Ferreira <susana@coder.com>
Date: Mon, 17 Aug 2026 09:51:14 +0000
Subject: [PATCH 1/6] feat: add source column to ai_model_prices

---
 coderd/aibridge/prices/prices.go              |   5 +-
 coderd/aibridge/prices/prices_test.go         | 166 ++++++++++++++++--
 .../aibridgedserver/aibridgedserver_test.go   |   4 +-
 coderd/database/dbauthz/dbauthz.go            |   4 +-
 coderd/database/dbauthz/dbauthz_test.go       |   5 +-
 coderd/database/dbmetrics/querymetrics.go     |   5 +-
 coderd/database/dbmock/dbmock.go              |   9 +-
 coderd/database/dump.sql                      |   8 +
 .../000574_ai_model_price_source.down.sql     |   3 +
 .../000574_ai_model_price_source.up.sql       |  14 ++
 coderd/database/models.go                     |  60 +++++++
 coderd/database/querier.go                    |  16 +-
 coderd/database/querier_test.go               |   2 +-
 coderd/database/queries.sql.go                |  70 +++++---
 coderd/database/queries/aicostcontrol.sql     |  49 ++++--
 enterprise/coderd/aimodelprices.go            |   5 +-
 16 files changed, 340 insertions(+), 85 deletions(-)
 create mode 100644 coderd/database/migrations/000574_ai_model_price_source.down.sql
 create mode 100644 coderd/database/migrations/000574_ai_model_price_source.up.sql

diff --git a/coderd/aibridge/prices/prices.go b/coderd/aibridge/prices/prices.go
index 5cac067c554..eff2551ae04 100644
--- a/coderd/aibridge/prices/prices.go
+++ b/coderd/aibridge/prices/prices.go
@@ -51,7 +51,10 @@ func SeedFromBytes(ctx context.Context, db database.Store, data []byte) error {
 	if len(rows) == 0 {
 		return xerrors.New("price seed is empty")
 	}
-	return db.UpsertAIModelPrices(ctx, data)
+	return db.UpsertAIModelPrices(ctx, database.UpsertAIModelPricesParams{
+		Seed:   data,
+		Source: database.AIModelPriceSourceDefault,
+	})
 }
 
 func parseSeed(data []byte) ([]seedRow, error) {
diff --git a/coderd/aibridge/prices/prices_test.go b/coderd/aibridge/prices/prices_test.go
index 4ecdabd2caa..2f50e5fe98e 100644
--- a/coderd/aibridge/prices/prices_test.go
+++ b/coderd/aibridge/prices/prices_test.go
@@ -75,6 +75,96 @@ func TestSeedFromBytes(t *testing.T) {
 		require.Zero(t, gpt.CacheWritePrice.Int64)
 	})
 
+	t.Run("SeededPricesAreDefault", func(t *testing.T) {
+		t.Parallel()
+		ctx := testutil.Context(t, testutil.WaitShort)
+		db, _ := dbtestutil.NewDB(t)
+
+		require.NoError(t, prices.SeedFromBytes(ctx, db, []byte(testSeedJSON)))
+
+		got, err := db.GetAIModelPriceByProviderModel(ctx, database.GetAIModelPriceByProviderModelParams{
+			Provider: "openai", Model: "gpt-4o",
+		})
+		require.NoError(t, err)
+		require.Equal(t, database.AIModelPriceSourceDefault, got.Source)
+		require.Equal(t, int64(2_500_000), got.InputPrice.Int64)
+		require.Equal(t, int64(10_000_000), got.OutputPrice.Int64)
+		require.Equal(t, int64(1_250_000), got.CacheReadPrice.Int64)
+		require.False(t, got.CacheWritePrice.Valid)
+	})
+
+	t.Run("OverwrittenPricesAreCustom", func(t *testing.T) {
+		t.Parallel()
+		ctx := testutil.Context(t, testutil.WaitShort)
+		db, _ := dbtestutil.NewDB(t)
+
+		require.NoError(t, prices.SeedFromBytes(ctx, db, []byte(testSeedJSON)))
+		seeded, err := db.GetAIModelPriceByProviderModel(ctx, database.GetAIModelPriceByProviderModelParams{
+			Provider: "openai", Model: "gpt-4o",
+		})
+		require.NoError(t, err)
+		require.Equal(t, database.AIModelPriceSourceDefault, seeded.Source)
+
+		// Update all prices.
+		require.NoError(t, db.UpsertAIModelPrices(ctx, database.UpsertAIModelPricesParams{
+			Seed: []byte(`[{
+				"provider": "openai",
+				"model": "gpt-4o",
+				"input_price": 5000000,
+				"output_price": 20000000,
+				"cache_read_price": 2000000,
+				"cache_write_price": 1000000
+			}]`),
+			Source: database.AIModelPriceSourceCustom,
+		}))
+
+		got, err := db.GetAIModelPriceByProviderModel(ctx, database.GetAIModelPriceByProviderModelParams{
+			Provider: "openai", Model: "gpt-4o",
+		})
+		require.NoError(t, err)
+		require.Equal(t, database.AIModelPriceSourceCustom, got.Source)
+		require.Equal(t, int64(5_000_000), got.InputPrice.Int64)
+		require.Equal(t, int64(20_000_000), got.OutputPrice.Int64)
+		require.Equal(t, int64(2_000_000), got.CacheReadPrice.Int64)
+		require.Equal(t, int64(1_000_000), got.CacheWritePrice.Int64)
+	})
+
+	t.Run("UnchangedPricesAreCustom", func(t *testing.T) {
+		t.Parallel()
+		ctx := testutil.Context(t, testutil.WaitShort)
+		db, _ := dbtestutil.NewDB(t)
+
+		require.NoError(t, prices.SeedFromBytes(ctx, db, []byte(testSeedJSON)))
+		seeded, err := db.GetAIModelPriceByProviderModel(ctx, database.GetAIModelPriceByProviderModelParams{
+			Provider: "openai", Model: "gpt-4o",
+		})
+		require.NoError(t, err)
+		require.Equal(t, database.AIModelPriceSourceDefault, seeded.Source)
+
+		// Declaring the prices the row already holds still marks it custom.
+		require.NoError(t, db.UpsertAIModelPrices(ctx, database.UpsertAIModelPricesParams{
+			Seed: []byte(`[{
+				"provider": "openai",
+				"model": "gpt-4o",
+				"input_price": 2500000,
+				"output_price": 10000000,
+				"cache_read_price": 1250000,
+				"cache_write_price": null
+			}]`),
+			Source: database.AIModelPriceSourceCustom,
+		}))
+
+		got, err := db.GetAIModelPriceByProviderModel(ctx, database.GetAIModelPriceByProviderModelParams{
+			Provider: "openai", Model: "gpt-4o",
+		})
+		require.NoError(t, err)
+		require.Equal(t, database.AIModelPriceSourceCustom, got.Source)
+		require.Equal(t, int64(2_500_000), got.InputPrice.Int64)
+		require.Equal(t, int64(10_000_000), got.OutputPrice.Int64)
+		require.Equal(t, int64(1_250_000), got.CacheReadPrice.Int64)
+		require.False(t, got.CacheWritePrice.Valid)
+	})
+
 	t.Run("Idempotent", func(t *testing.T) {
 		t.Parallel()
 		ctx := testutil.Context(t, testutil.WaitShort)
@@ -108,14 +198,17 @@ func TestSeedFromBytes(t *testing.T) {
 		// cache_write_price is set to a non-NULL value here even though the
 		// embedded seed leaves it NULL for OpenAI; Seed must replace it with
 		// NULL to keep the table in sync with the seed.
-		require.NoError(t, db.UpsertAIModelPrices(ctx, []byte(`[{
-			"provider": "openai",
-			"model": "gpt-4o",
-			"input_price": 1,
-			"output_price": 2,
-			"cache_read_price": 3,
-			"cache_write_price": 4
-		}]`)))
+		require.NoError(t, db.UpsertAIModelPrices(ctx, database.UpsertAIModelPricesParams{
+			Seed: []byte(`[{
+				"provider": "openai",
+				"model": "gpt-4o",
+				"input_price": 1,
+				"output_price": 2,
+				"cache_read_price": 3,
+				"cache_write_price": 4
+			}]`),
+			Source: database.AIModelPriceSourceDefault,
+		}))
 		before, err := db.GetAIModelPriceByProviderModel(ctx, database.GetAIModelPriceByProviderModelParams{
 			Provider: "openai", Model: "gpt-4o",
 		})
@@ -143,14 +236,17 @@ func TestSeedFromBytes(t *testing.T) {
 
 		// Insert a row for a (provider, model) the seed doesn't cover. After
 		// Seed it should still be there with its values intact.
-		require.NoError(t, db.UpsertAIModelPrices(ctx, []byte(`[{
-			"provider": "test-provider",
-			"model": "test-model-not-in-seed",
-			"input_price": 12345,
-			"output_price": 67890,
-			"cache_read_price": null,
-			"cache_write_price": null
-		}]`)))
+		require.NoError(t, db.UpsertAIModelPrices(ctx, database.UpsertAIModelPricesParams{
+			Seed: []byte(`[{
+				"provider": "test-provider",
+				"model": "test-model-not-in-seed",
+				"input_price": 12345,
+				"output_price": 67890,
+				"cache_read_price": null,
+				"cache_write_price": null
+			}]`),
+			Source: database.AIModelPriceSourceDefault,
+		}))
 
 		require.NoError(t, prices.SeedFromBytes(ctx, db, []byte(testSeedJSON)))
 
@@ -162,6 +258,44 @@ func TestSeedFromBytes(t *testing.T) {
 		require.Equal(t, int64(67890), got.OutputPrice.Int64)
 	})
 
+	t.Run("LeavesCustomPricesUntouched", func(t *testing.T) {
+		t.Parallel()
+		ctx := testutil.Context(t, testutil.WaitShort)
+		db, _ := dbtestutil.NewDB(t)
+
+		// Price a model the seed also covers.
+		require.NoError(t, db.UpsertAIModelPrices(ctx, database.UpsertAIModelPricesParams{
+			Seed: []byte(`[{
+				"provider": "openai",
+				"model": "gpt-4o",
+				"input_price": 1,
+				"output_price": 2,
+				"cache_read_price": 3,
+				"cache_write_price": 4
+			}]`),
+			Source: database.AIModelPriceSourceCustom,
+		}))
+		before, err := db.GetAIModelPriceByProviderModel(ctx, database.GetAIModelPriceByProviderModelParams{
+			Provider: "openai", Model: "gpt-4o",
+		})
+		require.NoError(t, err)
+		require.Equal(t, database.AIModelPriceSourceCustom, before.Source)
+
+		// Re-applying the price book skips custom rows.
+		require.NoError(t, prices.SeedFromBytes(ctx, db, []byte(testSeedJSON)))
+
+		got, err := db.GetAIModelPriceByProviderModel(ctx, database.GetAIModelPriceByProviderModelParams{
+			Provider: "openai", Model: "gpt-4o",
+		})
+		require.NoError(t, err)
+		require.Equal(t, int64(1), got.InputPrice.Int64)
+		require.Equal(t, int64(2), got.OutputPrice.Int64)
+		require.Equal(t, int64(3), got.CacheReadPrice.Int64)
+		require.Equal(t, int64(4), got.CacheWritePrice.Int64)
+		require.Equal(t, database.AIModelPriceSourceCustom, got.Source)
+		require.Equal(t, before.UpdatedAt, got.UpdatedAt)
+	})
+
 	// Verifies the chain: AsAIBridged context -> dbauthz wrapper auth check
 	// -> subjectAibridged's permission grant. A missing or wrong action on
 	// the subject would surface as "unauthorized: rbac: forbidden" here, even
diff --git a/coderd/aibridgedserver/aibridgedserver_test.go b/coderd/aibridgedserver/aibridgedserver_test.go
index 31fcbddd703..0cac462bab1 100644
--- a/coderd/aibridgedserver/aibridgedserver_test.go
+++ b/coderd/aibridgedserver/aibridgedserver_test.go
@@ -2506,7 +2506,7 @@ func TestRecordTokenUsageAuthorized(t *testing.T) {
 		"cache_write_price": 4_000_000,
 	}})
 	require.NoError(t, err)
-	require.NoError(t, rawDB.UpsertAIModelPrices(ctx, priceSeed), "seed model prices")
+	require.NoError(t, rawDB.UpsertAIModelPrices(ctx, database.UpsertAIModelPricesParams{Seed: priceSeed, Source: database.AIModelPriceSourceDefault}), "seed model prices")
 
 	// The interception's provider name resolves to this provider, whose type keys
 	// the price lookup.
@@ -2601,7 +2601,7 @@ func TestRecordTokenUsageProviderResolution(t *testing.T) {
 		{"provider": string(database.AIProviderTypeAzure), "model": gptModel, "input_price": azureInputPrice},
 	})
 	require.NoError(t, err)
-	require.NoError(t, rawDB.UpsertAIModelPrices(setupCtx, priceSeed), "seed model prices")
+	require.NoError(t, rawDB.UpsertAIModelPrices(setupCtx, database.UpsertAIModelPricesParams{Seed: priceSeed, Source: database.AIModelPriceSourceDefault}), "seed model prices")
 
 	srv, err := aibridgedserver.NewServer(setupCtx, aibridgedserver.Options{
 		Store:         authzDB,
diff --git a/coderd/database/dbauthz/dbauthz.go b/coderd/database/dbauthz/dbauthz.go
index 632efc4c8e1..104f188d951 100644
--- a/coderd/database/dbauthz/dbauthz.go
+++ b/coderd/database/dbauthz/dbauthz.go
@@ -8851,11 +8851,11 @@ func (q *querier) UpdateWorkspacesTTLByTemplateID(ctx context.Context, arg datab
 	return q.db.UpdateWorkspacesTTLByTemplateID(ctx, arg)
 }
 
-func (q *querier) UpsertAIModelPrices(ctx context.Context, seed json.RawMessage) error {
+func (q *querier) UpsertAIModelPrices(ctx context.Context, arg database.UpsertAIModelPricesParams) error {
 	if err := q.authorizeContext(ctx, policy.ActionUpdate, rbac.ResourceAiModelPrice); err != nil {
 		return err
 	}
-	return q.db.UpsertAIModelPrices(ctx, seed)
+	return q.db.UpsertAIModelPrices(ctx, arg)
 }
 
 func (q *querier) UpsertAISeatState(ctx context.Context, arg database.UpsertAISeatStateParams) (bool, error) {
diff --git a/coderd/database/dbauthz/dbauthz_test.go b/coderd/database/dbauthz/dbauthz_test.go
index 738b4a2de99..73b2c9654d2 100644
--- a/coderd/database/dbauthz/dbauthz_test.go
+++ b/coderd/database/dbauthz/dbauthz_test.go
@@ -6927,7 +6927,10 @@ func (s *MethodTestSuite) TestAIBridge() {
 
 	s.Run("UpsertAIModelPrices", s.Mocked(func(db *dbmock.MockStore, _ *gofakeit.Faker, check *expects) {
 		db.EXPECT().UpsertAIModelPrices(gomock.Any(), gomock.Any()).Return(nil).AnyTimes()
-		check.Args(json.RawMessage(`[]`)).Asserts(rbac.ResourceAiModelPrice, policy.ActionUpdate)
+		check.Args(database.UpsertAIModelPricesParams{
+			Seed:   json.RawMessage(`[]`),
+			Source: database.AIModelPriceSourceDefault,
+		}).Asserts(rbac.ResourceAiModelPrice, policy.ActionUpdate)
 	}))
 
 	s.Run("GetAIModelPriceByProviderModel", s.Mocked(func(db *dbmock.MockStore, _ *gofakeit.Faker, check *expects) {
diff --git a/coderd/database/dbmetrics/querymetrics.go b/coderd/database/dbmetrics/querymetrics.go
index 492686145ad..ca75449cca4 100644
--- a/coderd/database/dbmetrics/querymetrics.go
+++ b/coderd/database/dbmetrics/querymetrics.go
@@ -5,7 +5,6 @@ package dbmetrics
 
 import (
 	"context"
-	"encoding/json"
 	"slices"
 	"time"
 
@@ -6225,9 +6224,9 @@ func (m queryMetricsStore) UpdateWorkspacesTTLByTemplateID(ctx context.Context,
 	return r0
 }
 
-func (m queryMetricsStore) UpsertAIModelPrices(ctx context.Context, seed json.RawMessage) error {
+func (m queryMetricsStore) UpsertAIModelPrices(ctx context.Context, arg database.UpsertAIModelPricesParams) error {
 	start := time.Now()
-	r0 := m.s.UpsertAIModelPrices(ctx, seed)
+	r0 := m.s.UpsertAIModelPrices(ctx, arg)
 	m.queryLatencies.WithLabelValues("UpsertAIModelPrices").Observe(time.Since(start).Seconds())
 	m.queryCounts.WithLabelValues(httpmw.ExtractHTTPRoute(ctx), httpmw.ExtractHTTPMethod(ctx), "UpsertAIModelPrices").Inc()
 	return r0
diff --git a/coderd/database/dbmock/dbmock.go b/coderd/database/dbmock/dbmock.go
index 0c08caa669e..328c4d0e8cd 100644
--- a/coderd/database/dbmock/dbmock.go
+++ b/coderd/database/dbmock/dbmock.go
@@ -11,7 +11,6 @@ package dbmock
 
 import (
 	context "context"
-	json "encoding/json"
 	reflect "reflect"
 	time "time"
 
@@ -11716,17 +11715,17 @@ func (mr *MockStoreMockRecorder) UpdateWorkspacesTTLByTemplateID(ctx, arg any) *
 }
 
 // UpsertAIModelPrices mocks base method.
-func (m *MockStore) UpsertAIModelPrices(ctx context.Context, seed json.RawMessage) error {
+func (m *MockStore) UpsertAIModelPrices(ctx context.Context, arg database.UpsertAIModelPricesParams) error {
 	m.ctrl.T.Helper()
-	ret := m.ctrl.Call(m, "UpsertAIModelPrices", ctx, seed)
+	ret := m.ctrl.Call(m, "UpsertAIModelPrices", ctx, arg)
 	ret0, _ := ret[0].(error)
 	return ret0
 }
 
 // UpsertAIModelPrices indicates an expected call of UpsertAIModelPrices.
-func (mr *MockStoreMockRecorder) UpsertAIModelPrices(ctx, seed any) *gomock.Call {
+func (mr *MockStoreMockRecorder) UpsertAIModelPrices(ctx, arg any) *gomock.Call {
 	mr.mock.ctrl.T.Helper()
-	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "UpsertAIModelPrices", reflect.TypeOf((*MockStore)(nil).UpsertAIModelPrices), ctx, seed)
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "UpsertAIModelPrices", reflect.TypeOf((*MockStore)(nil).UpsertAIModelPrices), ctx, arg)
 }
 
 // UpsertAISeatState mocks base method.
diff --git a/coderd/database/dump.sql b/coderd/database/dump.sql
index 343491f73f0..1be2c4f37c3 100644
--- a/coderd/database/dump.sql
+++ b/coderd/database/dump.sql
@@ -10,6 +10,11 @@ CREATE TYPE agent_key_scope_enum AS ENUM (
     'no_user_data'
 );
 
+CREATE TYPE ai_model_price_source AS ENUM (
+    'default',
+    'custom'
+);
+
 CREATE TYPE ai_provider_type AS ENUM (
     'openai',
     'anthropic',
@@ -1573,6 +1578,7 @@ CREATE TABLE ai_model_prices (
     cache_write_price bigint,
     created_at timestamp with time zone DEFAULT now() NOT NULL,
     updated_at timestamp with time zone DEFAULT now() NOT NULL,
+    source ai_model_price_source NOT NULL,
     CONSTRAINT ai_model_prices_cache_read_price_check CHECK ((cache_read_price >= 0)),
     CONSTRAINT ai_model_prices_cache_write_price_check CHECK ((cache_write_price >= 0)),
     CONSTRAINT ai_model_prices_input_price_check CHECK ((input_price >= 0)),
@@ -1581,6 +1587,8 @@ CREATE TABLE ai_model_prices (
 
 COMMENT ON TABLE ai_model_prices IS 'Per-model token prices used by AI Bridge to compute interception cost.';
 
+COMMENT ON COLUMN ai_model_prices.source IS 'Where the price came from: default for the embedded price book, custom for a price an operator set. The startup seeder never overwrites a custom row.';
+
 CREATE TABLE ai_provider_keys (
     id uuid DEFAULT gen_random_uuid() NOT NULL,
     provider_id uuid NOT NULL,
diff --git a/coderd/database/migrations/000574_ai_model_price_source.down.sql b/coderd/database/migrations/000574_ai_model_price_source.down.sql
new file mode 100644
index 00000000000..f1e9a9b97c1
--- /dev/null
+++ b/coderd/database/migrations/000574_ai_model_price_source.down.sql
@@ -0,0 +1,3 @@
+ALTER TABLE ai_model_prices DROP COLUMN source;
+
+DROP TYPE ai_model_price_source;
diff --git a/coderd/database/migrations/000574_ai_model_price_source.up.sql b/coderd/database/migrations/000574_ai_model_price_source.up.sql
new file mode 100644
index 00000000000..4689db362f8
--- /dev/null
+++ b/coderd/database/migrations/000574_ai_model_price_source.up.sql
@@ -0,0 +1,14 @@
+-- source records where a price came from, either the price book embedded in
+-- each Coder release or a price set through the API. The startup seeder
+-- re-applies the book on every boot and reads source to leave custom rows
+-- alone.
+
+CREATE TYPE ai_model_price_source AS ENUM ('default', 'custom');
+
+ALTER TABLE ai_model_prices ADD COLUMN source ai_model_price_source;
+
+UPDATE ai_model_prices SET source = 'default' WHERE source IS NULL;
+
+ALTER TABLE ai_model_prices ALTER COLUMN source SET NOT NULL;
+
+COMMENT ON COLUMN ai_model_prices.source IS 'Where the price came from: default for the embedded price book, custom for a price an operator set. The startup seeder never overwrites a custom row.';
diff --git a/coderd/database/models.go b/coderd/database/models.go
index 68130280eb1..e6024549d3f 100644
--- a/coderd/database/models.go
+++ b/coderd/database/models.go
@@ -89,6 +89,64 @@ func AllAIBridgeInterceptionErrorTypeValues() []AIBridgeInterceptionErrorType {
 	}
 }
 
+type AIModelPriceSource string
+
+const (
+	AIModelPriceSourceDefault AIModelPriceSource = "default"
+	AIModelPriceSourceCustom  AIModelPriceSource = "custom"
+)
+
+func (e *AIModelPriceSource) Scan(src interface{}) error {
+	switch s := src.(type) {
+	case []byte:
+		*e = AIModelPriceSource(s)
+	case string:
+		*e = AIModelPriceSource(s)
+	default:
+		return fmt.Errorf("unsupported scan type for AIModelPriceSource: %T", src)
+	}
+	return nil
+}
+
+type NullAIModelPriceSource struct {
+	AIModelPriceSource AIModelPriceSource `json:"ai_model_price_source"`
+	Valid              bool               `json:"valid"` // Valid is true if AIModelPriceSource is not NULL
+}
+
+// Scan implements the Scanner interface.
+func (ns *NullAIModelPriceSource) Scan(value interface{}) error {
+	if value == nil {
+		ns.AIModelPriceSource, ns.Valid = "", false
+		return nil
+	}
+	ns.Valid = true
+	return ns.AIModelPriceSource.Scan(value)
+}
+
+// Value implements the driver Valuer interface.
+func (ns NullAIModelPriceSource) Value() (driver.Value, error) {
+	if !ns.Valid {
+		return nil, nil
+	}
+	return string(ns.AIModelPriceSource), nil
+}
+
+func (e AIModelPriceSource) Valid() bool {
+	switch e {
+	case AIModelPriceSourceDefault,
+		AIModelPriceSourceCustom:
+		return true
+	}
+	return false
+}
+
+func AllAIModelPriceSourceValues() []AIModelPriceSource {
+	return []AIModelPriceSource{
+		AIModelPriceSourceDefault,
+		AIModelPriceSourceCustom,
+	}
+}
+
 type AIProviderType string
 
 const (
@@ -4827,6 +4885,8 @@ type AIModelPrice struct {
 	CacheWritePrice sql.NullInt64 `db:"cache_write_price" json:"cache_write_price"`
 	CreatedAt       time.Time     `db:"created_at" json:"created_at"`
 	UpdatedAt       time.Time     `db:"updated_at" json:"updated_at"`
+	// Where the price came from: default for the embedded price book, custom for a price an operator set. The startup seeder never overwrites a custom row.
+	Source AIModelPriceSource `db:"source" json:"source"`
 }
 
 // Runtime configuration for AI providers. Authoritative source for the provider set served by aibridged. Replaces deployment-time CODER_AIBRIDGE_* environment variables.
diff --git a/coderd/database/querier.go b/coderd/database/querier.go
index ac1945f20f0..85a7204361a 100644
--- a/coderd/database/querier.go
+++ b/coderd/database/querier.go
@@ -6,7 +6,6 @@ package database
 
 import (
 	"context"
-	"encoding/json"
 	"time"
 
 	"github.com/google/uuid"
@@ -1616,13 +1615,14 @@ type sqlcQuerier interface {
 	UpdateWorkspaceTTL(ctx context.Context, arg UpdateWorkspaceTTLParams) error
 	UpdateWorkspacesDormantDeletingAtByTemplateID(ctx context.Context, arg UpdateWorkspacesDormantDeletingAtByTemplateIDParams) ([]WorkspaceTable, error)
 	UpdateWorkspacesTTLByTemplateID(ctx context.Context, arg UpdateWorkspacesTTLByTemplateIDParams) error
-	// 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
+	// Upsert a batch of (provider, model) rows from a JSON array, recording them
+	// under source. Each element must have provider, model, and the four price
+	// fields, and null prices are written as SQL NULL.
+	// A default write skips rows a custom price owns. Otherwise a conflicting row
+	// is only rewritten when a price or the source differs, so updated_at records
+	// when the row last changed. Prices are nullable and a NULL on either side
+	// counts as a difference.
+	UpsertAIModelPrices(ctx context.Context, arg UpsertAIModelPricesParams) error
 	// Returns true if a new rows was inserted, false otherwise.
 	UpsertAISeatState(ctx context.Context, arg UpsertAISeatStateParams) (bool, error)
 	UpsertAnnouncementBanners(ctx context.Context, value string) error
diff --git a/coderd/database/querier_test.go b/coderd/database/querier_test.go
index 2c4df715781..f66e50c2dc4 100644
--- a/coderd/database/querier_test.go
+++ b/coderd/database/querier_test.go
@@ -19112,7 +19112,7 @@ func TestGetAIModelPrices(t *testing.T) {
 			t.Parallel()
 			ctx := testutil.Context(t, testutil.WaitShort)
 			db, _ := dbtestutil.NewDB(t)
-			require.NoError(t, db.UpsertAIModelPrices(ctx, []byte(seed)))
+			require.NoError(t, db.UpsertAIModelPrices(ctx, database.UpsertAIModelPricesParams{Seed: []byte(seed), Source: database.AIModelPriceSourceDefault}))
 
 			prices, err := db.GetAIModelPrices(ctx, tt.params)
 			require.NoError(t, err)
diff --git a/coderd/database/queries.sql.go b/coderd/database/queries.sql.go
index e364a9400ea..cbab0fb7787 100644
--- a/coderd/database/queries.sql.go
+++ b/coderd/database/queries.sql.go
@@ -2862,7 +2862,7 @@ func (q *sqlQuerier) ExportOrganizationAISpend(ctx context.Context, arg ExportOr
 }
 
 const getAIModelPriceByProviderModel = `-- name: GetAIModelPriceByProviderModel :one
-SELECT provider, model, input_price, output_price, cache_read_price, cache_write_price, created_at, updated_at
+SELECT provider, model, input_price, output_price, cache_read_price, cache_write_price, created_at, updated_at, source
 FROM ai_model_prices
 WHERE provider = $1 AND model = $2
 `
@@ -2884,12 +2884,13 @@ func (q *sqlQuerier) GetAIModelPriceByProviderModel(ctx context.Context, arg Get
 		&i.CacheWritePrice,
 		&i.CreatedAt,
 		&i.UpdatedAt,
+		&i.Source,
 	)
 	return i, err
 }
 
 const getAIModelPrices = `-- name: GetAIModelPrices :many
-SELECT provider, model, input_price, output_price, cache_read_price, cache_write_price, created_at, updated_at
+SELECT provider, model, input_price, output_price, cache_read_price, cache_write_price, created_at, updated_at, source
 FROM ai_model_prices
     -- Filter by provider
 WHERE CASE
@@ -2929,6 +2930,7 @@ func (q *sqlQuerier) GetAIModelPrices(ctx context.Context, arg GetAIModelPricesP
 			&i.CacheWritePrice,
 			&i.CreatedAt,
 			&i.UpdatedAt,
+			&i.Source,
 		); err != nil {
 			return nil, err
 		}
@@ -3498,7 +3500,7 @@ func (q *sqlQuerier) IncrementUserAIDailySpend(ctx context.Context, arg Incremen
 
 const upsertAIModelPrices = `-- name: UpsertAIModelPrices :exec
 INSERT INTO ai_model_prices (
-	provider, model, input_price, output_price, cache_read_price, cache_write_price
+	provider, model, input_price, output_price, cache_read_price, cache_write_price, source
 )
 SELECT
 	elem->>'provider',
@@ -3506,35 +3508,51 @@ SELECT
 	(elem->>'input_price')::bigint,
 	(elem->>'output_price')::bigint,
 	(elem->>'cache_read_price')::bigint,
-	(elem->>'cache_write_price')::bigint
-FROM jsonb_array_elements($1::jsonb) AS elem
+	(elem->>'cache_write_price')::bigint,
+	$1::ai_model_price_source
+FROM jsonb_array_elements($2::jsonb) AS elem
 ON CONFLICT (provider, model) DO UPDATE SET
 	input_price       = EXCLUDED.input_price,
 	output_price      = EXCLUDED.output_price,
 	cache_read_price  = EXCLUDED.cache_read_price,
 	cache_write_price = EXCLUDED.cache_write_price,
+	source            = EXCLUDED.source,
 	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)
+WHERE CASE
+		-- A custom price claims any row, including one the price book wrote.
+		WHEN $1::ai_model_price_source = 'custom' THEN true
+		-- The price book leaves custom rows alone.
+		ELSE ai_model_prices.source <> 'custom'
+	END
+	AND (
+		ai_model_prices.input_price,
+		ai_model_prices.output_price,
+		ai_model_prices.cache_read_price,
+		ai_model_prices.cache_write_price,
+		ai_model_prices.source
+	) IS DISTINCT FROM (
+		EXCLUDED.input_price,
+		EXCLUDED.output_price,
+		EXCLUDED.cache_read_price,
+		EXCLUDED.cache_write_price,
+		EXCLUDED.source
+	)
+`
+
+type UpsertAIModelPricesParams struct {
+	Source AIModelPriceSource `db:"source" json:"source"`
+	Seed   json.RawMessage    `db:"seed" json:"seed"`
+}
+
+// Upsert a batch of (provider, model) rows from a JSON array, recording them
+// under source. Each element must have provider, model, and the four price
+// fields, and null prices are written as SQL NULL.
+// A default write skips rows a custom price owns. Otherwise a conflicting row
+// is only rewritten when a price or the source differs, so updated_at records
+// when the row last changed. Prices are nullable and a NULL on either side
+// counts as a difference.
+func (q *sqlQuerier) UpsertAIModelPrices(ctx context.Context, arg UpsertAIModelPricesParams) error {
+	_, err := q.db.ExecContext(ctx, upsertAIModelPrices, arg.Source, arg.Seed)
 	return err
 }
 
diff --git a/coderd/database/queries/aicostcontrol.sql b/coderd/database/queries/aicostcontrol.sql
index def9c0d316e..8b7649e39dc 100644
--- a/coderd/database/queries/aicostcontrol.sql
+++ b/coderd/database/queries/aicostcontrol.sql
@@ -1,12 +1,13 @@
 -- name: UpsertAIModelPrices :exec
--- 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.
+-- Upsert a batch of (provider, model) rows from a JSON array, recording them
+-- under source. Each element must have provider, model, and the four price
+-- fields, and null prices are written as SQL NULL.
+-- A default write skips rows a custom price owns. Otherwise a conflicting row
+-- is only rewritten when a price or the source differs, so updated_at records
+-- when the row 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
+	provider, model, input_price, output_price, cache_read_price, cache_write_price, source
 )
 SELECT
 	elem->>'provider',
@@ -14,25 +15,35 @@ SELECT
 	(elem->>'input_price')::bigint,
 	(elem->>'output_price')::bigint,
 	(elem->>'cache_read_price')::bigint,
-	(elem->>'cache_write_price')::bigint
+	(elem->>'cache_write_price')::bigint,
+	@source::ai_model_price_source
 FROM jsonb_array_elements(@seed::jsonb) AS elem
 ON CONFLICT (provider, model) DO UPDATE SET
 	input_price       = EXCLUDED.input_price,
 	output_price      = EXCLUDED.output_price,
 	cache_read_price  = EXCLUDED.cache_read_price,
 	cache_write_price = EXCLUDED.cache_write_price,
+	source            = EXCLUDED.source,
 	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
-);
+WHERE CASE
+		-- A custom price claims any row, including one the price book wrote.
+		WHEN @source::ai_model_price_source = 'custom' THEN true
+		-- The price book leaves custom rows alone.
+		ELSE ai_model_prices.source <> 'custom'
+	END
+	AND (
+		ai_model_prices.input_price,
+		ai_model_prices.output_price,
+		ai_model_prices.cache_read_price,
+		ai_model_prices.cache_write_price,
+		ai_model_prices.source
+	) IS DISTINCT FROM (
+		EXCLUDED.input_price,
+		EXCLUDED.output_price,
+		EXCLUDED.cache_read_price,
+		EXCLUDED.cache_write_price,
+		EXCLUDED.source
+	);
 
 -- name: GetAIModelPriceByProviderModel :one
 SELECT *
diff --git a/enterprise/coderd/aimodelprices.go b/enterprise/coderd/aimodelprices.go
index 78f2d068469..0499502f0a7 100644
--- a/enterprise/coderd/aimodelprices.go
+++ b/enterprise/coderd/aimodelprices.go
@@ -123,7 +123,10 @@ func (api *API) upsertAIModelPrices(rw http.ResponseWriter, r *http.Request) {
 		httpapi.InternalServerError(rw, err)
 		return
 	}
-	err = api.Database.UpsertAIModelPrices(ctx, seed)
+	err = api.Database.UpsertAIModelPrices(ctx, database.UpsertAIModelPricesParams{
+		Seed:   seed,
+		Source: database.AIModelPriceSourceCustom,
+	})
 	if dbauthz.IsNotAuthorizedError(err) {
 		httpapi.Forbidden(rw)
 		return

From 1c0e445903e2690b316b2827cf84e891833c4bab Mon Sep 17 00:00:00 2001
From: Susana Cardoso Ferreira <susana@coder.com>
Date: Wed, 19 Aug 2026 19:20:21 +0000
Subject: [PATCH 2/6] feat: keep one AI model price row per source

---
 coderd/aibridge/prices/prices_test.go         |  78 +----------
 .../aibridgedserver/aibridgedserver_test.go   | 130 +++++++++++++++++-
 coderd/database/dump.sql                      |   4 +-
 .../000574_ai_model_price_source.down.sql     |   8 ++
 .../000574_ai_model_price_source.up.sql       |  11 +-
 coderd/database/models.go                     |   2 +-
 coderd/database/querier.go                    |  16 ++-
 coderd/database/querier_test.go               | 116 +++++++++++++++-
 coderd/database/queries.sql.go                |  55 ++++----
 coderd/database/queries/aicostcontrol.sql     |  57 ++++----
 coderd/database/unique_constraint.go          |   2 +-
 11 files changed, 324 insertions(+), 155 deletions(-)

diff --git a/coderd/aibridge/prices/prices_test.go b/coderd/aibridge/prices/prices_test.go
index 2f50e5fe98e..d4d14cd5e6d 100644
--- a/coderd/aibridge/prices/prices_test.go
+++ b/coderd/aibridge/prices/prices_test.go
@@ -93,78 +93,6 @@ func TestSeedFromBytes(t *testing.T) {
 		require.False(t, got.CacheWritePrice.Valid)
 	})
 
-	t.Run("OverwrittenPricesAreCustom", func(t *testing.T) {
-		t.Parallel()
-		ctx := testutil.Context(t, testutil.WaitShort)
-		db, _ := dbtestutil.NewDB(t)
-
-		require.NoError(t, prices.SeedFromBytes(ctx, db, []byte(testSeedJSON)))
-		seeded, err := db.GetAIModelPriceByProviderModel(ctx, database.GetAIModelPriceByProviderModelParams{
-			Provider: "openai", Model: "gpt-4o",
-		})
-		require.NoError(t, err)
-		require.Equal(t, database.AIModelPriceSourceDefault, seeded.Source)
-
-		// Update all prices.
-		require.NoError(t, db.UpsertAIModelPrices(ctx, database.UpsertAIModelPricesParams{
-			Seed: []byte(`[{
-				"provider": "openai",
-				"model": "gpt-4o",
-				"input_price": 5000000,
-				"output_price": 20000000,
-				"cache_read_price": 2000000,
-				"cache_write_price": 1000000
-			}]`),
-			Source: database.AIModelPriceSourceCustom,
-		}))
-
-		got, err := db.GetAIModelPriceByProviderModel(ctx, database.GetAIModelPriceByProviderModelParams{
-			Provider: "openai", Model: "gpt-4o",
-		})
-		require.NoError(t, err)
-		require.Equal(t, database.AIModelPriceSourceCustom, got.Source)
-		require.Equal(t, int64(5_000_000), got.InputPrice.Int64)
-		require.Equal(t, int64(20_000_000), got.OutputPrice.Int64)
-		require.Equal(t, int64(2_000_000), got.CacheReadPrice.Int64)
-		require.Equal(t, int64(1_000_000), got.CacheWritePrice.Int64)
-	})
-
-	t.Run("UnchangedPricesAreCustom", func(t *testing.T) {
-		t.Parallel()
-		ctx := testutil.Context(t, testutil.WaitShort)
-		db, _ := dbtestutil.NewDB(t)
-
-		require.NoError(t, prices.SeedFromBytes(ctx, db, []byte(testSeedJSON)))
-		seeded, err := db.GetAIModelPriceByProviderModel(ctx, database.GetAIModelPriceByProviderModelParams{
-			Provider: "openai", Model: "gpt-4o",
-		})
-		require.NoError(t, err)
-		require.Equal(t, database.AIModelPriceSourceDefault, seeded.Source)
-
-		// Declaring the prices the row already holds still marks it custom.
-		require.NoError(t, db.UpsertAIModelPrices(ctx, database.UpsertAIModelPricesParams{
-			Seed: []byte(`[{
-				"provider": "openai",
-				"model": "gpt-4o",
-				"input_price": 2500000,
-				"output_price": 10000000,
-				"cache_read_price": 1250000,
-				"cache_write_price": null
-			}]`),
-			Source: database.AIModelPriceSourceCustom,
-		}))
-
-		got, err := db.GetAIModelPriceByProviderModel(ctx, database.GetAIModelPriceByProviderModelParams{
-			Provider: "openai", Model: "gpt-4o",
-		})
-		require.NoError(t, err)
-		require.Equal(t, database.AIModelPriceSourceCustom, got.Source)
-		require.Equal(t, int64(2_500_000), got.InputPrice.Int64)
-		require.Equal(t, int64(10_000_000), got.OutputPrice.Int64)
-		require.Equal(t, int64(1_250_000), got.CacheReadPrice.Int64)
-		require.False(t, got.CacheWritePrice.Valid)
-	})
-
 	t.Run("Idempotent", func(t *testing.T) {
 		t.Parallel()
 		ctx := testutil.Context(t, testutil.WaitShort)
@@ -280,8 +208,12 @@ func TestSeedFromBytes(t *testing.T) {
 		})
 		require.NoError(t, err)
 		require.Equal(t, database.AIModelPriceSourceCustom, before.Source)
+		require.Equal(t, int64(1), before.InputPrice.Int64)
+		require.Equal(t, int64(2), before.OutputPrice.Int64)
+		require.Equal(t, int64(3), before.CacheReadPrice.Int64)
+		require.Equal(t, int64(4), before.CacheWritePrice.Int64)
 
-		// Re-applying the price book skips custom rows.
+		// Re-applying the price book writes its own row and leaves this one be.
 		require.NoError(t, prices.SeedFromBytes(ctx, db, []byte(testSeedJSON)))
 
 		got, err := db.GetAIModelPriceByProviderModel(ctx, database.GetAIModelPriceByProviderModelParams{
diff --git a/coderd/aibridgedserver/aibridgedserver_test.go b/coderd/aibridgedserver/aibridgedserver_test.go
index 0cac462bab1..775e4a75c04 100644
--- a/coderd/aibridgedserver/aibridgedserver_test.go
+++ b/coderd/aibridgedserver/aibridgedserver_test.go
@@ -1731,7 +1731,6 @@ func TestRecordTokenUsage(t *testing.T) {
 					// No override
 					expectTokenUsageCostLookups(db, intc, nil, group, nil, price)
 
-					// input 300 + output 1200 + cache read 15 + cache write 40.
 					const wantCost int64 = 1555
 
 					db.EXPECT().InTx(gomock.Any(), nil).DoAndReturn(
@@ -2197,7 +2196,6 @@ func TestRecordTokenUsage(t *testing.T) {
 						func(fn func(database.Store) error, _ *database.TxOptions) error { return fn(db) },
 					)
 
-					// input 300 + output 1200 + cache read 15 + cache write 40.
 					const wantCost int64 = 1555
 
 					db.EXPECT().InsertAIBridgeTokenUsage(gomock.Any(), gomock.Cond(func(p database.InsertAIBridgeTokenUsageParams) bool {
@@ -2559,7 +2557,6 @@ func TestRecordTokenUsageAuthorized(t *testing.T) {
 	require.Equal(t, sql.NullInt64{Int64: 6_000_000, Valid: true}, tokenUsage.OutputPriceMicros, "output price")
 	require.Equal(t, sql.NullInt64{Int64: 300_000, Valid: true}, tokenUsage.CacheReadPriceMicros, "cache read price")
 	require.Equal(t, sql.NullInt64{Int64: 4_000_000, Valid: true}, tokenUsage.CacheWritePriceMicros, "cache write price")
-	// input 300 + output 1200 + cache read 15 + cache write 40.
 	const wantCost int64 = 1555
 	require.Equal(t, sql.NullInt64{Int64: wantCost, Valid: true}, tokenUsage.CostMicros, "cost")
 
@@ -2577,6 +2574,133 @@ func TestRecordTokenUsageAuthorized(t *testing.T) {
 	require.Equal(t, wantCost, spend.SpendMicros, "spend micros")
 }
 
+// TestRecordTokenUsageModelPriceResolution covers which price an interception
+// snapshots when a model carries a price from the embedded book, a price set
+// through the API, or both.
+func TestRecordTokenUsageModelPriceResolution(t *testing.T) {
+	t.Parallel()
+
+	const provider, model = "anthropic", "claude-sonnet-4-6"
+
+	priceSeed := func(input, output, cacheRead, cacheWrite int64) json.RawMessage {
+		seed, err := json.Marshal([]map[string]any{{
+			"provider":          provider,
+			"model":             model,
+			"input_price":       input,
+			"output_price":      output,
+			"cache_read_price":  cacheRead,
+			"cache_write_price": cacheWrite,
+		}})
+		require.NoError(t, err)
+		return seed
+	}
+
+	tests := []struct {
+		name        string
+		defaultSeed json.RawMessage
+		customSeed  json.RawMessage
+		want        database.AIBridgeTokenUsage
+	}{
+		{
+			name:        "DefaultOnly",
+			defaultSeed: priceSeed(3_000_000, 6_000_000, 300_000, 4_000_000),
+			want: database.AIBridgeTokenUsage{
+				InputPriceMicros:      sql.NullInt64{Int64: 3_000_000, Valid: true},
+				OutputPriceMicros:     sql.NullInt64{Int64: 6_000_000, Valid: true},
+				CacheReadPriceMicros:  sql.NullInt64{Int64: 300_000, Valid: true},
+				CacheWritePriceMicros: sql.NullInt64{Int64: 4_000_000, Valid: true},
+				// 100 input, 200 output, 50 cache read, and 10 cache write
+				// tokens, priced per million: 300 + 1200 + 15 + 40.
+				CostMicros: sql.NullInt64{Int64: 1555, Valid: true},
+			},
+		},
+		{
+			name:        "CustomWinsOverDefault",
+			defaultSeed: priceSeed(3_000_000, 6_000_000, 300_000, 4_000_000),
+			customSeed:  priceSeed(9_000_000, 12_000_000, 900_000, 8_000_000),
+			want: database.AIBridgeTokenUsage{
+				InputPriceMicros:      sql.NullInt64{Int64: 9_000_000, Valid: true},
+				OutputPriceMicros:     sql.NullInt64{Int64: 12_000_000, Valid: true},
+				CacheReadPriceMicros:  sql.NullInt64{Int64: 900_000, Valid: true},
+				CacheWritePriceMicros: sql.NullInt64{Int64: 8_000_000, Valid: true},
+				// 100 input, 200 output, 50 cache read, and 10 cache write
+				// tokens, priced per million: 900 + 2400 + 45 + 80.
+				CostMicros: sql.NullInt64{Int64: 3425, Valid: true},
+			},
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			t.Parallel()
+			ctx := testutil.Context(t, testutil.WaitLong)
+			logger := testutil.Logger(t)
+
+			rawDB, _ := dbtestutil.NewDB(t)
+			authzDB := dbauthz.New(rawDB, rbac.NewStrictAuthorizer(prometheus.NewRegistry()), logger, coderdtest.AccessControlStorePointer())
+
+			org := dbgen.Organization(t, rawDB, database.Organization{})
+			user := dbgen.User(t, rawDB, database.User{})
+			dbgen.OrganizationMember(t, rawDB, database.OrganizationMember{OrganizationID: org.ID, UserID: user.ID})
+
+			require.NoError(t, rawDB.UpsertAIModelPrices(ctx, database.UpsertAIModelPricesParams{
+				Seed:   tt.defaultSeed,
+				Source: database.AIModelPriceSourceDefault,
+			}), "seed model prices")
+			if tt.customSeed != nil {
+				require.NoError(t, rawDB.UpsertAIModelPrices(ctx, database.UpsertAIModelPricesParams{
+					Seed:   tt.customSeed,
+					Source: database.AIModelPriceSourceCustom,
+				}), "set custom model price")
+			}
+
+			aiProvider := dbgen.AIProvider(t, rawDB, database.AIProvider{
+				Name: "anthropic-eu",
+				Type: database.AIProviderTypeAnthropic,
+			})
+			intc := dbgen.AIBridgeInterception(t, rawDB, database.InsertAIBridgeInterceptionParams{
+				InitiatorID:  user.ID,
+				Provider:     provider,
+				ProviderName: aiProvider.Name,
+				Model:        model,
+			}, nil)
+
+			srv, err := aibridgedserver.NewServer(ctx, aibridgedserver.Options{
+				Store:         authzDB,
+				AISeatTracker: agplaiseats.Noop{},
+				AccessURL:     "/",
+				GatewayCfg:    codersdk.AIBridgeConfig{},
+				Experiments:   requiredExperiments,
+				Logger:        logger,
+				Clock:         quartz.NewReal(),
+			})
+			require.NoError(t, err)
+
+			_, err = srv.RecordTokenUsage(ctx, &proto.RecordTokenUsageRequest{
+				InterceptionId:        intc.ID.String(),
+				MsgId:                 "msg_price_resolution",
+				InputTokens:           100,
+				OutputTokens:          200,
+				CacheReadInputTokens:  50,
+				CacheWriteInputTokens: 10,
+				CreatedAt:             timestamppb.New(time.Date(2026, 6, 25, 14, 30, 0, 0, time.UTC)),
+			})
+			require.NoError(t, err, "record token usage")
+
+			tokenUsages, err := rawDB.GetAIBridgeTokenUsagesByInterceptionID(ctx, intc.ID)
+			require.NoError(t, err)
+			require.Len(t, tokenUsages, 1)
+
+			tokenUsage := tokenUsages[0]
+			require.Equal(t, tt.want.InputPriceMicros, tokenUsage.InputPriceMicros, "input price")
+			require.Equal(t, tt.want.OutputPriceMicros, tokenUsage.OutputPriceMicros, "output price")
+			require.Equal(t, tt.want.CacheReadPriceMicros, tokenUsage.CacheReadPriceMicros, "cache read price")
+			require.Equal(t, tt.want.CacheWritePriceMicros, tokenUsage.CacheWritePriceMicros, "cache write price")
+			require.Equal(t, tt.want.CostMicros, tokenUsage.CostMicros, "cost")
+		})
+	}
+}
+
 // TestRecordTokenUsageProviderResolution covers provider resolution against a real
 // database through dbauthz, where the live-row filter and name reuse apply.
 func TestRecordTokenUsageProviderResolution(t *testing.T) {
diff --git a/coderd/database/dump.sql b/coderd/database/dump.sql
index 1be2c4f37c3..f649f5e06af 100644
--- a/coderd/database/dump.sql
+++ b/coderd/database/dump.sql
@@ -1587,7 +1587,7 @@ CREATE TABLE ai_model_prices (
 
 COMMENT ON TABLE ai_model_prices IS 'Per-model token prices used by AI Bridge to compute interception cost.';
 
-COMMENT ON COLUMN ai_model_prices.source IS 'Where the price came from: default for the embedded price book, custom for a price an operator set. The startup seeder never overwrites a custom row.';
+COMMENT ON COLUMN ai_model_prices.source IS 'Where the price came from: default for the embedded price book, custom for a price set through the API. Both can exist for the same model.';
 
 CREATE TABLE ai_provider_keys (
     id uuid DEFAULT gen_random_uuid() NOT NULL,
@@ -4323,7 +4323,7 @@ ALTER TABLE ONLY ai_gateway_keys
     ADD CONSTRAINT ai_gateway_keys_pkey PRIMARY KEY (id);
 
 ALTER TABLE ONLY ai_model_prices
-    ADD CONSTRAINT ai_model_prices_pkey PRIMARY KEY (provider, model);
+    ADD CONSTRAINT ai_model_prices_pkey PRIMARY KEY (provider, model, source);
 
 ALTER TABLE ONLY ai_provider_keys
     ADD CONSTRAINT ai_provider_keys_pkey PRIMARY KEY (id);
diff --git a/coderd/database/migrations/000574_ai_model_price_source.down.sql b/coderd/database/migrations/000574_ai_model_price_source.down.sql
index f1e9a9b97c1..6bfdb71a99a 100644
--- a/coderd/database/migrations/000574_ai_model_price_source.down.sql
+++ b/coderd/database/migrations/000574_ai_model_price_source.down.sql
@@ -1,3 +1,11 @@
+-- Custom prices have no place in the single-row model, and a model may hold
+-- both rows, so they are dropped to restore uniqueness on (provider, model).
+DELETE FROM ai_model_prices WHERE source = 'custom';
+
+ALTER TABLE ai_model_prices DROP CONSTRAINT ai_model_prices_pkey;
+
+ALTER TABLE ai_model_prices ADD PRIMARY KEY (provider, model);
+
 ALTER TABLE ai_model_prices DROP COLUMN source;
 
 DROP TYPE ai_model_price_source;
diff --git a/coderd/database/migrations/000574_ai_model_price_source.up.sql b/coderd/database/migrations/000574_ai_model_price_source.up.sql
index 4689db362f8..61c86f631a7 100644
--- a/coderd/database/migrations/000574_ai_model_price_source.up.sql
+++ b/coderd/database/migrations/000574_ai_model_price_source.up.sql
@@ -1,7 +1,6 @@
 -- source records where a price came from, either the price book embedded in
--- each Coder release or a price set through the API. The startup seeder
--- re-applies the book on every boot and reads source to leave custom rows
--- alone.
+-- each Coder release or a price set through the API. Both can exist for the
+-- same model, so source joins the primary key.
 
 CREATE TYPE ai_model_price_source AS ENUM ('default', 'custom');
 
@@ -11,4 +10,8 @@ UPDATE ai_model_prices SET source = 'default' WHERE source IS NULL;
 
 ALTER TABLE ai_model_prices ALTER COLUMN source SET NOT NULL;
 
-COMMENT ON COLUMN ai_model_prices.source IS 'Where the price came from: default for the embedded price book, custom for a price an operator set. The startup seeder never overwrites a custom row.';
+ALTER TABLE ai_model_prices DROP CONSTRAINT ai_model_prices_pkey;
+
+ALTER TABLE ai_model_prices ADD PRIMARY KEY (provider, model, source);
+
+COMMENT ON COLUMN ai_model_prices.source IS 'Where the price came from: default for the embedded price book, custom for a price set through the API. Both can exist for the same model.';
diff --git a/coderd/database/models.go b/coderd/database/models.go
index e6024549d3f..6cea2c3f8a7 100644
--- a/coderd/database/models.go
+++ b/coderd/database/models.go
@@ -4885,7 +4885,7 @@ type AIModelPrice struct {
 	CacheWritePrice sql.NullInt64 `db:"cache_write_price" json:"cache_write_price"`
 	CreatedAt       time.Time     `db:"created_at" json:"created_at"`
 	UpdatedAt       time.Time     `db:"updated_at" json:"updated_at"`
-	// Where the price came from: default for the embedded price book, custom for a price an operator set. The startup seeder never overwrites a custom row.
+	// Where the price came from: default for the embedded price book, custom for a price set through the API. Both can exist for the same model.
 	Source AIModelPriceSource `db:"source" json:"source"`
 }
 
diff --git a/coderd/database/querier.go b/coderd/database/querier.go
index 85a7204361a..14c20ef75fe 100644
--- a/coderd/database/querier.go
+++ b/coderd/database/querier.go
@@ -323,7 +323,11 @@ type sqlcQuerier interface {
 	// returning the matched key. The lookup is an exact match on a unique index,
 	// so a returned row is itself proof the secret is valid.
 	GetAIGatewayKeyByHashedSecret(ctx context.Context, hashedSecret []byte) (AIGatewayKey, error)
+	// Returns the price in effect for the model, preferring a custom price over
+	// the price book.
 	GetAIModelPriceByProviderModel(ctx context.Context, arg GetAIModelPriceByProviderModelParams) (AIModelPrice, error)
+	// Returns the price in effect for each model, preferring a custom price over
+	// the price book.
 	GetAIModelPrices(ctx context.Context, arg GetAIModelPricesParams) ([]AIModelPrice, error)
 	GetAIProviderByID(ctx context.Context, id uuid.UUID) (AIProvider, error)
 	// Lock the provider row until the model-config write completes. The
@@ -1615,13 +1619,13 @@ type sqlcQuerier interface {
 	UpdateWorkspaceTTL(ctx context.Context, arg UpdateWorkspaceTTLParams) error
 	UpdateWorkspacesDormantDeletingAtByTemplateID(ctx context.Context, arg UpdateWorkspacesDormantDeletingAtByTemplateIDParams) ([]WorkspaceTable, error)
 	UpdateWorkspacesTTLByTemplateID(ctx context.Context, arg UpdateWorkspacesTTLByTemplateIDParams) error
-	// Upsert a batch of (provider, model) rows from a JSON array, recording them
-	// under source. Each element must have provider, model, and the four price
+	// Upsert a batch of model prices from a JSON array, all recorded under the
+	// given source. Each element must have provider, model, and the four price
 	// fields, and null prices are written as SQL NULL.
-	// A default write skips rows a custom price owns. Otherwise a conflicting row
-	// is only rewritten when a price or the source differs, so updated_at records
-	// when the row last changed. Prices are nullable and a NULL on either side
-	// counts as a difference.
+	// Each source keeps its own row, so the price book and a custom price never
+	// overwrite each other. 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, arg UpsertAIModelPricesParams) error
 	// Returns true if a new rows was inserted, false otherwise.
 	UpsertAISeatState(ctx context.Context, arg UpsertAISeatStateParams) (bool, error)
diff --git a/coderd/database/querier_test.go b/coderd/database/querier_test.go
index f66e50c2dc4..10b29923794 100644
--- a/coderd/database/querier_test.go
+++ b/coderd/database/querier_test.go
@@ -19059,6 +19059,85 @@ func TestOAuth2ProviderScopeNotEmpty(t *testing.T) {
 	})
 }
 
+func TestGetAIModelPriceByProviderModel(t *testing.T) {
+	t.Parallel()
+
+	const defaultSeed = `[{"provider":"anthropic","model":"model-a","input_price":1,"output_price":2,"cache_read_price":3,"cache_write_price":4}]`
+	const customSeed = `[{"provider":"anthropic","model":"model-a","input_price":5,"output_price":6,"cache_read_price":0,"cache_write_price":null}]`
+
+	defaultPrices := database.AIModelPrice{
+		InputPrice:      sql.NullInt64{Int64: 1, Valid: true},
+		OutputPrice:     sql.NullInt64{Int64: 2, Valid: true},
+		CacheReadPrice:  sql.NullInt64{Int64: 3, Valid: true},
+		CacheWritePrice: sql.NullInt64{Int64: 4, Valid: true},
+	}
+	customPrices := database.AIModelPrice{
+		InputPrice:      sql.NullInt64{Int64: 5, Valid: true},
+		OutputPrice:     sql.NullInt64{Int64: 6, Valid: true},
+		CacheReadPrice:  sql.NullInt64{Int64: 0, Valid: true},
+		CacheWritePrice: sql.NullInt64{},
+	}
+
+	tests := []struct {
+		name       string
+		seeds      []database.UpsertAIModelPricesParams
+		want       database.AIModelPrice
+		wantSource database.AIModelPriceSource
+	}{
+		{
+			name:       "DefaultOnly",
+			seeds:      []database.UpsertAIModelPricesParams{{Seed: []byte(defaultSeed), Source: database.AIModelPriceSourceDefault}},
+			want:       defaultPrices,
+			wantSource: database.AIModelPriceSourceDefault,
+		},
+		{
+			name:       "CustomOnly",
+			seeds:      []database.UpsertAIModelPricesParams{{Seed: []byte(customSeed), Source: database.AIModelPriceSourceCustom}},
+			want:       customPrices,
+			wantSource: database.AIModelPriceSourceCustom,
+		},
+		{
+			name: "CustomWinsOverDefault",
+			seeds: []database.UpsertAIModelPricesParams{
+				{Seed: []byte(defaultSeed), Source: database.AIModelPriceSourceDefault},
+				{Seed: []byte(customSeed), Source: database.AIModelPriceSourceCustom},
+			},
+			want:       customPrices,
+			wantSource: database.AIModelPriceSourceCustom,
+		},
+		{
+			name: "CustomWinsWhateverTheWriteOrder",
+			seeds: []database.UpsertAIModelPricesParams{
+				{Seed: []byte(customSeed), Source: database.AIModelPriceSourceCustom},
+				{Seed: []byte(defaultSeed), Source: database.AIModelPriceSourceDefault},
+			},
+			want:       customPrices,
+			wantSource: database.AIModelPriceSourceCustom,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			t.Parallel()
+			ctx := testutil.Context(t, testutil.WaitShort)
+			db, _ := dbtestutil.NewDB(t)
+			for _, seed := range tt.seeds {
+				require.NoError(t, db.UpsertAIModelPrices(ctx, seed))
+			}
+
+			got, err := db.GetAIModelPriceByProviderModel(ctx, database.GetAIModelPriceByProviderModelParams{
+				Provider: "anthropic", Model: "model-a",
+			})
+			require.NoError(t, err)
+			require.Equal(t, tt.wantSource, got.Source)
+			require.Equal(t, tt.want.InputPrice, got.InputPrice)
+			require.Equal(t, tt.want.OutputPrice, got.OutputPrice)
+			require.Equal(t, tt.want.CacheReadPrice, got.CacheReadPrice)
+			require.Equal(t, tt.want.CacheWritePrice, got.CacheWritePrice)
+		})
+	}
+}
+
 func TestGetAIModelPrices(t *testing.T) {
 	t.Parallel()
 
@@ -19070,15 +19149,21 @@ func TestGetAIModelPrices(t *testing.T) {
 		{"provider":"openai","model":"model-a","input_price":3,"output_price":null,"cache_read_price":null,"cache_write_price":null}
 	]`
 
+	// A custom price for anthropic/model-a, which the seed above also covers.
+	const customSeed = `[{"provider":"anthropic","model":"model-a","input_price":9,"output_price":null,"cache_read_price":null,"cache_write_price":null}]`
+
 	tests := []struct {
-		name   string
-		params database.GetAIModelPricesParams
-		want   []string
+		name       string
+		customSeed string
+		params     database.GetAIModelPricesParams
+		want       []string
+		wantPrices []int64
 	}{
 		{
-			name:   "NoFilterReturnsEveryPrice",
-			params: database.GetAIModelPricesParams{},
-			want:   []string{"anthropic/model-a", "anthropic/model-b", "openai/model-a"},
+			name:       "NoFilterReturnsEveryPrice",
+			params:     database.GetAIModelPricesParams{},
+			want:       []string{"anthropic/model-a", "anthropic/model-b", "openai/model-a"},
+			wantPrices: []int64{1, 2, 3},
 		},
 		{
 			name:   "ByProvider",
@@ -19105,6 +19190,14 @@ func TestGetAIModelPrices(t *testing.T) {
 			params: database.GetAIModelPricesParams{Provider: "openai", Model: "model-b"},
 			want:   nil,
 		},
+		{
+			// The anthropic/model-a is reported once, at the custom one.
+			name:       "ResolvesToTheCustomPrice",
+			customSeed: customSeed,
+			params:     database.GetAIModelPricesParams{},
+			want:       []string{"anthropic/model-a", "anthropic/model-b", "openai/model-a"},
+			wantPrices: []int64{9, 2, 3},
+		},
 	}
 
 	for _, tt := range tests {
@@ -19113,6 +19206,9 @@ func TestGetAIModelPrices(t *testing.T) {
 			ctx := testutil.Context(t, testutil.WaitShort)
 			db, _ := dbtestutil.NewDB(t)
 			require.NoError(t, db.UpsertAIModelPrices(ctx, database.UpsertAIModelPricesParams{Seed: []byte(seed), Source: database.AIModelPriceSourceDefault}))
+			if tt.customSeed != "" {
+				require.NoError(t, db.UpsertAIModelPrices(ctx, database.UpsertAIModelPricesParams{Seed: []byte(tt.customSeed), Source: database.AIModelPriceSourceCustom}))
+			}
 
 			prices, err := db.GetAIModelPrices(ctx, tt.params)
 			require.NoError(t, err)
@@ -19126,6 +19222,14 @@ func TestGetAIModelPrices(t *testing.T) {
 				return
 			}
 			require.Equal(t, tt.want, got)
+
+			if tt.wantPrices != nil {
+				gotPrices := make([]int64, 0, len(prices))
+				for _, price := range prices {
+					gotPrices = append(gotPrices, price.InputPrice.Int64)
+				}
+				require.Equal(t, tt.wantPrices, gotPrices)
+			}
 		})
 	}
 }
diff --git a/coderd/database/queries.sql.go b/coderd/database/queries.sql.go
index cbab0fb7787..be29c46a336 100644
--- a/coderd/database/queries.sql.go
+++ b/coderd/database/queries.sql.go
@@ -2865,6 +2865,8 @@ const getAIModelPriceByProviderModel = `-- name: GetAIModelPriceByProviderModel
 SELECT provider, model, input_price, output_price, cache_read_price, cache_write_price, created_at, updated_at, source
 FROM ai_model_prices
 WHERE provider = $1 AND model = $2
+ORDER BY CASE WHEN source = 'custom' THEN 0 ELSE 1 END
+LIMIT 1
 `
 
 type GetAIModelPriceByProviderModelParams struct {
@@ -2872,6 +2874,8 @@ type GetAIModelPriceByProviderModelParams struct {
 	Model    string `db:"model" json:"model"`
 }
 
+// Returns the price in effect for the model, preferring a custom price over
+// the price book.
 func (q *sqlQuerier) GetAIModelPriceByProviderModel(ctx context.Context, arg GetAIModelPriceByProviderModelParams) (AIModelPrice, error) {
 	row := q.db.QueryRowContext(ctx, getAIModelPriceByProviderModel, arg.Provider, arg.Model)
 	var i AIModelPrice
@@ -2890,7 +2894,7 @@ func (q *sqlQuerier) GetAIModelPriceByProviderModel(ctx context.Context, arg Get
 }
 
 const getAIModelPrices = `-- name: GetAIModelPrices :many
-SELECT provider, model, input_price, output_price, cache_read_price, cache_write_price, created_at, updated_at, source
+SELECT DISTINCT ON (provider, model) provider, model, input_price, output_price, cache_read_price, cache_write_price, created_at, updated_at, source
 FROM ai_model_prices
     -- Filter by provider
 WHERE CASE
@@ -2904,7 +2908,7 @@ WHERE CASE
             model = $2
         ELSE true
     END
-ORDER BY provider, model
+ORDER BY provider, model, CASE WHEN source = 'custom' THEN 0 ELSE 1 END
 `
 
 type GetAIModelPricesParams struct {
@@ -2912,6 +2916,8 @@ type GetAIModelPricesParams struct {
 	Model    string `db:"model" json:"model"`
 }
 
+// Returns the price in effect for each model, preferring a custom price over
+// the price book.
 func (q *sqlQuerier) GetAIModelPrices(ctx context.Context, arg GetAIModelPricesParams) ([]AIModelPrice, error) {
 	rows, err := q.db.QueryContext(ctx, getAIModelPrices, arg.Provider, arg.Model)
 	if err != nil {
@@ -3511,32 +3517,23 @@ SELECT
 	(elem->>'cache_write_price')::bigint,
 	$1::ai_model_price_source
 FROM jsonb_array_elements($2::jsonb) AS elem
-ON CONFLICT (provider, model) DO UPDATE SET
+ON CONFLICT (provider, model, source) DO UPDATE SET
 	input_price       = EXCLUDED.input_price,
 	output_price      = EXCLUDED.output_price,
 	cache_read_price  = EXCLUDED.cache_read_price,
 	cache_write_price = EXCLUDED.cache_write_price,
-	source            = EXCLUDED.source,
 	updated_at        = NOW()
-WHERE CASE
-		-- A custom price claims any row, including one the price book wrote.
-		WHEN $1::ai_model_price_source = 'custom' THEN true
-		-- The price book leaves custom rows alone.
-		ELSE ai_model_prices.source <> 'custom'
-	END
-	AND (
-		ai_model_prices.input_price,
-		ai_model_prices.output_price,
-		ai_model_prices.cache_read_price,
-		ai_model_prices.cache_write_price,
-		ai_model_prices.source
-	) IS DISTINCT FROM (
-		EXCLUDED.input_price,
-		EXCLUDED.output_price,
-		EXCLUDED.cache_read_price,
-		EXCLUDED.cache_write_price,
-		EXCLUDED.source
-	)
+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
+)
 `
 
 type UpsertAIModelPricesParams struct {
@@ -3544,13 +3541,13 @@ type UpsertAIModelPricesParams struct {
 	Seed   json.RawMessage    `db:"seed" json:"seed"`
 }
 
-// Upsert a batch of (provider, model) rows from a JSON array, recording them
-// under source. Each element must have provider, model, and the four price
+// Upsert a batch of model prices from a JSON array, all recorded under the
+// given source. Each element must have provider, model, and the four price
 // fields, and null prices are written as SQL NULL.
-// A default write skips rows a custom price owns. Otherwise a conflicting row
-// is only rewritten when a price or the source differs, so updated_at records
-// when the row last changed. Prices are nullable and a NULL on either side
-// counts as a difference.
+// Each source keeps its own row, so the price book and a custom price never
+// overwrite each other. 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, arg UpsertAIModelPricesParams) error {
 	_, err := q.db.ExecContext(ctx, upsertAIModelPrices, arg.Source, arg.Seed)
 	return err
diff --git a/coderd/database/queries/aicostcontrol.sql b/coderd/database/queries/aicostcontrol.sql
index 8b7649e39dc..1e9e0e668ed 100644
--- a/coderd/database/queries/aicostcontrol.sql
+++ b/coderd/database/queries/aicostcontrol.sql
@@ -1,11 +1,11 @@
 -- name: UpsertAIModelPrices :exec
--- Upsert a batch of (provider, model) rows from a JSON array, recording them
--- under source. Each element must have provider, model, and the four price
+-- Upsert a batch of model prices from a JSON array, all recorded under the
+-- given source. Each element must have provider, model, and the four price
 -- fields, and null prices are written as SQL NULL.
--- A default write skips rows a custom price owns. Otherwise a conflicting row
--- is only rewritten when a price or the source differs, so updated_at records
--- when the row last changed. Prices are nullable and a NULL on either side
--- counts as a difference.
+-- Each source keeps its own row, so the price book and a custom price never
+-- overwrite each other. 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, source
 )
@@ -18,40 +18,37 @@ SELECT
 	(elem->>'cache_write_price')::bigint,
 	@source::ai_model_price_source
 FROM jsonb_array_elements(@seed::jsonb) AS elem
-ON CONFLICT (provider, model) DO UPDATE SET
+ON CONFLICT (provider, model, source) DO UPDATE SET
 	input_price       = EXCLUDED.input_price,
 	output_price      = EXCLUDED.output_price,
 	cache_read_price  = EXCLUDED.cache_read_price,
 	cache_write_price = EXCLUDED.cache_write_price,
-	source            = EXCLUDED.source,
 	updated_at        = NOW()
-WHERE CASE
-		-- A custom price claims any row, including one the price book wrote.
-		WHEN @source::ai_model_price_source = 'custom' THEN true
-		-- The price book leaves custom rows alone.
-		ELSE ai_model_prices.source <> 'custom'
-	END
-	AND (
-		ai_model_prices.input_price,
-		ai_model_prices.output_price,
-		ai_model_prices.cache_read_price,
-		ai_model_prices.cache_write_price,
-		ai_model_prices.source
-	) IS DISTINCT FROM (
-		EXCLUDED.input_price,
-		EXCLUDED.output_price,
-		EXCLUDED.cache_read_price,
-		EXCLUDED.cache_write_price,
-		EXCLUDED.source
-	);
+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
+-- Returns the price in effect for the model, preferring a custom price over
+-- the price book.
 SELECT *
 FROM ai_model_prices
-WHERE provider = @provider AND model = @model;
+WHERE provider = @provider AND model = @model
+ORDER BY CASE WHEN source = 'custom' THEN 0 ELSE 1 END
+LIMIT 1;
 
 -- name: GetAIModelPrices :many
-SELECT *
+-- Returns the price in effect for each model, preferring a custom price over
+-- the price book.
+SELECT DISTINCT ON (provider, model) *
 FROM ai_model_prices
     -- Filter by provider
 WHERE CASE
@@ -65,7 +62,7 @@ WHERE CASE
             model = @model
         ELSE true
     END
-ORDER BY provider, model;
+ORDER BY provider, model, CASE WHEN source = 'custom' THEN 0 ELSE 1 END;
 
 -- name: GetGroupAIBudget :one
 SELECT *
diff --git a/coderd/database/unique_constraint.go b/coderd/database/unique_constraint.go
index 79ec48aa26d..87550acf8c0 100644
--- a/coderd/database/unique_constraint.go
+++ b/coderd/database/unique_constraint.go
@@ -8,7 +8,7 @@ type UniqueConstraint string
 const (
 	UniqueAgentStatsPkey                                      UniqueConstraint = "agent_stats_pkey"                                                // ALTER TABLE ONLY workspace_agent_stats ADD CONSTRAINT agent_stats_pkey PRIMARY KEY (id);
 	UniqueAIGatewayKeysPkey                                   UniqueConstraint = "ai_gateway_keys_pkey"                                            // ALTER TABLE ONLY ai_gateway_keys ADD CONSTRAINT ai_gateway_keys_pkey PRIMARY KEY (id);
-	UniqueAIModelPricesPkey                                   UniqueConstraint = "ai_model_prices_pkey"                                            // ALTER TABLE ONLY ai_model_prices ADD CONSTRAINT ai_model_prices_pkey PRIMARY KEY (provider, model);
+	UniqueAIModelPricesPkey                                   UniqueConstraint = "ai_model_prices_pkey"                                            // ALTER TABLE ONLY ai_model_prices ADD CONSTRAINT ai_model_prices_pkey PRIMARY KEY (provider, model, source);
 	UniqueAIProviderKeysPkey                                  UniqueConstraint = "ai_provider_keys_pkey"                                           // ALTER TABLE ONLY ai_provider_keys ADD CONSTRAINT ai_provider_keys_pkey PRIMARY KEY (id);
 	UniqueAIProvidersPkey                                     UniqueConstraint = "ai_providers_pkey"                                               // ALTER TABLE ONLY ai_providers ADD CONSTRAINT ai_providers_pkey PRIMARY KEY (id);
 	UniqueAISeatStatePkey                                     UniqueConstraint = "ai_seat_state_pkey"                                              // ALTER TABLE ONLY ai_seat_state ADD CONSTRAINT ai_seat_state_pkey PRIMARY KEY (user_id);

From 17c825729c1d1fabda65f789c45511ca0679facb Mon Sep 17 00:00:00 2001
From: Susana Cardoso Ferreira <susana@coder.com>
Date: Wed, 19 Aug 2026 19:35:36 +0000
Subject: [PATCH 3/6] chore: fix migration numbers

---
 ...rice_source.down.sql => 000577_ai_model_price_source.down.sql} | 0
 ...el_price_source.up.sql => 000577_ai_model_price_source.up.sql} | 0
 2 files changed, 0 insertions(+), 0 deletions(-)
 rename coderd/database/migrations/{000574_ai_model_price_source.down.sql => 000577_ai_model_price_source.down.sql} (100%)
 rename coderd/database/migrations/{000574_ai_model_price_source.up.sql => 000577_ai_model_price_source.up.sql} (100%)

diff --git a/coderd/database/migrations/000574_ai_model_price_source.down.sql b/coderd/database/migrations/000577_ai_model_price_source.down.sql
similarity index 100%
rename from coderd/database/migrations/000574_ai_model_price_source.down.sql
rename to coderd/database/migrations/000577_ai_model_price_source.down.sql
diff --git a/coderd/database/migrations/000574_ai_model_price_source.up.sql b/coderd/database/migrations/000577_ai_model_price_source.up.sql
similarity index 100%
rename from coderd/database/migrations/000574_ai_model_price_source.up.sql
rename to coderd/database/migrations/000577_ai_model_price_source.up.sql

From bbe9f357a26ef0c8266cbda373d13a8e13a1c3fe Mon Sep 17 00:00:00 2001
From: Susana Cardoso Ferreira <susana@coder.com>
Date: Wed, 19 Aug 2026 19:45:41 +0000
Subject: [PATCH 4/6] test: add
 TestRecordTokenUsageModelPriceResolution/CustomOnly

---
 .../aibridgedserver/aibridgedserver_test.go   | 24 +++++++++++++++----
 1 file changed, 20 insertions(+), 4 deletions(-)

diff --git a/coderd/aibridgedserver/aibridgedserver_test.go b/coderd/aibridgedserver/aibridgedserver_test.go
index 775e4a75c04..a68f2b18252 100644
--- a/coderd/aibridgedserver/aibridgedserver_test.go
+++ b/coderd/aibridgedserver/aibridgedserver_test.go
@@ -2614,6 +2614,20 @@ func TestRecordTokenUsageModelPriceResolution(t *testing.T) {
 				CostMicros: sql.NullInt64{Int64: 1555, Valid: true},
 			},
 		},
+		{
+			// A model the price book does not cover, priced through the API.
+			name:       "CustomOnly",
+			customSeed: priceSeed(2_000_000, 4_000_000, 200_000, 1_000_000),
+			want: database.AIBridgeTokenUsage{
+				InputPriceMicros:      sql.NullInt64{Int64: 2_000_000, Valid: true},
+				OutputPriceMicros:     sql.NullInt64{Int64: 4_000_000, Valid: true},
+				CacheReadPriceMicros:  sql.NullInt64{Int64: 200_000, Valid: true},
+				CacheWritePriceMicros: sql.NullInt64{Int64: 1_000_000, Valid: true},
+				// 100 input, 200 output, 50 cache read, and 10 cache write
+				// tokens, priced per million: 200 + 800 + 10 + 10.
+				CostMicros: sql.NullInt64{Int64: 1020, Valid: true},
+			},
+		},
 		{
 			name:        "CustomWinsOverDefault",
 			defaultSeed: priceSeed(3_000_000, 6_000_000, 300_000, 4_000_000),
@@ -2643,10 +2657,12 @@ func TestRecordTokenUsageModelPriceResolution(t *testing.T) {
 			user := dbgen.User(t, rawDB, database.User{})
 			dbgen.OrganizationMember(t, rawDB, database.OrganizationMember{OrganizationID: org.ID, UserID: user.ID})
 
-			require.NoError(t, rawDB.UpsertAIModelPrices(ctx, database.UpsertAIModelPricesParams{
-				Seed:   tt.defaultSeed,
-				Source: database.AIModelPriceSourceDefault,
-			}), "seed model prices")
+			if tt.defaultSeed != nil {
+				require.NoError(t, rawDB.UpsertAIModelPrices(ctx, database.UpsertAIModelPricesParams{
+					Seed:   tt.defaultSeed,
+					Source: database.AIModelPriceSourceDefault,
+				}), "seed model prices")
+			}
 			if tt.customSeed != nil {
 				require.NoError(t, rawDB.UpsertAIModelPrices(ctx, database.UpsertAIModelPricesParams{
 					Seed:   tt.customSeed,

From a246ba45ca2cd295b5de34eeb884984c7490e55e Mon Sep 17 00:00:00 2001
From: Susana Cardoso Ferreira <susana@coder.com>
Date: Thu, 20 Aug 2026 07:57:51 +0000
Subject: [PATCH 5/6] chore: address comments

---
 coderd/database/queries.sql.go            | 4 ++--
 coderd/database/queries/aicostcontrol.sql | 4 ++--
 2 files changed, 4 insertions(+), 4 deletions(-)

diff --git a/coderd/database/queries.sql.go b/coderd/database/queries.sql.go
index be29c46a336..59a65bea3c6 100644
--- a/coderd/database/queries.sql.go
+++ b/coderd/database/queries.sql.go
@@ -2865,7 +2865,7 @@ const getAIModelPriceByProviderModel = `-- name: GetAIModelPriceByProviderModel
 SELECT provider, model, input_price, output_price, cache_read_price, cache_write_price, created_at, updated_at, source
 FROM ai_model_prices
 WHERE provider = $1 AND model = $2
-ORDER BY CASE WHEN source = 'custom' THEN 0 ELSE 1 END
+ORDER BY CASE WHEN source = 'custom' THEN 0 ELSE 1 END ASC
 LIMIT 1
 `
 
@@ -2908,7 +2908,7 @@ WHERE CASE
             model = $2
         ELSE true
     END
-ORDER BY provider, model, CASE WHEN source = 'custom' THEN 0 ELSE 1 END
+ORDER BY provider ASC, model ASC, CASE WHEN source = 'custom' THEN 0 ELSE 1 END ASC
 `
 
 type GetAIModelPricesParams struct {
diff --git a/coderd/database/queries/aicostcontrol.sql b/coderd/database/queries/aicostcontrol.sql
index 1e9e0e668ed..f0ce38e524a 100644
--- a/coderd/database/queries/aicostcontrol.sql
+++ b/coderd/database/queries/aicostcontrol.sql
@@ -42,7 +42,7 @@ WHERE (
 SELECT *
 FROM ai_model_prices
 WHERE provider = @provider AND model = @model
-ORDER BY CASE WHEN source = 'custom' THEN 0 ELSE 1 END
+ORDER BY CASE WHEN source = 'custom' THEN 0 ELSE 1 END ASC
 LIMIT 1;
 
 -- name: GetAIModelPrices :many
@@ -62,7 +62,7 @@ WHERE CASE
             model = @model
         ELSE true
     END
-ORDER BY provider, model, CASE WHEN source = 'custom' THEN 0 ELSE 1 END;
+ORDER BY provider ASC, model ASC, CASE WHEN source = 'custom' THEN 0 ELSE 1 END ASC;
 
 -- name: GetGroupAIBudget :one
 SELECT *

From 8e531415bbcfbd5fb5902240acd46d2ce8475aa5 Mon Sep 17 00:00:00 2001
From: Susana Cardoso Ferreira <susana@coder.com>
Date: Thu, 20 Aug 2026 08:07:46 +0000
Subject: [PATCH 6/6] chore: fix migration numbers

---
 ...rice_source.down.sql => 000578_ai_model_price_source.down.sql} | 0
 ...el_price_source.up.sql => 000578_ai_model_price_source.up.sql} | 0
 2 files changed, 0 insertions(+), 0 deletions(-)
 rename coderd/database/migrations/{000577_ai_model_price_source.down.sql => 000578_ai_model_price_source.down.sql} (100%)
 rename coderd/database/migrations/{000577_ai_model_price_source.up.sql => 000578_ai_model_price_source.up.sql} (100%)

diff --git a/coderd/database/migrations/000577_ai_model_price_source.down.sql b/coderd/database/migrations/000578_ai_model_price_source.down.sql
similarity index 100%
rename from coderd/database/migrations/000577_ai_model_price_source.down.sql
rename to coderd/database/migrations/000578_ai_model_price_source.down.sql
diff --git a/coderd/database/migrations/000577_ai_model_price_source.up.sql b/coderd/database/migrations/000578_ai_model_price_source.up.sql
similarity index 100%
rename from coderd/database/migrations/000577_ai_model_price_source.up.sql
rename to coderd/database/migrations/000578_ai_model_price_source.up.sql