Thanks to visit codestin.com
Credit goes to github.com

Skip to content

Commit dae41eb

Browse files
evgeniy-scherbinaevgeniy-scherbina
authored
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 #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.
1 parent 9dcb75c commit dae41eb

4 files changed

Lines changed: 315 additions & 20 deletions

File tree

‎coderd/aibridgedserver/aibridgedserver.go‎

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -347,6 +347,17 @@ func (s *Server) RecordTokenUsage(ctx context.Context, in *proto.RecordTokenUsag
347347
)
348348
}
349349

350+
if err := validateTokenUsage(in); err != nil {
351+
s.logger.Error(ctx, "implausible token usage, discarding record",
352+
slog.F("interception_id", intcID),
353+
slog.F("input_tokens", in.GetInputTokens()),
354+
slog.F("output_tokens", in.GetOutputTokens()),
355+
slog.F("cache_read_input_tokens", in.GetCacheReadInputTokens()),
356+
slog.F("cache_write_input_tokens", in.GetCacheWriteInputTokens()),
357+
slog.Error(err))
358+
return nil, xerrors.Errorf("validate token usage for interception %q: %w", intcID, err)
359+
}
360+
350361
out, err := json.Marshal(metadata)
351362
if err != nil {
352363
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‎

Lines changed: 77 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,7 @@ import (
77
"database/sql"
88
"encoding/json"
99
"fmt"
10+
"math"
1011
"net"
1112
"net/url"
1213
"strconv"
@@ -1912,6 +1913,75 @@ func TestRecordTokenUsage(t *testing.T) {
19121913
prometheus.Labels{"provider": "anthropic", "model": "claude-sonnet-4-6"}))
19131914
},
19141915
},
1916+
{
1917+
// Implausible counts are rejected before any DB work, so no row
1918+
// is written. Persisting them would poison the organization
1919+
// spend export, whose SUM cast raises rather than wraps.
1920+
name: "token usage above the allowed range is rejected",
1921+
expectErrorLog: true,
1922+
request: &proto.RecordTokenUsageRequest{
1923+
InterceptionId: uuid.NewString(),
1924+
MsgId: "msg_123",
1925+
InputTokens: math.MaxInt64,
1926+
CreatedAt: timestamppb.Now(),
1927+
},
1928+
expectedErr: "reported token usage is out of range",
1929+
},
1930+
{
1931+
name: "negative token usage is rejected",
1932+
expectErrorLog: true,
1933+
request: &proto.RecordTokenUsageRequest{
1934+
InterceptionId: uuid.NewString(),
1935+
MsgId: "msg_123",
1936+
InputTokens: -1_000_000,
1937+
OutputTokens: 2_000_000,
1938+
CreatedAt: timestamppb.Now(),
1939+
},
1940+
expectedErr: "reported token usage is out of range",
1941+
},
1942+
{
1943+
// Plausible token counts against a price row six orders of
1944+
// magnitude too high. The record is written anyway, with prices
1945+
// snapshotted and cost NULL.
1946+
name: "valid token usage with cost out of range",
1947+
expectErrorLog: true,
1948+
request: &proto.RecordTokenUsageRequest{
1949+
InterceptionId: uuid.NewString(),
1950+
MsgId: "msg_123",
1951+
InputTokens: 1_000_000,
1952+
CreatedAt: timestamppb.Now(),
1953+
},
1954+
setupMocks: func(t *testing.T, db *dbmock.MockStore, req *proto.RecordTokenUsageRequest) {
1955+
interceptionID, err := uuid.Parse(req.GetInterceptionId())
1956+
assert.NoError(t, err, "parse interception UUID")
1957+
1958+
intc := newTestInterception(interceptionID)
1959+
groupID := uuid.New()
1960+
group := &database.GetHighestGroupAIBudgetByUserRow{GroupID: groupID, SpendLimitMicros: 1_000_000_000}
1961+
// $20M per million tokens puts a 1M-token request well past
1962+
// the per-interception cost bound.
1963+
price := &database.AIModelPrice{InputPrice: sql.NullInt64{Int64: 20_000_000_000_000, Valid: true}}
1964+
expectTokenUsageCostLookups(db, intc, nil, group, nil, price)
1965+
1966+
db.EXPECT().InTx(gomock.Any(), nil).DoAndReturn(
1967+
func(fn func(database.Store) error, _ *database.TxOptions) error { return fn(db) },
1968+
)
1969+
1970+
db.EXPECT().InsertAIBridgeTokenUsage(gomock.Any(), gomock.Cond(func(p database.InsertAIBridgeTokenUsageParams) bool {
1971+
// Prices and tokens are populated even though cost is NULL.
1972+
if !assert.Equal(t, uuid.NullUUID{UUID: groupID, Valid: true}, p.EffectiveGroupID, "effective group ID") ||
1973+
!assert.True(t, p.InputPriceMicros.Valid, "input price populated") ||
1974+
!assert.False(t, p.CostMicros.Valid, "cost null") ||
1975+
!assert.Equal(t, int64(1_000_000), p.InputTokens, "input tokens recorded") {
1976+
return false
1977+
}
1978+
return true
1979+
})).Return(database.AIBridgeTokenUsage{ID: uuid.New(), InterceptionID: interceptionID}, nil)
1980+
1981+
// Spend update is skipped because cost is NULL.
1982+
db.EXPECT().IncrementUserAIDailySpend(gomock.Any(), gomock.Any()).Times(0)
1983+
},
1984+
},
19151985
{
19161986
// Price row exists with NULL columns, so cost is 0 (Valid).
19171987
name: "valid token usage with effective group and NULL prices",
@@ -3534,6 +3604,10 @@ type testRecordMethodCase[Req any] struct {
35343604
// assertMetrics, when set, is called after the method returns to assert
35353605
// the metrics recorded on the server's registry.
35363606
assertMetrics func(t *testing.T, reg *prometheus.Registry)
3607+
// expectErrorLog tolerates ERROR-level logs, which slogtest otherwise
3608+
// treats as a test failure. Set it only for cases whose expected behavior
3609+
// includes logging an error, so every other case stays strict.
3610+
expectErrorLog bool
35373611
}
35383612

35393613
// testRecordMethod is a helper that abstracts the common testing pattern for all Record* methods.
@@ -3551,6 +3625,9 @@ func testRecordMethod[Req any, Resp any](
35513625
ctrl := gomock.NewController(t)
35523626
db := dbmock.NewMockStore(ctrl)
35533627
logger := testutil.Logger(t)
3628+
if tc.expectErrorLog {
3629+
logger = slogtest.Make(t, &slogtest.Options{IgnoreErrors: true}).Leveled(slog.LevelDebug)
3630+
}
35543631

35553632
if tc.setupMocks != nil {
35563633
tc.setupMocks(t, db, tc.request)

‎coderd/aibridgedserver/cost.go‎

Lines changed: 92 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,7 @@ import (
66
"errors"
77

88
"github.com/google/uuid"
9+
"github.com/shopspring/decimal"
910
"golang.org/x/xerrors"
1011

1112
"cdr.dev/slog/v3"
@@ -15,9 +16,46 @@ import (
1516
"github.com/coder/coder/v2/codersdk"
1617
)
1718

18-
// tokensPerMillion is the divisor for prices, which are quoted per million
19-
// tokens.
20-
const tokensPerMillion = 1_000_000
19+
// maxAllowedTokenUsage bounds the token count an interception may report per
20+
// category. A 1M-token context is the current frontier, so this leaves six
21+
// orders of magnitude of headroom.
22+
const maxAllowedTokenUsage int64 = 1_000_000_000_000
23+
24+
var (
25+
// tokensPerMillion is the divisor for prices, which are quoted per million
26+
// tokens.
27+
tokensPerMillion = decimal.NewFromInt(1_000_000)
28+
// maxCostMicros bounds one interception's cost at $10M.
29+
maxCostMicros = decimal.NewFromInt(10_000_000_000_000)
30+
)
31+
32+
// errTokenUsageOutOfRange reports a token count outside [0, maxAllowedTokenUsage].
33+
var errTokenUsageOutOfRange = xerrors.New("reported token usage is out of range")
34+
35+
// errCostOutOfRange reports a cost outside [0, maxCostMicros]. Real
36+
// usage cannot reach it, so it means a wrong price row or implausible
37+
// provider-reported token counts.
38+
var errCostOutOfRange = xerrors.New("computed cost is out of range")
39+
40+
// validateTokenUsage rejects an interception whose reported token counts fall
41+
// outside [0, maxAllowedTokenUsage].
42+
func validateTokenUsage(in *proto.RecordTokenUsageRequest) error {
43+
for _, category := range []struct {
44+
name string
45+
count int64
46+
}{
47+
{"input_tokens", in.GetInputTokens()},
48+
{"output_tokens", in.GetOutputTokens()},
49+
{"cache_read_input_tokens", in.GetCacheReadInputTokens()},
50+
{"cache_write_input_tokens", in.GetCacheWriteInputTokens()},
51+
} {
52+
if category.count < 0 || category.count > maxAllowedTokenUsage {
53+
return xerrors.Errorf("%s is %d, outside [0, %d]: %w",
54+
category.name, category.count, maxAllowedTokenUsage, errTokenUsageOutOfRange)
55+
}
56+
}
57+
return nil
58+
}
2159

2260
// tokenUsageCost holds the cost-attribution columns snapshotted onto a token
2361
// usage record. A field left unset (Valid == false) is recorded as SQL NULL; a
@@ -34,11 +72,11 @@ type tokenUsageCost struct {
3472
}
3573

3674
// resolveTokenUsageCost resolves the effective group and per-token prices for an
37-
// interception and computes its cost. Three independent conditions yield a NULL
75+
// interception and computes its cost. Four independent conditions yield a NULL
3876
// column rather than an error: an unresolved effective group (the user has no
3977
// org membership), an interception whose provider name matches no configured
40-
// provider, and a model absent from the price table. The latter two leave prices
41-
// and cost NULL (a NULL cost unambiguously means "model not priced").
78+
// provider, a model absent from the price table, and a cost outside the
79+
// maxCostMicros range. A NULL cost means the cost is unknown.
4280
// Any other error is returned.
4381
func (s *Server) resolveTokenUsageCost(ctx context.Context, intc database.AIBridgeInterception, in *proto.RecordTokenUsageRequest) (tokenUsageCost, error) {
4482
var result tokenUsageCost
@@ -104,30 +142,65 @@ func (s *Server) resolveTokenUsageCost(ctx context.Context, intc database.AIBrid
104142
result.outputPriceMicros = price.OutputPrice
105143
result.cacheReadPriceMicros = price.CacheReadPrice
106144
result.cacheWritePriceMicros = price.CacheWritePrice
107-
result.costMicros = sql.NullInt64{
108-
Int64: computeCost(price,
109-
in.GetInputTokens(), in.GetOutputTokens(),
110-
in.GetCacheReadInputTokens(), in.GetCacheWriteInputTokens()),
111-
Valid: true,
145+
146+
costMicros, err := computeCost(price,
147+
in.GetInputTokens(), in.GetOutputTokens(),
148+
in.GetCacheReadInputTokens(), in.GetCacheWriteInputTokens())
149+
if err != nil {
150+
// No trustworthy cost exists, so record it as unknown rather than
151+
// storing a figure derived from bad inputs.
152+
s.logger.Error(ctx, "cost out of range, recording token usage with NULL cost",
153+
slog.F("interception_id", intc.ID),
154+
slog.F("initiator_id", intc.InitiatorID),
155+
slog.F("provider", intc.Provider), slog.F("model", intc.Model),
156+
slog.F("input_tokens", in.GetInputTokens()),
157+
slog.F("output_tokens", in.GetOutputTokens()),
158+
slog.F("cache_read_input_tokens", in.GetCacheReadInputTokens()),
159+
slog.F("cache_write_input_tokens", in.GetCacheWriteInputTokens()),
160+
slog.Error(err))
161+
return result, nil
112162
}
163+
result.costMicros = sql.NullInt64{Int64: costMicros, Valid: true}
113164
return result, nil
114165
}
115166

116167
// computeCost returns the cost of an interception in micro-units, snapshotting
117168
// the per-token prices from the price table. Prices are expressed per million
118169
// tokens; a NULL price column is treated as zero (e.g. providers that do not
119170
// charge for cache writes).
120-
func computeCost(price database.AIModelPrice, inputTokens, outputTokens, cacheReadTokens, cacheWriteTokens int64) int64 {
121-
return tokenCost(inputTokens, price.InputPrice) +
122-
tokenCost(outputTokens, price.OutputPrice) +
123-
tokenCost(cacheReadTokens, price.CacheReadPrice) +
124-
tokenCost(cacheWriteTokens, price.CacheWritePrice)
171+
func computeCost(price database.AIModelPrice, inputTokens, outputTokens, cacheReadTokens, cacheWriteTokens int64) (int64, error) {
172+
total := tokenCost(inputTokens, price.InputPrice).
173+
Add(tokenCost(outputTokens, price.OutputPrice)).
174+
Add(tokenCost(cacheReadTokens, price.CacheReadPrice)).
175+
Add(tokenCost(cacheWriteTokens, price.CacheWritePrice))
176+
177+
if err := validateTotalCost(total); err != nil {
178+
return 0, err
179+
}
180+
return total.IntPart(), nil
181+
}
182+
183+
// validateTotalCost rejects a computed cost outside [0, maxCostMicros].
184+
//
185+
// Rejecting the negative case early keeps it from reaching the
186+
// cost_micros >= 0 check constraint, which would discard the whole record.
187+
func validateTotalCost(total decimal.Decimal) error {
188+
if total.IsNegative() || total.GreaterThan(maxCostMicros) {
189+
return xerrors.Errorf("cost %s micro-units: %w", total.String(), errCostOutOfRange)
190+
}
191+
return nil
125192
}
126193

127194
// tokenCost returns tokens * price / 1,000,000, treating a NULL price as zero.
128-
func tokenCost(tokens int64, pricePerMillion sql.NullInt64) int64 {
195+
//
196+
// Each category is divided and truncated on its own, which makes a per-category breakdown
197+
// recomputed from the snapshotted price columns add up to the stored cost.
198+
func tokenCost(tokens int64, pricePerMillion sql.NullInt64) decimal.Decimal {
129199
if !pricePerMillion.Valid {
130-
return 0
200+
return decimal.Zero
131201
}
132-
return tokens * pricePerMillion.Int64 / tokensPerMillion
202+
quotient, _ := decimal.NewFromInt(tokens).
203+
Mul(decimal.NewFromInt(pricePerMillion.Int64)).
204+
QuoRem(tokensPerMillion, 0)
205+
return quotient
133206
}

0 commit comments

Comments
 (0)