From ca32a2e8371e82e7cfeb0aa5b674311afb440c50 Mon Sep 17 00:00:00 2001 From: Yevhenii Shcherbina Date: Tue, 28 Jul 2026 17:45:08 +0000 Subject: [PATCH 01/10] fix: detect out-of-range AI Gateway costs instead of wrapping silently --- .../aibridgedserver/aibridgedserver_test.go | 51 +++++++++++++ coderd/aibridgedserver/cost.go | 73 ++++++++++++++----- coderd/aibridgedserver/cost_internal_test.go | 53 +++++++++++++- 3 files changed, 156 insertions(+), 21 deletions(-) diff --git a/coderd/aibridgedserver/aibridgedserver_test.go b/coderd/aibridgedserver/aibridgedserver_test.go index f5641a0531c..6c8dfe8a5ca 100644 --- a/coderd/aibridgedserver/aibridgedserver_test.go +++ b/coderd/aibridgedserver/aibridgedserver_test.go @@ -7,6 +7,7 @@ import ( "database/sql" "encoding/json" "fmt" + "math" "net" "net/url" "strconv" @@ -1901,6 +1902,49 @@ func TestRecordTokenUsage(t *testing.T) { prometheus.Labels{"provider": "anthropic", "model": "claude-sonnet-4-6"})) }, }, + { + // Token counts large enough that the cost cannot be stored. The + // record is written anyway, with prices snapshotted and cost + // NULL. + name: "valid token usage with cost out of range", + expectErrorLog: true, + request: &proto.RecordTokenUsageRequest{ + InterceptionId: uuid.NewString(), + MsgId: "msg_123", + InputTokens: math.MaxInt64, + CreatedAt: timestamppb.Now(), + }, + setupMocks: func(t *testing.T, db *dbmock.MockStore, req *proto.RecordTokenUsageRequest) { + interceptionID, err := uuid.Parse(req.GetInterceptionId()) + assert.NoError(t, err, "parse interception UUID") + + intc := newTestInterception(interceptionID) + groupID := uuid.New() + group := &database.GetHighestGroupAIBudgetByUserRow{GroupID: groupID, SpendLimitMicros: 1_000_000_000} + // 2 micro-units per token against int64 max tokens exceeds + // the column, so the cost cannot be represented. + price := &database.AIModelPrice{InputPrice: sql.NullInt64{Int64: 2_000_000, Valid: true}} + expectTokenUsageCostLookups(db, intc, nil, group, nil, price) + + db.EXPECT().InTx(gomock.Any(), nil).DoAndReturn( + func(fn func(database.Store) error, _ *database.TxOptions) error { return fn(db) }, + ) + + db.EXPECT().InsertAIBridgeTokenUsage(gomock.Any(), gomock.Cond(func(p database.InsertAIBridgeTokenUsageParams) bool { + // Prices and tokens are populated even though cost is NULL. + if !assert.Equal(t, uuid.NullUUID{UUID: groupID, Valid: true}, p.EffectiveGroupID, "effective group ID") || + !assert.True(t, p.InputPriceMicros.Valid, "input price populated") || + !assert.False(t, p.CostMicros.Valid, "cost null") || + !assert.Equal(t, int64(math.MaxInt64), p.InputTokens, "input tokens recorded") { + return false + } + return true + })).Return(database.AIBridgeTokenUsage{ID: uuid.New(), InterceptionID: interceptionID}, nil) + + // Spend update is skipped because cost is NULL. + db.EXPECT().IncrementUserAIDailySpend(gomock.Any(), gomock.Any()).Times(0) + }, + }, { // Price row exists with NULL columns, so cost is 0 (Valid). name: "valid token usage with effective group and NULL prices", @@ -3220,6 +3264,10 @@ type testRecordMethodCase[Req any] struct { // assertMetrics, when set, is called after the method returns to assert // the metrics recorded on the server's registry. assertMetrics func(t *testing.T, reg *prometheus.Registry) + // expectErrorLog tolerates ERROR-level logs, which slogtest otherwise + // treats as a test failure. Set it only for cases whose expected behavior + // includes logging an error, so every other case stays strict. + expectErrorLog bool } // testRecordMethod is a helper that abstracts the common testing pattern for all Record* methods. @@ -3237,6 +3285,9 @@ func testRecordMethod[Req any, Resp any]( ctrl := gomock.NewController(t) db := dbmock.NewMockStore(ctrl) logger := testutil.Logger(t) + if tc.expectErrorLog { + logger = slogtest.Make(t, &slogtest.Options{IgnoreErrors: true}).Leveled(slog.LevelDebug) + } if tc.setupMocks != nil { tc.setupMocks(t, db, tc.request) diff --git a/coderd/aibridgedserver/cost.go b/coderd/aibridgedserver/cost.go index 9bd280b53c5..bdc06553981 100644 --- a/coderd/aibridgedserver/cost.go +++ b/coderd/aibridgedserver/cost.go @@ -4,8 +4,10 @@ import ( "context" "database/sql" "errors" + "math" "github.com/google/uuid" + "github.com/shopspring/decimal" "golang.org/x/xerrors" "cdr.dev/slog/v3" @@ -15,9 +17,18 @@ import ( "github.com/coder/coder/v2/codersdk" ) -// tokensPerMillion is the divisor for prices, which are quoted per million -// tokens. -const tokensPerMillion = 1_000_000 +var ( + // tokensPerMillion is the divisor for prices, which are quoted per million + // tokens. + tokensPerMillion = decimal.NewFromInt(1_000_000) + // maxCostMicros is the largest cost the database column can hold. + maxCostMicros = decimal.NewFromInt(math.MaxInt64) +) + +// errCostOutOfRange reports a cost that cannot be stored in the database column. +// Real usage cannot reach it, so it means a wrong price row or implausible +// provider-reported token counts. +var errCostOutOfRange = xerrors.New("computed cost is out of range") // tokenUsageCost holds the cost-attribution columns snapshotted onto a token // usage record. A field left unset (Valid == false) is recorded as SQL NULL; a @@ -34,10 +45,10 @@ type tokenUsageCost struct { } // resolveTokenUsageCost resolves the effective group and per-token prices for an -// interception and computes its cost. Two independent conditions yield a NULL -// column rather than an error: an unresolved effective group (the user has no -// org membership), and a model absent from the price table leaves prices and -// cost NULL (a NULL cost unambiguously means "model not priced"). +// interception and computes its cost. Three conditions yield a NULL column +// rather than an error: an unresolved effective group (the user has no org +// membership), a model absent from the price table, and a cost outside the +// maxCostMicros range. A NULL cost means the cost is unknown. // Any other error is returned. func (s *Server) resolveTokenUsageCost(ctx context.Context, intc database.AIBridgeInterception, in *proto.RecordTokenUsageRequest) (tokenUsageCost, error) { var result tokenUsageCost @@ -85,12 +96,24 @@ func (s *Server) resolveTokenUsageCost(ctx context.Context, intc database.AIBrid result.outputPriceMicros = price.OutputPrice result.cacheReadPriceMicros = price.CacheReadPrice result.cacheWritePriceMicros = price.CacheWritePrice - result.costMicros = sql.NullInt64{ - Int64: computeCost(price, - in.GetInputTokens(), in.GetOutputTokens(), - in.GetCacheReadInputTokens(), in.GetCacheWriteInputTokens()), - Valid: true, + + costMicros, err := computeCost(price, + in.GetInputTokens(), in.GetOutputTokens(), + in.GetCacheReadInputTokens(), in.GetCacheWriteInputTokens()) + if err != nil { + // No trustworthy cost exists, so record it as unknown rather than + // storing a figure derived from bad inputs. The token counts are logged + // because the range error alone does not say which input was wrong. + s.logger.Error(ctx, "cost out of range, recording token usage with NULL cost", + slog.F("provider", intc.Provider), slog.F("model", intc.Model), + slog.F("input_tokens", in.GetInputTokens()), + slog.F("output_tokens", in.GetOutputTokens()), + slog.F("cache_read_input_tokens", in.GetCacheReadInputTokens()), + slog.F("cache_write_input_tokens", in.GetCacheWriteInputTokens()), + slog.Error(err)) + return result, nil } + result.costMicros = sql.NullInt64{Int64: costMicros, Valid: true} return result, nil } @@ -98,17 +121,27 @@ func (s *Server) resolveTokenUsageCost(ctx context.Context, intc database.AIBrid // the per-token prices from the price table. Prices are expressed per million // tokens; a NULL price column is treated as zero (e.g. providers that do not // charge for cache writes). -func computeCost(price database.AIModelPrice, inputTokens, outputTokens, cacheReadTokens, cacheWriteTokens int64) int64 { - return tokenCost(inputTokens, price.InputPrice) + - tokenCost(outputTokens, price.OutputPrice) + - tokenCost(cacheReadTokens, price.CacheReadPrice) + - tokenCost(cacheWriteTokens, price.CacheWritePrice) +func computeCost(price database.AIModelPrice, inputTokens, outputTokens, cacheReadTokens, cacheWriteTokens int64) (int64, error) { + total := tokenCost(inputTokens, price.InputPrice). + Add(tokenCost(outputTokens, price.OutputPrice)). + Add(tokenCost(cacheReadTokens, price.CacheReadPrice)). + Add(tokenCost(cacheWriteTokens, price.CacheWritePrice)) + + // Rejecting the negative case here keeps it from reaching the + // cost_micros >= 0 check constraint, which would discard the whole record. + if total.IsNegative() || total.GreaterThan(maxCostMicros) { + return 0, xerrors.Errorf("cost %s micro-units: %w", total.String(), errCostOutOfRange) + } + return total.IntPart(), nil } // tokenCost returns tokens * price / 1,000,000, treating a NULL price as zero. -func tokenCost(tokens int64, pricePerMillion sql.NullInt64) int64 { +func tokenCost(tokens int64, pricePerMillion sql.NullInt64) decimal.Decimal { if !pricePerMillion.Valid { - return 0 + return decimal.Zero } - return tokens * pricePerMillion.Int64 / tokensPerMillion + return decimal.NewFromInt(tokens). + Mul(decimal.NewFromInt(pricePerMillion.Int64)). + Div(tokensPerMillion). + Truncate(0) } diff --git a/coderd/aibridgedserver/cost_internal_test.go b/coderd/aibridgedserver/cost_internal_test.go index a36d24e16e1..5285eb06047 100644 --- a/coderd/aibridgedserver/cost_internal_test.go +++ b/coderd/aibridgedserver/cost_internal_test.go @@ -2,6 +2,8 @@ package aibridgedserver import ( "database/sql" + "errors" + "math" "testing" "github.com/coder/coder/v2/coderd/database" @@ -17,6 +19,7 @@ func TestComputeCost(t *testing.T) { price database.AIModelPrice inputTokens, outputTokens, cacheReadTokens, cacheWriteTokens int64 want int64 + wantOutOfRange bool }{ { name: "all priced", @@ -109,13 +112,61 @@ func TestComputeCost(t *testing.T) { inputTokens: 122_000_000_000, // 122e9 * 75e6 = 9.15e18 < int64 max want: 9_150_000_000_000, }, + { + name: "cost exactly at the column ceiling is in range", + price: database.AIModelPrice{ + // 1 micro-unit per token, so cost equals the token count. + InputPrice: nullInt64(1_000_000), + }, + inputTokens: math.MaxInt64, + want: math.MaxInt64, + }, + { + name: "cost above the column ceiling is out of range", + price: database.AIModelPrice{ + InputPrice: nullInt64(2_000_000), // 2 micro-units per token + }, + inputTokens: math.MaxInt64, // 2 * int64 max + wantOutOfRange: true, + }, + { + // Each category fits on its own; only their sum exceeds the column, + // so the range check has to run on the total. + name: "sum of in-range categories above the ceiling is out of range", + price: database.AIModelPrice{ + InputPrice: nullInt64(1_000_000), + OutputPrice: nullInt64(1_000_000), + }, + inputTokens: math.MaxInt64/2 + 1, + outputTokens: math.MaxInt64/2 + 1, + wantOutOfRange: true, + }, + { + // The cost column forbids negatives, so an implausible token count + // is rejected here rather than failing the insert. + name: "negative cost is out of range", + price: database.AIModelPrice{ + InputPrice: nullInt64(3_000_000), + }, + inputTokens: -1_000_000, + wantOutOfRange: true, + }, } for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { t.Parallel() - got := computeCost(tt.price, tt.inputTokens, tt.outputTokens, tt.cacheReadTokens, tt.cacheWriteTokens) + got, err := computeCost(tt.price, tt.inputTokens, tt.outputTokens, tt.cacheReadTokens, tt.cacheWriteTokens) + if tt.wantOutOfRange { + if !errors.Is(err, errCostOutOfRange) { + t.Fatalf("computeCost error = %v, want errCostOutOfRange", err) + } + return + } + if err != nil { + t.Fatalf("computeCost error = %v, want nil", err) + } if got != tt.want { t.Fatalf("computeCost = %d, want %d", got, tt.want) } From 69b0fa04d9c9bb91163d288df1c884d3ffcd621e Mon Sep 17 00:00:00 2001 From: Yevhenii Shcherbina Date: Wed, 29 Jul 2026 13:38:59 -0400 Subject: [PATCH 02/10] fix: cap interception cost at $10M (#27651) --- coderd/aibridgedserver/cost.go | 9 ++--- coderd/aibridgedserver/cost_internal_test.go | 40 +++++++++++--------- 2 files changed, 26 insertions(+), 23 deletions(-) diff --git a/coderd/aibridgedserver/cost.go b/coderd/aibridgedserver/cost.go index bdc06553981..3a22bd35a82 100644 --- a/coderd/aibridgedserver/cost.go +++ b/coderd/aibridgedserver/cost.go @@ -4,7 +4,6 @@ import ( "context" "database/sql" "errors" - "math" "github.com/google/uuid" "github.com/shopspring/decimal" @@ -21,12 +20,12 @@ var ( // tokensPerMillion is the divisor for prices, which are quoted per million // tokens. tokensPerMillion = decimal.NewFromInt(1_000_000) - // maxCostMicros is the largest cost the database column can hold. - maxCostMicros = decimal.NewFromInt(math.MaxInt64) + // maxCostMicros bounds one interception's cost at $10M. + maxCostMicros = decimal.NewFromInt(10_000_000_000_000) ) -// errCostOutOfRange reports a cost that cannot be stored in the database column. -// Real usage cannot reach it, so it means a wrong price row or implausible +// errCostOutOfRange reports a cost outside [0, maxCostMicros]. Real +// usage cannot reach it, so it means a wrong price row or implausible // provider-reported token counts. var errCostOutOfRange = xerrors.New("computed cost is out of range") diff --git a/coderd/aibridgedserver/cost_internal_test.go b/coderd/aibridgedserver/cost_internal_test.go index 5285eb06047..28d76eb93e9 100644 --- a/coderd/aibridgedserver/cost_internal_test.go +++ b/coderd/aibridgedserver/cost_internal_test.go @@ -14,6 +14,9 @@ func TestComputeCost(t *testing.T) { nullInt64 := func(v int64) sql.NullInt64 { return sql.NullInt64{Int64: v, Valid: true} } + const oneMicroPerToken = 1_000_000 + bound := maxCostMicros.IntPart() + tests := []struct { name string price database.AIModelPrice @@ -113,32 +116,33 @@ func TestComputeCost(t *testing.T) { want: 9_150_000_000_000, }, { - name: "cost exactly at the column ceiling is in range", - price: database.AIModelPrice{ - // 1 micro-unit per token, so cost equals the token count. - InputPrice: nullInt64(1_000_000), - }, - inputTokens: math.MaxInt64, - want: math.MaxInt64, + name: "cost exactly at the bound is in range", + price: database.AIModelPrice{InputPrice: nullInt64(oneMicroPerToken)}, + inputTokens: bound, + want: bound, }, { - name: "cost above the column ceiling is out of range", - price: database.AIModelPrice{ - InputPrice: nullInt64(2_000_000), // 2 micro-units per token - }, - inputTokens: math.MaxInt64, // 2 * int64 max + name: "cost one micro-unit above the bound is out of range", + price: database.AIModelPrice{InputPrice: nullInt64(oneMicroPerToken)}, + inputTokens: bound + 1, + wantOutOfRange: true, + }, + { + name: "cost of int64 max is out of range", + price: database.AIModelPrice{InputPrice: nullInt64(oneMicroPerToken)}, + inputTokens: math.MaxInt64, wantOutOfRange: true, }, { - // Each category fits on its own; only their sum exceeds the column, + // Each category fits on its own; only their sum exceeds the bound, // so the range check has to run on the total. - name: "sum of in-range categories above the ceiling is out of range", + name: "sum of in-range categories above the bound is out of range", price: database.AIModelPrice{ - InputPrice: nullInt64(1_000_000), - OutputPrice: nullInt64(1_000_000), + InputPrice: nullInt64(oneMicroPerToken), + OutputPrice: nullInt64(oneMicroPerToken), }, - inputTokens: math.MaxInt64/2 + 1, - outputTokens: math.MaxInt64/2 + 1, + inputTokens: bound/2 + 1, + outputTokens: bound/2 + 1, wantOutOfRange: true, }, { From 6bd034c337d9f41805d86be7a75dffec1370d077 Mon Sep 17 00:00:00 2001 From: Yevhenii Shcherbina Date: Wed, 29 Jul 2026 17:54:44 +0000 Subject: [PATCH 03/10] test: pin per-category cost truncation --- coderd/aibridgedserver/cost.go | 4 ++++ coderd/aibridgedserver/cost_internal_test.go | 12 ++++++++++++ 2 files changed, 16 insertions(+) diff --git a/coderd/aibridgedserver/cost.go b/coderd/aibridgedserver/cost.go index 3a22bd35a82..f667b8b87ae 100644 --- a/coderd/aibridgedserver/cost.go +++ b/coderd/aibridgedserver/cost.go @@ -135,6 +135,10 @@ func computeCost(price database.AIModelPrice, inputTokens, outputTokens, cacheRe } // tokenCost returns tokens * price / 1,000,000, treating a NULL price as zero. +// +// Truncate(0) belongs here rather than on the summed total: truncating each +// category is what makes a per-category breakdown recomputed from the +// snapshotted price columns add up to the stored cost. func tokenCost(tokens int64, pricePerMillion sql.NullInt64) decimal.Decimal { if !pricePerMillion.Valid { return decimal.Zero diff --git a/coderd/aibridgedserver/cost_internal_test.go b/coderd/aibridgedserver/cost_internal_test.go index 28d76eb93e9..627c4f73eba 100644 --- a/coderd/aibridgedserver/cost_internal_test.go +++ b/coderd/aibridgedserver/cost_internal_test.go @@ -115,6 +115,18 @@ func TestComputeCost(t *testing.T) { inputTokens: 122_000_000_000, // 122e9 * 75e6 = 9.15e18 < int64 max want: 9_150_000_000_000, }, + { + // Each category costs 37.5 micro-units, so truncating per category + // gives 37 + 37 = 74. + name: "each category truncates before the sum", + price: database.AIModelPrice{ + InputPrice: nullInt64(37_500_000), + OutputPrice: nullInt64(37_500_000), + }, + inputTokens: 1, + outputTokens: 1, + want: 74, + }, { name: "cost exactly at the bound is in range", price: database.AIModelPrice{InputPrice: nullInt64(oneMicroPerToken)}, From 9286dbefecf91dbc27fee03d8665093b5c8ca8b9 Mon Sep 17 00:00:00 2001 From: Yevhenii Shcherbina Date: Wed, 29 Jul 2026 18:44:12 +0000 Subject: [PATCH 04/10] fix: identify the record in the out-of-range cost log --- coderd/aibridgedserver/cost.go | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/coderd/aibridgedserver/cost.go b/coderd/aibridgedserver/cost.go index f667b8b87ae..e614d330cc8 100644 --- a/coderd/aibridgedserver/cost.go +++ b/coderd/aibridgedserver/cost.go @@ -101,9 +101,10 @@ func (s *Server) resolveTokenUsageCost(ctx context.Context, intc database.AIBrid in.GetCacheReadInputTokens(), in.GetCacheWriteInputTokens()) if err != nil { // No trustworthy cost exists, so record it as unknown rather than - // storing a figure derived from bad inputs. The token counts are logged - // because the range error alone does not say which input was wrong. + // storing a figure derived from bad inputs. s.logger.Error(ctx, "cost out of range, recording token usage with NULL cost", + slog.F("interception_id", intc.ID), + slog.F("initiator_id", intc.InitiatorID), slog.F("provider", intc.Provider), slog.F("model", intc.Model), slog.F("input_tokens", in.GetInputTokens()), slog.F("output_tokens", in.GetOutputTokens()), From 1149fef8b61d0d52e236077f6a26dee5810ba8f2 Mon Sep 17 00:00:00 2001 From: Yevhenii Shcherbina Date: Wed, 29 Jul 2026 19:04:39 +0000 Subject: [PATCH 05/10] refactor: divide token cost with QuoRem --- coderd/aibridgedserver/cost.go | 11 +++++------ 1 file changed, 5 insertions(+), 6 deletions(-) diff --git a/coderd/aibridgedserver/cost.go b/coderd/aibridgedserver/cost.go index e614d330cc8..71dbed863fb 100644 --- a/coderd/aibridgedserver/cost.go +++ b/coderd/aibridgedserver/cost.go @@ -137,15 +137,14 @@ func computeCost(price database.AIModelPrice, inputTokens, outputTokens, cacheRe // tokenCost returns tokens * price / 1,000,000, treating a NULL price as zero. // -// Truncate(0) belongs here rather than on the summed total: truncating each -// category is what makes a per-category breakdown recomputed from the -// snapshotted price columns add up to the stored cost. +// Each category is divided and truncated on its own, which makes a per-category breakdown +// recomputed from the snapshotted price columns add up to the stored cost. func tokenCost(tokens int64, pricePerMillion sql.NullInt64) decimal.Decimal { if !pricePerMillion.Valid { return decimal.Zero } - return decimal.NewFromInt(tokens). + quotient, _ := decimal.NewFromInt(tokens). Mul(decimal.NewFromInt(pricePerMillion.Int64)). - Div(tokensPerMillion). - Truncate(0) + QuoRem(tokensPerMillion, 0) + return quotient } From 0acc49a564d59a596428cfbbfe04488dcef53ddc Mon Sep 17 00:00:00 2001 From: Yevhenii Shcherbina Date: Wed, 29 Jul 2026 16:32:20 -0400 Subject: [PATCH 06/10] fix: validate reported token usage (#27660) --- coderd/aibridgedserver/aibridgedserver.go | 4 ++ .../aibridgedserver/aibridgedserver_test.go | 40 +++++++++++++++---- coderd/aibridgedserver/cost.go | 28 +++++++++++++ 3 files changed, 64 insertions(+), 8 deletions(-) diff --git a/coderd/aibridgedserver/aibridgedserver.go b/coderd/aibridgedserver/aibridgedserver.go index c9906bd35aa..7d450d436ff 100644 --- a/coderd/aibridgedserver/aibridgedserver.go +++ b/coderd/aibridgedserver/aibridgedserver.go @@ -330,6 +330,10 @@ func (s *Server) RecordTokenUsage(ctx context.Context, in *proto.RecordTokenUsag return nil, xerrors.Errorf("failed to parse interception_id %q: %w", in.GetInterceptionId(), err) } + if err := validateTokenUsage(in); err != nil { + return nil, xerrors.Errorf("validate token usage for interception %q: %w", intcID, err) + } + metadata := metadataToMap(in.GetMetadata()) if s.structuredLogging { diff --git a/coderd/aibridgedserver/aibridgedserver_test.go b/coderd/aibridgedserver/aibridgedserver_test.go index 6c8dfe8a5ca..ec8b2004d85 100644 --- a/coderd/aibridgedserver/aibridgedserver_test.go +++ b/coderd/aibridgedserver/aibridgedserver_test.go @@ -1903,15 +1903,39 @@ func TestRecordTokenUsage(t *testing.T) { }, }, { - // Token counts large enough that the cost cannot be stored. The - // record is written anyway, with prices snapshotted and cost - // NULL. + // Implausible counts are rejected before any DB work, so no row + // is written. Persisting them would poison the organization + // spend export, whose SUM cast raises rather than wraps. + name: "token usage above the allowed range is rejected", + request: &proto.RecordTokenUsageRequest{ + InterceptionId: uuid.NewString(), + MsgId: "msg_123", + InputTokens: math.MaxInt64, + CreatedAt: timestamppb.Now(), + }, + expectedErr: "reported token usage is out of range", + }, + { + name: "negative token usage is rejected", + request: &proto.RecordTokenUsageRequest{ + InterceptionId: uuid.NewString(), + MsgId: "msg_123", + InputTokens: -1_000_000, + OutputTokens: 2_000_000, + CreatedAt: timestamppb.Now(), + }, + expectedErr: "reported token usage is out of range", + }, + { + // Plausible token counts against a price row six orders of + // magnitude too high. The record is written anyway, with prices + // snapshotted and cost NULL. name: "valid token usage with cost out of range", expectErrorLog: true, request: &proto.RecordTokenUsageRequest{ InterceptionId: uuid.NewString(), MsgId: "msg_123", - InputTokens: math.MaxInt64, + InputTokens: 1_000_000, CreatedAt: timestamppb.Now(), }, setupMocks: func(t *testing.T, db *dbmock.MockStore, req *proto.RecordTokenUsageRequest) { @@ -1921,9 +1945,9 @@ func TestRecordTokenUsage(t *testing.T) { intc := newTestInterception(interceptionID) groupID := uuid.New() group := &database.GetHighestGroupAIBudgetByUserRow{GroupID: groupID, SpendLimitMicros: 1_000_000_000} - // 2 micro-units per token against int64 max tokens exceeds - // the column, so the cost cannot be represented. - price := &database.AIModelPrice{InputPrice: sql.NullInt64{Int64: 2_000_000, Valid: true}} + // $20M per million tokens puts a 1M-token request well past + // the per-interception cost bound. + price := &database.AIModelPrice{InputPrice: sql.NullInt64{Int64: 20_000_000_000_000, Valid: true}} expectTokenUsageCostLookups(db, intc, nil, group, nil, price) db.EXPECT().InTx(gomock.Any(), nil).DoAndReturn( @@ -1935,7 +1959,7 @@ func TestRecordTokenUsage(t *testing.T) { if !assert.Equal(t, uuid.NullUUID{UUID: groupID, Valid: true}, p.EffectiveGroupID, "effective group ID") || !assert.True(t, p.InputPriceMicros.Valid, "input price populated") || !assert.False(t, p.CostMicros.Valid, "cost null") || - !assert.Equal(t, int64(math.MaxInt64), p.InputTokens, "input tokens recorded") { + !assert.Equal(t, int64(1_000_000), p.InputTokens, "input tokens recorded") { return false } return true diff --git a/coderd/aibridgedserver/cost.go b/coderd/aibridgedserver/cost.go index 71dbed863fb..859bac86d6d 100644 --- a/coderd/aibridgedserver/cost.go +++ b/coderd/aibridgedserver/cost.go @@ -16,6 +16,11 @@ import ( "github.com/coder/coder/v2/codersdk" ) +// maxAllowedTokenUsage bounds the token count an interception may report per +// category. A 1M-token context is the current frontier, so this leaves six +// orders of magnitude of headroom. +const maxAllowedTokenUsage = 1_000_000_000_000 + var ( // tokensPerMillion is the divisor for prices, which are quoted per million // tokens. @@ -24,11 +29,34 @@ var ( maxCostMicros = decimal.NewFromInt(10_000_000_000_000) ) +// errTokenUsageOutOfRange reports a token count outside [0, maxAllowedTokenUsage]. +var errTokenUsageOutOfRange = xerrors.New("reported token usage is out of range") + // errCostOutOfRange reports a cost outside [0, maxCostMicros]. Real // usage cannot reach it, so it means a wrong price row or implausible // provider-reported token counts. var errCostOutOfRange = xerrors.New("computed cost is out of range") +// validateTokenUsage rejects an interception whose reported token counts fall +// outside [0, maxAllowedTokenUsage]. +func validateTokenUsage(in *proto.RecordTokenUsageRequest) error { + for _, category := range []struct { + name string + count int64 + }{ + {"input_tokens", in.GetInputTokens()}, + {"output_tokens", in.GetOutputTokens()}, + {"cache_read_input_tokens", in.GetCacheReadInputTokens()}, + {"cache_write_input_tokens", in.GetCacheWriteInputTokens()}, + } { + if category.count < 0 || category.count > maxAllowedTokenUsage { + return xerrors.Errorf("%s is %d, outside [0, %d]: %w", + category.name, category.count, maxAllowedTokenUsage, errTokenUsageOutOfRange) + } + } + return nil +} + // tokenUsageCost holds the cost-attribution columns snapshotted onto a token // usage record. A field left unset (Valid == false) is recorded as SQL NULL; a // price or cost of 0 is recorded as 0, which is distinct from NULL. From 4c2dc74f343c4d2505e04244a0d79a96cc4b61b4 Mon Sep 17 00:00:00 2001 From: Yevhenii Shcherbina Date: Wed, 29 Jul 2026 20:49:18 +0000 Subject: [PATCH 07/10] fix: log discarded token usage records --- coderd/aibridgedserver/aibridgedserver.go | 15 +++-- .../aibridgedserver/aibridgedserver_test.go | 6 +- coderd/aibridgedserver/cost_internal_test.go | 67 +++++++++++++++++++ 3 files changed, 82 insertions(+), 6 deletions(-) diff --git a/coderd/aibridgedserver/aibridgedserver.go b/coderd/aibridgedserver/aibridgedserver.go index 7d450d436ff..7b9895cbb69 100644 --- a/coderd/aibridgedserver/aibridgedserver.go +++ b/coderd/aibridgedserver/aibridgedserver.go @@ -330,10 +330,6 @@ func (s *Server) RecordTokenUsage(ctx context.Context, in *proto.RecordTokenUsag return nil, xerrors.Errorf("failed to parse interception_id %q: %w", in.GetInterceptionId(), err) } - if err := validateTokenUsage(in); err != nil { - return nil, xerrors.Errorf("validate token usage for interception %q: %w", intcID, err) - } - metadata := metadataToMap(in.GetMetadata()) if s.structuredLogging { @@ -350,6 +346,17 @@ func (s *Server) RecordTokenUsage(ctx context.Context, in *proto.RecordTokenUsag ) } + if err := validateTokenUsage(in); err != nil { + s.logger.Error(ctx, "implausible token usage, discarding record", + slog.F("interception_id", intcID), + slog.F("input_tokens", in.GetInputTokens()), + slog.F("output_tokens", in.GetOutputTokens()), + slog.F("cache_read_input_tokens", in.GetCacheReadInputTokens()), + slog.F("cache_write_input_tokens", in.GetCacheWriteInputTokens()), + slog.Error(err)) + return nil, xerrors.Errorf("validate token usage for interception %q: %w", intcID, err) + } + out, err := json.Marshal(metadata) if err != nil { s.logger.Warn(ctx, "failed to marshal aibridge metadata from proto to JSON", slog.F("metadata", in), slog.Error(err)) diff --git a/coderd/aibridgedserver/aibridgedserver_test.go b/coderd/aibridgedserver/aibridgedserver_test.go index ec8b2004d85..f25312cdf25 100644 --- a/coderd/aibridgedserver/aibridgedserver_test.go +++ b/coderd/aibridgedserver/aibridgedserver_test.go @@ -1906,7 +1906,8 @@ func TestRecordTokenUsage(t *testing.T) { // Implausible counts are rejected before any DB work, so no row // is written. Persisting them would poison the organization // spend export, whose SUM cast raises rather than wraps. - name: "token usage above the allowed range is rejected", + name: "token usage above the allowed range is rejected", + expectErrorLog: true, request: &proto.RecordTokenUsageRequest{ InterceptionId: uuid.NewString(), MsgId: "msg_123", @@ -1916,7 +1917,8 @@ func TestRecordTokenUsage(t *testing.T) { expectedErr: "reported token usage is out of range", }, { - name: "negative token usage is rejected", + name: "negative token usage is rejected", + expectErrorLog: true, request: &proto.RecordTokenUsageRequest{ InterceptionId: uuid.NewString(), MsgId: "msg_123", diff --git a/coderd/aibridgedserver/cost_internal_test.go b/coderd/aibridgedserver/cost_internal_test.go index 627c4f73eba..ba387c38754 100644 --- a/coderd/aibridgedserver/cost_internal_test.go +++ b/coderd/aibridgedserver/cost_internal_test.go @@ -6,6 +6,7 @@ import ( "math" "testing" + "github.com/coder/coder/v2/coderd/aibridged/proto" "github.com/coder/coder/v2/coderd/database" ) @@ -189,3 +190,69 @@ func TestComputeCost(t *testing.T) { }) } } + +func TestValidateTokenUsage(t *testing.T) { + t.Parallel() + + bound := int64(maxAllowedTokenUsage) + + tests := []struct { + name string + request *proto.RecordTokenUsageRequest + wantOutOfRange bool + }{ + { + // A frontier-sized request, with one category at zero. + name: "plausible counts", + request: &proto.RecordTokenUsageRequest{ + InputTokens: 1_000_000, OutputTokens: 128_000, + CacheReadInputTokens: 500_000, + }, + }, + { + // The bound is inclusive. + name: "every category exactly at the bound", + request: &proto.RecordTokenUsageRequest{ + InputTokens: bound, OutputTokens: bound, + CacheReadInputTokens: bound, CacheWriteInputTokens: bound, + }, + }, + { + name: "above the bound", + request: &proto.RecordTokenUsageRequest{InputTokens: bound + 1}, + wantOutOfRange: true, + }, + { + // The last category is checked too, not just the first. + name: "negative cache write", + request: &proto.RecordTokenUsageRequest{CacheWriteInputTokens: -1}, + wantOutOfRange: true, + }, + { + // A negative offset by a larger positive still totals in range, so + // the check has to run per category rather than on the sum. + name: "negative input offset by positive output", + request: &proto.RecordTokenUsageRequest{ + InputTokens: -1_000_000, OutputTokens: 2_000_000, + }, + wantOutOfRange: true, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + + err := validateTokenUsage(tt.request) + if tt.wantOutOfRange { + if !errors.Is(err, errTokenUsageOutOfRange) { + t.Fatalf("validateTokenUsage error = %v, want errTokenUsageOutOfRange", err) + } + return + } + if err != nil { + t.Fatalf("validateTokenUsage error = %v, want nil", err) + } + }) + } +} From 5b8c0e0d27a0be91db2e6a0a390f7426ee6bb79d Mon Sep 17 00:00:00 2001 From: Yevhenii Shcherbina Date: Wed, 29 Jul 2026 22:14:45 +0000 Subject: [PATCH 08/10] fix: type maxAllowedTokenUsage as int64 --- coderd/aibridgedserver/cost.go | 2 +- coderd/aibridgedserver/cost_internal_test.go | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/coderd/aibridgedserver/cost.go b/coderd/aibridgedserver/cost.go index 859bac86d6d..134a27f0b1c 100644 --- a/coderd/aibridgedserver/cost.go +++ b/coderd/aibridgedserver/cost.go @@ -19,7 +19,7 @@ import ( // maxAllowedTokenUsage bounds the token count an interception may report per // category. A 1M-token context is the current frontier, so this leaves six // orders of magnitude of headroom. -const maxAllowedTokenUsage = 1_000_000_000_000 +const maxAllowedTokenUsage int64 = 1_000_000_000_000 var ( // tokensPerMillion is the divisor for prices, which are quoted per million diff --git a/coderd/aibridgedserver/cost_internal_test.go b/coderd/aibridgedserver/cost_internal_test.go index ba387c38754..be33059e23d 100644 --- a/coderd/aibridgedserver/cost_internal_test.go +++ b/coderd/aibridgedserver/cost_internal_test.go @@ -194,7 +194,7 @@ func TestComputeCost(t *testing.T) { func TestValidateTokenUsage(t *testing.T) { t.Parallel() - bound := int64(maxAllowedTokenUsage) + bound := maxAllowedTokenUsage tests := []struct { name string From 5aa541cd2bde3a93f5890797c7fa654b2a892f4e Mon Sep 17 00:00:00 2001 From: Yevhenii Shcherbina Date: Wed, 5 Aug 2026 17:50:15 +0000 Subject: [PATCH 09/10] doc: update docs for resolveTokenUsageCost func --- coderd/aibridgedserver/cost.go | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/coderd/aibridgedserver/cost.go b/coderd/aibridgedserver/cost.go index 562a874b08f..31ec1dbf70f 100644 --- a/coderd/aibridgedserver/cost.go +++ b/coderd/aibridgedserver/cost.go @@ -72,11 +72,11 @@ type tokenUsageCost struct { } // resolveTokenUsageCost resolves the effective group and per-token prices for an -// interception and computes its cost. Three independent conditions yield a NULL +// interception and computes its cost. Four independent conditions yield a NULL // column rather than an error: an unresolved effective group (the user has no // org membership), an interception whose provider name matches no configured -// provider, and a model absent from the price table. The latter two leave prices -// and cost NULL (a NULL cost unambiguously means "model not priced"). +// provider, a model absent from the price table, and a cost outside the +// maxCostMicros range. A NULL cost means the cost is unknown. // Any other error is returned. func (s *Server) resolveTokenUsageCost(ctx context.Context, intc database.AIBridgeInterception, in *proto.RecordTokenUsageRequest) (tokenUsageCost, error) { var result tokenUsageCost From 1792455b191a95bafb87866ef814fb35c05b4fc4 Mon Sep 17 00:00:00 2001 From: Yevhenii Shcherbina Date: Wed, 5 Aug 2026 18:18:04 +0000 Subject: [PATCH 10/10] refactor(coderd/aibridgedserver): extract validateTotalCost --- coderd/aibridgedserver/cost.go | 17 +++++++++++++---- 1 file changed, 13 insertions(+), 4 deletions(-) diff --git a/coderd/aibridgedserver/cost.go b/coderd/aibridgedserver/cost.go index 31ec1dbf70f..08b84aef6c8 100644 --- a/coderd/aibridgedserver/cost.go +++ b/coderd/aibridgedserver/cost.go @@ -174,14 +174,23 @@ func computeCost(price database.AIModelPrice, inputTokens, outputTokens, cacheRe Add(tokenCost(cacheReadTokens, price.CacheReadPrice)). Add(tokenCost(cacheWriteTokens, price.CacheWritePrice)) - // Rejecting the negative case here keeps it from reaching the - // cost_micros >= 0 check constraint, which would discard the whole record. - if total.IsNegative() || total.GreaterThan(maxCostMicros) { - return 0, xerrors.Errorf("cost %s micro-units: %w", total.String(), errCostOutOfRange) + if err := validateTotalCost(total); err != nil { + return 0, err } return total.IntPart(), nil } +// validateTotalCost rejects a computed cost outside [0, maxCostMicros]. +// +// Rejecting the negative case early keeps it from reaching the +// cost_micros >= 0 check constraint, which would discard the whole record. +func validateTotalCost(total decimal.Decimal) error { + if total.IsNegative() || total.GreaterThan(maxCostMicros) { + return xerrors.Errorf("cost %s micro-units: %w", total.String(), errCostOutOfRange) + } + return nil +} + // tokenCost returns tokens * price / 1,000,000, treating a NULL price as zero. // // Each category is divided and truncated on its own, which makes a per-category breakdown