diff --git a/coderd/aibridgedserver/aibridgedserver.go b/coderd/aibridgedserver/aibridgedserver.go index 038286555f7..bf1beb83048 100644 --- a/coderd/aibridgedserver/aibridgedserver.go +++ b/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)) diff --git a/coderd/aibridgedserver/aibridgedserver_test.go b/coderd/aibridgedserver/aibridgedserver_test.go index 8fa4776fb92..42915b95892 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" @@ -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) diff --git a/coderd/aibridgedserver/cost.go b/coderd/aibridgedserver/cost.go index ba56412a455..08b84aef6c8 100644 --- a/coderd/aibridgedserver/cost.go +++ b/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 } diff --git a/coderd/aibridgedserver/cost_internal_test.go b/coderd/aibridgedserver/cost_internal_test.go index a36d24e16e1..be33059e23d 100644 --- a/coderd/aibridgedserver/cost_internal_test.go +++ b/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) + } + }) + } +}