feat: record cost on aibridge token usages - #26229
Conversation
57c5717 to
a65787b
Compare
a65787b to
3792e85
Compare
|
/coder-agents-review |
|
Chat: Review in progress | View chat deep-review v0.7.1 | Round 1 | Last posted: Round 1, 8 findings (4 P3, 4 Nit), COMMENT. Review Finding inventoryFindings
Round logRound 1Panel: Bisky, Hisoka, Mafu-san, Mafuuu, Pariston, Chopper, Ging-Go, Gon, Knuckle, Kurapika, Leorio, Meruem, Knov. 4 P3, 4 Nit, 2 Note new. 2 Notes dropped. Reviewed against ad127d0..06993fe. About deep-reviewCRF = Coder Review Finding (P0-P4, Nit, Note)
|
There was a problem hiding this comment.
The cost-attribution design is solid. Point-in-time snapshots are the right call for financial audit records (query-time JOINs would break when prices update or groups are deleted). The budget.Store interface extraction is clean, the test density is good (66%), and the overflow analysis is thorough.
Severity count: 4 P3, 4 Nit.
The main structural concern is that tokenCost inherits signed int64 token counts from the proto with no validation, and calculateActualInputTokenUsage actively subtracts CachedTokens from PromptTokens. Before this PR, a negative count was harmlessly stored; after this PR, it triggers the cost_micros CHECK (>= 0) constraint, and the callers (5 of 6 call sites) discard the error with _ =. Silent data loss on a path that was previously safe.
The test matrix covers three of the four (budget, price) corners well but misses two combinations that anchor the independence of budget resolution and price lookup.
"If someone reordered the code to bail on missing price before resolving the budget, no test catches the regression." (Bisky)
coderd/database/dbgen/dbgen.go:2028
P3 [CRF-3] dbgen.AIBridgeTokenUsage does not forward the six new cost-attribution columns from the seed.
The InsertAIBridgeTokenUsageParams literal here omits EffectiveGroupID, InputPriceMicros, OutputPriceMicros, CacheReadPriceMicros, CacheWritePriceMicros, and CostMicros. The existing CacheReadInputTokens and CacheWriteInputTokens are forwarded via seed.CacheReadInputTokens, so the pattern is established.
Nothing breaks today (zero values map to SQL NULL, which is valid), but any future integration test that needs cost data on a fixture will silently get NULLs.
EffectiveGroupID: seed.EffectiveGroupID,
InputPriceMicros: seed.InputPriceMicros,
OutputPriceMicros: seed.OutputPriceMicros,
CacheReadPriceMicros: seed.CacheReadPriceMicros,
CacheWritePriceMicros: seed.CacheWritePriceMicros,
CostMicros: seed.CostMicros,(Knuckle)
🤖
🤖 This review was automatically generated with Coder Agents.
| []testRecordMethodCase[*proto.RecordTokenUsageRequest]{ | ||
| { | ||
| name: "valid token usage", | ||
| name: "valid token usage with null cost", |
There was a problem hiding this comment.
P3 [CRF-2] Test matrix for (budget, price) cross-product is missing two corners.
The test suite covers:
- (no budget, no price) -> NULL group, NULL cost ("valid token usage with null cost")
- (group, price) -> group set, cost computed ("valid token usage with cost")
- (override, price) -> override group, cost computed ("valid token usage with user override")
Missing:
- (budget, no price): verifies
effectiveGroupIDsurvives the early-return whenGetAIModelPriceByProviderModelreturnssql.ErrNoRows. Setup:expectTokenUsageCostLookups(db, intc, nil, group, nil). Assert:EffectiveGroupID.Valid == true,CostMicros.Valid == false, all price columnsValid == false. - (no budget, price): verifies cost is computed even without a budget. Setup:
expectTokenUsageCostLookups(db, intc, nil, nil, price). Assert:EffectiveGroupID.Valid == false,CostMicros.Valid == truewith expected value.
These anchor the stated independence of budget resolution and price lookup. If someone refactored to gate cost computation on budget presence, no test catches the regression. (Bisky P3, Chopper P3)
🤖
| ADD COLUMN effective_group_id UUID, | ||
| -- Snapshotted prices at interception time, in micro-units per million | ||
| -- tokens. NULL if the model is not present in ai_model_prices. | ||
| ADD COLUMN input_price_micros BIGINT CHECK (input_price_micros >= 0), |
There was a problem hiding this comment.
P3 [CRF-4] Snapshot columns add a _micros suffix that the source columns in ai_model_prices lack, for the same value.
ai_model_prices stores input_price; the snapshot here stores input_price_micros. Both hold the same micro-unit value with no conversion (result.inputPriceMicros = price.InputPrice at cost.go:67). A reader writing a JOIN or tracing a value between the two tables will assume _micros signals a unit difference.
The snapshot naming is arguably more correct (it declares the unit), and the source naming is the one missing context. This can't be fixed in this PR. Worth noting for consistency if the source table is ever touched. (Gon)
🤖
There was a problem hiding this comment.
Yeah, it's unfortunate that ai_model_prices is the only place where we don't follow the micros convention.
We can add a migration, but I'm not sure it's worth it.
| return result, nil | ||
| } | ||
|
|
||
| // computeCost returns the cost of an interception in micro-units, snapshotting |
There was a problem hiding this comment.
Nit [CRF-5] Doc comment says computeCost is "snapshotting the per-token prices from the price table." It does not snapshot anything; it multiplies token counts by prices that were already fetched. The snapshotting happens in the caller (resolveTokenUsageCost). Replace "snapshotting the per-token prices from the price table" with language that describes the computation, e.g. "given per-million-token prices and token counts." (Leorio)
🤖
|
|
||
| // NewAIBudgetPolicyFromString converts s to an AIBudgetPolicy, falling back to | ||
| // AIBudgetPolicyHighest when s is empty or not a recognized policy. | ||
| func NewAIBudgetPolicyFromString(s string) AIBudgetPolicy { |
There was a problem hiding this comment.
Nit [CRF-6] NewAIBudgetPolicyFromString is named like a constructor (New), but it's a string-to-typed-value conversion with a default fallback. Go convention uses Parse for this pattern (time.ParseDuration, url.Parse). ParseAIBudgetPolicy or AIBudgetPolicyFrom would match better. (Gon)
🤖
| "github.com/coder/coder/v2/coderd/database" | ||
| ) | ||
|
|
||
| // tokensPerMillion is the divisor for prices, which are quoted per million |
There was a problem hiding this comment.
Nit [CRF-7] Several comments in this file restate what the code already shows. The constant doc here restates the name and value ("tokensPerMillion is the divisor for prices"); the tokenCost doc at line 91 restates the 3-line function body; the inline comment at line 59 duplicates the debug log on the next line. The project mandates substantive comments that describe behavior not already visible from code. Consider trimming to the domain context that the names alone don't convey, e.g. this constant's doc could be just: // Prices are quoted per million tokens. (Gon)
🤖
| } | ||
| } | ||
|
|
||
| // expectTokenUsageCostLookups mocks the store lookups made by resolveTokenUsageCost |
There was a problem hiding this comment.
Nit [CRF-8] (budget resolution and the price lookup), A nil override... has a comma before "A" where a period belongs. Run-on between independent clauses. (Leorio)
🤖
| // 300 + 1200 + 15 + 40. | ||
| const wantCost int64 = 1555 |
There was a problem hiding this comment.
Nice 👍 could you add the token type on the comment, for additional clarity?
| }, | ||
| }, | ||
| { | ||
| name: "valid token usage with budget but no price", |
There was a problem hiding this comment.
We should probably add a case where some of the prices are 0...and maybe a case where the total cost is 0 (all prices are 0), IIRC, this was a valid scenario, right? 👀
| usages, err := rawDB.GetAIBridgeTokenUsagesByInterceptionID(ctx, intc.ID) | ||
| require.NoError(t, err) | ||
| require.Len(t, usages, 1) | ||
| got := usages[0] |
There was a problem hiding this comment.
Hum 🤔 is it possible that an interception has multiple token usages? I guess so in an inner agentic loop...So an interception's cost is the result of the sum of their token usage. I think I forgot about this on the RFC...but this is not a problem, right?
nit: adding a test with multiple token usages
There was a problem hiding this comment.
Hum 🤔 is it possible that an interception has multiple token usages?
it was discussed, yeah it's possible, one example is Streaming Anthropic LLM Request.
nit: adding a test with multiple token usages
I think it doesn't make sense in context of TestRecordTokenUsage[Authorized]? TestRecordTokenUsage is designed to have one token_usage record.
Multiple token usages maybe useful in context of aggregation, but we don't do it here.
| }{ | ||
| {name: "supported", in: "highest", want: codersdk.AIBudgetPolicyHighest}, | ||
| {name: "empty falls back to highest", in: "", want: codersdk.AIBudgetPolicyHighest}, | ||
| {name: "unknown falls back to highest", in: "unsupported", want: codersdk.AIBudgetPolicyHighest}, |
There was a problem hiding this comment.
🤔 shouldn't we fail in this case? This is a configuration error
There was a problem hiding this comment.
it should be unreachable case, because it's checked on CLI level
| -- user has no effective group (no budget configured). Intentionally not a | ||
| -- foreign key: this is an immutable historical attribution that must | ||
| -- survive group deletion, so the id is retained even after the group is gone. | ||
| ADD COLUMN effective_group_id UUID, |
There was a problem hiding this comment.
We should probably add an index here, right? I mean, we will be using mostly the daily aggregated table, but we've seen that there are edge cases where we might have to go straight to the source. Wdyt?
There was a problem hiding this comment.
We discussed it, and decided to postpone this until we actually need an index here.
Implements https://linear.app/codercom/issue/AIGOV-286/add-interception-cost-calculation-to-aibridge-token-usages
Adds spend attribution to AI Gateway. After the upstream response, each token-usage record now captures the user's effective group, the per-token prices in effect at that moment, and a computed cost — so spend is recorded as an immutable, point-in-time snapshot.
Concretely,
aibridge_token_usagesgainseffective_group_id,input_price_micros,output_price_micros,cache_read_price_micros,cache_write_price_micros, andcost_micros. When a usage record is written, the effective group is resolved (per-user override, else the deployment budget policy), the(provider, model)price is looked up and snapshotted onto the row, and cost is computed from the provider-reported token counts. A model that isn't in the price table records its tokens with aNULLcost; any other resolution failure fails the write, so aNULLcost unambiguously means "model not priced" rather than "lookup errored."All values are stored in micro-units (1 unit = 1,000,000 micro-units; Phase 1 assumes USD, so 1 micro-unit = $0.000001). Prices are quoted per million tokens.
This also grants the AI Bridge RBAC subject
readonai_model_prices(the per-interception price lookup needs it; it previously only hadupdatefor the startup seeder).Cost precision
Cost is computed per token category as
tokens × price / 1_000_000with integer division, then the four categories are summed. The division is done per category (not once over the summed numerator) on purpose: it keeps the per-category line items summing exactly to the stored total — no "the parts don't add up to the whole" in reporting).Integer division truncates sub-micro-unit fractions. For example, a cheap model at $0.10 per million tokens is a price of
100_000; 9 tokens cost9 × 100_000 / 1_000_000 = 900_000 / 1_000_000 = 0(the true 0.9 micro-units floors to 0). At real list prices this rarely bites — $3/M input is a price of3_000_000, so even a single token is 3 micro-units. The per-record under-count is bounded below 1 micro-unit per category, so under $0.000004 total across the four categories, which is acceptable for list-price-based cost approximation.Overflow safety
cost_microsis aBIGINT(int64), and the largest intermediate value is a single category'stokens × pricebefore division. int64's ceiling is ≈9.223e18.75_000_000), overflow would require ~123 billion tokens in one response:123e9 × 75e6 = 9.225e18, just over the limit.122e9stays under at9.15e18.1.5e13— roughly six orders of magnitude below the ceiling.So overflow is unreachable from real token counts.
Multi-currency support
In the future, we may encounter issues with multi-currency support, especially when dealing with currencies that have very large exchange rates relative to USD, for example:
IRR: ~1,300,000 IRR ≈ 1 USD
VND: ~26,000 VND ≈ 1 USD
For currencies with such large denominations, numeric overflow is technically possible, considering that we have only about six orders of magnitude of headroom before reaching the limit (see above).
effective_group_idhas no foreign keyeffective_group_idrecords the group a spend was attributed to, as an immutable historical fact. It is intentionally not a foreign key, so the record survives deletion of the group.Alternatives were considered and rejected:
ON DELETE SET NULLwould mutate an "immutable" record — deleting a group silently erases that interception's attribution and under-counts the group's historical spend.RESTRICT/NO ACTIONwould block group deletion entirely (groups are hard-deleted).CASCADEwould delete spend history when a group is deleted — the worst outcome for an audit record.There is also no insert-time check that the group still exists: the id comes from a budget that was just resolved, meaning it was valid at some point.
Open question: group name snapshotting
Should we also snapshot the group name onto each record? Two options:
Leaning toward postponing until a concrete reporting need settles the drift question.