fix: detect out-of-range AI Gateway costs instead of wrapping silently (#27602)

Implements:
https://linear.app/codercom/issue/AIGOV-448/use-decimal-for-cost-computation
Follow-up to https://github.com/coder/coder/pull/26229

Follow-up to the AI Gateway cost-control work. Cost is computed per
token category as `tokens × price / 1_000_000` in `int64`, then summed.
This change makes an unrepresentable result a defined outcome instead of
an accident of integer wrap-around.

## Motivation

The intermediate `tokens × price` can exceed `int64`. Real usage cannot
get there: at a $75/M model the product overflows at roughly 123 billion
tokens in a single response, about six orders of magnitude above a
maxed-out Opus request, so this is not a live incident. The problem is
what happens if it ever does, because the sign of the wrapped value
silently selects between two different failure modes, neither of which
was chosen:

1. **Wraps positive.** A plausible-looking cost is stored, incremented
into the user's daily spend, and enforced against their AI budget. No
error, no signal, wrong number.
2. **Wraps negative.** The value violates `CHECK (cost_micros >= 0)`,
the insert fails, the surrounding transaction rolls back, and
`RecordTokenUsage` returns a Postgres constraint error that says nothing
about overflow. The token usage record is lost entirely, along with its
token counts.

So the same class of bad input either corrupts budget accounting or
discards an audit record, depending on arithmetic that nobody reasoned
about. That is the undefined behaviour.

## Decision

**An unrepresentable cost is treated as bad input, not a large bill.**
Since real usage cannot produce one, it can only mean a wrong price row
or implausible provider-reported token counts. In both cases the true
cost is unknowable, so no number is stored.

**Detect rather than avoid.** `computeCost` now evaluates in `decimal`,
so nothing wraps, and range-checks the total against `[0, MaxInt64]`
before converting back. Out of range returns `errCostOutOfRange`.
Rejecting negatives in the same check also keeps them away from the
non-negative column constraint, which would otherwise discard the
record.

**Log, do not block.** The error is swallowed at the call site: the
record is written with token counts intact and `cost_micros` NULL, the
spend update is skipped, and the condition is logged at ERROR.

**Per-category truncation is unchanged.** Each category is still
truncated independently rather than the total being rounded once, so a
per-category breakdown recomputed from the snapshotted price columns
sums exactly to the stored total. Every existing `computeCost` test case
passes unmodified.

coderd/aibridgedserver/aibridgedserver.go

@@ -347,6 +347,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))

coderd/aibridgedserver/aibridgedserver_test.go

@@ -7,6 +7,7 @@ import (
 	"database/sql"
 	"encoding/json"
 	"fmt"
+	"math"
 	"net"
 	"net/url"
 	"strconv"
@@ -1912,6 +1913,75 @@ func TestRecordTokenUsage(t *testing.T) {
 						prometheus.Labels{"provider": "anthropic", "model": "claude-sonnet-4-6"}))
 				},
 			},
+			{
+				// 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",
+				expectErrorLog: true,
+				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",
+				expectErrorLog: true,
+				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:    1_000_000,
+					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}
+					// $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(
+						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(1_000_000), 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",
@@ -3534,6 +3604,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.
@@ -3551,6 +3625,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)

coderd/aibridgedserver/cost.go

@@ -6,6 +6,7 @@ import (
 	"errors"
 
 	"github.com/google/uuid"
+	"github.com/shopspring/decimal"
 	"golang.org/x/xerrors"
 
 	"cdr.dev/slog/v3"
@@ -15,9 +16,46 @@ import (
 	"github.com/coder/coder/v2/codersdk"
 )
 
-// tokensPerMillion is the divisor for prices, which are quoted per million
-// tokens.
-const tokensPerMillion = 1_000_000
+// 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 int64 = 1_000_000_000_000
+
+var (
+	// tokensPerMillion is the divisor for prices, which are quoted per million
+	// tokens.
+	tokensPerMillion = decimal.NewFromInt(1_000_000)
+	// maxCostMicros bounds one interception's cost at $10M.
+	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
@@ -34,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
@@ -104,12 +142,25 @@ 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.
+		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()),
+			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
 }
 
@@ -117,17 +168,39 @@ 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))
+
+	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.
-func tokenCost(tokens int64, pricePerMillion sql.NullInt64) int64 {
+//
+// 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 0
+		return decimal.Zero
 	}
-	return tokens * pricePerMillion.Int64 / tokensPerMillion
+	quotient, _ := decimal.NewFromInt(tokens).
+		Mul(decimal.NewFromInt(pricePerMillion.Int64)).
+		QuoRem(tokensPerMillion, 0)
+	return quotient
 }

coderd/aibridgedserver/cost_internal_test.go

@@ -2,8 +2,11 @@ package aibridgedserver
 
 import (
 	"database/sql"
+	"errors"
+	"math"
 	"testing"
 
+	"github.com/coder/coder/v2/coderd/aibridged/proto"
 	"github.com/coder/coder/v2/coderd/database"
 )
 
@@ -12,11 +15,15 @@ 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
 		inputTokens, outputTokens, cacheReadTokens, cacheWriteTokens int64
 		want                                                         int64
+		wantOutOfRange                                               bool
 	}{
 		{
 			name: "all priced",
@@ -109,16 +116,143 @@ 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)},
+			inputTokens: bound,
+			want:        bound,
+		},
+		{
+			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 bound,
+			// so the range check has to run on the total.
+			name: "sum of in-range categories above the bound is out of range",
+			price: database.AIModelPrice{
+				InputPrice:  nullInt64(oneMicroPerToken),
+				OutputPrice: nullInt64(oneMicroPerToken),
+			},
+			inputTokens:    bound/2 + 1,
+			outputTokens:   bound/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)
 			}
 		})
 	}
 }
+
+func TestValidateTokenUsage(t *testing.T) {
+	t.Parallel()
+
+	bound := 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)
+			}
+		})
+	}
+}