Fix legacy token billing projection

This commit is contained in:
Bryan Helmkamp 2026-07-22 21:37:13 -04:00
parent 3c4ed41d57
commit b39e3767a3
No known key found for this signature in database
4 changed files with 142 additions and 40 deletions

View file

@ -27,6 +27,15 @@ pub(crate) fn parse_tool_arguments_or_empty(raw_arguments: &str) -> serde_json::
serde_json::from_str(raw_arguments).unwrap_or_else(|_| serde_json::json!({}))
}
/// Split an inclusive provider token total into disjoint base and detail
/// buckets. Provider detail counts are advisory and occasionally exceed their
/// parent total, so bound both values while preserving the nonnegative total.
pub(crate) fn split_inclusive_token_total(total: i64, detail: i64) -> (i64, i64) {
let total = total.max(0);
let detail = detail.clamp(0, total);
(total - detail, detail)
}
/// Merge `provider_options.<provider_name>` fields into an encoded request
/// body. Used by codecs whose provider-options namespace is adapter-name keyed
/// rather than a single fixed provider.

View file

@ -1,5 +1,6 @@
//! Serde types mirroring the OpenAI Chat Completions wire shapes.
use crate::codec::split_inclusive_token_total;
use crate::types::{ReasoningEffort, TokenCounts};
#[derive(serde::Serialize)]
@ -139,29 +140,32 @@ impl ApiUsage {
/// reasoning tokens out of `output_tokens`, mirroring the
/// `openai_responses` convention.
pub(super) fn token_counts(&self) -> TokenCounts {
let cached = self
let cached_detail = self
.prompt_tokens_details
.as_ref()
.and_then(|d| d.cached_tokens)
.unwrap_or(0);
let cache_write = self
let cache_write_detail = self
.prompt_tokens_details
.as_ref()
.and_then(|d| d.cache_write_tokens)
.unwrap_or(0);
let reasoning = self
let reasoning_detail = self
.completion_tokens_details
.as_ref()
.and_then(|d| d.reasoning_tokens)
.unwrap_or(0);
let (uncached_input, cached) =
split_inclusive_token_total(self.prompt_tokens, cached_detail);
let (input_tokens, cache_write) =
split_inclusive_token_total(uncached_input, cache_write_detail);
let (output_tokens, reasoning) =
split_inclusive_token_total(self.completion_tokens, reasoning_detail);
TokenCounts {
input_tokens: self
.prompt_tokens
.saturating_sub(cached)
.saturating_sub(cache_write),
output_tokens: self.completion_tokens.saturating_sub(reasoning),
reasoning_tokens: reasoning,
cache_read_tokens: cached,
input_tokens,
output_tokens,
reasoning_tokens: reasoning,
cache_read_tokens: cached,
cache_write_tokens: cache_write,
}
}
@ -225,7 +229,25 @@ pub(super) struct AccumulatedToolCall {
#[cfg(test)]
mod tests {
use super::{ApiResponse, StreamChunk};
use super::{ApiResponse, ApiUsage, StreamChunk};
use crate::types::TokenCounts;
#[test]
fn token_counts_bound_detail_to_parent_totals() {
let usage: ApiUsage = serde_json::from_value(serde_json::json!({
"prompt_tokens": 53,
"completion_tokens": 59,
"completion_tokens_details": {"reasoning_tokens": 66}
}))
.unwrap();
assert_eq!(usage.token_counts(), TokenCounts {
input_tokens: 53,
output_tokens: 0,
reasoning_tokens: 59,
..TokenCounts::default()
});
}
#[test]
fn reasoning_accepts_provider_and_openrouter_spellings() {

View file

@ -3,7 +3,7 @@
use serde::Deserialize;
use super::wire::{ApiResponse, ApiUsage, InputTokensResponse};
use crate::codec::{CodecCtx, parse_tool_arguments_or_empty};
use crate::codec::{CodecCtx, parse_tool_arguments_or_empty, split_inclusive_token_total};
use crate::error::Error;
use crate::types::{
ContentPart, FinishReason, Message, RateLimitInfo, Response, Role, TokenCounts, ToolCall,
@ -11,19 +11,23 @@ use crate::types::{
pub(super) fn token_counts_from_api_usage(usage: Option<&ApiUsage>) -> TokenCounts {
usage.map_or_else(TokenCounts::default, |u| {
let cached_tokens = u
let cached_detail = u
.input_tokens_details
.as_ref()
.and_then(|d| d.cached_tokens)
.unwrap_or(0);
let reasoning_tokens = u
let reasoning_detail = u
.output_tokens_details
.as_ref()
.and_then(|d| d.reasoning_tokens)
.unwrap_or(0);
let (input_tokens, cached_tokens) =
split_inclusive_token_total(u.input_tokens, cached_detail);
let (output_tokens, reasoning_tokens) =
split_inclusive_token_total(u.output_tokens, reasoning_detail);
TokenCounts {
input_tokens: u.input_tokens.saturating_sub(cached_tokens),
output_tokens: u.output_tokens.saturating_sub(reasoning_tokens),
input_tokens,
output_tokens,
reasoning_tokens,
cache_read_tokens: cached_tokens,
..TokenCounts::default()
@ -205,6 +209,24 @@ mod tests {
use super::super::encode;
use super::*;
#[test]
fn token_counts_bound_detail_to_parent_totals() {
let usage: ApiUsage = serde_json::from_value(serde_json::json!({
"input_tokens": 53,
"output_tokens": 59,
"input_tokens_details": null,
"output_tokens_details": {"reasoning_tokens": 66}
}))
.unwrap();
assert_eq!(token_counts_from_api_usage(Some(&usage)), TokenCounts {
input_tokens: 53,
output_tokens: 0,
reasoning_tokens: 59,
..TokenCounts::default()
});
}
#[test]
fn parse_output_preserves_both_ids_on_function_call() {
let output = vec![serde_json::json!({

View file

@ -3,7 +3,7 @@ use std::fmt::Write as _;
use std::sync::LazyLock;
use chrono::{DateTime, Utc};
use fabro_types::{Run, RunId, RunSize, RunStatusKind, RunTiming};
use fabro_types::{BilledTokenCounts, Run, RunId, RunSize, RunStatusKind, RunTiming};
use sqlx::sqlite::{SqliteConnection, SqliteRow};
use sqlx::{QueryBuilder, Row as _, Sqlite, SqlitePool};
use strum::VariantArray as _;
@ -313,7 +313,7 @@ impl ProjectedRunSummary {
.unwrap_or(run.timestamps.created_at);
run.timing = entry.projection.live_run_timing(at);
}
let billing = projected_billing(&entry.projection);
let billing = normalize_billing_for_read_model(projected_billing(&entry.projection));
let workflow_name = run.workflow.display_name().map(str::to_string);
let repository_name = run
.repository
@ -335,6 +335,33 @@ impl ProjectedRunSummary {
}
}
/// Older provider codecs could persist a negative disjoint bucket when a
/// detail count exceeded its inclusive parent total. The SQLite summary is a
/// rebuildable, nonnegative read model, so normalize those legacy values here
/// without rewriting the authoritative run events.
fn normalize_billing_for_read_model(mut billing: BilledTokenCounts) -> BilledTokenCounts {
let input_total = billing
.input_tokens
.saturating_add(billing.cache_read_tokens)
.saturating_add(billing.cache_write_tokens)
.max(0);
billing.cache_read_tokens = billing.cache_read_tokens.clamp(0, input_total);
billing.cache_write_tokens = billing
.cache_write_tokens
.clamp(0, input_total - billing.cache_read_tokens);
billing.input_tokens = input_total - billing.cache_read_tokens - billing.cache_write_tokens;
let output_total = billing
.output_tokens
.saturating_add(billing.reasoning_tokens)
.max(0);
billing.reasoning_tokens = billing.reasoning_tokens.clamp(0, output_total);
billing.output_tokens = output_total - billing.reasoning_tokens;
billing.total_tokens = input_total.saturating_add(output_total);
billing.total_usd_micros = billing.total_usd_micros.map(|value| value.max(0));
billing
}
async fn upsert_run(connection: &mut SqliteConnection, record: &ProjectedRunSummary) -> Result<()> {
let run = &record.run;
let diff = run.diff.unwrap_or_default();
@ -803,6 +830,47 @@ mod tests {
assert_eq!(run.size, RunSize::S);
}
#[tokio::test]
async fn projection_normalizes_legacy_overlapping_reasoning_tokens() {
let (_directory, store) = store().await;
let created_at = dt("2026-07-11T12:00:00Z");
let run_id = run_id(created_at.timestamp_millis().cast_unsigned(), 1);
let mut projection = projection(run_id, "legacy billing", created_at);
projection.conclusion = Some(Conclusion {
timestamp: created_at,
status: StageOutcome::Succeeded,
timing: RunTiming::default(),
failure: None,
final_git_commit_sha: None,
stages: Vec::new(),
billing: Some(BilledTokenCounts {
input_tokens: 53,
output_tokens: -7,
total_tokens: 112,
reasoning_tokens: 66,
..BilledTokenCounts::default()
}),
total_retries: 0,
diff: RunDiff::default(),
});
store
.upsert_projection(&entry(projection, 1))
.await
.unwrap();
let row = sqlx::query(
"SELECT input_tokens, output_tokens, reasoning_tokens FROM runs WHERE id = ?",
)
.bind(run_id.to_string())
.fetch_one(&store.pool)
.await
.unwrap();
assert_eq!(sqlx::Row::get::<i64, _>(&row, "input_tokens"), 53);
assert_eq!(sqlx::Row::get::<i64, _>(&row, "output_tokens"), 0);
assert_eq!(sqlx::Row::get::<i64, _>(&row, "reasoning_tokens"), 59);
}
#[tokio::test]
async fn reconcile_removes_rows_absent_from_authoritative_entries() {
let (_directory, store) = store().await;
@ -833,33 +901,14 @@ mod tests {
.unwrap();
let good = entry(projection(good_id, "good", created_at), 1);
let mut invalid_projection = projection(recovered_id, "recovered", created_at);
invalid_projection.conclusion = Some(Conclusion {
timestamp: created_at,
status: StageOutcome::Succeeded,
timing: RunTiming::default(),
failure: None,
final_git_commit_sha: None,
stages: Vec::new(),
billing: Some(BilledTokenCounts {
input_tokens: -1,
..BilledTokenCounts::default()
}),
total_retries: 0,
diff: RunDiff::default(),
});
let invalid = entry(invalid_projection.clone(), 1);
let recovered_projection = projection(recovered_id, "recovered", created_at);
let invalid = entry(recovered_projection.clone(), 0);
assert!(store.reconcile(&[good.clone(), invalid]).await.is_err());
assert!(store.get(&stale_id, created_at).await.unwrap().is_some());
assert!(store.get(&good_id, created_at).await.unwrap().is_none());
invalid_projection.conclusion.as_mut().unwrap().billing = Some(BilledTokenCounts {
input_tokens: 1,
total_tokens: 1,
..BilledTokenCounts::default()
});
let recovered = entry(invalid_projection, 1);
let recovered = entry(recovered_projection, 1);
store.reconcile(&[good, recovered]).await.unwrap();
assert!(store.get(&stale_id, created_at).await.unwrap().is_none());