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

Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
11 changes: 11 additions & 0 deletions coderd/aibridgedserver/aibridgedserver.go
Original file line number Diff line number Diff line change
Expand Up @@ -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()),

Copy link
Copy Markdown
Contributor

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.

@evgeniy-scherbina evgeniy-scherbina Aug 5, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

  1. The provider and model variables 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?
  2. I'll create a follow-up

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))
Expand Down
77 changes: 77 additions & 0 deletions coderd/aibridgedserver/aibridgedserver_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,7 @@ import (
"database/sql"
"encoding/json"
"fmt"
"math"
"net"
"net/url"
"strconv"
Expand Down Expand Up @@ -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",
Expand Down Expand Up @@ -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

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Nit [CRF-17] expectErrorLog names an assertion the field never makes. (Gon Nit, Leorio Nit)

The field's only effect is IgnoreErrors: true at line 3289. Nothing verifies an error log was emitted, and setting it on a case that logs nothing passes silently. The field's own doc gets the verb right on the first word, "tolerates", while the identifier says "expect", which in test-helper vocabulary means the test fails if it does not happen. The next person to write expectErrorLog: true will believe they have coverage of the log.

ignoreErrorLogs or tolerateErrorLogs describes what it does, and the doc then shortens to match the name instead of correcting it. Distinct from CRF-2: renaming fixes the mismatch without adding an assertion, and adding the assertion would not fix the name.

🤖

}

// testRecordMethod is a helper that abstracts the common testing pattern for all Record* methods.
Expand All @@ -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 {

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The 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)

expectErrorLog swaps in slogtest.Make(t, &slogtest.Options{IgnoreErrors: true}), which suppresses the failure slogtest would otherwise raise. It does not assert. Delete the entire s.logger.Error(...) block at cost.go:107-113 and the new case at line 1905 still passes: it checks CostMicros.Valid == false, that prices are populated, and that IncrementUserAIDailySpend is not called. Nothing observes the log. The token-count fields, which cost.go:105-106 names as the point of logging, are entirely untested.

This file already has the pattern: slog.Make(slogjson.Sink(buf)) plus parseLogLines and getLogLinesWithMessage at lines 3626 to 3663. Route the out-of-range case through a buffer sink and assert the message and the fields.

On the blanket IgnoreErrors: true substitution itself, Mafu-san and Zoro both checked before flagging it and did not: it appears in roughly forty test files across cli/, provisionersdk/ and scaletest/, twice already in this file, so it is the established pattern rather than a divergence. The cost is that it also drops the targeted flaky-error suppression testutil.Logger provides.

🤖

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The 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)
Expand Down
111 changes: 92 additions & 19 deletions coderd/aibridgedserver/cost.go
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,7 @@ import (
"errors"

"github.com/google/uuid"
"github.com/shopspring/decimal"
"golang.org/x/xerrors"

"cdr.dev/slog/v3"
Expand All @@ -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].
Comment thread
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
Expand All @@ -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
Expand Down Expand Up @@ -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 {

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P3 [CRF-12] The caller converts every computeCost error into "cost unknown, store NULL", so the sentinel is discriminated nowhere outside the test. (Chopper P3, Meruem P3, Zoro P3, Knov P3)

errCostOutOfRange exists so a caller can tell "unrepresentable, keep going" apart from "something else went wrong". The only errors.Is against it is cost_internal_test.go:162. Production does if err != nil and then logs a message that asserts a specific cause.

Today computeCost has exactly one failure mode, so nothing is misreported. What is wrong is the contract: the next error added to computeCost, for any reason, is silently downgraded to a NULL cost under a log line naming the wrong diagnosis, with no code change and no review signal. Chopper called it the unhandled middle inverted, handling all errors as one known error.

The file already has the shape twelve lines up, where the price lookup separates the known case from the rest. errors is already imported:

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.

🤖

@evgeniy-scherbina evgeniy-scherbina Jul 29, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The 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

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The 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",

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P2 [CRF-11] The errCostOutOfRange doc names two suspects; the log records only one of them, and identifies no record. (Leorio P2, Chopper P3, Ryosuke P3)

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 errCostOutOfRange's own doc says the condition means "a wrong price row or implausible provider-reported token counts", and only one of the two is in the log.

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. price is in scope from line 78.

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 aibridge_token_usages row, the interception, or the user whose spend is now understated; on a busy deployment the row has to be recovered by matching provider, model and token counts inside a timestamp window. Every other diagnostic in this package carries one: the skipped-spend Debug logs interception_id and initiator_id (aibridgedserver.go:409-410), the threshold-detection Error logs both (:435-436), the unresolved-group Warn logs user_id (cost.go:66).

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. intc.ID, intc.InitiatorID and the four price fields are all in scope at this call.

🤖

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The 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 prices.json and embedded into the binary, so it's unlikely to be incorrect. If needed, it can always be derived from prices.json by provider and model.

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The 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 (cost.go:87-89, s.metrics.UnpricedTokenUsageRecords). Two consequences. unpriced_token_usage_records_total, documented in docs/admin/integrations/prometheus.md:119 as counting records "for which no model price was found", no longer accounts for every NULL cost_micros, so a dashboard reconciling "records with NULL cost" against that counter silently disagrees. And the new condition has no numeric signal at all.

metrics.go:25-34 already carries provider/model-labelled counters, so this is one field plus one Inc() with the same labels. s.metrics is nil-checked at cost.go:87, so the same guard applies here.

🤖

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The 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),

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The 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.

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The 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).
Comment thread
ssncferreira marked this conversation as resolved.

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The 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. chatcost.CalculateTotalCostMicros:63 computes the same four-category sum and applies Ceil() once to the total; computeCost truncates each category. Identical token counts at identical prices therefore produce different *_cost_micros values depending on which product recorded them, by up to four micro-units.

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.

🤖

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The 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: input_price_micros and its siblings are write-only, referenced solely by InsertAIBridgeTokenUsage (queries/aibridge.sql:46).

When someone writes it, the natural form is input_tokens * input_price_micros / 1000000, and bigint * bigint raises exactly where the old Go code wrapped (verified: select 9223372036854775807::bigint * 2::bigint returns ERROR: bigint out of range). Whoever writes that query needs ::numeric on the multiply. One line in the column comment, while the reasoning is fresh, is cheaper than rebuilding the invariant in the one language that still fails at the same boundary.

🤖

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
}
Loading
Loading