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
31 changes: 31 additions & 0 deletions crates/tokscale-cli/src/main.rs
Original file line number Diff line number Diff line change
Expand Up @@ -5481,6 +5481,32 @@ fn report_excluded_tokenless_rows(excluded: &[ExcludedTokenlessRow]) {
println!();
}

fn report_unpriced_submission_exclusions(
excluded: &[tokscale_core::UnpricedSubmissionExclusion],
has_remaining_usage: bool,
) {
use colored::Colorize;

for row in excluded {
let remaining_usage_message = has_remaining_usage
.then_some(" Remaining priced usage will be submitted.")
.unwrap_or_default();
println!(
"{}",
format!(
" Warning: excluded {} unpriced {}/{} message(s) ({} tokens): {}.{}",
row.message_count,
row.provider_id,
row.model_id,
format_tokens_with_commas(row.total_tokens),
row.reason,
remaining_usage_message,
)
.yellow()
);
}
}

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
enum SubmitMode {
Interactive,
Expand Down Expand Up @@ -5636,6 +5662,10 @@ fn run_submit_command(
// left out, so a single legacy charge can't block the whole submission.
let excluded_rows = exclude_tokenless_cost_contributions(&mut graph_result);
report_excluded_tokenless_rows(&excluded_rows);
report_unpriced_submission_exclusions(
&graph_result.unpriced_submission_exclusions,
graph_result.summary.total_tokens > 0,
);

println!("{}", " Data to submit:".white());
println!(
Expand Down Expand Up @@ -6568,6 +6598,7 @@ mod tests {
years: calculate_years(&contributions),
contributions,
time_metrics: None,
unpriced_submission_exclusions: Vec::new(),
}
}

Expand Down
33 changes: 33 additions & 0 deletions crates/tokscale-cli/tests/cli_tests.rs
Original file line number Diff line number Diff line change
Expand Up @@ -3560,6 +3560,39 @@ fn test_submit_offline_without_pricing_cache_fails() {
"stderr should contain a pricing/network error: {stderr}"
);
}

#[test]
fn test_submit_excluding_only_generic_gemini_usage_does_not_promise_submission() {
let tmp = create_empty_fixture_dir();
let message_dir = tmp
.path()
.join(".local/share/opencode/storage/message/gemini-default");
fs::create_dir_all(&message_dir).unwrap();
fs::write(
message_dir.join("gemini-default.json"),
r#"{
"id": "gemini-default",
"sessionID": "gemini-default",
"role": "assistant",
"modelID": "gemini-default",
"providerID": "google",
"tokens": { "input": 1, "output": 0, "reasoning": 0, "cache": { "read": 0, "write": 0 } },
"time": { "created": 1736510400000.0 }
}"#,
)
.unwrap();

cmd_with_home(tmp.path())
.env("TOKSCALE_API_TOKEN", "test-token")
.args(["submit", "--client", "opencode", "--dry-run"])
.assert()
.success()
.stdout(predicate::str::contains(
"excluded 1 unpriced google/gemini-default message(s) (1 tokens)",
))
.stdout(predicate::str::contains("Remaining priced usage will be submitted.").not())
.stdout(predicate::str::contains("No usage data found to submit."));
}
// ── gjc client filter tests ────────────────────────────────────────────────

/// Write a gjc session JSONL file at
Expand Down
1 change: 1 addition & 0 deletions crates/tokscale-core/src/aggregator.rs
Original file line number Diff line number Diff line change
Expand Up @@ -220,6 +220,7 @@ pub fn generate_graph_result(
years,
contributions,
time_metrics: None,
unpriced_submission_exclusions: Vec::new(),
}
}

Expand Down
169 changes: 165 additions & 4 deletions crates/tokscale-core/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -496,6 +496,19 @@ pub struct GraphResult {
pub contributions: Vec<DailyContribution>,
#[serde(skip_serializing_if = "Option::is_none")]
pub time_metrics: Option<sessionize::TimeMetrics>,
#[serde(skip)]
pub unpriced_submission_exclusions: Vec<UnpricedSubmissionExclusion>,
}

/// Token-bearing usage excluded only from a submission because its generic
/// routing label has no authoritative model-to-price mapping.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct UnpricedSubmissionExclusion {
pub provider_id: String,
pub model_id: String,
pub message_count: usize,
pub total_tokens: i64,
pub reason: &'static str,
}

#[derive(Debug, Clone, Default)]
Expand Down Expand Up @@ -2927,9 +2940,15 @@ fn build_graph_from_messages(
start: Instant,
bucket_timezone: &bucket_tz::BucketTimezone,
) -> Result<GraphResult, String> {
if matches!(pricing_requirement, GraphPricingRequirement::Submission) {
validate_priced_messages(&filtered, pricing)?;
}
let (filtered, unpriced_submission_exclusions) = match pricing_requirement {
GraphPricingRequirement::Lenient => (filtered, Vec::new()),
GraphPricingRequirement::Submission => {
let (submitted, exclusions) =
exclude_generic_unpriced_submission_messages(filtered, pricing);
validate_priced_messages(&submitted, pricing)?;
(submitted, exclusions)
}
};

let intervals = sessionize::sessionize(&filtered, sessionize::DEFAULT_IDLE_GAP_MS);
let time_metrics =
Expand All @@ -2944,6 +2963,7 @@ fn build_graph_from_messages(
let processing_time_ms = start.elapsed().as_millis() as u32;
let mut result = aggregator::generate_graph_result(contributions, processing_time_ms);
result.time_metrics = Some(time_metrics);
result.unpriced_submission_exclusions = unpriced_submission_exclusions;

for contribution in &mut result.contributions {
if let Some(&ms) = daily_active_time.get(&contribution.date) {
Expand All @@ -2954,6 +2974,58 @@ fn build_graph_from_messages(
Ok(result)
}

const GEMINI_DEFAULT_UNPRICED_REASON: &str =
"generic routing label has no authoritative model-to-price mapping";

fn exclude_generic_unpriced_submission_messages(
messages: Vec<UnifiedMessage>,
pricing: Option<&pricing::PricingService>,
) -> (Vec<UnifiedMessage>, Vec<UnpricedSubmissionExclusion>) {
let Some(pricing) = pricing else {
return (messages, Vec::new());
};

let mut submitted = Vec::with_capacity(messages.len());
let mut exclusions: std::collections::BTreeMap<(String, String), (usize, i64)> =
std::collections::BTreeMap::new();

for message in messages {
let is_unpriced_gemini_default = message.tokens.total() > 0
&& !message.has_authoritative_cost()
&& message.provider_id.eq_ignore_ascii_case("google")
&& message.model_id.eq_ignore_ascii_case("gemini-default")
&& !pricing.covers_usage_with_provider(
&message.model_id,
Some(&message.provider_id),
&message.tokens,
);

if is_unpriced_gemini_default {
let entry = exclusions
.entry((message.provider_id.clone(), message.model_id.clone()))
.or_default();
entry.0 += 1;
entry.1 = entry.1.saturating_add(message.tokens.total());
} else {
submitted.push(message);
}
}

let exclusions = exclusions
.into_iter()
.map(|((provider_id, model_id), (message_count, total_tokens))| {
UnpricedSubmissionExclusion {
provider_id,
model_id,
message_count,
total_tokens,
reason: GEMINI_DEFAULT_UNPRICED_REASON,
}
})
.collect();
(submitted, exclusions)
}

#[derive(Debug, Clone, serde::Serialize)]
pub struct TimeMetricsReport {
pub metrics: sessionize::TimeMetrics,
Expand Down Expand Up @@ -4150,7 +4222,8 @@ mod tests {
parse_local_clients, parsed_to_unified, paths, pricing, retain_for_requested_clients,
scanner, select_local_parse_pricing, unified_to_parsed, validate_priced_messages, ClientId,
GraphPricingRequirement, GroupBy, LocalParseOptions, ReportOptions, TokenBreakdown,
UnifiedMessage, UNKNOWN_WORKSPACE_LABEL,
UnifiedMessage, UnpricedSubmissionExclusion, GEMINI_DEFAULT_UNPRICED_REASON,
UNKNOWN_WORKSPACE_LABEL,
};
use serial_test::serial;
use std::collections::{HashMap, HashSet};
Expand Down Expand Up @@ -7842,6 +7915,94 @@ mod tests {
assert!(submission_error.contains("unknown-provider/genuinely-unpriced-model"));
}

#[test]
fn submission_excludes_unpriced_generic_gemini_default_but_keeps_priceable_usage() {
let mut litellm = HashMap::new();
litellm.insert(
"gpt-4o".to_string(),
pricing::ModelPricing {
input_cost_per_token: Some(1e-6),
..Default::default()
},
);
let pricing = pricing::PricingService::new(litellm, HashMap::new());
let generic = UnifiedMessage::new(
"antigravity-cli",
"gemini-default",
"google",
"generic",
1_736_510_400_000,
TokenBreakdown {
input: 7,
cache_read: 11,
..Default::default()
},
0.0,
);
let concrete = UnifiedMessage::new(
"synthetic",
"gpt-4o",
"openai",
"concrete",
1_736_510_400_000,
TokenBreakdown {
input: 13,
..Default::default()
},
0.0,
);

let graph = build_graph_from_messages(
vec![generic, concrete],
Some(&pricing),
GraphPricingRequirement::Submission,
std::time::Instant::now(),
)
.expect("generic routing label must not block fully priced submission usage");

assert_eq!(graph.summary.total_tokens, 13);
assert_eq!(graph.contributions[0].clients.len(), 1);
assert_eq!(graph.contributions[0].clients[0].model_id, "gpt-4o");
assert_eq!(graph.unpriced_submission_exclusions.len(), 1);
assert_eq!(
graph.unpriced_submission_exclusions[0],
UnpricedSubmissionExclusion {
provider_id: "google".to_string(),
model_id: "gemini-default".to_string(),
message_count: 1,
total_tokens: 18,
reason: GEMINI_DEFAULT_UNPRICED_REASON,
}
);
}

#[test]
fn submission_still_rejects_unpriced_concrete_models() {
let concrete = UnifiedMessage::new(
"synthetic",
"gemini-3.5-pro",
"google",
"concrete",
1_736_510_400_000,
TokenBreakdown {
input: 1,
..Default::default()
},
0.0,
);
let pricing = pricing::PricingService::new(HashMap::new(), HashMap::new());

let error = build_graph_from_messages(
vec![concrete],
Some(&pricing),
GraphPricingRequirement::Submission,
std::time::Instant::now(),
)
.expect_err("concrete unpriced models must remain a submission error");

assert!(error.contains("google/gemini-3.5-pro"));
}

#[test]
fn strict_pricing_validation_accepts_bundled_pricing() {
let pricing = pricing::PricingService::new(HashMap::new(), HashMap::new());
Expand Down
44 changes: 40 additions & 4 deletions crates/tokscale-core/src/pricing/lookup.rs
Original file line number Diff line number Diff line change
Expand Up @@ -47,10 +47,11 @@ const RESELLER_PROVIDER_PREFIXES: &[&str] = &[
"orcarouter/",
];

// Bare brand tokens ("claude", "anthropic") are blocked because they contain
// no model information: a fuzzy hit from them can land on any model of the
// brand (e.g. retired `claude-2.1` eroding to `claude` and billing at an
// opus-fast key), so such a match is never trustworthy.
// Bare brand tokens ("claude", "anthropic", "gemini") are blocked because they
// contain no model information: a fuzzy hit from them can land on any model of
// the brand (e.g. retired `claude-2.1` eroding to `claude` and billing at an
// opus-fast key, or `gemini-default` eroding to `gemini` and landing on a
// native-audio preview key), so such a match is never trustworthy.
//
// Generic English words ("model", "router") are blocked for the same reason:
// they carry no model identity, yet substring-match real priced keys
Expand All @@ -65,6 +66,7 @@ const FUZZY_BLOCKLIST: &[&str] = &[
"base",
"claude",
"anthropic",
"gemini",
"model",
"router",
];
Expand Down Expand Up @@ -3181,6 +3183,40 @@ mod tests {
);
}

#[test]
fn incomplete_unhinted_result_does_not_replace_provider_pricing() {
let mut litellm = HashMap::new();
litellm.insert(
"azure/gpt-fallback-guard".into(),
ModelPricing {
input_cost_per_token: Some(1.0),
..Default::default()
},
);
litellm.insert(
"gpt-fallback-guard".into(),
ModelPricing {
output_cost_per_token: Some(2.0),
..Default::default()
},
);
let lookup = PricingLookup::new(litellm, HashMap::new(), HashMap::new());
let usage = TokenBreakdown {
input: 1,
output: 1,
cache_read: 0,
cache_write: 0,
reasoning: 0,
};

// Neither row covers both populated buckets. Retain the provider row
// rather than replacing it with an unhinted row that silently prices
// the input bucket at zero.
assert_eq!(
lookup.calculate_cost_with_provider("gpt-fallback-guard", Some("azure"), &usage),
1.0
);
}
#[test]
fn test_provider_hint_normalizes_openai_codex_alias() {
let mut litellm = HashMap::new();
Expand Down
13 changes: 13 additions & 0 deletions crates/tokscale-core/src/sessions/antigravity_cli.rs
Original file line number Diff line number Diff line change
Expand Up @@ -526,6 +526,19 @@ mod tests {
assert_eq!(message.provider_id, "google");
}

// The generic routing label is preserved verbatim. It is not a concrete
// billable model id, so submit can exclude it instead of inventing a cost.
#[test]
fn gemini_default_response_model_is_preserved() {
let blob = build_gen_metadata_with_model("gemini-default");
let mut seen = HashSet::new();

let message = parse_gen_metadata(&blob, "session", 1_000, &mut seen).unwrap();

assert_eq!(message.model_id, "gemini-default");
assert_eq!(message.provider_id, "google");
}

#[test]
fn per_generation_timestamp_overrides_session_fallback() {
// chatModel.#9.#4 = {#1: seconds, #2: nanos} is the per-turn wall-clock
Expand Down
Loading