Skip to content

Commit 2ced786

Browse files
ryoppippikkoomen
andcommitted
fix(pricing): migrate merged adapters to timestamp-aware costs
Use the current timestamp-aware cost API in the newly merged ZCode and Antigravity adapters so each event can resolve scheduled pricing without retaining the removed helper. Co-authored-by: Kim Koomen <[email protected]>
1 parent fad8134 commit 2ced786

2 files changed

Lines changed: 15 additions & 5 deletions

File tree

‎rust/adapters/antigravity/src/parser.rs‎

Lines changed: 12 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -4,7 +4,7 @@ use jiff::tz::TimeZone as JiffTimeZone;
44

55
use crate::{
66
LoadedEntry, PricingMap, Result, TimestampMs, TokenUsageRaw, UsageEntry, UsageMessage,
7-
calculate_cost_for_usage, cli::CostMode, cli_error, format_date_tz,
7+
calculate_cost_for_usage_at, cli::CostMode, cli_error, format_date_tz,
88
missing_pricing_model_for_candidates,
99
};
1010

@@ -904,7 +904,14 @@ pub(super) fn event_to_loaded(
904904
let extra_total_tokens = event
905905
.total_output_tokens
906906
.saturating_sub(event.output_tokens);
907-
let cost = calculate_antigravity_cost(&event.model, event.provider, cost_usage, mode, pricing);
907+
let cost = calculate_antigravity_cost(
908+
&event.model,
909+
event.provider,
910+
cost_usage,
911+
event.timestamp,
912+
mode,
913+
pricing,
914+
);
908915
let missing_pricing_model =
909916
missing_antigravity_pricing(&event.model, event.provider, cost_usage, mode, pricing);
910917
let data = UsageEntry {
@@ -942,6 +949,7 @@ fn calculate_antigravity_cost(
942949
model: &str,
943950
provider: Option<u64>,
944951
usage: TokenUsageRaw,
952+
timestamp: TimestampMs,
945953
mode: CostMode,
946954
pricing: &PricingMap,
947955
) -> f64 {
@@ -951,10 +959,11 @@ fn calculate_antigravity_cost(
951959
.into_iter()
952960
.find_map(|candidate| {
953961
pricing.find(&candidate).map(|_| {
954-
calculate_cost_for_usage(
962+
calculate_cost_for_usage_at(
955963
Some(&candidate),
956964
usage,
957965
None,
966+
Some(timestamp),
958967
CostMode::Calculate,
959968
Some(pricing),
960969
)

‎rust/adapters/zcode/src/parser.rs‎

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -4,7 +4,7 @@ use jiff::tz::TimeZone as JiffTimeZone;
44

55
use crate::{
66
LoadedEntry, PricingMap, TimestampMs, TokenUsageRaw, UsageEntry, UsageMessage,
7-
apply_total_token_fallback, calculate_cost_for_usage,
7+
apply_total_token_fallback, calculate_cost_for_usage_at,
88
cli::{CostMode, PricingOverride},
99
format_date_tz, format_rfc3339_millis, total_usage_tokens,
1010
};
@@ -94,10 +94,11 @@ pub(super) fn row_to_entry(
9494
let cost = if mode == CostMode::Display {
9595
0.0
9696
} else {
97-
calculate_cost_for_usage(
97+
calculate_cost_for_usage_at(
9898
cost_model.as_deref(),
9999
cost_usage,
100100
None,
101+
Some(timestamp),
101102
CostMode::Calculate,
102103
Some(pricing),
103104
)

0 commit comments

Comments
 (0)