mirror of
https://github.com/fabro-sh/fabro.git
synced 2026-08-28 05:27:41 +00:00
Fix legacy token billing projection
This commit is contained in:
parent
3c4ed41d57
commit
b39e3767a3
4 changed files with 142 additions and 40 deletions
|
|
@ -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.
|
||||
|
|
|
|||
|
|
@ -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() {
|
||||
|
|
|
|||
|
|
@ -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!({
|
||||
|
|
|
|||
|
|
@ -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());
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue