-
Notifications
You must be signed in to change notification settings - Fork 1.5k
fix: detect out-of-range AI Gateway costs instead of wrapping silently #27602
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Changes from all commits
ca32a2e
69b0fa0
6bd034c
9286dbe
1149fef
0acc49a
4c2dc74
5b8c0e0
68609be
5aa541c
1792455
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -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 | ||
|
Contributor
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Nit [CRF-17] The field's only effect is
|
||
| } | ||
|
|
||
| // 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 { | ||
|
Contributor
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. P3 [CRF-2] The ERROR log is the only operator-facing output of the new behavior, and no test asserts it exists. (Netero)
This file already has the pattern: On the blanket
Contributor
Author
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. I think it's minor. |
||
| logger = slogtest.Make(t, &slogtest.Options{IgnoreErrors: true}).Leveled(slog.LevelDebug) | ||
| } | ||
|
|
||
| if tc.setupMocks != nil { | ||
| tc.setupMocks(t, db, tc.request) | ||
|
|
||
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -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]. | ||
|
evgeniy-scherbina marked this conversation as resolved.
|
||
| 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,30 +142,65 @@ 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 { | ||
|
Contributor
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. P3 [CRF-12] The caller converts every
Today The file already has the shape twelve lines up, where the price lookup separates the known case from the rest. switch {
case errors.Is(err, errCostOutOfRange):
// 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", ...)
return result, nil
case err != nil:
return tokenUsageCost{}, xerrors.Errorf("compute cost for %s/%s: %w", intc.Provider, intc.Model, err)
}Meruem's alternative is equally valid and cheaper: if discrimination is not wanted, drop the sentinel, because an unchecked sentinel is a distinction the program does not make.
Contributor
Author
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Makes sense, but I don't expect computeCost to change much. In any case, if computeCost fails - there is nothing better than set cost to NULL, but don't fail recording. |
||
| // No trustworthy cost exists, so record it as unknown rather than | ||
|
Contributor
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Nit [CRF-22] The second sentence narrates the six log fields directly beneath it. (Gon P2) "The token counts are logged because the range error alone does not say which input was wrong" tells the reader that lines 109 to 112 log token counts, which those lines already show, wrapped in a rationale that is the obvious purpose of a diagnostic log. The first sentence carries the actual why-not-what, the decision to record NULL rather than return an error or store a figure. // No trustworthy cost exists, so record it as unknown rather than storing a
// figure derived from bad inputs.(CRF-11 argues the fields themselves are incomplete, which is the substantive half of this statement.)
|
||
| // storing a figure derived from bad inputs. | ||
| s.logger.Error(ctx, "cost out of range, recording token usage with NULL cost", | ||
|
Contributor
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. P2 [CRF-11] The The comment two lines up says token counts are logged "because the range error alone does not say which input was wrong", which is the right instinct. But Leorio traced the case that matters: a price row six orders of magnitude high makes every token count in the log entry look completely normal. The operator sees plausible tokens, a 20-digit total, and no reason to suspect the price table. Chopper and Ryosuke add the other half. No field identifies the affected record, so an operator holding "cost out of range for anthropic/claude-sonnet-4-6" cannot reach the This matters more here than for the sibling unpriced Debug log at line 85, which omits the ID too, because that path is a benign recurring condition with a metric behind it while this one is a per-record anomaly whose entire purpose is to be investigated.
Contributor
Author
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. I added interception_id and initiator_id (9286dbe). The pricing is defined in
Contributor
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. P3 [CRF-1] This is the only NULL-cost outcome with no metric, so the condition the PR exists to surface is detectable only by log scraping. (Netero) The sibling NULL-cost condition four blocks up increments a counter (
Contributor
Author
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. I don't think we need a metric for this extremely rare scenario. |
||
| slog.F("interception_id", intc.ID), | ||
| slog.F("initiator_id", intc.InitiatorID), | ||
| slog.F("provider", intc.Provider), slog.F("model", intc.Model), | ||
|
Contributor
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Another metric here as well. In fact, thinking about it better, we should have a metric for all the NULL cases 🤔 we already have for the unpriced model, but we should have for the other 2 cases.
Contributor
Author
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. I'll create a follow-up |
||
| 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 | ||
| } | ||
|
|
||
| // computeCost returns the cost of an interception in micro-units, snapshotting | ||
| // 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). | ||
|
ssncferreira marked this conversation as resolved.
Contributor
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Note [CRF-19] This subsystem and chatd now agree on the arithmetic and disagree on the rounding. (Zoro) Not CRF-3, which is the missing overflow guard in chatd. This is the rounding policy. Both are defensible in isolation. Worth knowing before anyone reconciles AI Gateway spend against chat spend, or unifies the two functions on the assumption that they are the same computation.
Contributor
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Note [CRF-20] The per-category reconciliation the description justifies has no query yet, and the natural form of it overflows in the one language you just left. (Knuckle) The PR justifies per-category truncation by saying a breakdown recomputed from the snapshotted price columns sums exactly to the stored total. No such query exists: When someone writes it, the natural form is
|
||
| 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 | ||
| } | ||
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
It might be useful to include provider and model here as well.
Additionally, it might make sense to create a metric for this, with provider and model as well. Can be done in a follow-up PR.
Uh oh!
There was an error while loading. Please reload this page.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
providerandmodelvariables are not in scope. It's possible to refactor this, but it turns out that would require changes to both the code and the tests. I guess it's not worth it?