diff --git a/rust/src/providers/antigravity/cli_fallback.rs b/rust/src/providers/antigravity/cli_fallback.rs index c4474bc8ca..f2fbdaca80 100644 --- a/rust/src/providers/antigravity/cli_fallback.rs +++ b/rust/src/providers/antigravity/cli_fallback.rs @@ -144,13 +144,15 @@ pub(super) async fn managed_spawn_is_csrf_gated(binary: Option) -> bool let Some(binary) = binary else { return false; }; - let Ok(version) = - run_cli_command(&binary, &VERSION_ARGS, VERSION_TIMEOUT, VERSION_TOO_LARGE).await - else { - return false; - }; - let version = String::from_utf8_lossy(&version); - is_csrf_gated_version(version.trim()) + agy_version(&binary, VERSION_TIMEOUT) + .await + .is_ok_and(|version| is_csrf_gated_version(&version)) +} + +/// `agy --version` output, trimmed. +async fn agy_version(binary: &Path, timeout: Duration) -> Result { + let version = run_cli_command(binary, &VERSION_ARGS, timeout, VERSION_TOO_LARGE).await?; + Ok(String::from_utf8_lossy(&version).trim().to_string()) } async fn fetch_print_usage(binary: &Path) -> Result { @@ -161,10 +163,7 @@ async fn fetch_print_usage_with_version_timeout( binary: &Path, version_timeout: Duration, ) -> Result { - let version = - run_cli_command(binary, &VERSION_ARGS, version_timeout, VERSION_TOO_LARGE).await?; - let version = String::from_utf8_lossy(&version); - if !is_supported_version(version.trim()) { + if !is_supported_version(&agy_version(binary, version_timeout).await?) { return Err(ProviderError::Parse( "Antigravity CLI usage reports require agy 1.1.11 or later".into(), ) diff --git a/rust/src/providers/antigravity/legacy_status.rs b/rust/src/providers/antigravity/legacy_status.rs index 852561a399..a93366f129 100644 --- a/rust/src/providers/antigravity/legacy_status.rs +++ b/rust/src/providers/antigravity/legacy_status.rs @@ -13,10 +13,6 @@ pub(crate) struct UserStatusResponse { #[derive(Debug, Deserialize)] #[serde(rename_all = "camelCase")] pub(crate) struct UserStatus { - #[allow( - dead_code, - reason = "field mirrors the Antigravity API user payload; deserialized for round-trip fidelity but not read yet" - )] email: Option, plan_status: Option, user_tier: Option, @@ -224,9 +220,14 @@ fn model_window_id(config: &ModelConfig) -> String { .as_deref() .or(config.id.as_deref()) .unwrap_or_else(|| model_label(config)); - let raw = canonical_model_id(raw); - let slug = raw - .chars() + let slug = slug(canonical_model_id(raw)); + format!("model-{}", if slug.is_empty() { "unknown" } else { &slug }) +} + +/// Lowercase ASCII alphanumerics; every other character becomes `-`, and +/// leading or trailing dashes are dropped. +pub(super) fn slug(raw: &str) -> String { + raw.chars() .map(|ch| { if ch.is_ascii_alphanumeric() { ch.to_ascii_lowercase() @@ -236,8 +237,7 @@ fn model_window_id(config: &ModelConfig) -> String { }) .collect::() .trim_matches('-') - .to_string(); - format!("model-{}", if slug.is_empty() { "unknown" } else { &slug }) + .to_string() } fn rate_window_from_quota(quota: &QuotaInfo) -> RateWindow { diff --git a/rust/src/providers/antigravity/local_bot_id.rs b/rust/src/providers/antigravity/local_bot_id.rs index ac732ee84b..d9d9c26da2 100644 --- a/rust/src/providers/antigravity/local_bot_id.rs +++ b/rust/src/providers/antigravity/local_bot_id.rs @@ -116,14 +116,7 @@ pub(super) fn append_recovered_events( continue; } if bot_id_uses.get(bot_id) == Some(&1) { - let mut turn = pending.turn.clone(); - turn.timestamp_ms = Some(exact.timestamp_ms); - events.push(Event { - session: session.to_string(), - row: pending.row, - turn, - total: pending.total, - }); + events.push(recovered_event(session, pending, exact.timestamp_ms)); recovered += 1; continue; } @@ -141,16 +134,20 @@ pub(super) fn append_recovered_events( let Some(timestamp_ms) = timestamps.get(*offset) else { continue; }; - let mut turn = pending.turn.clone(); - turn.timestamp_ms = Some(*timestamp_ms); - events.push(Event { - session: session.to_string(), - row: pending.row, - turn, - total: pending.total, - }); + events.push(recovered_event(session, pending, *timestamp_ms)); recovered += 1; } events.sort_by_key(|event| event.row); recovered } + +fn recovered_event(session: &str, pending: &PendingTimestampRow, timestamp_ms: i64) -> Event { + let mut turn = pending.turn.clone(); + turn.timestamp_ms = Some(timestamp_ms); + Event { + session: session.to_string(), + row: pending.row, + turn, + total: pending.total, + } +} diff --git a/rust/src/providers/antigravity/local_history.rs b/rust/src/providers/antigravity/local_history.rs index 5581c65618..3bcc6e712d 100644 --- a/rust/src/providers/antigravity/local_history.rs +++ b/rust/src/providers/antigravity/local_history.rs @@ -9,13 +9,6 @@ use chrono::{DateTime, Utc}; use crate::spend_contract::{LocalHistoryCoverage, LocalTokenHistorySummary}; -fn clean_env_path(value: Option<&str>) -> Option { - value - .map(str::trim) - .filter(|value| !value.is_empty()) - .map(PathBuf::from) -} - fn configured_database_roots(home: &Path) -> [PathBuf; 3] { let gemini_cli_home = std::env::var("GEMINI_CLI_HOME").ok(); configured_database_roots_from_values(home, gemini_cli_home.as_deref()) @@ -25,7 +18,8 @@ fn configured_database_roots_from_values( home: &Path, gemini_cli_home: Option<&str>, ) -> [PathBuf; 3] { - let gemini_base = clean_env_path(gemini_cli_home).unwrap_or_else(|| home.join(".gemini")); + let gemini_base = + local_sessions::clean_env_path(gemini_cli_home).unwrap_or_else(|| home.join(".gemini")); local_sqlite::database_roots(&gemini_base) } diff --git a/rust/src/providers/antigravity/local_proto.rs b/rust/src/providers/antigravity/local_proto.rs index 851b80a2fa..cffe7893ef 100644 --- a/rust/src/providers/antigravity/local_proto.rs +++ b/rust/src/providers/antigravity/local_proto.rs @@ -127,11 +127,6 @@ fn identity_text(field: Field<'_>) -> Option> { Some((!value.trim().is_empty()).then(|| value.to_string())) } -fn auxiliary_text(field: Field<'_>) -> Result, ()> { - let value = std::str::from_utf8(message(field).ok_or(())?).map_err(|_| ())?; - Ok((!value.trim().is_empty()).then(|| value.to_string())) -} - #[derive(Debug, Default)] struct AuxiliaryIdentifier { value: Option, @@ -150,12 +145,9 @@ impl AuxiliaryIdentifier { if !self.valid { return; } - match auxiliary_text(field) { - Ok(value) => self.value = value, - Err(()) => { - self.value = None; - self.valid = false; - } + match identity_text(field) { + Some(value) => self.value = value, + None => self.invalidate(), } } @@ -165,11 +157,20 @@ impl AuxiliaryIdentifier { } } +/// Positive and no later than 9999-12-31T23:59:59Z. +fn valid_seconds(seconds: u64) -> bool { + seconds != 0 && seconds <= 253_402_300_799 +} + +fn valid_nanos(nanos: u64) -> bool { + nanos <= 999_999_999 +} + fn timestamp_millis(seconds: Option, nanos: u64) -> Option> { let Some(seconds) = seconds else { return Some(None); }; - if seconds == 0 || seconds > 253_402_300_799 || nanos > 999_999_999 { + if !valid_seconds(seconds) || !valid_nanos(nanos) { return None; } let seconds = i64::try_from(seconds).ok()?; @@ -302,29 +303,17 @@ fn parse_generation(bytes: &[u8], seconds: &mut Option, nanos: &mut u64) -> fn parse_timestamp_field(bytes: &[u8], seconds: &mut Option, nanos: &mut u64) -> Option<()> { fields(bytes, |stamp| { match stamp.number { - 1 => { - let value = integer(stamp)?; - if value == 0 || value > 253_402_300_799 { - return None; - } - *seconds = Some(value); - } - 2 => { - let value = integer(stamp)?; - if value > 999_999_999 { - return None; - } - *nanos = value; - } + 1 => *seconds = Some(integer(stamp).filter(|value| valid_seconds(*value))?), + 2 => *nanos = integer(stamp).filter(|value| valid_nanos(*value))?, _ => {} } Some(()) }) } +/// Protobuf encoders for building test blobs. #[cfg(test)] -mod tests { - use super::*; +pub(super) mod test_fixtures { fn varint(mut value: u64) -> Vec { let mut bytes = Vec::new(); loop { @@ -339,17 +328,25 @@ mod tests { } } } - fn field_varint(number: u64, value: u64) -> Vec { - let mut b = varint(number << 3); - b.extend(varint(value)); - b + + pub(crate) fn field_varint(number: u64, value: u64) -> Vec { + let mut bytes = varint(number << 3); + bytes.extend(varint(value)); + bytes } - fn field_bytes(number: u64, value: &[u8]) -> Vec { - let mut b = varint((number << 3) | 2); - b.extend(varint(value.len() as u64)); - b.extend(value); - b + + pub(crate) fn field_bytes(number: u64, value: &[u8]) -> Vec { + let mut bytes = varint((number << 3) | 2); + bytes.extend(varint(value.len() as u64)); + bytes.extend(value); + bytes } +} + +#[cfg(test)] +mod tests { + use super::test_fixtures::{field_bytes, field_varint}; + use super::*; #[test] fn decodes_generation_usage_and_timestamp() { let mut usage = Vec::new(); diff --git a/rust/src/providers/antigravity/local_sessions_reader.rs b/rust/src/providers/antigravity/local_sessions_reader.rs index 6b494f628e..487ede9382 100644 --- a/rust/src/providers/antigravity/local_sessions_reader.rs +++ b/rust/src/providers/antigravity/local_sessions_reader.rs @@ -37,7 +37,7 @@ pub(super) fn configured_tokscale_sessions(home: &Path) -> PathBuf { tokscale_sessions_from_values(home, tokscale.as_deref()) } -fn clean_env_path(value: Option<&str>) -> Option { +pub(super) fn clean_env_path(value: Option<&str>) -> Option { value .map(str::trim) .filter(|value| !value.is_empty()) @@ -413,6 +413,18 @@ fn has_valid_token_fields(value: &Value) -> bool { mod tests { use super::*; + fn now() -> DateTime { + Utc.timestamp_millis_opt(1787576400000).single().unwrap() + } + + /// Summarize one session file holding `text` over the last 7 days. + fn summarize_lines(text: impl AsRef<[u8]>) -> LocalTokenHistorySummary { + let dir = tempfile::tempdir().unwrap(); + let path = dir.path().join("session-a.jsonl"); + fs::write(&path, text).unwrap(); + summarize_paths(&[path], now(), 7, false) + } + #[test] fn tokscale_discovery_bounds_entries_and_retained_paths() { let dir = tempfile::tempdir().unwrap(); @@ -452,9 +464,8 @@ mod tests { ), ) .unwrap(); - let now = Utc.timestamp_millis_opt(1787576400000).single().unwrap(); - let summary = summarize_paths(&[known, unknown], now, 7, false); + let summary = summarize_paths(&[known, unknown], now(), 7, false); assert_eq!(summary.cost_estimate.coverage.estimated, 1); assert_eq!(summary.cost_estimate.coverage.unpriced, 1); @@ -463,15 +474,11 @@ mod tests { } #[test] fn summarizes_tokscale_jsonl_and_deduplicates_response_ids() { - let dir = tempfile::tempdir().unwrap(); - let path = dir.path().join("session-a.jsonl"); - fs::write(&path, concat!( + let summary = summarize_lines(concat!( "{\"type\":\"session_meta\",\"modelId\":\"test-model-antigravity-a\"}\n", "{\"type\":\"usage\",\"responseId\":\"r1\",\"timestamp\":1787572800000,\"input\":100,\"output\":20,\"cacheRead\":10,\"cacheWrite\":5}\n", "{\"type\":\"usage\",\"response_id\":\"r1\",\"timestamp\":1787572800000,\"input\":100,\"output\":20}\n" - )).unwrap(); - let now = Utc.timestamp_millis_opt(1787576400000).single().unwrap(); - let summary = summarize_paths(&[path], now, 7, false); + )); assert_eq!(summary.total_tokens, 135); assert_eq!(summary.session_count, 1); assert_eq!(summary.coverage, LocalHistoryCoverage::Complete); @@ -486,35 +493,24 @@ mod tests { b"{\"type\":\"usage\",\"timestamp\":1787572800000,\"input\":10}\n", ) .unwrap(); - let now = Utc.timestamp_millis_opt(1787576400000).single().unwrap(); - let truncated = summarize_paths(std::slice::from_ref(&path), now, 7, true); + let truncated = summarize_paths(std::slice::from_ref(&path), now(), 7, true); assert_eq!(truncated.coverage, LocalHistoryCoverage::Partial); - let missing = summarize_paths(&[dir.path().join("missing.jsonl")], now, 7, false); + let missing = summarize_paths(&[dir.path().join("missing.jsonl")], now(), 7, false); assert_eq!(missing.coverage, LocalHistoryCoverage::Partial); } #[test] fn excludes_usage_outside_requested_window() { - let dir = tempfile::tempdir().unwrap(); - let path = dir.path().join("session-a.jsonl"); - fs::write( - &path, - concat!( - "{\"type\":\"usage\",\"timestamp\":1787572800000,\"input\":10,\"output\":5}\n", - "{\"type\":\"usage\",\"timestamp\":1784894400000,\"input\":99,\"output\":99}\n" - ), - ) - .unwrap(); - let now = Utc.timestamp_millis_opt(1787576400000).single().unwrap(); - let summary = summarize_paths(&[path], now, 7, false); + let summary = summarize_lines(concat!( + "{\"type\":\"usage\",\"timestamp\":1787572800000,\"input\":10,\"output\":5}\n", + "{\"type\":\"usage\",\"timestamp\":1784894400000,\"input\":99,\"output\":99}\n" + )); assert_eq!(summary.total_tokens, 15); assert_eq!(summary.session_count, 1); } #[test] fn oversized_jsonl_line_marks_coverage_partial_and_next_row_is_counted() { - let dir = tempfile::tempdir().unwrap(); - let path = dir.path().join("session-a.jsonl"); let mut text = format!( r#"{{"type":"usage","padding":"{}"}}"#, "x".repeat(MAX_JSONL_LINE_BYTES + 32) @@ -522,9 +518,7 @@ mod tests { text.push('\n'); text.push_str(r#"{"type":"usage","timestamp":1787572800000,"input":10,"output":5}"#); text.push('\n'); - fs::write(&path, text).unwrap(); - let now = Utc.timestamp_millis_opt(1787576400000).single().unwrap(); - let summary = summarize_paths(&[path], now, 7, false); + let summary = summarize_lines(text); assert_eq!(summary.total_tokens, 15); assert_eq!(summary.session_count, 1); assert_eq!(summary.coverage, LocalHistoryCoverage::Partial); @@ -532,21 +526,12 @@ mod tests { #[test] fn malformed_jsonl_record_marks_coverage_partial_and_next_row_is_counted() { - let dir = tempfile::tempdir().unwrap(); - let path = dir.path().join("session-a.jsonl"); - fs::write( - &path, - concat!( - "{malformed json}\n", - "{\"type\":\"usage\",\"timestamp\":1787572800000,\"input\":\"invalid\"}\n", - "{\"type\":\"usage\",\"input\":10}\n", - "{\"type\":\"usage\",\"timestamp\":1787572800000,\"input\":10,\"output\":5}\n" - ), - ) - .unwrap(); - let now = Utc.timestamp_millis_opt(1787576400000).single().unwrap(); - - let summary = summarize_paths(&[path], now, 7, false); + let summary = summarize_lines(concat!( + "{malformed json}\n", + "{\"type\":\"usage\",\"timestamp\":1787572800000,\"input\":\"invalid\"}\n", + "{\"type\":\"usage\",\"input\":10}\n", + "{\"type\":\"usage\",\"timestamp\":1787572800000,\"input\":10,\"output\":5}\n" + )); assert_eq!(summary.total_tokens, 15); assert_eq!(summary.coverage, LocalHistoryCoverage::Partial); @@ -567,10 +552,9 @@ mod tests { "{\"type\":\"usage\",\"timestamp\":1787572800000,\"input\":20}\n", ) .unwrap(); - let now = Utc.timestamp_millis_opt(1787576400000).single().unwrap(); let summary = - summarize_paths_with_budget(&[first_path, second_path], now, 7, false, first.len()); + summarize_paths_with_budget(&[first_path, second_path], now(), 7, false, first.len()); // A hard budget stop is withheld: no total is published from it. assert_eq!(summary.total_tokens, 0); @@ -581,20 +565,11 @@ mod tests { #[test] fn token_sum_overflow_marks_coverage_partial_without_saturation() { - let dir = tempfile::tempdir().unwrap(); - let path = dir.path().join("session-a.jsonl"); - fs::write( - &path, - concat!( - "{\"type\":\"usage\",\"timestamp\":1787572800000,\"input\":18446744073709551615,\"output\":1}\n", - "{\"type\":\"usage\",\"timestamp\":1787572800000,\"input\":18446744073709551615}\n", - "{\"type\":\"usage\",\"timestamp\":1787572800000,\"input\":1}\n" - ), - ) - .unwrap(); - let now = Utc.timestamp_millis_opt(1787576400000).single().unwrap(); - - let summary = summarize_paths(&[path], now, 7, false); + let summary = summarize_lines(concat!( + "{\"type\":\"usage\",\"timestamp\":1787572800000,\"input\":18446744073709551615,\"output\":1}\n", + "{\"type\":\"usage\",\"timestamp\":1787572800000,\"input\":18446744073709551615}\n", + "{\"type\":\"usage\",\"timestamp\":1787572800000,\"input\":1}\n" + )); assert_eq!(summary.total_tokens, u64::MAX); assert_eq!(summary.session_count, 1); diff --git a/rust/src/providers/antigravity/local_sqlite.rs b/rust/src/providers/antigravity/local_sqlite.rs index e820538cb2..e89a41728c 100644 --- a/rust/src/providers/antigravity/local_sqlite.rs +++ b/rust/src/providers/antigravity/local_sqlite.rs @@ -46,6 +46,27 @@ enum DatabaseScan { Unsupported, } +impl DatabaseScan { + fn partial(events: Vec) -> Self { + Self::Supported { + events, + complete: false, + } + } +} + +/// Outcome of the per-row size checks shared by both table readers. +enum CheckedBlob<'r> { + Blob { + idx: i64, + bytes: &'r [u8], + }, + /// Unusable row: skip it and mark the read incomplete. + Invalid, + /// A byte limit ran out (already recorded on the budget): stop the read. + Exhausted, +} + #[derive(Debug, Clone, Copy, PartialEq, Eq)] enum SchemaInspection { Supported, @@ -406,21 +427,15 @@ fn discover_databases(roots: &[PathBuf], budget: &mut Budget) -> (Vec, fn read_database(path: &Path, budget: &mut Budget) -> rusqlite::Result { if !budget.check() { - return Ok(DatabaseScan::Supported { - events: Vec::new(), - complete: false, - }); + return Ok(DatabaseScan::partial(Vec::new())); } let mut conn = Connection::open_with_flags(path, OpenFlags::SQLITE_OPEN_READ_ONLY)?; let tx = conn.transaction_with_behavior(TransactionBehavior::Deferred)?; - match supported_schema(&tx, budget)? { + match table_schema(&tx, budget, "gen_metadata", &["idx", "data"])? { SchemaInspection::Supported => {} SchemaInspection::Unsupported => return Ok(DatabaseScan::Unsupported), SchemaInspection::Incomplete => { - return Ok(DatabaseScan::Supported { - events: Vec::new(), - complete: false, - }); + return Ok(DatabaseScan::partial(Vec::new())); } } @@ -439,20 +454,14 @@ fn read_database(path: &Path, budget: &mut Budget) -> rusqlite::Result rusqlite::Result scan, Err(_) => { - return Ok(DatabaseScan::Supported { - events: rows.events, - complete: false, - }); + return Ok(DatabaseScan::partial(rows.events)); } }; if !step_scan.complete { - return Ok(DatabaseScan::Supported { - events: rows.events, - complete: false, - }); + return Ok(DatabaseScan::partial(rows.events)); } if !embedded_timestamps_agree(&rows.occurrences, &step_scan, &rows.bot_id_uses) { - return Ok(DatabaseScan::Supported { - events: rows.events, - complete: false, - }); + return Ok(DatabaseScan::partial(rows.events)); } let resolved = resolve_step_timestamps( @@ -551,46 +551,17 @@ fn read_generation_rows( break; } - let idx: i64 = match row.get(0) { - Ok(value) if value >= 0 => value, - _ => { + let (idx, blob) = match checked_blob(row, &mut database_bytes, budget)? { + CheckedBlob::Blob { idx, bytes } => (idx, bytes), + CheckedBlob::Invalid => { complete = false; continue; } - }; - let declared: Option = row.get(1).ok(); - let Some(declared) = declared.and_then(|value| usize::try_from(value).ok()) else { - complete = false; - continue; - }; - database_bytes = match database_bytes.checked_add(declared) { - Some(value) if value <= MAX_DATABASE_BYTES => value, - _ => { + CheckedBlob::Exhausted => { complete = false; - budget.exhausted = true; break; } }; - budget.bytes = match budget.bytes.checked_add(declared) { - Some(value) if value <= MAX_TOTAL_BYTES => value, - _ => { - complete = false; - budget.exhausted = true; - break; - } - }; - if declared == 0 || declared > MAX_BLOB_BYTES { - complete = false; - continue; - } - - let blob = match row.get_ref(2)? { - ValueRef::Blob(bytes) if bytes.len() == declared => bytes, - _ => { - complete = false; - continue; - } - }; let Some(turn) = parse_turn(blob) else { complete = false; continue; @@ -684,46 +655,17 @@ fn read_step_timestamps( break; } - let idx: i64 = match row.get(0) { - Ok(value) if value >= 0 => value, - _ => { + let (idx, blob) = match checked_blob(row, database_bytes, budget)? { + CheckedBlob::Blob { idx, bytes } => (idx, bytes), + CheckedBlob::Invalid => { rows_are_valid = false; continue; } - }; - let declared: Option = row.get(1).ok(); - let Some(declared) = declared.and_then(|value| usize::try_from(value).ok()) else { - rows_are_valid = false; - continue; - }; - *database_bytes = match (*database_bytes).checked_add(declared) { - Some(value) if value <= MAX_DATABASE_BYTES => value, - _ => { - complete = false; - budget.exhausted = true; - break; - } - }; - budget.bytes = match budget.bytes.checked_add(declared) { - Some(value) if value <= MAX_TOTAL_BYTES => value, - _ => { + CheckedBlob::Exhausted => { complete = false; - budget.exhausted = true; break; } }; - if declared == 0 || declared > MAX_BLOB_BYTES { - rows_are_valid = false; - continue; - } - - let blob = match row.get_ref(2)? { - ValueRef::Blob(bytes) if bytes.len() == declared => bytes, - _ => { - rows_are_valid = false; - continue; - } - }; let Some(metadata) = parse_step_metadata(blob) else { rows_are_valid = false; continue; @@ -763,44 +705,54 @@ fn read_step_timestamps( }) } -fn supported_schema(conn: &Connection, budget: &mut Budget) -> rusqlite::Result { - let mut statement = - conn.prepare("SELECT name, type, rootpage FROM main.sqlite_master LIMIT ?1")?; - let mut rows = statement.query([i64::try_from(MAX_SCHEMA_ENTRIES + 1).unwrap_or(i64::MAX)])?; - let mut found = false; - let mut schema_entries = 0usize; - while let Some(row) = rows.next()? { - if !budget.check() { - return Ok(SchemaInspection::Incomplete); - } - schema_entries += 1; - if schema_entries > MAX_SCHEMA_ENTRIES { - return Ok(SchemaInspection::Incomplete); - } - let name: String = row.get(0)?; - let kind: String = row.get(1)?; - if !budget.charge_schema_text(&name) || !budget.charge_schema_text(&kind) { - return Ok(SchemaInspection::Incomplete); - } - if !name.eq_ignore_ascii_case("gen_metadata") { - continue; - } - let rootpage: i64 = row.get(2)?; - if kind != "table" || rootpage <= 0 || found { - return Ok(SchemaInspection::Unsupported); - } - found = true; - } - if !found { - return Ok(SchemaInspection::Unsupported); +/// Size-check one `(idx, declared length, blob)` row. Rows with a bad index or +/// length are skipped before any bytes are charged; out-of-range lengths are +/// charged first, then rejected. +fn checked_blob<'r>( + row: &'r rusqlite::Row<'_>, + database_bytes: &mut usize, + budget: &mut Budget, +) -> rusqlite::Result> { + let idx: i64 = match row.get(0) { + Ok(value) if value >= 0 => value, + _ => return Ok(CheckedBlob::Invalid), + }; + let declared: Option = row.get(1).ok(); + let Some(declared) = declared.and_then(|value| usize::try_from(value).ok()) else { + return Ok(CheckedBlob::Invalid); + }; + let Some(next_database) = database_bytes + .checked_add(declared) + .filter(|value| *value <= MAX_DATABASE_BYTES) + else { + budget.exhausted = true; + return Ok(CheckedBlob::Exhausted); + }; + *database_bytes = next_database; + let Some(next_total) = budget + .bytes + .checked_add(declared) + .filter(|value| *value <= MAX_TOTAL_BYTES) + else { + budget.exhausted = true; + return Ok(CheckedBlob::Exhausted); + }; + budget.bytes = next_total; + if declared == 0 || declared > MAX_BLOB_BYTES { + return Ok(CheckedBlob::Invalid); } - - has_stored_columns(conn, "gen_metadata", &["idx", "data"], budget) + Ok(match row.get_ref(2)? { + ValueRef::Blob(bytes) if bytes.len() == declared => CheckedBlob::Blob { idx, bytes }, + _ => CheckedBlob::Invalid, + }) } -fn supported_steps_schema( +/// Exactly one real `table` in `main` with the `required` stored columns. +fn table_schema( conn: &Connection, budget: &mut Budget, + table: &str, + required: &[&str], ) -> rusqlite::Result { let mut statement = conn.prepare("SELECT name, type, rootpage FROM main.sqlite_master LIMIT ?1")?; @@ -820,7 +772,7 @@ fn supported_steps_schema( if !budget.charge_schema_text(&name) || !budget.charge_schema_text(&kind) { return Ok(SchemaInspection::Incomplete); } - if !name.eq_ignore_ascii_case("steps") { + if !name.eq_ignore_ascii_case(table) { continue; } let rootpage: i64 = row.get(2)?; @@ -833,7 +785,7 @@ fn supported_steps_schema( return Ok(SchemaInspection::Unsupported); } - has_stored_columns(conn, "steps", &["idx", "metadata"], budget) + has_stored_columns(conn, table, required, budget) } fn has_stored_columns( diff --git a/rust/src/providers/antigravity/local_sqlite_synthetic_tests.rs b/rust/src/providers/antigravity/local_sqlite_synthetic_tests.rs index 6259052d82..71dcd1ada2 100644 --- a/rust/src/providers/antigravity/local_sqlite_synthetic_tests.rs +++ b/rust/src/providers/antigravity/local_sqlite_synthetic_tests.rs @@ -6,6 +6,7 @@ use rusqlite::{Connection, params}; use tempfile::TempDir; use super::*; +use crate::providers::antigravity::local_proto::test_fixtures::{field_bytes, field_varint}; const NOW_SECONDS: u64 = 1_800_000_000; @@ -15,34 +16,6 @@ fn now() -> DateTime { .unwrap() } -fn varint(mut value: u64) -> Vec { - let mut bytes = Vec::new(); - loop { - let mut byte = (value & 0x7f) as u8; - value >>= 7; - if value != 0 { - byte |= 0x80; - } - bytes.push(byte); - if value == 0 { - return bytes; - } - } -} - -fn field_varint(number: u64, value: u64) -> Vec { - let mut bytes = varint(number << 3); - bytes.extend(varint(value)); - bytes -} - -fn field_bytes(number: u64, value: &[u8]) -> Vec { - let mut bytes = varint((number << 3) | 2); - bytes.extend(varint(value.len() as u64)); - bytes.extend(value); - bytes -} - fn turn_blob(step_uuid: Option<&str>, input: u64, timestamp: Option) -> Vec { turn_blob_with_bot_id(step_uuid, input, timestamp, None) } @@ -145,60 +118,60 @@ fn database( path } -fn summary(dir: &TempDir) -> LocalSessionSummary { +/// Build one session database in a fresh directory, scan it, and check the +/// coverage and token total. +fn expect_summary( + session: &str, + generation_rows: &[(i64, Vec)], + step_rows: Option>, + coverage: LocalHistoryCoverage, + total_tokens: u64, +) -> LocalSessionSummary { + let dir = tempfile::tempdir().unwrap(); + database(&dir, session, generation_rows, step_rows); let root = dir.path().join("conversations"); let SQLiteScan::Summary(summary) = summarize(&[root], now(), 30) else { panic!("database should be attempted"); }; + assert_eq!(summary.coverage, coverage, "{session}"); + assert_eq!(summary.total_tokens, total_tokens, "{session}"); summary } #[test] fn newer_step_timestamp_recovery_preserves_legacy_and_new_totals() { - let dir = tempfile::tempdir().unwrap(); let uuid = "new-step"; - database( - &dir, + let summary = expect_summary( "mixed", &[ (0, turn_blob(None, 100, Some(NOW_SECONDS - 120))), (1, turn_blob(Some(uuid), 200, None)), ], Some(&[(10, Some(step_metadata(Some(uuid), Some(NOW_SECONDS - 60))))]), + LocalHistoryCoverage::Complete, + 496, ); - - let summary = summary(&dir); - - assert_eq!(summary.coverage, LocalHistoryCoverage::Complete); - assert_eq!(summary.total_tokens, 496); assert_eq!(summary.session_count, 1); } #[test] fn absent_steps_table_withholds_newer_rows_but_keeps_embedded_totals() { - let dir = tempfile::tempdir().unwrap(); - database( - &dir, + expect_summary( "missing-table", &[ (0, turn_blob(None, 100, Some(NOW_SECONDS - 120))), (1, turn_blob(Some("missing-table"), 200, None)), ], None, + LocalHistoryCoverage::Partial, + 198, ); - - let summary = summary(&dir); - - assert_eq!(summary.coverage, LocalHistoryCoverage::Partial); - assert_eq!(summary.total_tokens, 198); } #[test] fn reused_step_uuid_follows_stored_idx_order() { - let dir = tempfile::tempdir().unwrap(); let uuid = "reused-step"; - database( - &dir, + expect_summary( "ordered", &[ (0, turn_blob(Some(uuid), 100, None)), @@ -208,32 +181,24 @@ fn reused_step_uuid_follows_stored_idx_order() { (20, Some(step_metadata(Some(uuid), Some(NOW_SECONDS - 60)))), (10, Some(step_metadata(Some(uuid), Some(NOW_SECONDS - 120)))), ]), + LocalHistoryCoverage::Complete, + 496, ); - - let summary = summary(&dir); - - assert_eq!(summary.coverage, LocalHistoryCoverage::Complete); - assert_eq!(summary.total_tokens, 496); } #[test] fn missing_lowest_step_timestamp_withholds_pending_tokens() { - let dir = tempfile::tempdir().unwrap(); let uuid = "missing-lowest"; - database( - &dir, + expect_summary( "missing", &[(0, turn_blob(Some(uuid), 100, None))], Some(&[ (10, Some(step_metadata(Some(uuid), None))), (20, Some(step_metadata(Some(uuid), Some(NOW_SECONDS - 60)))), ]), + LocalHistoryCoverage::Partial, + 0, ); - - let summary = summary(&dir); - - assert_eq!(summary.coverage, LocalHistoryCoverage::Partial); - assert_eq!(summary.total_tokens, 0); } #[test] @@ -257,35 +222,25 @@ fn duplicate_or_malformed_step_rows_fail_closed() { vec![(10, Some(malformed_step_metadata("malformed")))], ), ] { - let dir = tempfile::tempdir().unwrap(); - database( - &dir, + expect_summary( session, &[(0, turn_blob(Some(session), 100, None))], Some(&step_rows), + LocalHistoryCoverage::Partial, + 0, ); - - let summary = summary(&dir); - - assert_eq!(summary.coverage, LocalHistoryCoverage::Partial, "{session}"); - assert_eq!(summary.total_tokens, 0, "{session}"); } } #[test] fn null_step_rows_fail_closed() { - let dir = tempfile::tempdir().unwrap(); - database( - &dir, + expect_summary( "null-step", &[(0, turn_blob(Some("null-step"), 100, None))], Some(&[(10, None)]), + LocalHistoryCoverage::Partial, + 0, ); - - let summary = summary(&dir); - - assert_eq!(summary.coverage, LocalHistoryCoverage::Partial); - assert_eq!(summary.total_tokens, 0); } #[test] @@ -307,22 +262,13 @@ fn unidentified_step_rows_do_not_invalidate_timestamp_recovery() { ], ), ] { - let dir = tempfile::tempdir().unwrap(); - database( - &dir, + expect_summary( session, &[(0, turn_blob(Some(uuid), 100, None))], Some(&step_rows), - ); - - let summary = summary(&dir); - - assert_eq!( - summary.coverage, LocalHistoryCoverage::Complete, - "{session}" + 198, ); - assert_eq!(summary.total_tokens, 198, "{session}"); } } @@ -331,7 +277,6 @@ fn exact_bot_id_recovery_survives_an_unidentified_step_row() { let uuid = "exact-with-unidentified"; let bot_id = "exact-bot"; for unidentified_first in [false, true] { - let dir = tempfile::tempdir().unwrap(); let identified = ( 20, Some(step_metadata_with_bot_id( @@ -346,29 +291,23 @@ fn exact_bot_id_recovery_survives_an_unidentified_step_row() { } else { vec![identified, unidentified] }; - database( - &dir, + expect_summary( "exact-with-unidentified", &[( 0, turn_blob_with_bot_id(Some(uuid), 100, None, Some(bot_id)), )], Some(&step_rows), + LocalHistoryCoverage::Complete, + 198, ); - - let summary = summary(&dir); - - assert_eq!(summary.coverage, LocalHistoryCoverage::Complete); - assert_eq!(summary.total_tokens, 198); } } #[test] fn embedded_and_step_timestamps_must_agree() { - let dir = tempfile::tempdir().unwrap(); let uuid = "conflicting"; - database( - &dir, + expect_summary( "conflict", &[ (0, turn_blob(Some(uuid), 100, Some(NOW_SECONDS - 120))), @@ -378,20 +317,15 @@ fn embedded_and_step_timestamps_must_agree() { (10, Some(step_metadata(Some(uuid), Some(NOW_SECONDS - 60)))), (20, Some(step_metadata(Some(uuid), Some(NOW_SECONDS - 30)))), ]), + LocalHistoryCoverage::Partial, + 198, ); - - let summary = summary(&dir); - - assert_eq!(summary.coverage, LocalHistoryCoverage::Partial); - assert_eq!(summary.total_tokens, 198); } #[test] fn unique_bot_ids_recover_each_turn_despite_auxiliary_step_rows() { - let dir = tempfile::tempdir().unwrap(); let uuid = "per-turn"; - database( - &dir, + expect_summary( "bot-correlation", &[ ( @@ -429,19 +363,15 @@ fn unique_bot_ids_recover_each_turn_despite_auxiliary_step_rows() { )), ), ]), + LocalHistoryCoverage::Complete, + 496, ); - - let summary = summary(&dir); - - assert_eq!(summary.coverage, LocalHistoryCoverage::Complete); - assert_eq!(summary.total_tokens, 496); } #[test] fn reused_step_uuid_recovers_distinct_timestamps_with_an_unidentified_row() { let uuid = "reused-with-unidentified"; for unidentified_position in 0..=2 { - let dir = tempfile::tempdir().unwrap(); let mut step_rows = vec![ (10, Some(step_metadata(Some(uuid), Some(NOW_SECONDS - 120)))), (20, Some(step_metadata(Some(uuid), Some(NOW_SECONDS - 60)))), @@ -450,29 +380,23 @@ fn reused_step_uuid_recovers_distinct_timestamps_with_an_unidentified_row() { unidentified_position, (15, Some(step_metadata(None, Some(NOW_SECONDS - 300)))), ); - database( - &dir, + expect_summary( "reused-with-unidentified", &[ (0, turn_blob(Some(uuid), 100, None)), (1, turn_blob(Some(uuid), 200, None)), ], Some(&step_rows), + LocalHistoryCoverage::Complete, + 496, ); - - let summary = summary(&dir); - - assert_eq!(summary.coverage, LocalHistoryCoverage::Complete); - assert_eq!(summary.total_tokens, 496); } } #[test] fn reused_step_uuid_with_one_identified_row_and_an_unidentified_row_stays_partial() { - let dir = tempfile::tempdir().unwrap(); let uuid = "single-identified-with-unidentified"; - database( - &dir, + expect_summary( "single-identified-with-unidentified", &[ (0, turn_blob(Some(uuid), 100, None)), @@ -482,12 +406,9 @@ fn reused_step_uuid_with_one_identified_row_and_an_unidentified_row_stays_partia (10, Some(step_metadata(Some(uuid), Some(NOW_SECONDS - 120)))), (20, Some(step_metadata(None, Some(NOW_SECONDS - 60)))), ]), + LocalHistoryCoverage::Partial, + 0, ); - - let summary = summary(&dir); - - assert_eq!(summary.coverage, LocalHistoryCoverage::Partial); - assert_eq!(summary.total_tokens, 0); } #[test] @@ -496,7 +417,6 @@ fn unidentified_duplicate_bot_id_withholds_exact_and_positional_recovery() { let bot_id = "shared-bot"; for unidentified_first in [false, true] { for unidentified_timestamp in [Some(NOW_SECONDS - 60), None] { - let dir = tempfile::tempdir().unwrap(); let identified = ( 10, Some(step_metadata_with_bot_id( @@ -518,30 +438,24 @@ fn unidentified_duplicate_bot_id_withholds_exact_and_positional_recovery() { } else { vec![identified, unidentified] }; - database( - &dir, + expect_summary( "duplicate-bot-with-unidentified", &[( 0, turn_blob_with_bot_id(Some(uuid), 100, None, Some(bot_id)), )], Some(&step_rows), + LocalHistoryCoverage::Partial, + 0, ); - - let summary = summary(&dir); - - assert_eq!(summary.coverage, LocalHistoryCoverage::Partial); - assert_eq!(summary.total_tokens, 0); } } } #[test] fn conflicting_bot_id_evidence_is_withheld() { - let dir = tempfile::tempdir().unwrap(); let uuid = "ambiguous-bot"; - database( - &dir, + expect_summary( "ambiguous-bot", &[( (0), @@ -565,19 +479,14 @@ fn conflicting_bot_id_evidence_is_withheld() { )), ), ]), + LocalHistoryCoverage::Partial, + 0, ); - - let summary = summary(&dir); - - assert_eq!(summary.coverage, LocalHistoryCoverage::Partial); - assert_eq!(summary.total_tokens, 0); } #[test] fn cross_step_bot_id_evidence_is_withheld() { - let dir = tempfile::tempdir().unwrap(); - database( - &dir, + expect_summary( "cross-step-bot", &[( (0), @@ -601,10 +510,7 @@ fn cross_step_bot_id_evidence_is_withheld() { )), ), ]), + LocalHistoryCoverage::Partial, + 0, ); - - let summary = summary(&dir); - - assert_eq!(summary.coverage, LocalHistoryCoverage::Partial); - assert_eq!(summary.total_tokens, 0); } diff --git a/rust/src/providers/antigravity/local_sqlite_tests.rs b/rust/src/providers/antigravity/local_sqlite_tests.rs index a8a93a1103..ffdc79b747 100644 --- a/rust/src/providers/antigravity/local_sqlite_tests.rs +++ b/rust/src/providers/antigravity/local_sqlite_tests.rs @@ -1,34 +1,7 @@ use super::*; +use crate::providers::antigravity::local_proto::test_fixtures::{field_bytes, field_varint}; use rusqlite::params; -fn varint(mut value: u64) -> Vec { - let mut bytes = Vec::new(); - loop { - let mut byte = (value & 0x7f) as u8; - value >>= 7; - if value != 0 { - byte |= 0x80; - } - bytes.push(byte); - if value == 0 { - return bytes; - } - } -} - -fn field_varint(number: u64, value: u64) -> Vec { - let mut bytes = varint(number << 3); - bytes.extend(varint(value)); - bytes -} - -fn field_bytes(number: u64, value: &[u8]) -> Vec { - let mut bytes = varint((number << 3) | 2); - bytes.extend(varint(value.len() as u64)); - bytes.extend(value); - bytes -} - fn valid_turn_blob(input: u64, timestamp_seconds: u64) -> Vec { valid_turn_blob_with_model(input, timestamp_seconds, None) } @@ -58,6 +31,38 @@ fn zero_token_turn_blob(timestamp_seconds: u64, model: &str) -> Vec { field_bytes(1, &chat) } +/// The CLI conversations directory under `dir`. +fn cli_root(dir: &tempfile::TempDir) -> PathBuf { + dir.path().join(".gemini/antigravity-cli/conversations") +} + +/// Open `root/name`, creating `root`, with an empty `gen_metadata` table. +fn database_at(root: &Path, name: &str) -> Connection { + fs::create_dir_all(root).unwrap(); + let conn = Connection::open(root.join(name)).unwrap(); + conn.execute("CREATE TABLE gen_metadata(idx INTEGER, data BLOB)", []) + .unwrap(); + conn +} + +fn insert_row(conn: &Connection, idx: i64, data: impl rusqlite::ToSql) { + conn.execute( + "INSERT INTO gen_metadata(idx, data) VALUES(?1, ?2)", + params![idx, data], + ) + .unwrap(); +} + +/// Scan the default `.gemini` roots under `dir`. +fn expect_summary(dir: &tempfile::TempDir) -> LocalSessionSummary { + let SQLiteScan::Summary(summary) = + summarize(&database_roots(&dir.path().join(".gemini")), Utc::now(), 30) + else { + panic!("supported database should produce coverage"); + }; + summary +} + #[test] fn missing_databases_falls_through() { let dir = tempfile::tempdir().unwrap(); @@ -70,7 +75,7 @@ fn missing_databases_falls_through() { #[test] fn foreign_database_is_non_authoritative() { let dir = tempfile::tempdir().unwrap(); - let root = dir.path().join(".gemini/antigravity-cli/conversations"); + let root = cli_root(&dir); fs::create_dir_all(&root).unwrap(); let conn = Connection::open(root.join("one.db")).unwrap(); conn.execute("CREATE TABLE wrong(idx INTEGER, data BLOB)", []) @@ -85,17 +90,8 @@ fn foreign_database_is_non_authoritative() { #[test] fn empty_supported_database_is_confirmed_zero() { let dir = tempfile::tempdir().unwrap(); - let root = dir.path().join(".gemini/antigravity-cli/conversations"); - fs::create_dir_all(&root).unwrap(); - let conn = Connection::open(root.join("one.db")).unwrap(); - conn.execute("CREATE TABLE gen_metadata(idx INTEGER, data BLOB)", []) - .unwrap(); - drop(conn); - let SQLiteScan::Summary(summary) = - summarize(&database_roots(&dir.path().join(".gemini")), Utc::now(), 30) - else { - panic!("supported database should produce coverage"); - }; + drop(database_at(&cli_root(&dir), "one.db")); + let summary = expect_summary(&dir); assert_eq!(summary.coverage, LocalHistoryCoverage::Complete); assert_eq!(summary.total_tokens, 0); assert_eq!(summary.session_count, 0); @@ -106,19 +102,11 @@ fn same_named_databases_in_separate_roots_keep_distinct_rows_and_sessions() { let dir = tempfile::tempdir().unwrap(); let first_root = dir.path().join("first"); let second_root = dir.path().join("second"); - fs::create_dir_all(&first_root).unwrap(); - fs::create_dir_all(&second_root).unwrap(); let timestamp = u64::try_from(Utc::now().timestamp()).unwrap(); for (root, input) in [(&first_root, 100_u64), (&second_root, 200_u64)] { - let conn = Connection::open(root.join("session.db")).unwrap(); - conn.execute("CREATE TABLE gen_metadata(idx INTEGER, data BLOB)", []) - .unwrap(); - conn.execute( - "INSERT INTO gen_metadata(idx, data) VALUES(1, ?1)", - [valid_turn_blob(input, timestamp)], - ) - .unwrap(); + let conn = database_at(root, "session.db"); + insert_row(&conn, 1, valid_turn_blob(input, timestamp)); } let SQLiteScan::Summary(summary) = summarize(&[first_root, second_root], Utc::now(), 30) else { @@ -133,34 +121,18 @@ fn same_named_databases_in_separate_roots_keep_distinct_rows_and_sessions() { #[test] fn zero_token_unknown_model_does_not_poison_priced_history() { let dir = tempfile::tempdir().unwrap(); - let root = dir.path().join(".gemini/antigravity-cli/conversations"); - fs::create_dir_all(&root).unwrap(); + let root = cli_root(&dir); let now = Utc::now(); let timestamp = u64::try_from(now.timestamp()).unwrap(); - let priced = Connection::open(root.join("priced.db")).unwrap(); - priced - .execute("CREATE TABLE gen_metadata(idx INTEGER, data BLOB)", []) - .unwrap(); - priced - .execute( - "INSERT INTO gen_metadata(idx, data) VALUES(1, ?1)", - [valid_turn_blob_with_model( - 100, - timestamp, - Some("claude-sonnet-4-6"), - )], - ) - .unwrap(); - - let zero = Connection::open(root.join("zero.db")).unwrap(); - zero.execute("CREATE TABLE gen_metadata(idx INTEGER, data BLOB)", []) - .unwrap(); - zero.execute( - "INSERT INTO gen_metadata(idx, data) VALUES(1, ?1)", - [zero_token_turn_blob(timestamp, "unknown-model")], - ) - .unwrap(); + let priced = database_at(&root, "priced.db"); + insert_row( + &priced, + 1, + valid_turn_blob_with_model(100, timestamp, Some("claude-sonnet-4-6")), + ); + let zero = database_at(&root, "zero.db"); + insert_row(&zero, 1, zero_token_turn_blob(timestamp, "unknown-model")); let SQLiteScan::Summary(summary) = summarize(&[root], now, 30) else { panic!("supported databases should produce coverage"); @@ -186,23 +158,10 @@ fn zero_token_unknown_model_does_not_poison_priced_history() { #[test] fn non_blob_rows_make_coverage_partial() { let dir = tempfile::tempdir().unwrap(); - let root = dir.path().join(".gemini/antigravity-cli/conversations"); - fs::create_dir_all(&root).unwrap(); - let conn = Connection::open(root.join("one.db")).unwrap(); - conn.execute("CREATE TABLE gen_metadata(idx INTEGER, data BLOB)", []) - .unwrap(); - conn.execute( - "INSERT INTO gen_metadata(idx,data) VALUES(?1,?2)", - params![1_i64, "not-a-blob"], - ) - .unwrap(); + let conn = database_at(&cli_root(&dir), "one.db"); + insert_row(&conn, 1, "not-a-blob"); drop(conn); - let SQLiteScan::Summary(summary) = - summarize(&database_roots(&dir.path().join(".gemini")), Utc::now(), 30) - else { - panic!("supported database should produce coverage"); - }; - assert_eq!(summary.coverage, LocalHistoryCoverage::Partial); + assert_eq!(expect_summary(&dir).coverage, LocalHistoryCoverage::Partial); } #[test] fn discovery_allows_exactly_500_databases_but_marks_501_partial() { @@ -245,7 +204,7 @@ fn extra_columns_and_without_rowid_schema_is_supported() { .unwrap(); let mut budget = Budget::new(); assert_eq!( - supported_schema(&conn, &mut budget).unwrap(), + table_schema(&conn, &mut budget, "gen_metadata", &["idx", "data"]).unwrap(), SchemaInspection::Supported ); } @@ -260,7 +219,7 @@ fn generated_columns_are_rejected() { .unwrap(); let mut budget = Budget::new(); assert_eq!( - supported_schema(&conn, &mut budget).unwrap(), + table_schema(&conn, &mut budget, "gen_metadata", &["idx", "data"]).unwrap(), SchemaInspection::Unsupported ); } @@ -274,7 +233,7 @@ fn schema_entry_budget_is_incomplete_not_foreign() { } let mut budget = Budget::new(); assert_eq!( - supported_schema(&conn, &mut budget).unwrap(), + table_schema(&conn, &mut budget, "gen_metadata", &["idx", "data"]).unwrap(), SchemaInspection::Incomplete ); } @@ -282,29 +241,13 @@ fn schema_entry_budget_is_incomplete_not_foreign() { #[test] fn undecodable_row_beside_valid_rows_yields_a_lower_bound() { let dir = tempfile::tempdir().unwrap(); - let root = dir.path().join(".gemini/antigravity-cli/conversations"); - fs::create_dir_all(&root).unwrap(); let timestamp = u64::try_from(Utc::now().timestamp()).unwrap(); - let conn = Connection::open(root.join("one.db")).unwrap(); - conn.execute("CREATE TABLE gen_metadata(idx INTEGER, data BLOB)", []) - .unwrap(); - conn.execute( - "INSERT INTO gen_metadata(idx, data) VALUES(1, ?1)", - [valid_turn_blob(100, timestamp)], - ) - .unwrap(); - conn.execute( - "INSERT INTO gen_metadata(idx, data) VALUES(2, ?1)", - params!["not-a-blob"], - ) - .unwrap(); + let conn = database_at(&cli_root(&dir), "one.db"); + insert_row(&conn, 1, valid_turn_blob(100, timestamp)); + insert_row(&conn, 2, "not-a-blob"); drop(conn); - let SQLiteScan::Summary(summary) = - summarize(&database_roots(&dir.path().join(".gemini")), Utc::now(), 30) - else { - panic!("supported database should produce coverage"); - }; + let summary = expect_summary(&dir); assert_eq!(summary.coverage, LocalHistoryCoverage::Partial); assert!(summary.total_tokens > 0); @@ -316,26 +259,14 @@ fn undecodable_row_beside_valid_rows_yields_a_lower_bound() { #[test] fn contradicting_rows_for_one_index_are_withheld_not_published() { let dir = tempfile::tempdir().unwrap(); - let root = dir.path().join(".gemini/antigravity-cli/conversations"); - fs::create_dir_all(&root).unwrap(); let timestamp = u64::try_from(Utc::now().timestamp()).unwrap(); - let conn = Connection::open(root.join("one.db")).unwrap(); - conn.execute("CREATE TABLE gen_metadata(idx INTEGER, data BLOB)", []) - .unwrap(); + let conn = database_at(&cli_root(&dir), "one.db"); for input in [100_u64, 900_u64] { - conn.execute( - "INSERT INTO gen_metadata(idx, data) VALUES(1, ?1)", - [valid_turn_blob(input, timestamp)], - ) - .unwrap(); + insert_row(&conn, 1, valid_turn_blob(input, timestamp)); } drop(conn); - let SQLiteScan::Summary(summary) = - summarize(&database_roots(&dir.path().join(".gemini")), Utc::now(), 30) - else { - panic!("supported database should produce coverage"); - }; + let summary = expect_summary(&dir); assert_eq!(summary.coverage, LocalHistoryCoverage::Partial); assert_eq!(summary.total_tokens, 0); @@ -348,22 +279,19 @@ fn list_price_uses_prompt_plus_input_and_output_plus_reasoning() { // Upstream AntigravityLocalReaderTests: a known model gets a list-price estimate, a routing // variant prices from its base model, and an unknown model stays unpriced. let dir = tempfile::tempdir().unwrap(); - let root = dir.path().join(".gemini/antigravity-cli/conversations"); - fs::create_dir_all(&root).unwrap(); + let root = cli_root(&dir); let now = Utc::now(); let timestamp = u64::try_from(now.timestamp()).unwrap(); for (name, model) in [ ("direct.db", "claude-sonnet-4-6"), ("routed.db", "claude-sonnet-4-6-Thinking"), ] { - let conn = Connection::open(root.join(name)).unwrap(); - conn.execute("CREATE TABLE gen_metadata(idx INTEGER, data BLOB)", []) - .unwrap(); - conn.execute( - "INSERT INTO gen_metadata(idx, data) VALUES(1, ?1)", - [valid_turn_blob_with_model(100, timestamp, Some(model))], - ) - .unwrap(); + let conn = database_at(&root, name); + insert_row( + &conn, + 1, + valid_turn_blob_with_model(100, timestamp, Some(model)), + ); } let SQLiteScan::Summary(summary) = summarize(std::slice::from_ref(&root), now, 30) else { @@ -378,18 +306,12 @@ fn list_price_uses_prompt_plus_input_and_output_plus_reasoning() { assert_eq!(summary.cost_estimate.coverage.unpriced, 0); assert_eq!(summary.total_usd(), Some(per_request * 2.0)); - let conn = Connection::open(root.join("unknown.db")).unwrap(); - conn.execute("CREATE TABLE gen_metadata(idx INTEGER, data BLOB)", []) - .unwrap(); - conn.execute( - "INSERT INTO gen_metadata(idx, data) VALUES(1, ?1)", - [valid_turn_blob_with_model( - 100, - timestamp, - Some("fixture-unpriced"), - )], - ) - .unwrap(); + let conn = database_at(&root, "unknown.db"); + insert_row( + &conn, + 1, + valid_turn_blob_with_model(100, timestamp, Some("fixture-unpriced")), + ); drop(conn); let SQLiteScan::Summary(summary) = summarize(std::slice::from_ref(&root), now, 30) else { @@ -426,18 +348,16 @@ fn routine_read_offers_background_pricing_only_for_unpriced_history() { // rows; `Some(Some(model))`: one recorded request for `model`. let routine_read = |database: Option>| { let dir = tempfile::tempdir().unwrap(); - let root = dir.path().join(".gemini/antigravity-cli/conversations"); + let root = cli_root(&dir); fs::create_dir_all(&root).unwrap(); if let Some(model) = database { - let conn = Connection::open(root.join("one.db")).unwrap(); - conn.execute("CREATE TABLE gen_metadata(idx INTEGER, data BLOB)", []) - .unwrap(); + let conn = database_at(&root, "one.db"); if let Some(model) = model { - conn.execute( - "INSERT INTO gen_metadata(idx, data) VALUES(1, ?1)", - [valid_turn_blob_with_model(100, timestamp, Some(model))], - ) - .unwrap(); + insert_row( + &conn, + 1, + valid_turn_blob_with_model(100, timestamp, Some(model)), + ); } } match summarize(&[root], now, 30) { @@ -465,3 +385,202 @@ fn routine_read_offers_background_pricing_only_for_unpriced_history() { assert_eq!(unknown.total_usd(), None); assert!(background_pricing_refresh(&unknown).is_some()); } + +fn inspect_schema( + table: &str, + conn: &Connection, + budget: &mut Budget, +) -> rusqlite::Result { + let required: &[&str] = match table { + "gen_metadata" => &["idx", "data"], + _ => &["idx", "metadata"], + }; + table_schema(conn, budget, table, required) +} + +#[test] +fn schema_inspection_requires_one_stored_table_with_the_reader_columns() { + use SchemaInspection::{Incomplete, Supported, Unsupported}; + for (table, column) in [("gen_metadata", "data"), ("steps", "metadata")] { + let upper = (table.to_ascii_uppercase(), column.to_ascii_uppercase()); + let unrelated = (0..=MAX_SCHEMA_ENTRIES) + .map(|index| format!("CREATE TABLE unrelated_{index}(value TEXT);")) + .collect::(); + let cases = [ + ( + format!("CREATE TABLE {table}(idx INTEGER, {column} BLOB);"), + false, + Supported, + ), + ( + format!("CREATE TABLE {}(IDX INTEGER, {} BLOB);", upper.0, upper.1), + false, + Supported, + ), + ( + format!("CREATE TABLE {table}(idx INTEGER);"), + false, + Unsupported, + ), + ( + format!("CREATE VIEW {table} AS SELECT 1 AS idx, x'00' AS {column};"), + false, + Unsupported, + ), + ( + format!("CREATE TABLE other(idx INTEGER, {column} BLOB);"), + false, + Unsupported, + ), + (String::new(), false, Unsupported), + ( + format!( + "CREATE TABLE {table}(idx INTEGER, {column} BLOB, derived TEXT GENERATED ALWAYS AS (idx || 'x') VIRTUAL);" + ), + false, + Unsupported, + ), + ( + format!("CREATE TABLE {table}(idx INTEGER, {column} BLOB);"), + true, + Incomplete, + ), + (unrelated, false, Incomplete), + ]; + for (ddl, expired, expected) in cases { + let conn = Connection::open_in_memory().unwrap(); + conn.execute_batch(&ddl).unwrap(); + let mut budget = if expired { + Budget::with_deadline(Instant::now()) + } else { + Budget::new() + }; + assert_eq!( + inspect_schema(table, &conn, &mut budget).unwrap(), + expected, + "{table}: {ddl:.80}" + ); + } + } +} + +/// (complete, exhausted, budget rows, budget bytes, database bytes) after one read. +type RowCharge = (bool, bool, usize, usize, usize); +/// (rows as (idx, SQL value), preset (rows, bytes, database bytes), expected charge). +type RowCase<'a> = (&'a [(i64, &'a str)], (usize, usize, usize), RowCharge); + +fn charge_rows(steps: bool, rows: &[(i64, &str)], preset: (usize, usize, usize)) -> RowCharge { + let (table, column) = if steps { + ("steps", "metadata") + } else { + ("gen_metadata", "data") + }; + let conn = Connection::open_in_memory().unwrap(); + conn.execute( + &format!("CREATE TABLE {table}(idx INTEGER, {column} BLOB)"), + [], + ) + .unwrap(); + for (idx, value) in rows { + conn.execute(&format!("INSERT INTO {table} VALUES (?1, {value})"), [idx]) + .unwrap(); + } + let (preset_rows, preset_bytes, preset_database_bytes) = preset; + let mut budget = Budget::new(); + budget.rows = preset_rows; + budget.bytes = preset_bytes; + let (complete, database_bytes) = if steps { + let mut database_bytes = preset_database_bytes; + let scan = + read_step_timestamps(&conn, &HashMap::new(), &mut budget, &mut database_bytes).unwrap(); + (scan.complete, database_bytes) + } else { + assert_eq!( + preset_database_bytes, 0, + "the generation reader starts at zero" + ); + let parsed = read_generation_rows(&conn, "session", &mut budget).unwrap(); + (parsed.complete, parsed.database_bytes) + }; + ( + complete, + budget.exhausted, + budget.rows, + budget.bytes, + database_bytes, + ) +} + +#[test] +fn row_readers_skip_unusable_rows_before_charging_bytes() { + let oversized = format!("zeroblob({})", MAX_BLOB_BYTES + 1); + let near_full = MAX_TOTAL_BYTES - 5; + for steps in [false, true] { + let cases: [RowCase<'_>; 8] = [ + ( + &[(-1, "zeroblob(10)")], + (0, MAX_TOTAL_BYTES, 0), + (false, false, 1, MAX_TOTAL_BYTES, 0), + ), + ( + &[(0, "NULL")], + (0, MAX_TOTAL_BYTES, 0), + (false, false, 1, MAX_TOTAL_BYTES, 0), + ), + ( + &[(0, "'text'")], + (0, MAX_TOTAL_BYTES, 0), + (false, false, 1, MAX_TOTAL_BYTES, 0), + ), + (&[(0, "zeroblob(0)")], (0, 0, 0), (false, false, 1, 0, 0)), + ( + &[(0, oversized.as_str())], + (0, 0, 0), + (false, false, 1, MAX_BLOB_BYTES + 1, MAX_BLOB_BYTES + 1), + ), + ( + &[(0, "zeroblob(10)")], + (0, near_full, 0), + (false, true, 1, near_full, 10), + ), + ( + &[(0, "zeroblob(4)"), (1, "zeroblob(10)")], + (0, MAX_TOTAL_BYTES - 9, 0), + (false, true, 2, MAX_TOTAL_BYTES - 5, 14), + ), + ( + &[(0, "zeroblob(10)")], + (MAX_ROWS, 0, 0), + (false, true, MAX_ROWS + 1, 0, 0), + ), + ]; + for (rows, preset, expected) in cases { + assert_eq!( + charge_rows(steps, rows, preset), + expected, + "steps={steps} {rows:?} {preset:?}" + ); + } + } + // The steps reader carries the database total from the generation scan. + let full = MAX_DATABASE_BYTES - 5; + assert_eq!( + charge_rows(true, &[(0, "zeroblob(10)")], (0, 0, full)), + (false, true, 1, 0, full) + ); + assert_eq!( + charge_rows(true, &[(0, "zeroblob(5)")], (0, 0, full)), + (false, false, 1, 5, MAX_DATABASE_BYTES) + ); + + let valid = valid_turn_blob(100, 1_785_700_000); + let hex = valid + .iter() + .map(|byte| format!("{byte:02x}")) + .collect::(); + let literal = format!("x'{hex}'"); + assert_eq!( + charge_rows(false, &[(0, literal.as_str())], (0, 0, 0)), + (true, false, 1, valid.len(), valid.len()) + ); +} diff --git a/rust/src/providers/antigravity/local_step_resolver.rs b/rust/src/providers/antigravity/local_step_resolver.rs index 52b49fd72f..c2b54dad48 100644 --- a/rust/src/providers/antigravity/local_step_resolver.rs +++ b/rust/src/providers/antigravity/local_step_resolver.rs @@ -96,41 +96,30 @@ fn has_duplicate_rows(rows: impl Iterator) -> bool { mod tests { use super::*; - fn occurrences(rows: &[(i64, Option)]) -> Vec { - rows.iter() - .map(|&(row, timestamp_ms)| StepOccurrence { - row, - timestamp_ms, - bot_id: None, - }) - .collect() - } - - fn timestamps(rows: &[(i64, Option)]) -> Vec { - rows.iter() + /// `(row, timestamp)` pairs without bot ids, keyed under one `step`. + fn step(rows: &[(i64, Option)]) -> HashMap> { + let rows = rows + .iter() .map(|&(row, timestamp_ms)| StepOccurrence { row, timestamp_ms, bot_id: None, }) - .collect() + .collect(); + HashMap::from([("step".to_string(), rows)]) } #[test] fn orders_rows_and_repeats_one_shared_timestamp() { - let occurrences = - HashMap::from([("step".to_string(), occurrences(&[(0, None), (1, None)]))]); - let timestamp_map = HashMap::from([( - "step".to_string(), - timestamps(&[(20, Some(200)), (10, Some(100))]), - )]); + let occurrences = step(&[(0, None), (1, None)]); + let timestamp_map = step(&[(20, Some(200)), (10, Some(100))]); assert_eq!( resolve_step_timestamps(×tamp_map, &occurrences, &HashSet::new(), false)["step"], vec![100, 200] ); - let timestamps = HashMap::from([("step".to_string(), timestamps(&[(10, Some(100))]))]); + let timestamps = step(&[(10, Some(100))]); assert_eq!( resolve_step_timestamps(×tamps, &occurrences, &HashSet::new(), false)["step"], vec![100, 100] @@ -139,13 +128,11 @@ mod tests { #[test] fn withholds_missing_duplicate_and_conflicting_evidence() { - let occurrences = - HashMap::from([("step".to_string(), occurrences(&[(0, None), (1, None)]))]); - for timestamps in [ - timestamps(&[(10, None), (20, Some(200))]), - timestamps(&[(10, Some(100)), (10, Some(200))]), + let occurrences = step(&[(0, None), (1, None)]); + for evidence in [ + step(&[(10, None), (20, Some(200))]), + step(&[(10, Some(100)), (10, Some(200))]), ] { - let evidence = HashMap::from([("step".to_string(), timestamps)]); assert!( !resolve_step_timestamps(&evidence, &occurrences, &HashSet::new(), false) .contains_key("step") @@ -155,14 +142,8 @@ mod tests { #[test] fn embedded_timestamp_must_agree_with_aligned_step() { - let occurrences = HashMap::from([( - "step".to_string(), - occurrences(&[(0, Some(100)), (1, None)]), - )]); - let timestamps = HashMap::from([( - "step".to_string(), - timestamps(&[(10, Some(101)), (20, Some(200))]), - )]); + let occurrences = step(&[(0, Some(100)), (1, None)]); + let timestamps = step(&[(10, Some(101)), (20, Some(200))]); assert!( !resolve_step_timestamps(×tamps, &occurrences, &HashSet::new(), false) @@ -172,9 +153,8 @@ mod tests { #[test] fn unidentified_rows_do_not_reuse_one_timestamp_for_repeated_occurrences() { - let occurrences = - HashMap::from([("step".to_string(), occurrences(&[(0, None), (1, None)]))]); - let timestamps = HashMap::from([("step".to_string(), timestamps(&[(10, Some(100))]))]); + let occurrences = step(&[(0, None), (1, None)]); + let timestamps = step(&[(10, Some(100))]); assert!( !resolve_step_timestamps(×tamps, &occurrences, &HashSet::new(), true) @@ -184,8 +164,7 @@ mod tests { #[test] fn ambiguous_bot_timestamp_keeps_its_positional_slot() { - let occurrences = - HashMap::from([("step".to_string(), occurrences(&[(0, None), (1, None)]))]); + let occurrences = step(&[(0, None), (1, None)]); let timestamps = HashMap::from([( "step".to_string(), vec![ diff --git a/rust/src/providers/antigravity/mod.rs b/rust/src/providers/antigravity/mod.rs index 8c6c3fb577..011dfd2f94 100755 --- a/rust/src/providers/antigravity/mod.rs +++ b/rust/src/providers/antigravity/mod.rs @@ -17,7 +17,7 @@ mod local_step_resolver; mod offline_reason; mod quota_summary; -use legacy_status::{UserStatus, UserStatusResponse}; +use legacy_status::UserStatusResponse; use offline_reason::LiveFailure; #[cfg(windows)] @@ -106,6 +106,59 @@ fn flag_re(flag: &str) -> Regex { Regex::new(&format!("--{f}(?:\\s+|\\s*=\\s*)(\\S+)", f = flag)).expect("valid flag pattern") } +fn flag_value<'a>(re: &Regex, line: &'a str) -> Option<&'a str> { + re.captures(line).and_then(|c| c.get(1)).map(|m| m.as_str()) +} + +/// Client for the loopback language server. +/// +/// SECURITY: TLS verification is disabled because the local language server +/// uses a self-signed certificate; this client only talks to 127.0.0.1. It also +/// bypasses the app-wide outbound proxy. +fn loopback_client(timeout: std::time::Duration) -> Result { + crate::core::credentialed_http_client_builder() + .no_proxy() + .timeout(timeout) + .danger_accept_invalid_certs(true) + .redirect(reqwest::redirect::Policy::none()) + .build() + .map_err(|e| ProviderError::Other(e.to_string())) +} + +fn user_status_body() -> serde_json::Value { + serde_json::json!({ + "metadata": { + "ideName": "antigravity", + "extensionName": "antigravity", + "ideVersion": "unknown", + "locale": "en" + } + }) +} + +/// A Connect-protocol JSON POST to the local language server. +fn local_post( + client: &reqwest::Client, + url: &str, + body: &serde_json::Value, + timeout: std::time::Duration, +) -> reqwest::RequestBuilder { + client + .post(url) + .header("Content-Type", "application/json") + .header("Connect-Protocol-Version", "1") + .timeout(timeout) + .json(body) +} + +async fn response_bytes(response: reqwest::Response) -> Result, LiveFailure> { + response + .bytes() + .await + .map(|bytes| bytes.to_vec()) + .map_err(|e| LiveFailure::request("Failed to read response", &e)) +} + /// The kind of local Antigravity process a `ProcessInfo` was derived from. /// /// The desktop IDE/app language server authenticates local requests with a @@ -200,25 +253,12 @@ impl AntigravityProvider { None => (None, line), }; - let csrf_token = csrf_re - .captures(line) - .and_then(|c| c.get(1)) - .map(|m| m.as_str().to_string()); - - let ext_csrf_token = ext_csrf_re - .captures(line) - .and_then(|c| c.get(1)) - .map(|m| m.as_str().to_string()); - - let port = port_re - .captures(line) - .and_then(|c| c.get(1)) - .and_then(|m| m.as_str().parse::().ok()) + let csrf_token = flag_value(&csrf_re, line).map(str::to_string); + let ext_csrf_token = flag_value(&ext_csrf_re, line).map(str::to_string); + let port = flag_value(&port_re, line) + .and_then(|value| value.parse::().ok()) .or_else(|| { - https_port_re - .captures(line) - .and_then(|c| c.get(1)) - .and_then(|m| m.as_str().parse::().ok()) + flag_value(&https_port_re, line).and_then(|value| value.parse::().ok()) }); // Desktop IDE/app language server: requires --csrf_token. @@ -258,20 +298,9 @@ impl AntigravityProvider { // real gRPC/Connect API port is not guaranteed to be within a small window above it. // Mirror the macOS/Linux probe (which uses `lsof`) by enumerating the language-server // process's own listening ports first, then fall back to a heuristic window above the - // extension port and a few historically-seen ports. - // - // SECURITY: TLS verification is disabled because the local language server uses a - // self-signed certificate. This is scoped to 127.0.0.1 only; we confirm a port by - // checking that it answers the expected gRPC endpoint. - // The language server is a local loopback endpoint. Do not route it - // through the app-wide outbound proxy. - let client = crate::core::credentialed_http_client_builder() - .no_proxy() - .timeout(std::time::Duration::from_secs(2)) - .danger_accept_invalid_certs(true) - .redirect(reqwest::redirect::Policy::none()) - .build() - .map_err(|e| ProviderError::Other(e.to_string()))?; + // extension port and a few historically-seen ports. A port is confirmed only when it + // answers the expected gRPC endpoint. + let client = loopback_client(std::time::Duration::from_secs(2))?; // Ordered candidate ports: the process's real listening ports first (Windows // equivalent of `lsof`), then the heuristic window above the extension port, then a @@ -383,14 +412,7 @@ impl AntigravityProvider { process_info: &ProcessInfo, api_port: u16, ) -> Result { - // SECURITY: TLS verification disabled only for this loopback language server. - let client = crate::core::credentialed_http_client_builder() - .no_proxy() - .timeout(std::time::Duration::from_secs(8)) - .danger_accept_invalid_certs(true) - .redirect(reqwest::redirect::Policy::none()) - .build() - .map_err(|e| ProviderError::Other(e.to_string()))?; + let client = loopback_client(std::time::Duration::from_secs(8))?; let quota_body = serde_json::json!({ "forceRefresh": true }); match Self::fetch_local_payload( @@ -407,20 +429,12 @@ impl AntigravityProvider { Ok(mut snapshot) => { // Identity is best-effort enrichment and must not displace a // successful quota-summary result. - let identity_body = serde_json::json!({ - "metadata": { - "ideName": "antigravity", - "extensionName": "antigravity", - "ideVersion": "unknown", - "locale": "en" - } - }); if let Ok(identity_bytes) = Self::fetch_local_payload( &client, process_info, api_port, GET_USER_STATUS_PATH, - &identity_body, + &user_status_body(), std::time::Duration::from_secs(1), ) .await @@ -442,26 +456,18 @@ impl AntigravityProvider { ), } - let body = serde_json::json!({ - "metadata": { - "ideName": "antigravity", - "extensionName": "antigravity", - "ideVersion": "unknown", - "locale": "en" - } - }); let bytes = Self::fetch_local_payload( &client, process_info, api_port, GET_USER_STATUS_PATH, - &body, + &user_status_body(), std::time::Duration::from_secs(8), ) .await?; let response: UserStatusResponse = serde_json::from_slice(&bytes) .map_err(|e| ProviderError::Parse(format!("Failed to parse response: {e}")))?; - let usage = self.parse_user_status(response)?; + let usage = legacy_status::parse_user_status(response)?; Ok(Self::fetch_result(usage, AntigravityStrategyId::Local)) } @@ -507,13 +513,7 @@ impl AntigravityProvider { let Some(binary) = cli_resolution::locate_agy_binary()? else { return Ok(ManagedAgyOutcome::Missing); }; - let probe_client = crate::core::credentialed_http_client_builder() - .no_proxy() - .timeout(AGY_PROBE_TIMEOUT) - .danger_accept_invalid_certs(true) - .redirect(reqwest::redirect::Policy::none()) - .build() - .map_err(|error| ProviderError::Other(error.to_string()))?; + let probe_client = loopback_client(AGY_PROBE_TIMEOUT)?; let config = ManagedProcessConfig { program: binary, args: Vec::new(), @@ -797,12 +797,7 @@ impl AntigravityProvider { .extension_server_csrf_token .as_deref() .unwrap_or(&process_info.csrf_token); - let mut request = client - .post(&url) - .header("Content-Type", "application/json") - .header("Connect-Protocol-Version", "1") - .timeout(timeout) - .json(body); + let mut request = local_post(client, &url, body, timeout); if requires_csrf { request = request.header("X-Codeium-Csrf-Token", csrf_token); } @@ -811,42 +806,30 @@ impl AntigravityProvider { .await .map_err(|e| LiveFailure::request("API request failed", &e))?; if response.status().is_success() { - return response - .bytes() - .await - .map(|bytes| bytes.to_vec()) - .map_err(|e| LiveFailure::request("Failed to read response", &e)); + return response_bytes(response).await; } let status = response.status(); let text = response.text().await.unwrap_or_default(); if requires_csrf && process_info.extension_server_csrf_token.is_some() { - let retry = client - .post(&url) - .header("Content-Type", "application/json") - .header("Connect-Protocol-Version", "1") + let retry = local_post(client, &url, body, timeout) .header("X-Codeium-Csrf-Token", &process_info.csrf_token) - .timeout(timeout) - .json(body) .send() .await; if let Ok(retry) = retry && retry.status().is_success() { - return retry - .bytes() - .await - .map(|bytes| bytes.to_vec()) - .map_err(|e| LiveFailure::request("Failed to read response", &e)); + return response_bytes(retry).await; } } + let lower = text.to_ascii_lowercase(); if process_info.source == ProcessSource::Cli && (status == reqwest::StatusCode::UNAUTHORIZED || status == reqwest::StatusCode::FORBIDDEN - || text.to_ascii_lowercase().contains("not logged") - || text.to_ascii_lowercase().contains("login method") - || text.to_ascii_lowercase().contains("keyring")) + || ["not logged", "login method", "keyring"] + .iter() + .any(|needle| lower.contains(needle))) { return Err(ProviderError::AuthRequired.into()); } @@ -855,17 +838,6 @@ impl AntigravityProvider { ProviderError::Other(format!("API error {status}: {text}")), )) } - - fn resolve_plan_name(status: &UserStatus) -> Option { - legacy_status::resolve_plan_name(status) - } - - fn parse_user_status( - &self, - response: UserStatusResponse, - ) -> Result { - legacy_status::parse_user_status(response) - } } #[async_trait] diff --git a/rust/src/providers/antigravity/offline_reason_tests.rs b/rust/src/providers/antigravity/offline_reason_tests.rs index 00e12c9196..c9b2d5cf2f 100644 --- a/rust/src/providers/antigravity/offline_reason_tests.rs +++ b/rust/src/providers/antigravity/offline_reason_tests.rs @@ -3,7 +3,7 @@ use super::super::AntigravityProvider; use super::super::cli_print_failure::ExitReason; -use super::super::tests::offline_result; +use super::super::tests::{offline_result, run_fallback}; use super::*; use crate::core::{ProviderFetchResult, RateWindow, UsageSnapshot}; @@ -18,6 +18,14 @@ fn offline_detail_value(result: &ProviderFetchResult) -> String { details[0].value().to_string() } +/// Resolve `failure` against offline history and return its explanation. +fn offline_detail_for(failure: impl Into) -> String { + let resolved = AntigravityProvider::resolve_probe_failure(failure, Some(offline_result())) + .expect("offline history is preserved"); + assert_eq!(resolved.source_label, "offline"); + offline_detail_value(&resolved) +} + #[test] fn offline_fallback_explains_each_typed_failure() { let other = |message: &str| ProviderError::Other(message.to_string()); @@ -62,11 +70,8 @@ fn offline_fallback_explains_each_typed_failure() { (ProviderError::NoCookies.into(), HINT), ]; for (failure, reason) in cases { - let resolved = AntigravityProvider::resolve_probe_failure(failure, Some(offline_result())) - .expect("offline history is preserved"); - assert_eq!(resolved.source_label, "offline"); assert_eq!( - offline_detail_value(&resolved), + offline_detail_for(failure), format!("{OFFLINE_DETAIL_PREFIX} {reason}") ); } @@ -74,29 +79,19 @@ fn offline_fallback_explains_each_typed_failure() { #[test] fn offline_explanation_keeps_only_the_http_status() { - let body = AntigravityProvider::resolve_probe_failure( - LiveFailure::http_status( - 500, - ProviderError::Other( - r#"API error 500 Internal Server Error: {"secret":"token-value"}"#.to_string(), - ), + let body = offline_detail_for(LiveFailure::http_status( + 500, + ProviderError::Other( + r#"API error 500 Internal Server Error: {"secret":"token-value"}"#.to_string(), ), - Some(offline_result()), - ) - .expect("offline history is preserved"); - let body = offline_detail_value(&body); + )); assert!(body.contains("HTTP 500"), "{body}"); assert!(!body.contains("token-value"), "{body}"); - let expired = AntigravityProvider::resolve_probe_failure( - LiveFailure::http_status( - 403, - ProviderError::Other(r#"API error 403 Forbidden: {"detail":"private"}"#.to_string()), - ), - Some(offline_result()), - ) - .expect("offline history is preserved"); - let expired = offline_detail_value(&expired); + let expired = offline_detail_for(LiveFailure::http_status( + 403, + ProviderError::Other(r#"API error 403 Forbidden: {"detail":"private"}"#.to_string()), + )); assert!(expired.contains("session expired"), "{expired}"); assert!(!expired.contains("private"), "{expired}"); } @@ -118,19 +113,16 @@ fn http_failure_without_history_surfaces_the_original_error() { #[tokio::test] async fn offline_fallback_explains_the_masked_cli_failure() { - let result = AntigravityProvider::new() - .resolve_runtime_fallback_with_offline( - Err(LiveFailure::not_running()), - || async { - Err(LiveFailure::cli_report(CliPrintFailure::Exited { - code: 1, - reason: ExitReason::EligibilityNetwork, - })) - }, - Some(offline_result()), - ) - .await - .expect("offline history should survive a failed CLI probe"); + let (result, _) = run_fallback( + Err(LiveFailure::not_running()), + Err(LiveFailure::cli_report(CliPrintFailure::Exited { + code: 1, + reason: ExitReason::EligibilityNetwork, + })), + Some(offline_result()), + ) + .await; + let result = result.expect("offline history should survive a failed CLI probe"); assert_eq!(result.source_label, "offline"); assert_eq!( offline_detail_value(&result), @@ -140,17 +132,16 @@ async fn offline_fallback_explains_the_masked_cli_failure() { #[tokio::test] async fn unavailable_cli_keeps_the_local_failure_reason() { - let result = AntigravityProvider::new() - .resolve_runtime_fallback_with_offline( - Err(LiveFailure::http_status( - 500, - ProviderError::Other("API error 500 Internal Server Error: busy".to_string()), - )), - || async { Ok(None) }, - Some(offline_result()), - ) - .await - .expect("offline history is preserved"); + let (result, _) = run_fallback( + Err(LiveFailure::http_status( + 500, + ProviderError::Other("API error 500 Internal Server Error: busy".to_string()), + )), + Ok(None), + Some(offline_result()), + ) + .await; + let result = result.expect("offline history is preserved"); assert_eq!( offline_detail_value(&result), format!("{OFFLINE_DETAIL_PREFIX} the usage request failed (HTTP 500)") @@ -240,9 +231,7 @@ fn offline_explanation_never_echoes_error_text() { ProviderError::OAuth(leaky), ]; for error in errors { - let resolved = AntigravityProvider::resolve_probe_failure(error, Some(offline_result())) - .expect("offline history is preserved"); - let value = offline_detail_value(&resolved); + let value = offline_detail_for(error); for secret in secrets { assert!(!value.contains(secret), "{secret} leaked into {value}"); } @@ -266,28 +255,26 @@ fn auth_required_adds_no_offline_explanation() { #[tokio::test] async fn successful_live_fallback_carries_no_offline_explanation() { let live = ProviderFetchResult::new(UsageSnapshot::new(RateWindow::new(10.0)), "cli"); - let result = AntigravityProvider::new() - .resolve_runtime_fallback_with_offline( - Err(ProviderError::Timeout.into()), - || async { Ok(Some(live)) }, - Some(offline_result()), - ) - .await - .expect("live CLI fallback wins over offline history"); + let (result, _) = run_fallback( + Err(ProviderError::Timeout.into()), + Ok(Some(live)), + Some(offline_result()), + ) + .await; + let result = result.expect("live CLI fallback wins over offline history"); assert_eq!(result.source_label, "cli"); assert!(result.display_details().is_empty()); } #[tokio::test] async fn failed_cli_fallback_offline_result_explains_failure() { - let result = AntigravityProvider::new() - .resolve_runtime_fallback_with_offline( - Err(ProviderError::AuthRequired.into()), - || async { Err(LiveFailure::from(ProviderError::Timeout)) }, - Some(offline_result()), - ) - .await - .expect("offline history should survive a failed CLI probe"); + let (result, _) = run_fallback( + Err(ProviderError::AuthRequired.into()), + Err(LiveFailure::from(ProviderError::Timeout)), + Some(offline_result()), + ) + .await; + let result = result.expect("offline history should survive a failed CLI probe"); assert_eq!( offline_detail_value(&result), format!("{OFFLINE_DETAIL_PREFIX} Antigravity quota request timed out.") diff --git a/rust/src/providers/antigravity/quota_summary.rs b/rust/src/providers/antigravity/quota_summary.rs index 079f74a029..b23ad03a22 100644 --- a/rust/src/providers/antigravity/quota_summary.rs +++ b/rust/src/providers/antigravity/quota_summary.rs @@ -2,6 +2,7 @@ use chrono::{DateTime, TimeZone, Utc}; use serde::Deserialize; use std::collections::{HashMap, HashSet}; +use super::legacy_status::slug; use crate::core::{NamedRateWindow, ProviderError, RateWindow, UsageSnapshot}; const WINDOW_ID_PREFIX: &str = "antigravity-quota-summary-"; @@ -287,49 +288,35 @@ fn most_constrained_named(windows: &[NamedRateWindow], minutes: u32) -> Option<& }) } -fn group_rank(group: &QuotaSummaryGroup) -> u8 { - let title = group - .display_name - .as_deref() - .or(group.name.as_deref()) - .unwrap_or_default() - .trim() - .to_ascii_lowercase(); - if title.contains("gemini") { +/// The group's trimmed name and its rank: 0 for Gemini, 1 for Claude/GPT, +/// 2 for everything else. +fn group_family(group: &QuotaSummaryGroup) -> (Option<&str>, u8) { + let title = non_empty(group.display_name.as_deref().or(group.name.as_deref())); + let lower = title.unwrap_or_default().to_ascii_lowercase(); + let rank = if lower.contains("gemini") { 0 - } else if title.contains("claude") || title.contains("gpt") { + } else if lower.contains("claude") || lower.contains("gpt") { 1 } else { 2 - } + }; + (title, rank) +} + +fn group_rank(group: &QuotaSummaryGroup) -> u8 { + group_family(group).1 } fn group_title(group: &QuotaSummaryGroup) -> String { - let title = - non_empty(group.display_name.as_deref().or(group.name.as_deref())).unwrap_or("Quota"); - let lower = title.to_ascii_lowercase(); - if lower.contains("gemini") { - "Gemini".into() - } else if lower.contains("claude") || lower.contains("gpt") { - "Claude/GPT".into() - } else { - title.to_string() + match group_family(group) { + (_, 0) => "Gemini".into(), + (_, 1) => "Claude/GPT".into(), + (title, _) => title.unwrap_or("Quota").to_string(), } } fn group_scope(group: &QuotaSummaryGroup) -> String { - group_title(group) - .chars() - .map(|character| { - if character.is_ascii_alphanumeric() { - character.to_ascii_lowercase() - } else { - '-' - } - }) - .collect::() - .trim_matches('-') - .to_string() + slug(&group_title(group)) } fn bucket_kind(bucket: &QuotaSummaryBucket) -> BucketKind { @@ -603,6 +590,41 @@ mod tests { ); } + #[test] + fn group_rank_title_and_scope_follow_the_group_family() { + // (displayName, name, rank, title, scope) + let cases = [ + (Some("Gemini Pro"), None, 0, "Gemini", "gemini"), + (None, Some("claude & GPT"), 1, "Claude/GPT", "claude-gpt"), + ( + Some("OpenAI gpt-oss"), + Some("Gemini"), + 1, + "Claude/GPT", + "claude-gpt", + ), + (Some(" "), Some("Gemini"), 2, "Quota", "quota"), + ( + Some(" Third Party! "), + None, + 2, + "Third Party!", + "third-party", + ), + (None, None, 2, "Quota", "quota"), + ]; + for (display_name, name, rank, title, scope) in cases { + let group: QuotaSummaryGroup = serde_json::from_value(serde_json::json!({ + "displayName": display_name, + "name": name, + })) + .unwrap(); + assert_eq!(group_rank(&group), rank, "{display_name:?} {name:?}"); + assert_eq!(group_title(&group), title, "{display_name:?} {name:?}"); + assert_eq!(group_scope(&group), scope, "{display_name:?} {name:?}"); + } + } + fn single_bucket_snapshot(bucket_id: &str, display_name: &str, window: &str) -> UsageSnapshot { let data = format!( r#"{{"groups":[{{"displayName":"Gemini Models","buckets":[ diff --git a/rust/src/providers/antigravity/tests.rs b/rust/src/providers/antigravity/tests.rs index e248151f6d..1dfa8a4f34 100644 --- a/rust/src/providers/antigravity/tests.rs +++ b/rust/src/providers/antigravity/tests.rs @@ -1,4 +1,7 @@ -use super::legacy_status::{ModelFamily, canonical_model_id, classify_model}; +use super::legacy_status::{ + ModelFamily, UserStatus, canonical_model_id, classify_model, parse_user_status, + resolve_plan_name, +}; use super::*; use std::sync::{ Arc, @@ -27,24 +30,22 @@ fn explicit_quota_summary_cadence_label_survives_normalization() { #[test] fn test_classify_model_families() { - assert_eq!(classify_model("Claude 3.5 Sonnet"), ModelFamily::Claude); - assert_eq!(classify_model("claude-4-opus"), ModelFamily::Claude); - assert_eq!( - classify_model("Claude Thinking"), - ModelFamily::ClaudeThinking - ); - assert_eq!( - classify_model("claude-3.5-sonnet-thinking"), - ModelFamily::ClaudeThinking - ); - assert_eq!(classify_model("Gemini 2.5 Pro Low"), ModelFamily::GeminiPro); - assert_eq!(classify_model("gemini-pro-low"), ModelFamily::GeminiPro); - assert_eq!(classify_model("Pro Low Latency"), ModelFamily::GeminiPro); - assert_eq!(classify_model("Gemini 2.5 Flash"), ModelFamily::GeminiFlash); - assert_eq!(classify_model("gemini-flash"), ModelFamily::GeminiFlash); - assert_eq!(classify_model("Flash Model"), ModelFamily::GeminiFlash); - assert_eq!(classify_model("GPT-4o"), ModelFamily::Other); - assert_eq!(classify_model("unknown-model"), ModelFamily::Other); + for (label, family) in [ + ("Claude 3.5 Sonnet", ModelFamily::Claude), + ("claude-4-opus", ModelFamily::Claude), + ("Claude Thinking", ModelFamily::ClaudeThinking), + ("claude-3.5-sonnet-thinking", ModelFamily::ClaudeThinking), + ("Gemini 2.5 Pro Low", ModelFamily::GeminiPro), + ("gemini-pro-low", ModelFamily::GeminiPro), + ("Pro Low Latency", ModelFamily::GeminiPro), + ("Gemini 2.5 Flash", ModelFamily::GeminiFlash), + ("gemini-flash", ModelFamily::GeminiFlash), + ("Flash Model", ModelFamily::GeminiFlash), + ("GPT-4o", ModelFamily::Other), + ("unknown-model", ModelFamily::Other), + ] { + assert_eq!(classify_model(label), family, "{label}"); + } } #[test] @@ -61,54 +62,112 @@ fn retired_flash_ids_collapse_to_current_wire_id() { } #[test] -fn parses_current_language_server_process() { - let output = r"4242 C:\Users\test\AppData\Local\Programs\Antigravity\resources\bin\language_server.exe --csrf_token 11111111-2222-3333-4444-555555555555 --extension_server_port 54123"; - - let process = AntigravityProvider::parse_process_info(output).expect("process info"); - - assert_eq!(process.pid, Some(4242)); - assert_eq!(process.extension_port, Some(54123)); - assert_eq!(process.csrf_token, "11111111-2222-3333-4444-555555555555"); - assert_eq!(process.source, ProcessSource::Ide); -} - -#[test] -fn parses_language_server_without_extension_server_port() { - let output = "34564\tC:\\Users\\test\\AppData\\Local\\Programs\\Antigravity\\resources\\bin\\language_server.exe --standalone --override_ide_name antigravity --subclient_type hub --override_ide_version 2.0.11 --https_server_port 0 --csrf_token 68dda2fb-6b26-40c0-aeef-b9a628615714 --app_data_dir antigravity"; - - let process = - AntigravityProvider::parse_process_info(output).expect("process info should be detected"); - - assert_eq!(process.pid, Some(34564)); - assert_eq!(process.extension_port, Some(0)); - assert_eq!(process.csrf_token, "68dda2fb-6b26-40c0-aeef-b9a628615714"); - assert_eq!(process.source, ProcessSource::Ide); +fn unselected_model_window_ids_are_lowercase_ascii_slugs() { + let response: UserStatusResponse = serde_json::from_value(serde_json::json!({ + "userStatus": { + "cascadeModelConfigData": { + "clientModelConfigs": [ + {"label": "Claude Sonnet", "modelId": "claude-sonnet", "quotaInfo": {"remainingFraction": 0.5}}, + {"label": "Other One", "modelId": " GPT-OSS 120B (Medium) ", "quotaInfo": {"remainingFraction": 0.5}}, + {"label": "Other Two", "modelId": "!!!", "quotaInfo": {"remainingFraction": 0.5}}, + {"label": "Other Three", "id": "Mod\u{e8}le_A", "quotaInfo": {"remainingFraction": 0.5}}, + {"label": "Plain Label", "quotaInfo": {"remainingFraction": 0.5}} + ] + } + } + })) + .unwrap(); + let snap = parse_user_status(response).unwrap(); + let mut ids: Vec<_> = snap + .extra_rate_windows + .iter() + .map(|window| window.id.as_str()) + .collect(); + ids.sort_unstable(); + assert_eq!( + ids, + [ + "model-gpt-oss-120b--medium", + "model-mod-le-a", + "model-plain-label", + "model-unknown", + ] + ); } #[test] -fn parses_language_server_without_any_port_arg() { - let output = "34564\tC:\\Users\\test\\AppData\\Local\\Programs\\Antigravity\\resources\\bin\\language_server.exe --standalone --csrf_token aabbccdd-1122-3344-5566-778899001122 --app_data_dir antigravity"; - - let process = - AntigravityProvider::parse_process_info(output).expect("process info should be detected"); +fn local_post_sends_connect_json_with_the_request_timeout() { + let client = reqwest::Client::new(); + let request = local_post( + &client, + "https://127.0.0.1:1/exa.language_server_pb.LanguageServerService/GetUserStatus", + &user_status_body(), + std::time::Duration::from_secs(4), + ) + .header("X-Codeium-Csrf-Token", "tok") + .build() + .unwrap(); - assert_eq!(process.pid, Some(34564)); - assert_eq!(process.extension_port, None); - assert_eq!(process.csrf_token, "aabbccdd-1122-3344-5566-778899001122"); - assert_eq!(process.source, ProcessSource::Ide); + assert_eq!(request.method(), reqwest::Method::POST); + let headers: Vec<_> = request + .headers() + .iter() + .map(|(name, value)| (name.as_str(), value.to_str().unwrap())) + .collect(); + assert_eq!( + headers, + [ + ("content-type", "application/json"), + ("connect-protocol-version", "1"), + ("x-codeium-csrf-token", "tok"), + ] + ); + assert_eq!(request.timeout(), Some(&std::time::Duration::from_secs(4))); + let body: serde_json::Value = + serde_json::from_slice(request.body().unwrap().as_bytes().unwrap()).unwrap(); + assert_eq!( + body, + serde_json::json!({ + "metadata": { + "ideName": "antigravity", + "extensionName": "antigravity", + "ideVersion": "unknown", + "locale": "en" + } + }) + ); } #[test] -fn parses_equals_form_args() { - let output = "34564\tC:\\Users\\test\\AppData\\Local\\Programs\\Antigravity\\resources\\bin\\language_server.exe --csrf_token=68dda2fb-6b26-40c0-aeef-b9a628615714 --https_server_port=61999"; - - let process = - AntigravityProvider::parse_process_info(output).expect("process info should be detected"); - - assert_eq!(process.pid, Some(34564)); - assert_eq!(process.extension_port, Some(61999)); - assert_eq!(process.csrf_token, "68dda2fb-6b26-40c0-aeef-b9a628615714"); - assert_eq!(process.source, ProcessSource::Ide); +fn ide_flags_capture_extension_token_and_prefer_the_extension_port() { + // (command line, extension token, port) + let cases = [ + ( + "1 ls.exe --extension_server_csrf_token ext-tok --csrf_token main --extension_server_port 54123 --https_server_port 61999", + Some("ext-tok"), + Some(54123), + ), + ( + "1 ls.exe --csrf_token=main --extension_server_port=abc --https_server_port=61999", + None, + Some(61999), + ), + ( + "1 ls.exe --csrf_token main --extension_server_port 70000", + None, + None, + ), + ]; + for (line, ext_token, port) in cases { + let process = AntigravityProvider::parse_process_info(line).expect(line); + assert_eq!(process.csrf_token, "main", "{line}"); + assert_eq!( + process.extension_server_csrf_token.as_deref(), + ext_token, + "{line}" + ); + assert_eq!(process.extension_port, port, "{line}"); + } } fn make_response(models: Vec<(&str, f64)>) -> UserStatusResponse { @@ -156,7 +215,7 @@ fn antigravity_extra_windows_preserve_usage_known() { } }); let resp: UserStatusResponse = serde_json::from_value(json).unwrap(); - let snap = AntigravityProvider::new().parse_user_status(resp).unwrap(); + let snap = parse_user_status(resp).unwrap(); let claude = snap .extra_rate_windows .iter() @@ -166,98 +225,135 @@ fn antigravity_extra_windows_preserve_usage_known() { assert_eq!(claude.window.used_percent, 0.0); } -#[test] -fn test_parse_user_status_standard() { - let resp = make_response(vec![ - ("Claude 3.5 Sonnet", 0.8), - ("Gemini 2.5 Pro Low", 0.5), - ("Gemini 2.5 Flash", 0.9), - ]); - let provider = AntigravityProvider::new(); - let snap = provider.parse_user_status(resp).unwrap(); - - assert!((snap.primary.used_percent - 20.0).abs() < 0.1); - let sec = snap.secondary.unwrap(); - assert!((sec.used_percent - 50.0).abs() < 0.1); - let ter = snap.model_specific.unwrap(); - assert!((ter.used_percent - 10.0).abs() < 0.1); - assert!(snap.extra_rate_windows.is_empty()); -} - -#[test] -fn antigravity_extra_windows_preserve_all_unselected_configs() { - let resp = make_response(vec![ - ("Claude 4 Sonnet", 0.8), - ("GPT-4o", 0.8), - ("Mistral Large", 0.6), - ("Qwen Max", 0.6), - ]); - let provider = AntigravityProvider::new(); - let snap = provider.parse_user_status(resp).unwrap(); - - assert_eq!(snap.extra_rate_windows.len(), 3); - assert!( - snap.extra_rate_windows - .iter() - .any(|window| window.title == "GPT-4o") - ); - assert!( - snap.extra_rate_windows - .iter() - .any(|window| window.title == "Mistral Large") - ); - assert!( - snap.extra_rate_windows - .iter() - .any(|window| window.title == "Qwen Max") - ); -} - -#[test] -fn test_parse_user_status_thinking_skipped() { - let resp = make_response(vec![ - ("Claude Thinking", 0.6), - ("Claude 3.5 Sonnet", 0.7), - ("Gemini 2.5 Flash", 0.5), - ]); - let provider = AntigravityProvider::new(); - let snap = provider.parse_user_status(resp).unwrap(); - - assert!((snap.primary.used_percent - 30.0).abs() < 0.1); -} - -#[test] -fn test_parse_user_status_fallback_first() { - let resp = make_response(vec![("GPT-4o", 0.4), ("Mistral Large", 0.6)]); - let provider = AntigravityProvider::new(); - let snap = provider.parse_user_status(resp).unwrap(); - - assert!((snap.primary.used_percent - 60.0).abs() < 0.1); - assert!(snap.secondary.is_none()); - assert!(snap.model_specific.is_none()); -} - -#[test] -fn test_noisy_models_do_not_drive_summary_windows() { - let resp = make_response(vec![ - ("Gemini 2.5 Flash Image", 0.01), - ("Gemini 2.5 Pro Lite", 0.02), - ("Gemini autocomplete internal", 0.03), - ("Claude 4 Sonnet", 0.8), - ("Gemini 2.5 Pro Low", 0.6), - ("Gemini 2.5 Flash", 0.7), - ]); - let provider = AntigravityProvider::new(); - let snap = provider.parse_user_status(resp).unwrap(); - - assert!((snap.primary.used_percent - 20.0).abs() < 0.1); - assert!((snap.secondary.unwrap().used_percent - 40.0).abs() < 0.1); - assert!((snap.model_specific.unwrap().used_percent - 30.0).abs() < 0.1); - assert!( - snap.extra_rate_windows +/// (models, primary, secondary, model-specific used %, sorted extra titles) +type WindowCase<'a> = ( + &'a [(&'a str, f64)], + f64, + Option, + Option, + &'a [&'a str], +); + +#[test] +fn parse_user_status_selects_summary_windows_and_keeps_the_rest() { + let cases: [WindowCase<'_>; 7] = [ + ( + &[ + ("Claude 3.5 Sonnet", 0.8), + ("Gemini 2.5 Pro Low", 0.5), + ("Gemini 2.5 Flash", 0.9), + ], + 20.0, + Some(50.0), + Some(10.0), + &[], + ), + // Upstream 0.50.1 #2963: every unselected config stays visible. + ( + &[ + ("Claude 4 Sonnet", 0.8), + ("GPT-4o", 0.8), + ("Mistral Large", 0.6), + ("Qwen Max", 0.6), + ], + 20.0, + None, + None, + &["GPT-4o", "Mistral Large", "Qwen Max"], + ), + // Thinking variants never drive the Claude window. + ( + &[ + ("Claude Thinking", 0.6), + ("Claude 3.5 Sonnet", 0.7), + ("Gemini 2.5 Flash", 0.5), + ], + 30.0, + None, + Some(50.0), + &["Claude Thinking"], + ), + // Without a known family, the first model is primary. + ( + &[("GPT-4o", 0.4), ("Mistral Large", 0.6)], + 60.0, + None, + None, + &["Mistral Large"], + ), + // Noisy Gemini variants stay visible but do not drive summary windows. + ( + &[ + ("Gemini 2.5 Flash Image", 0.01), + ("Gemini 2.5 Pro Lite", 0.02), + ("Gemini autocomplete internal", 0.03), + ("Claude 4 Sonnet", 0.8), + ("Gemini 2.5 Pro Low", 0.6), + ("Gemini 2.5 Flash", 0.7), + ], + 20.0, + Some(40.0), + Some(30.0), + &[ + "Gemini 2.5 Flash Image", + "Gemini 2.5 Pro Lite", + "Gemini autocomplete internal", + ], + ), + // Equal readings are not a pool identity: both unselected configs stay + // visible even though the canonical Claude and Gemini configs are selected. + ( + &[ + ("Claude 4 Sonnet", 0.8), + ("Gemini 2.5 Pro Low", 0.5), + ("Mistral Large", 0.8), + ("Qwen Max", 0.8), + ], + 20.0, + Some(50.0), + None, + &["Mistral Large", "Qwen Max"], + ), + // Models in distinct quota buckets keep separate lanes. + ( + &[ + ("Claude 3.5 Sonnet", 0.8), + ("Claude 4 Sonnet", 0.7), + ("Gemini 2.5 Pro Low", 0.5), + ], + 30.0, + Some(50.0), + None, + &["Claude 3.5 Sonnet"], + ), + ]; + let near = |actual: Option, expected: Option| match (actual, expected) { + (Some(actual), Some(expected)) => (actual - expected).abs() < 0.1, + (actual, expected) => actual.is_none() && expected.is_none(), + }; + for (models, primary, secondary, model_specific, extras) in cases { + let snap = parse_user_status(make_response(models.to_vec())).unwrap(); + let mut titles: Vec<_> = snap + .extra_rate_windows .iter() - .any(|window| window.title == "Gemini 2.5 Flash Image") - ); + .map(|window| window.title.as_str()) + .collect(); + titles.sort_unstable(); + let actual = ( + snap.primary.used_percent, + snap.secondary.as_ref().map(|window| window.used_percent), + snap.model_specific + .as_ref() + .map(|window| window.used_percent), + ); + assert!( + near(Some(actual.0), Some(primary)) + && near(actual.1, secondary) + && near(actual.2, model_specific), + "{models:?}: {actual:?}" + ); + assert_eq!(titles, extras, "{models:?}"); + } } #[test] @@ -318,166 +414,149 @@ fn managed_auth_required_surfaces_instead_of_offline() { // ── agy CLI process matching ─────────────────────────────────────── #[test] -fn detects_agy_exe_cli_process_with_empty_csrf() { - // agy.exe hosts the language server in-process with no --csrf_token. - let output = - "7777\tC:\\Users\\test\\AppData\\Local\\agy\\bin\\agy.exe session --model gemini-2.5-pro"; - - let process = AntigravityProvider::parse_process_info(output) - .expect("agy CLI process should be detected"); - - assert_eq!(process.pid, Some(7777)); - assert_eq!(process.source, ProcessSource::Cli); - assert_eq!(process.csrf_token, ""); - assert!( - process.csrf_token.is_empty(), - "agy CLI requires no CSRF token" - ); - assert_eq!(process.extension_server_csrf_token, None); - assert_eq!(process.extension_port, None); -} - -#[test] -fn detects_quoted_agy_exe_cli_process() { - // Windows CIM quotes an executable path that contains path separators. - let output = "7777\t\"C:\\Users\\user\\AppData\\Local\\agy\\bin\\agy.exe\" --model gemini-3.7-flash-high"; - - let process = AntigravityProvider::parse_process_info(output) - .expect("quoted agy CLI process should be detected"); - - assert_eq!(process.pid, Some(7777)); - assert_eq!(process.source, ProcessSource::Cli); - assert!(process.csrf_token.is_empty()); -} - -#[test] -fn detects_bare_agy_command() { - // The CLI may appear under the bare `agy` name (no .exe suffix). - let output = "8888\tagy serve"; - - let process = AntigravityProvider::parse_process_info(output) - .expect("bare agy command should be detected"); - - assert_eq!(process.pid, Some(8888)); - assert_eq!(process.source, ProcessSource::Cli); - assert!(process.csrf_token.is_empty()); -} - -#[test] -fn detects_antigravity_cli_command() { - // Upstream also matches antigravity-cli / antigravity_cli. - let output = "9999\t/opt/homebrew/bin/antigravity-cli status"; - - let process = AntigravityProvider::parse_process_info(output) - .expect("antigravity-cli command should be detected"); - - assert_eq!(process.pid, Some(9999)); - assert_eq!(process.source, ProcessSource::Cli); - assert!(process.csrf_token.is_empty()); -} - -#[test] -fn ide_match_preferred_over_agy_cli_when_both_running() { - // When the desktop IDE server and the agy CLI are both running, the - // CSRF-protected IDE match wins (mirrors upstream process-kind precedence). - let output = "4242\tC:\\Antigravity\\language_server.exe --csrf_token deadbeef-aaaa-bbbb-cccc-dddddddddddd --extension_server_port 54123\n\ - 7777\tC:\\Users\\test\\AppData\\Local\\agy\\bin\\agy.exe session"; - - let process = - AntigravityProvider::parse_process_info(output).expect("a process should be detected"); - - assert_eq!(process.pid, Some(4242)); - assert_eq!(process.source, ProcessSource::Ide); - assert_eq!(process.csrf_token, "deadbeef-aaaa-bbbb-cccc-dddddddddddd"); -} - -#[test] -fn agy_cli_matches_when_only_cli_running() { - // No --csrf_token anywhere: only the agy CLI line should match. - let output = "7777\tC:\\Users\\test\\AppData\\Local\\agy\\bin\\agy.exe"; - - let process = AntigravityProvider::parse_process_info(output) - .expect("agy CLI should be detected when it is the only match"); - - assert_eq!(process.source, ProcessSource::Cli); - assert!(process.csrf_token.is_empty()); -} - -#[test] -fn non_antigravity_process_without_csrf_is_not_matched() { - // An unrelated tokenless process must not be mistaken for the agy CLI. - let output = "1234\tC:\\Windows\\System32\\notepad.exe"; - - let process = AntigravityProvider::parse_process_info(output); - - assert!(process.is_none(), "unrelated process must not match"); -} - -#[test] -fn is_agy_cli_command_matches_known_names() { - assert!(is_agy_cli_command("agy serve")); - assert!(is_agy_cli_command( - "C:\\Users\\test\\AppData\\Local\\agy\\bin\\agy.exe session" - )); - assert!(is_agy_cli_command( - "C:\\Users\\user\\AppData\\Local\\agy\\bin\\AGY.EXE --model gemini-3.7-flash-high" - )); - assert!(is_agy_cli_command( - "\"C:\\Users\\user\\AppData\\Local\\agy\\bin\\agy.exe\" --model gemini-3.7-flash-high" - )); - assert!(is_agy_cli_command("/usr/local/bin/antigravity-cli status")); - assert!(is_agy_cli_command("/opt/antigravity_cli run")); - assert!(is_agy_cli_command("\"C:\\Tools\\antigravity-cli\" status")); - assert!(is_agy_cli_command("\"C:\\Tools\\antigravity_cli\" run")); -} - -#[test] -fn is_agy_cli_command_rejects_unrelated_names() { - // A leading path separator prevents `notantigravity-cli` from matching. - assert!(!is_agy_cli_command( - "notagy.exe --model gemini-3.7-flash-high" - )); - assert!(!is_agy_cli_command( - "C:\\Tools\\someagy.exe --model gemini-3.7-flash-high" - )); - assert!(!is_agy_cli_command("notantigravity-cli status")); - assert!(!is_agy_cli_command("C:\\Tools\\notantigravity-cli status")); - assert!(!is_agy_cli_command("C:\\Windows\\System32\\notepad.exe")); - assert!(!is_agy_cli_command("language_server.exe --csrf_token abc")); - assert!(!is_agy_cli_command("")); -} - -// ── Upstream 0.50.1 #2963: preserve unselected quota configs ──────────────── - -#[test] -fn multiple_unselected_models_with_same_reading_remain_visible() { - // Equal readings are not a pool identity. Both unselected configs remain - // visible even though the canonical Claude and Gemini configs are selected. - let resp = make_response(vec![ - ("Claude 4 Sonnet", 0.8), - ("Gemini 2.5 Pro Low", 0.5), - ("Mistral Large", 0.8), - ("Qwen Max", 0.8), - ]); - let provider = AntigravityProvider::new(); - let snap = provider.parse_user_status(resp).unwrap(); - assert_eq!( - snap.extra_rate_windows.len(), - 2, - "distinct unselected configs with equal readings remain visible" - ); +fn parse_process_info_matches_ide_servers_and_the_agy_cli() { + // (output, expected (pid, source, csrf token, extension-server csrf token, port)) + let cases = [ + ( + r"4242 C:\Users\test\AppData\Local\Programs\Antigravity\resources\bin\language_server.exe --csrf_token 11111111-2222-3333-4444-555555555555 --extension_server_port 54123", + Some(( + Some(4242), + ProcessSource::Ide, + "11111111-2222-3333-4444-555555555555", + None, + Some(54123), + )), + ), + // No --extension_server_port: the https port is used. + ( + "34564\tC:\\Users\\test\\AppData\\Local\\Programs\\Antigravity\\resources\\bin\\language_server.exe --standalone --override_ide_name antigravity --subclient_type hub --override_ide_version 2.0.11 --https_server_port 0 --csrf_token 68dda2fb-6b26-40c0-aeef-b9a628615714 --app_data_dir antigravity", + Some(( + Some(34564), + ProcessSource::Ide, + "68dda2fb-6b26-40c0-aeef-b9a628615714", + None, + Some(0), + )), + ), + ( + "34564\tC:\\Users\\test\\AppData\\Local\\Programs\\Antigravity\\resources\\bin\\language_server.exe --standalone --csrf_token aabbccdd-1122-3344-5566-778899001122 --app_data_dir antigravity", + Some(( + Some(34564), + ProcessSource::Ide, + "aabbccdd-1122-3344-5566-778899001122", + None, + None, + )), + ), + // `--flag=value` form. + ( + "34564\tC:\\Users\\test\\AppData\\Local\\Programs\\Antigravity\\resources\\bin\\language_server.exe --csrf_token=68dda2fb-6b26-40c0-aeef-b9a628615714 --https_server_port=61999", + Some(( + Some(34564), + ProcessSource::Ide, + "68dda2fb-6b26-40c0-aeef-b9a628615714", + None, + Some(61999), + )), + ), + // agy.exe hosts the language server in-process with no --csrf_token. + ( + "7777\tC:\\Users\\test\\AppData\\Local\\agy\\bin\\agy.exe session --model gemini-2.5-pro", + Some((Some(7777), ProcessSource::Cli, "", None, None)), + ), + // Windows CIM quotes an executable path that contains path separators. + ( + "7777\t\"C:\\Users\\user\\AppData\\Local\\agy\\bin\\agy.exe\" --model gemini-3.7-flash-high", + Some((Some(7777), ProcessSource::Cli, "", None, None)), + ), + // The CLI may appear under the bare `agy` name (no .exe suffix). + ( + "8888\tagy serve", + Some((Some(8888), ProcessSource::Cli, "", None, None)), + ), + // Upstream also matches antigravity-cli / antigravity_cli. + ( + "9999\t/opt/homebrew/bin/antigravity-cli status", + Some((Some(9999), ProcessSource::Cli, "", None, None)), + ), + // When the desktop IDE server and the agy CLI are both running, the + // CSRF-protected IDE match wins (mirrors upstream process-kind precedence). + ( + "4242\tC:\\Antigravity\\language_server.exe --csrf_token deadbeef-aaaa-bbbb-cccc-dddddddddddd --extension_server_port 54123\n\ + 7777\tC:\\Users\\test\\AppData\\Local\\agy\\bin\\agy.exe session", + Some(( + Some(4242), + ProcessSource::Ide, + "deadbeef-aaaa-bbbb-cccc-dddddddddddd", + None, + Some(54123), + )), + ), + // No --csrf_token anywhere: only the agy CLI line matches. + ( + "7777\tC:\\Users\\test\\AppData\\Local\\agy\\bin\\agy.exe", + Some((Some(7777), ProcessSource::Cli, "", None, None)), + ), + // An unrelated tokenless process must not be mistaken for the agy CLI. + ("1234\tC:\\Windows\\System32\\notepad.exe", None), + ]; + for (output, expected) in cases { + let actual = AntigravityProvider::parse_process_info(output).map(|process| { + ( + process.pid, + process.source, + process.csrf_token, + process.extension_server_csrf_token, + process.extension_port, + ) + }); + let expected = expected.map(|(pid, source, csrf, extension_csrf, port)| { + ( + pid, + source, + csrf.to_string(), + extension_csrf.map(str::to_string), + port, + ) + }); + assert_eq!(actual, expected, "{output}"); + } } #[test] -fn models_in_distinct_quota_buckets_keep_separate_lanes() { - let resp = make_response(vec![ - ("Claude 3.5 Sonnet", 0.8), - ("Claude 4 Sonnet", 0.7), - ("Gemini 2.5 Pro Low", 0.5), - ]); - let provider = AntigravityProvider::new(); - let snap = provider.parse_user_status(resp).unwrap(); - assert_eq!(snap.extra_rate_windows.len(), 1); +fn is_agy_cli_command_matches_only_cli_names() { + for (command, expected) in [ + ("agy serve", true), + ( + "C:\\Users\\test\\AppData\\Local\\agy\\bin\\agy.exe session", + true, + ), + ( + "C:\\Users\\user\\AppData\\Local\\agy\\bin\\AGY.EXE --model gemini-3.7-flash-high", + true, + ), + ( + "\"C:\\Users\\user\\AppData\\Local\\agy\\bin\\agy.exe\" --model gemini-3.7-flash-high", + true, + ), + ("/usr/local/bin/antigravity-cli status", true), + ("/opt/antigravity_cli run", true), + ("\"C:\\Tools\\antigravity-cli\" status", true), + ("\"C:\\Tools\\antigravity_cli\" run", true), + // A leading path separator prevents `notantigravity-cli` from matching. + ("notagy.exe --model gemini-3.7-flash-high", false), + ( + "C:\\Tools\\someagy.exe --model gemini-3.7-flash-high", + false, + ), + ("notantigravity-cli status", false), + ("C:\\Tools\\notantigravity-cli status", false), + ("C:\\Windows\\System32\\notepad.exe", false), + ("language_server.exe --csrf_token abc", false), + ("", false), + ] { + assert_eq!(is_agy_cli_command(command), expected, "{command}"); + } } #[test] @@ -590,15 +669,39 @@ fn non_auth_failure_without_history_surfaces_error() { assert!(matches!(resolved, Err(ProviderError::Other(_)))); } +fn local_result() -> ProviderFetchResult { + ProviderFetchResult::new(UsageSnapshot::new(RateWindow::new(10.0)), "local") +} + +/// Resolve `local` with the given offline history and a CLI fallback that +/// yields `cli`. Also reports whether the CLI fallback ran. +pub(super) async fn run_fallback( + local: Result, LiveFailure>, + cli: Result, LiveFailure>, + offline: Option, +) -> (Result, bool) { + let ran = Arc::new(AtomicBool::new(false)); + let marker = Arc::clone(&ran); + let result = AntigravityProvider::new() + .resolve_runtime_fallback_with_offline( + local, + move || async move { + marker.store(true, Ordering::SeqCst); + cli + }, + offline, + ) + .await; + (result, ran.load(Ordering::SeqCst)) +} + #[tokio::test] async fn local_probe_success_does_not_run_structured_cli_fallback() { let provider = AntigravityProvider::new(); let fallback_called = Arc::new(AtomicBool::new(false)); let marker = Arc::clone(&fallback_called); - let local = ProviderFetchResult::new(UsageSnapshot::new(RateWindow::new(10.0)), "local"); - let result = provider - .resolve_runtime_fallback(Ok(Some(local)), move || async move { + .resolve_runtime_fallback(Ok(Some(local_result())), move || async move { marker.store(true, Ordering::SeqCst); Ok(Some(structured_cli_result())) }) @@ -612,28 +715,19 @@ async fn local_probe_success_does_not_run_structured_cli_fallback() { #[tokio::test] async fn local_probe_success_wins_over_an_invalid_cli_override() { - let provider = AntigravityProvider::new(); - let fallback_called = Arc::new(AtomicBool::new(false)); - let marker = Arc::clone(&fallback_called); - let local = ProviderFetchResult::new(UsageSnapshot::new(RateWindow::new(10.0)), "local"); - - let result = provider - .resolve_runtime_fallback_with_offline( - Ok(Some(local)), - move || async move { - marker.store(true, Ordering::SeqCst); - Err(LiveFailure::from(ProviderError::NotInstalled( - "ANTIGRAVITY_CLI_PATH is set but unusable".to_string(), - ))) - }, - Some(offline_result()), - ) - .await - .expect("successful local desktop probe must remain authoritative"); + let (result, fallback_ran) = run_fallback( + Ok(Some(local_result())), + Err(LiveFailure::from(ProviderError::NotInstalled( + "ANTIGRAVITY_CLI_PATH is set but unusable".to_string(), + ))), + Some(offline_result()), + ) + .await; + let result = result.expect("successful local desktop probe must remain authoritative"); assert_eq!(result.source_label, "local"); assert_eq!(result.usage.primary.used_percent, 10.0); - assert!(!fallback_called.load(Ordering::SeqCst)); + assert!(!fallback_ran); } #[tokio::test] @@ -676,35 +770,29 @@ async fn unauthenticated_local_and_unavailable_cli_paths_remain_auth_required() #[tokio::test] async fn malformed_structured_cli_json_from_fallback_is_a_parse_error() { - let result = AntigravityProvider::new() - .resolve_runtime_fallback_with_offline( - Err(ProviderError::AuthRequired.into()), - || async { - let error = quota_summary::parse_cli_usage_report(br#"{"status":"SUCCESS""#) - .expect_err("malformed JSON must fail parsing"); - Err(LiveFailure::from(error)) - }, - None, - ) - .await; + let error = quota_summary::parse_cli_usage_report(br#"{"status":"SUCCESS""#) + .expect_err("malformed JSON must fail parsing"); + let (result, _) = run_fallback( + Err(ProviderError::AuthRequired.into()), + Err(LiveFailure::from(error)), + None, + ) + .await; assert!(matches!(result, Err(ProviderError::Parse(_)))); } #[tokio::test] async fn cli_fallback_error_prefers_offline_history() { - let result = AntigravityProvider::new() - .resolve_runtime_fallback_with_offline( - Err(ProviderError::AuthRequired.into()), - || async { - Err(LiveFailure::from(ProviderError::Parse( - "Antigravity CLI usage report: malformed JSON".to_string(), - ))) - }, - Some(offline_result()), - ) - .await - .expect("offline history should survive a failed CLI probe"); + let (result, _) = run_fallback( + Err(ProviderError::AuthRequired.into()), + Err(LiveFailure::from(ProviderError::Parse( + "Antigravity CLI usage report: malformed JSON".to_string(), + ))), + Some(offline_result()), + ) + .await; + let result = result.expect("offline history should survive a failed CLI probe"); assert_eq!(result.source_label, "offline"); assert_eq!(result.usage.login_method.as_deref(), Some("offline")); @@ -712,18 +800,15 @@ async fn cli_fallback_error_prefers_offline_history() { #[tokio::test] async fn unusable_cli_override_prefers_offline_history() { - let result = AntigravityProvider::new() - .resolve_runtime_fallback_with_offline( - Err(ProviderError::AuthRequired.into()), - || async { - Err(LiveFailure::from(ProviderError::NotInstalled( - "ANTIGRAVITY_CLI_PATH is set but does not point to a usable agy file".into(), - ))) - }, - Some(offline_result()), - ) - .await - .expect("offline history should survive an unusable CLI override"); + let (result, _) = run_fallback( + Err(ProviderError::AuthRequired.into()), + Err(LiveFailure::from(ProviderError::NotInstalled( + "ANTIGRAVITY_CLI_PATH is set but does not point to a usable agy file".into(), + ))), + Some(offline_result()), + ) + .await; + let result = result.expect("offline history should survive an unusable CLI override"); assert_eq!(result.source_label, "offline"); assert_eq!(result.usage.login_method.as_deref(), Some("offline")); @@ -731,17 +816,14 @@ async fn unusable_cli_override_prefers_offline_history() { #[tokio::test] async fn unusable_cli_override_without_history_reports_the_override() { - let result = AntigravityProvider::new() - .resolve_runtime_fallback_with_offline( - Err(ProviderError::Other("local API unavailable".to_string()).into()), - || async { - Err(LiveFailure::from(ProviderError::NotInstalled( - "ANTIGRAVITY_CLI_PATH is set but does not point to a usable agy file".into(), - ))) - }, - None, - ) - .await; + let (result, _) = run_fallback( + Err(ProviderError::Other("local API unavailable".to_string()).into()), + Err(LiveFailure::from(ProviderError::NotInstalled( + "ANTIGRAVITY_CLI_PATH is set but does not point to a usable agy file".into(), + ))), + None, + ) + .await; match result { Err(ProviderError::NotInstalled(message)) => { @@ -752,91 +834,56 @@ async fn unusable_cli_override_without_history_reports_the_override() { } #[test] -fn user_tier_resolves_google_ai_ultra_plan_name() { - let json = serde_json::json!({ - "userStatus": { - "email": "user@example.com", - "planStatus": { - "planInfo": { - "planName": "Pro" - } - }, - "userTier": { +fn user_tier_name_wins_over_the_plan_status_name() { + // (userTier, login method) + for (user_tier, plan) in [ + ( + Some(serde_json::json!({ "id": "g1-ultra-tier", "name": "Google AI Ultra", "description": "Google AI Ultra" - }, - "cascadeModelConfigData": { - "clientModelConfigs": [ - { - "label": "Gemini 2.5 Pro", - "quotaInfo": {"remainingFraction": 0.8} - } - ] - } - } - }); - let resp: UserStatusResponse = serde_json::from_value(json).unwrap(); - let snap = AntigravityProvider::new().parse_user_status(resp).unwrap(); - assert_eq!(snap.login_method.as_deref(), Some("Google AI Ultra")); - assert_eq!(snap.account_email.as_deref(), Some("user@example.com")); -} - -#[test] -fn user_status_falls_back_to_plan_status_when_user_tier_is_absent() { - let json = serde_json::json!({ - "userStatus": { + })), + "Google AI Ultra", + ), + (None, "Pro"), + ] { + let mut status = serde_json::json!({ "email": "user@example.com", - "planStatus": { - "planInfo": { - "planName": "Pro" - } - }, + "planStatus": {"planInfo": {"planName": "Pro"}}, "cascadeModelConfigData": { "clientModelConfigs": [ - { - "label": "Gemini 2.5 Pro", - "quotaInfo": {"remainingFraction": 0.8} - } + {"label": "Gemini 2.5 Pro", "quotaInfo": {"remainingFraction": 0.8}} ] } + }); + if let Some(user_tier) = user_tier { + status["userTier"] = user_tier; } - }); - let resp: UserStatusResponse = serde_json::from_value(json).unwrap(); - let snap = AntigravityProvider::new().parse_user_status(resp).unwrap(); - assert_eq!(snap.login_method.as_deref(), Some("Pro")); -} - -#[test] -fn plan_name_fallback_skips_blank_user_tier_name() { - let status: UserStatus = serde_json::from_value(serde_json::json!({ - "userTier": { - "name": " ", - "description": " Google AI Ultra " - } - })) - .unwrap(); - - assert_eq!( - AntigravityProvider::resolve_plan_name(&status).as_deref(), - Some("Google AI Ultra") - ); + let resp: UserStatusResponse = + serde_json::from_value(serde_json::json!({ "userStatus": status })).unwrap(); + let snap = parse_user_status(resp).unwrap(); + assert_eq!(snap.login_method.as_deref(), Some(plan)); + assert_eq!(snap.account_email.as_deref(), Some("user@example.com")); + } } #[test] -fn plan_name_fallback_skips_blank_display_name() { - let status: UserStatus = serde_json::from_value(serde_json::json!({ - "planStatus": { - "planInfo": { - "planDisplayName": " ", - "planName": " Pro " - } - } - })) - .unwrap(); - - assert_eq!( - AntigravityProvider::resolve_plan_name(&status).as_deref(), - Some("Pro") - ); +fn plan_name_fallback_skips_blank_names() { + for (status, plan) in [ + ( + serde_json::json!({ + "userTier": {"name": " ", "description": " Google AI Ultra "} + }), + "Google AI Ultra", + ), + ( + serde_json::json!({ + "planStatus": {"planInfo": {"planDisplayName": " ", "planName": " Pro "}} + }), + "Pro", + ), + ] { + let status: UserStatus = serde_json::from_value(status).unwrap(); + assert_eq!(resolve_plan_name(&status).as_deref(), Some(plan)); + } } diff --git a/rust/src/providers/kimi/code_api.rs b/rust/src/providers/kimi/code_api.rs index 624f35fa72..3cda29ee0f 100644 --- a/rust/src/providers/kimi/code_api.rs +++ b/rust/src/providers/kimi/code_api.rs @@ -56,10 +56,7 @@ pub(crate) async fn fetch_via_code_api( let api_key = code_api_key(api_key_override.or(ctx.api_key.as_deref()))?; let base_url = code_api_base_url(region)?; let endpoint = code_api_usage_endpoint(&base_url)?; - let client = crate::core::credentialed_http_client_builder() - .timeout(std::time::Duration::from_secs(30)) - .build() - .map_err(|e| ProviderError::Other(e.to_string()))?; + let client = web::client()?; let mut request = client .get(endpoint) diff --git a/rust/src/providers/kimi/desktop_token.rs b/rust/src/providers/kimi/desktop_token.rs index ac40f793ba..f0ab4a1e1a 100644 --- a/rust/src/providers/kimi/desktop_token.rs +++ b/rust/src/providers/kimi/desktop_token.rs @@ -20,7 +20,6 @@ use base64::Engine; use std::path::{Path, PathBuf}; -use std::time::{SystemTime, UNIX_EPOCH}; pub struct KimiDesktopAuthToken; @@ -44,10 +43,6 @@ impl KimiDesktopAuthToken { /// Most recently accessed `kimi-auth` token from the signed-in Kimi /// Desktop session, or `None` when the app/database/cookie is absent or /// unreadable. Production entry point. - pub fn load() -> Option { - Self::load_for_region(super::KimiRegion::China) - } - pub fn load_for_region(region: super::KimiRegion) -> Option { let data_root = dirs::data_dir()?; Self::load_from_region(&data_root, region) @@ -101,7 +96,7 @@ impl KimiDesktopAuthToken { }) .ok() .and_then(|row| decode_cookie_value(row, aes_key)) - .filter(|token| !is_expired_jwt(token, unix_now_secs())) + .filter(|token| !is_expired_jwt(token, super::unix_now_secs())) } } @@ -133,13 +128,6 @@ fn is_expired_jwt(token: &str, now_unix: f64) -> bool { expiry.is_finite() && expiry <= now_unix } -fn unix_now_secs() -> f64 { - SystemTime::now() - .duration_since(UNIX_EPOCH) - .map(|duration| duration.as_secs_f64()) - .unwrap_or(0.0) -} - /// Decode a `(value, encrypted_value)` pair: plaintext first, AES-256-GCM /// fallback (upstream reads `value` only; Windows Chromium rows are usually /// encrypted). diff --git a/rust/src/providers/kimi/mod.rs b/rust/src/providers/kimi/mod.rs index 4e6c27c292..3f140a2c0d 100755 --- a/rust/src/providers/kimi/mod.rs +++ b/rust/src/providers/kimi/mod.rs @@ -594,19 +594,23 @@ fn ascii_header_value(raw: &str) -> String { } } -fn format_usage_amount(value: f64) -> String { - if value.is_finite() +/// `value` as an i64 when it is finite, integral and inside the i64 range. +fn exact_i64(value: f64) -> Option { + // The strict upper bound excludes 2^63, which is representable as f64 + // but has no exact i64 representation. + let integral = value.is_finite() && value.fract() == 0.0 && value >= i64::MIN as f64 - && value < i64::MAX as f64 - { - // The strict upper bound excludes 2^63, which is representable as f64 - // but has no exact i64 representation. - #[allow( - clippy::cast_possible_truncation, - reason = "finite integral value is bounded to the i64 range above" - )] - let integral = value as i64; + && value < i64::MAX as f64; + #[allow( + clippy::cast_possible_truncation, + reason = "finite integral value is bounded to the i64 range above" + )] + integral.then_some(value as i64) +} + +fn format_usage_amount(value: f64) -> String { + if let Some(integral) = exact_i64(value) { format!("{integral}") } else if value.is_finite() && value.fract() == 0.0 { // Preserve a large integral value without saturating it to i64::MAX. @@ -619,7 +623,21 @@ fn format_usage_amount(value: f64) -> String { #[cfg(test)] mod tests { use super::*; - use serde_json::json; + use serde_json::{Value, json}; + + fn web_snapshot(usage: Value, subscription: Value) -> UsageSnapshot { + let usage: KimiWebUsageResponse = serde_json::from_value(usage).unwrap(); + let subscription: KimiSubscriptionStatsResponse = + serde_json::from_value(subscription).unwrap(); + web::snapshot_from_web_usage_response(usage, Some(subscription)).unwrap() + } + + fn has_extra(snapshot: &UsageSnapshot, id: &str) -> bool { + snapshot + .extra_rate_windows + .iter() + .any(|window| window.id == id) + } #[test] fn auth_token_search_skips_unrelated_cookie_headers() { @@ -716,31 +734,29 @@ mod tests { #[test] fn parses_web_usage_with_subscription_windows() { - let usage: KimiWebUsageResponse = serde_json::from_value(json!({ - "usages": [{ - "scope": "FEATURE_CODING", - "detail": { "limit": "2048", "used": "375", "resetTime": "2026-01-09T15:23:13Z" }, - "limits": [{ - "window": { "duration": 300, "timeUnit": "TIME_UNIT_MINUTE" }, - "detail": { "limit": "100", "used": "25" } + let snapshot = web_snapshot( + json!({ + "usages": [{ + "scope": "FEATURE_CODING", + "detail": { "limit": "2048", "used": "375", "resetTime": "2026-01-09T15:23:13Z" }, + "limits": [{ + "window": { "duration": 300, "timeUnit": "TIME_UNIT_MINUTE" }, + "detail": { "limit": "100", "used": "25" } + }] }] - }] - })) - .unwrap(); - let subscription: KimiSubscriptionStatsResponse = serde_json::from_value(json!({ - "subscriptionBalance": { - "amountUsedRatio": 0.7716, - "expireTime": "2026-07-23T00:00:00Z" - }, - "ratelimitCode7d": { - "ratio": 0.0946, - "enabled": true, - "resetTime": "2026-07-13T15:28:00Z" - } - })) - .unwrap(); - - let snapshot = web::snapshot_from_web_usage_response(usage, Some(subscription)).unwrap(); + }), + json!({ + "subscriptionBalance": { + "amountUsedRatio": 0.7716, + "expireTime": "2026-07-23T00:00:00Z" + }, + "ratelimitCode7d": { + "ratio": 0.0946, + "enabled": true, + "resetTime": "2026-07-13T15:28:00Z" + } + }), + ); assert!((snapshot.primary.used_percent - 18.310546875).abs() < f64::EPSILON); assert_eq!( @@ -770,29 +786,22 @@ mod tests { fn feature_scoped_balance_is_not_the_total_usage_lane() { // Upstream 0.49.0 #2741: only the omni/subscription pool maps to the // "Total usage" lane; feature-scoped balances must not. - let usage: KimiWebUsageResponse = serde_json::from_value(json!({ - "usages": [{ - "scope": "FEATURE_CODING", - "detail": { "limit": "2048", "used": "375" } - }] - })) - .unwrap(); - let subscription: KimiSubscriptionStatsResponse = serde_json::from_value(json!({ - "subscriptionBalance": { - "amountUsedRatio": 0.5, - "feature": "FEATURE_CODING", - "type": "SUBSCRIPTION" - } - })) - .unwrap(); - - let snapshot = web::snapshot_from_web_usage_response(usage, Some(subscription)).unwrap(); - assert!( - snapshot - .extra_rate_windows - .iter() - .all(|window| window.id != "kimi-monthly") + let snapshot = web_snapshot( + json!({ + "usages": [{ + "scope": "FEATURE_CODING", + "detail": { "limit": "2048", "used": "375" } + }] + }), + json!({ + "subscriptionBalance": { + "amountUsedRatio": 0.5, + "feature": "FEATURE_CODING", + "type": "SUBSCRIPTION" + } + }), ); + assert!(!has_extra(&snapshot, "kimi-monthly")); } #[test] @@ -800,7 +809,7 @@ mod tests { // Upstream 0.49.0 #2741: when the membership Code 7-day ratio and the // primary weekly window agree (percent within 1 point, resets within // 5 minutes, weekly counter reliable), the extra row is suppressed. - let usage: KimiWebUsageResponse = serde_json::from_value(json!({ + let usage = json!({ "usages": [{ "scope": "FEATURE_CODING", "detail": { @@ -809,58 +818,27 @@ mod tests { "resetTime": "2026-08-13T15:28:00Z" } }] - })) - .unwrap(); - let subscription_matching: KimiSubscriptionStatsResponse = serde_json::from_value(json!({ - "ratelimitCode7d": { - "ratio": 0.421, - "enabled": true, - "resetTime": "2026-08-13T15:30:00Z" - } - })) - .unwrap(); - - let snapshot = - web::snapshot_from_web_usage_response(usage, Some(subscription_matching)).unwrap(); + }); + let code_7d = |ratio: f64| { + json!({ + "ratelimitCode7d": { + "ratio": ratio, + "enabled": true, + "resetTime": "2026-08-13T15:30:00Z" + } + }) + }; + let snapshot = web_snapshot(usage.clone(), code_7d(0.421)); assert!((snapshot.primary.used_percent - 42.0).abs() < f64::EPSILON); assert!( - snapshot - .extra_rate_windows - .iter() - .all(|window| window.id != "kimi-code-7d"), + !has_extra(&snapshot, "kimi-code-7d"), "matching Code 7-day row should be suppressed" ); // Diverging ratio (or missing reset evidence) keeps the row. - let subscription_diverging: KimiSubscriptionStatsResponse = serde_json::from_value(json!({ - "ratelimitCode7d": { - "ratio": 0.9, - "enabled": true, - "resetTime": "2026-08-13T15:30:00Z" - } - })) - .unwrap(); - let usage_diverging: KimiWebUsageResponse = serde_json::from_value(json!({ - "usages": [{ - "scope": "FEATURE_CODING", - "detail": { - "limit": "1000", - "used": "420", - "resetTime": "2026-08-13T15:28:00Z" - } - }] - })) - .unwrap(); - let snapshot = - web::snapshot_from_web_usage_response(usage_diverging, Some(subscription_diverging)) - .unwrap(); - assert!( - snapshot - .extra_rate_windows - .iter() - .any(|window| window.id == "kimi-code-7d") - ); + let snapshot = web_snapshot(usage, code_7d(0.9)); + assert!(has_extra(&snapshot, "kimi-code-7d")); } #[test] @@ -871,12 +849,21 @@ mod tests { } #[test] - fn oversized_integral_usage_amount_is_not_saturated_to_i64_max() { - let value = 2_f64.powi(63); - let formatted = format_usage_amount(value); - - assert!(formatted.starts_with("9223372036854775808")); - assert_ne!(formatted, i64::MAX.to_string()); + fn usage_amount_formats_integral_fractional_and_non_finite_values() { + let cases = [ + (42.0, "42"), + (-0.0, "0"), + (1.5, "1.50"), + (-2_f64.powi(63), "-9223372036854775808"), + // Not saturated to i64::MAX. + (2_f64.powi(63), "9223372036854775808"), + (1e20, "100000000000000000000"), + (f64::NAN, "NaN"), + (f64::INFINITY, "inf"), + ]; + for (value, expected) in cases { + assert_eq!(format_usage_amount(value), expected, "{value}"); + } } #[test] diff --git a/rust/src/providers/kimi/ratio_pool.rs b/rust/src/providers/kimi/ratio_pool.rs index 941f394765..e9c3f9d43e 100644 --- a/rust/src/providers/kimi/ratio_pool.rs +++ b/rust/src/providers/kimi/ratio_pool.rs @@ -57,18 +57,9 @@ fn usage_counts(detail: &KimiUsageDetail) -> Option { /// strings do not. fn integer_counter(value: Option<&serde_json::Value>) -> Option { match value? { - serde_json::Value::Number(number) => number.as_i64().or_else(|| { - let value = number.as_f64()?; - let integral = value.is_finite() - && value.fract() == 0.0 - && value >= i64::MIN as f64 - && value < i64::MAX as f64; - #[allow( - clippy::cast_possible_truncation, - reason = "finite integral value is bounded to the i64 range above" - )] - integral.then_some(value as i64) - }), + serde_json::Value::Number(number) => number + .as_i64() + .or_else(|| number.as_f64().and_then(super::exact_i64)), serde_json::Value::String(text) => text.parse().ok(), _ => None, } @@ -158,11 +149,31 @@ mod tests { use super::super::code_api::{MISSING_WEEKLY_DESCRIPTION, snapshot_from_code_api_response}; use super::super::{ KimiCodeApiUsageResponse, KimiSubscriptionStatsResponse, MONTHLY_WINDOW_ID, ProviderError, - UsageSnapshot, apply_subscription_windows, + RateWindow, UsageSnapshot, apply_subscription_windows, }; use chrono::{DateTime, Utc}; use serde_json::{Value, json}; + #[test] + fn integer_counter_accepts_exact_integers_only() { + let cases = [ + (json!(3), Some(3)), + (json!(3.0), Some(3)), + (json!(3.5), None), + (json!(i64::MIN as f64), Some(i64::MIN)), + (json!(-(i64::MIN as f64)), None), + (json!(u64::MAX), None), + (json!("12"), Some(12)), + (json!(" 12"), None), + (json!("1,000"), None), + (json!(null), None), + ]; + for (value, expected) in cases { + assert_eq!(super::integer_counter(Some(&value)), expected, "{value}"); + } + assert_eq!(super::integer_counter(None), None); + } + fn try_parse(value: Value) -> Result { let response: KimiCodeApiUsageResponse = serde_json::from_value(value).expect("fixture parses"); @@ -197,6 +208,14 @@ mod tests { ); } + /// The 5-hour rate-limit lane, asserting its percent and duration. + fn assert_rate_limit(snapshot: &UsageSnapshot, percent: f64) -> &RateWindow { + let rate_limit = snapshot.secondary.as_ref().expect("rate-limit lane"); + assert_eq!(rate_limit.used_percent, percent); + assert_eq!(rate_limit.window_minutes, Some(300)); + rate_limit + } + fn rate_limit_percent(snapshot: &UsageSnapshot) -> f64 { snapshot .secondary @@ -256,9 +275,7 @@ mod tests { } })); assert_weekly_absent(&snapshot); - let rate_limit = snapshot.secondary.as_ref().expect("rate-limit lane"); - assert_eq!(rate_limit.used_percent, 0.0); - assert_eq!(rate_limit.window_minutes, Some(300)); + let rate_limit = assert_rate_limit(&snapshot, 0.0); assert_eq!(rate_limit.resets_at, at("2026-09-16T20:15:44Z")); assert_eq!(rate_limit.reset_description, None); assert!(snapshot.tertiary.is_none()); @@ -287,9 +304,7 @@ mod tests { assert_eq!(snapshot.primary.used_percent, 12.5); assert_eq!(snapshot.primary.window_minutes, Some(10_080)); assert_eq!(snapshot.primary.resets_at, at("2026-09-20T00:00:00Z")); - let rate_limit = snapshot.secondary.as_ref().expect("rate-limit lane"); - assert_eq!(rate_limit.used_percent, 62.5); - assert_eq!(rate_limit.window_minutes, Some(300)); + assert_rate_limit(&snapshot, 62.5); assert!(snapshot.extra_rate_windows.is_empty()); } @@ -315,9 +330,7 @@ mod tests { }] })); assert_weekly_absent(&snapshot); - let rate_limit = snapshot.secondary.as_ref().expect("rate-limit lane"); - assert_eq!(rate_limit.used_percent, 25.0); - assert_eq!(rate_limit.window_minutes, Some(300)); + let rate_limit = assert_rate_limit(&snapshot, 25.0); assert_eq!( rate_limit.reset_description.as_deref(), Some("25/100 credits") @@ -496,9 +509,7 @@ mod tests { snapshot.primary.resets_at, at("2026-09-19T16:45:59.449979Z") ); - let rate_limit = snapshot.secondary.as_ref().expect("rate-limit lane"); - assert_eq!(rate_limit.used_percent, 1.0); - assert_eq!(rate_limit.window_minutes, Some(300)); + let rate_limit = assert_rate_limit(&snapshot, 1.0); assert_eq!(rate_limit.resets_at, at("2026-09-19T14:45:59.449979Z")); assert!(snapshot.tertiary.is_none()); assert!(snapshot.extra_rate_windows.is_empty()); @@ -564,9 +575,7 @@ mod tests { response["limits"][0]["window"]["duration"] = json!(120); let snapshot = parse(response); assert_eq!(snapshot.primary.used_percent, 19.0); - let rate_limit = snapshot.secondary.as_ref().expect("rate-limit lane"); - assert_eq!(rate_limit.used_percent, 0.0); - assert_eq!(rate_limit.window_minutes, Some(300)); + assert_rate_limit(&snapshot, 0.0); } // Upstream `usageCounts`: an invalid `used` falls back to a valid @@ -607,9 +616,7 @@ mod tests { let snapshot = parse(response); assert_eq!(snapshot.primary.used_percent, 0.0); assert_eq!(snapshot.primary.window_minutes, Some(10_080)); - let rate_limit = snapshot.secondary.as_ref().expect("rate-limit lane"); - assert_eq!(rate_limit.used_percent, 0.0); - assert_eq!(rate_limit.window_minutes, Some(300)); + assert_rate_limit(&snapshot, 0.0); } } @@ -650,9 +657,7 @@ mod tests { #[test] fn missing_legacy_window_does_not_override_zero_session_ratio() { let snapshot = zero_session_ratio_with_legacy_window(None); - let rate_limit = snapshot.secondary.as_ref().expect("rate-limit lane"); - assert_eq!(rate_limit.window_minutes, Some(300)); - assert_eq!(rate_limit.used_percent, 0.0); + assert_rate_limit(&snapshot, 0.0); } #[test] @@ -661,8 +666,6 @@ mod tests { "duration": 300, "timeUnit": "TIME_UNIT_FORTNIGHT" }))); - let rate_limit = snapshot.secondary.as_ref().expect("rate-limit lane"); - assert_eq!(rate_limit.window_minutes, Some(300)); - assert_eq!(rate_limit.used_percent, 0.0); + assert_rate_limit(&snapshot, 0.0); } } diff --git a/rust/src/providers/kimi/web.rs b/rust/src/providers/kimi/web.rs index c6c1b8f59f..02d6976b93 100644 --- a/rust/src/providers/kimi/web.rs +++ b/rust/src/providers/kimi/web.rs @@ -53,6 +53,13 @@ pub(super) struct WebFetchFailure { } impl WebFetchFailure { + fn before_token(error: ProviderError) -> Self { + Self { + error, + had_token: false, + } + } + fn after_token(error: ProviderError) -> Self { Self { error, @@ -61,6 +68,18 @@ impl WebFetchFailure { } } +/// One automatic token attempt. `None` means the server rejected the token +/// and the chain moves on; anything else ends it. +fn settle_attempt( + result: Result, +) -> Option> { + match result { + Ok(usage) => Some(Ok(usage)), + Err(ProviderError::AuthRequired) => None, + Err(error) => Some(Err(WebFetchFailure::after_token(error))), + } +} + /// Web auth token chain for both the web fetch and the Code-API enrichment /// (upstream `KimiWebEnrichmentTokenResolver.resolve`): /// 1. Manual cookie header (its `kimi-auth`/auth cookie), source-independent. @@ -194,24 +213,9 @@ pub(super) async fn fetch_web_session_isolated( // its own session cookie and never fall back to another ambient credential // or browser account (upstream per-provider token account routing). if account_isolated { - let token = match selected_account_auth_token(cookie_header) { - Ok(token) => token, - Err(error) => { - return Err(WebFetchFailure { - error, - had_token: false, - }); - } - }; - let http = match client() { - Ok(http) => http, - Err(error) => { - return Err(WebFetchFailure { - error, - had_token: false, - }); - } - }; + let token = + selected_account_auth_token(cookie_header).map_err(WebFetchFailure::before_token)?; + let http = client().map_err(WebFetchFailure::before_token)?; return fetch_via_web_token(&http, &token, region) .await .map_err(WebFetchFailure::after_token); @@ -267,50 +271,41 @@ where } if !browser_import_allowed(input.cookie_source) { - return Err(WebFetchFailure { - error: browser_import_error(input.cookie_source), - had_token: false, - }); + return Err(WebFetchFailure::before_token(browser_import_error( + input.cookie_source, + ))); } // Read and try the desktop session first. Browser cookies are intentionally // read only after the server rejects this automatic session, so a healthy // desktop account never causes another credential store to be touched. let desktop_token = (input.desktop_token)(input.region); - if let Some(token) = desktop_token.clone() { - match fetch(token).await { - Ok(usage) => return Ok(usage), - Err(ProviderError::AuthRequired) => {} - Err(error) => return Err(WebFetchFailure::after_token(error)), - } + if let Some(token) = desktop_token.clone() + && let Some(outcome) = settle_attempt(fetch(token).await) + { + return outcome; } let browser_token = (input.browser_token)(input.region).filter(|token| desktop_token.as_ref() != Some(token)); - if let Some(token) = browser_token.clone() { - match fetch(token).await { - Ok(usage) => return Ok(usage), - Err(ProviderError::AuthRequired) => {} - Err(error) => return Err(WebFetchFailure::after_token(error)), - } + if let Some(token) = browser_token.clone() + && let Some(outcome) = settle_attempt(fetch(token).await) + { + return outcome; } // Local-storage tokens are read last, only after every cookie source // was rejected. - let mut seen = std::collections::HashSet::new(); - if let Some(token) = desktop_token.clone() { - seen.insert(token); - } - if let Some(token) = browser_token.clone() { - seen.insert(token); - } + let mut seen: std::collections::HashSet = desktop_token + .iter() + .chain(&browser_token) + .cloned() + .collect(); for candidate in resolve_web_tokens(input).into_iter().filter(|candidate| { candidate.source == WebTokenSource::LocalStorage && seen.insert(candidate.token.clone()) }) { - match fetch(candidate.token).await { - Ok(usage) => return Ok(usage), - Err(ProviderError::AuthRequired) => {} - Err(error) => return Err(WebFetchFailure::after_token(error)), + if let Some(outcome) = settle_attempt(fetch(candidate.token).await) { + return outcome; } } @@ -330,7 +325,7 @@ fn selected_account_auth_token(cookie_header: Option<&str>) -> Result Result { +pub(super) fn client() -> Result { crate::core::credentialed_http_client_builder() .timeout(std::time::Duration::from_secs(30)) .build() @@ -381,14 +376,17 @@ async fn fetch_subscription_details( // completed statistics response when the plan endpoint is slow or absent. let stats = tokio::time::timeout( SUBSCRIPTION_ENRICHMENT_TIMEOUT, - fetch_subscription_for_enrichment(client, token, region), + fetch_subscription_for_enrichment_result(client, token, region), ); let plan = tokio::time::timeout( SUBSCRIPTION_ENRICHMENT_TIMEOUT, fetch_subscription_plan(client, token, region), ); let (stats, plan) = tokio::join!(stats, plan); - (stats.ok().flatten(), plan.ok().flatten()) + ( + stats.ok().and_then(Result::ok).flatten(), + plan.ok().flatten(), + ) } pub(super) fn snapshot_from_web_usage_response( @@ -444,19 +442,6 @@ pub(super) async fn fetch_subscription_plan( } } -// Kept for `code_api`: resolve the subscription stats snapshot with a web -// token; any failure means "no enrichment", never an error. -pub(super) async fn fetch_subscription_for_enrichment( - client: &Client, - token: &str, - region: KimiRegion, -) -> Option { - fetch_subscription_for_enrichment_result(client, token, region) - .await - .ok() - .flatten() -} - pub(super) async fn fetch_subscription_for_enrichment_result( client: &Client, token: &str, @@ -648,6 +633,120 @@ mod tests { Vec::new() } + thread_local! { + static READS: std::cell::Cell<[u32; 3]> = const { std::cell::Cell::new([0; 3]) }; + } + + fn count_read(source: usize) { + READS.with(|reads| { + let mut counts = reads.get(); + counts[source] += 1; + reads.set(counts); + }); + } + + /// Desktop, browser and local-storage reads since the last call. + fn take_reads() -> [u32; 3] { + READS.with(|reads| reads.replace([0; 3])) + } + + fn counted_desktop(_: KimiRegion) -> Option { + count_read(0); + Some("desktop-token".to_string()) + } + + fn counted_browser(_: KimiRegion) -> Option { + count_read(1); + Some("browser-token".to_string()) + } + + fn counted_no_desktop(_: KimiRegion) -> Option { + count_read(0); + None + } + + fn counted_no_browser(_: KimiRegion) -> Option { + count_read(1); + None + } + + fn counted_local_storage(_: KimiRegion) -> Vec { + count_read(2); + vec!["browser-token".to_string(), "storage-token".to_string()] + } + + fn accept_storage_only(token: &str) -> Result { + if token == "storage-token" { + Ok(usage(40.0)) + } else { + Err(ProviderError::AuthRequired) + } + } + + #[tokio::test] + async fn rejected_local_storage_only_token_reports_none_tried() { + take_reads(); + let (result, sent) = run_chain( + WebTokenInput { + local_storage_tokens: counted_local_storage, + ..input(None, "auto", counted_no_desktop, counted_no_browser) + }, + reject_all, + ) + .await; + let failure = result.expect_err("local-storage token rejected"); + // Kept quirk: a rejected local-storage-only token reads as "no token". + assert!(!failure.had_token); + assert!(matches!(failure.error, ProviderError::AuthRequired)); + assert_eq!(sent, ["browser-token", "storage-token"]); + assert_eq!(take_reads(), [2, 2, 2]); + } + + #[tokio::test] + async fn local_storage_tokens_follow_rejected_cookie_tokens() { + take_reads(); + let (result, sent) = run_chain( + WebTokenInput { + local_storage_tokens: counted_local_storage, + ..input(None, "auto", counted_desktop, counted_browser) + }, + accept_storage_only, + ) + .await; + assert_eq!( + result.expect("storage token accepted").primary.used_percent, + 40.0 + ); + assert_eq!(sent, ["desktop-token", "browser-token", "storage-token"]); + assert_eq!(take_reads(), [2, 2, 1]); + + let (result, sent) = run_chain( + WebTokenInput { + local_storage_tokens: counted_local_storage, + ..input(None, "auto", counted_desktop, counted_browser) + }, + reject_all, + ) + .await; + assert!(result.expect_err("every token rejected").had_token); + assert_eq!(sent, ["desktop-token", "browser-token", "storage-token"]); + assert_eq!(take_reads(), [2, 2, 1]); + + let (result, sent) = run_chain( + WebTokenInput { + local_storage_tokens: counted_local_storage, + ..input(None, "auto", counted_no_desktop, counted_no_browser) + }, + server_error, + ) + .await; + let failure = result.expect_err("server error"); + assert!(failure.had_token); + assert!(matches!(failure.error, ProviderError::Other(message) if message.contains("500"))); + assert_eq!(sent, ["browser-token"]); + assert_eq!(take_reads(), [2, 2, 1]); + } + fn static_local_storage(_: KimiRegion) -> Vec { vec!["browser-token".to_string(), "storage-token".to_string()] } diff --git a/rust/src/providers/minimax/billing.rs b/rust/src/providers/minimax/billing.rs new file mode 100644 index 0000000000..5098c2c6b9 --- /dev/null +++ b/rust/src/providers/minimax/billing.rs @@ -0,0 +1,278 @@ +//! MiniMax billing-history summary: parse, aggregate and attach as extra windows. + +use chrono::{DateTime, Duration, TimeZone, Utc}; +use serde::Deserialize; +use std::collections::HashMap; + +use super::{format_count, scalar_string, value_i64}; +use crate::core::{CostSnapshot, ProviderError, ProviderFetchResult, RateWindow}; +use crate::providers::json; + +#[derive(Debug, Deserialize)] +struct MiniMaxBillingHistoryPayload { + #[serde(default)] + base_resp: Option, + #[serde(default)] + charge_records: Vec, +} + +#[derive(Debug, Deserialize)] +struct MiniMaxBaseResponse { + status_code: Option, + status_msg: Option, +} + +#[derive(Debug, Deserialize)] +pub(super) struct MiniMaxBillingRecord { + pub(super) consume_token: Option, + pub(super) consume_input_token: Option, + pub(super) consume_output_token: Option, + pub(super) consume_cash: Option, + pub(super) consume_cash_after_voucher: Option, + pub(super) created_at: Option, + pub(super) ymd: Option, + pub(super) consume_time: Option, + pub(super) method: Option, + pub(super) model: Option, + pub(super) result: Option, + pub(super) status: Option, +} + +#[derive(Debug, Clone)] +pub(super) struct MiniMaxBillingSummary { + pub(super) today_tokens: i64, + pub(super) last_30_days_tokens: i64, + pub(super) today_cash: Option, + pub(super) last_30_days_cash: Option, + pub(super) top_methods: Vec, + pub(super) top_models: Vec, +} + +#[derive(Debug, Clone)] +pub(super) struct MiniMaxBillingBreakdown { + pub(super) name: String, + pub(super) tokens: i64, + pub(super) cash: Option, +} + +pub(super) fn parse_billing_summary( + json: &serde_json::Value, +) -> Result { + let payload: MiniMaxBillingHistoryPayload = serde_json::from_value(json.clone()) + .map_err(|e| ProviderError::Parse(format!("Failed to parse MiniMax billing: {e}")))?; + if let Some(base) = payload.base_resp + && let Some(status) = base.status_code + && status != 0 + { + return Err(ProviderError::Other( + base.status_msg + .unwrap_or_else(|| format!("MiniMax billing status {status}")), + )); + } + if payload.charge_records.is_empty() { + return Err(ProviderError::Parse( + "MiniMax billing records not present".to_string(), + )); + } + Ok(aggregate_billing(&payload.charge_records, Utc::now())) +} + +pub(super) fn aggregate_billing( + records: &[MiniMaxBillingRecord], + now: DateTime, +) -> MiniMaxBillingSummary { + let today_start = now.date_naive().and_hms_opt(0, 0, 0).unwrap().and_utc(); + let window_start = today_start - Duration::days(29); + let mut today_tokens = 0; + let mut last_30_days_tokens = 0; + let mut today_cash = 0.0; + let mut today_has_cash = false; + let mut last_30_days_cash = 0.0; + let mut last_30_has_cash = false; + let mut method_totals: HashMap = HashMap::new(); + let mut model_totals: HashMap = HashMap::new(); + + for record in records { + if !billing_record_succeeded(record) { + continue; + } + let Some(date) = record_date(record) else { + continue; + }; + if date < window_start || date > now { + continue; + } + let tokens = record_token_count(record); + let cash = record_cash(record); + last_30_days_tokens += tokens; + if let Some(cash) = cash { + last_30_days_cash += cash; + last_30_has_cash = true; + } + if date >= today_start { + today_tokens += tokens; + if let Some(cash) = cash { + today_cash += cash; + today_has_cash = true; + } + } + add_breakdown(&mut method_totals, record.method.as_deref(), tokens, cash); + add_breakdown(&mut model_totals, record.model.as_deref(), tokens, cash); + } + + MiniMaxBillingSummary { + today_tokens, + last_30_days_tokens, + today_cash: today_has_cash.then_some(today_cash), + last_30_days_cash: last_30_has_cash.then_some(last_30_days_cash), + top_methods: top_breakdowns(method_totals), + top_models: top_breakdowns(model_totals), + } +} + +fn billing_record_succeeded(record: &MiniMaxBillingRecord) -> bool { + let status = + scalar_string(record.result.as_ref()).or_else(|| scalar_string(record.status.as_ref())); + match status + .as_deref() + .map(str::trim) + .filter(|value| !value.is_empty()) + { + None => true, + Some(value) => value.eq_ignore_ascii_case("success"), + } +} + +pub(super) fn attach_billing_summary( + mut result: ProviderFetchResult, + summary: MiniMaxBillingSummary, +) -> ProviderFetchResult { + let mut rows = vec![ + ( + "billing-tokens-today".to_string(), + "Tokens today".to_string(), + format_count(summary.today_tokens), + ), + ( + "billing-tokens-30d".to_string(), + "Tokens (30 days)".to_string(), + format_count(summary.last_30_days_tokens), + ), + ]; + if let Some(cash) = summary.today_cash { + rows.push(( + "billing-cash-today".to_string(), + "Spend today".to_string(), + format!("${cash:.2}"), + )); + } + if let Some(cash) = summary.last_30_days_cash { + rows.push(( + "billing-cash-30d".to_string(), + "Spend (30 days)".to_string(), + format!("${cash:.2}"), + )); + result.cost = Some(CostSnapshot::new(cash, "USD", "Last 30 days")); + } + for (kind, label, items) in [ + ("method", "Method", &summary.top_methods), + ("model", "Model", &summary.top_models), + ] { + for (idx, item) in items.iter().enumerate() { + rows.push(( + format!("billing-{kind}-{idx}"), + format!("{label}: {}", item.name), + breakdown_description(item), + )); + } + } + for (id, title, detail) in rows { + result.usage = result.usage.with_extra_rate_window( + id, + title, + RateWindow::with_details(0.0, None, None, Some(detail)), + ); + } + result +} + +fn add_breakdown( + totals: &mut HashMap, + raw_name: Option<&str>, + tokens: i64, + cash: Option, +) { + let Some(name) = raw_name.map(str::trim).filter(|s| !s.is_empty()) else { + return; + }; + let total = totals.entry(name.to_string()).or_default(); + total.0 += tokens; + if let Some(cash) = cash { + total.1 += cash; + total.2 = true; + } +} + +fn top_breakdowns(totals: HashMap) -> Vec { + let mut items: Vec<_> = totals + .into_iter() + .map(|(name, (tokens, cash, has_cash))| MiniMaxBillingBreakdown { + name, + tokens, + cash: has_cash.then_some(cash), + }) + .collect(); + items.sort_by(|a, b| b.tokens.cmp(&a.tokens).then_with(|| a.name.cmp(&b.name))); + items.truncate(3); + items +} + +fn breakdown_description(item: &MiniMaxBillingBreakdown) -> String { + match item.cash { + Some(cash) => format!("{} tokens / ${cash:.2}", format_count(item.tokens)), + None => format!("{} tokens", format_count(item.tokens)), + } +} + +fn record_token_count(record: &MiniMaxBillingRecord) -> i64 { + let direct = value_i64(record.consume_token.as_ref()).unwrap_or(0); + if direct > 0 { + return direct; + } + value_i64(record.consume_input_token.as_ref()).unwrap_or(0) + + value_i64(record.consume_output_token.as_ref()).unwrap_or(0) +} + +fn record_cash(record: &MiniMaxBillingRecord) -> Option { + json::lenient_f64(record.consume_cash_after_voucher.as_ref()) + .or_else(|| json::lenient_f64(record.consume_cash.as_ref())) +} + +fn record_date(record: &MiniMaxBillingRecord) -> Option> { + if let Some(created_at) = value_i64(record.created_at.as_ref()) { + let seconds = if created_at > 1_000_000_000_000 { + created_at / 1000 + } else { + created_at + }; + return Utc.timestamp_opt(seconds, 0).single(); + } + if let Some(ymd) = record.ymd.as_deref() { + for format in ["%Y-%m-%d", "%Y%m%d", "%Y/%m/%d"] { + if let Ok(date) = chrono::NaiveDate::parse_from_str(ymd.trim(), format) { + return Some(date.and_hms_opt(0, 0, 0)?.and_utc()); + } + } + } + if let Some(text) = record.consume_time.as_deref() { + for format in ["%Y-%m-%d %H:%M:%S", "%Y/%m/%d %H:%M:%S"] { + if let Ok(date) = chrono::NaiveDateTime::parse_from_str(text.trim(), format) { + return Some(date.and_utc()); + } + } + if let Ok(date) = DateTime::parse_from_rfc3339(text.trim()) { + return Some(date.with_timezone(&Utc)); + } + } + None +} diff --git a/rust/src/providers/minimax/coding_plan.rs b/rust/src/providers/minimax/coding_plan.rs index 25694854c7..fc3a3bbf05 100644 --- a/rust/src/providers/minimax/coding_plan.rs +++ b/rust/src/providers/minimax/coding_plan.rs @@ -65,12 +65,98 @@ pub(super) struct RemainsRow { // for non-weekly rows, IANA timezone names, host/env URL overrides // (`MiniMaxSettingsReader`). +/// MiniMax console times are China Standard Time. +const CN_TZ: FixedOffset = FixedOffset::east_opt(8 * 3600).unwrap(); + /// Look up a value trying snake_case then camelCase (upstream normalizes the /// camelCase aliases onto the snake_case keys). fn get_field<'a>(obj: &'a Value, snake: &str, camel: &str) -> Option<&'a Value> { obj.get(snake).or_else(|| obj.get(camel)) } +/// Snake/camel key pairs for one quota lane of a `model_remains[]` entry. +struct LaneKeys { + total: [&'static str; 2], + usage: [&'static str; 2], + remaining_percent: [&'static str; 2], + status: [&'static str; 2], + start: [&'static str; 2], + end: [&'static str; 2], + remains: [&'static str; 2], +} + +const INTERVAL_LANE: LaneKeys = LaneKeys { + total: ["current_interval_total_count", "currentIntervalTotalCount"], + usage: ["current_interval_usage_count", "currentIntervalUsageCount"], + remaining_percent: [ + "current_interval_remaining_percent", + "currentIntervalRemainingPercent", + ], + status: ["current_interval_status", "currentIntervalStatus"], + start: ["start_time", "startTime"], + end: ["end_time", "endTime"], + remains: ["remains_time", "remainsTime"], +}; + +const WEEKLY_LANE: LaneKeys = LaneKeys { + total: ["current_weekly_total_count", "currentWeeklyTotalCount"], + usage: ["current_weekly_usage_count", "currentWeeklyUsageCount"], + remaining_percent: [ + "current_weekly_remaining_percent", + "currentWeeklyRemainingPercent", + ], + status: ["current_weekly_status", "currentWeeklyStatus"], + start: ["weekly_start_time", "weeklyStartTime"], + end: ["weekly_end_time", "weeklyEndTime"], + remains: ["weekly_remains_time", "weeklyRemainsTime"], +}; + +/// Raw lane fields. MiniMax's `usage_count` is the remaining count. +#[derive(Clone, Copy)] +struct LaneFields { + total: Option, + remaining: Option, + remaining_percent: Option, + status: Option, + start: Option, + end: Option, + remains: Option, +} + +impl LaneKeys { + fn read(&self, entry: &Value) -> LaneFields { + let field = |[snake, camel]: [&str; 2]| get_field(entry, snake, camel); + LaneFields { + total: value_i64(field(self.total)), + remaining: value_i64(field(self.usage)), + remaining_percent: json::lenient_f64(field(self.remaining_percent)), + status: value_i64(field(self.status)), + start: value_i64(field(self.start)), + end: value_i64(field(self.end)), + remains: value_i64(field(self.remains)), + } + } +} + +/// The next `time` today in `tz`, or a day from now once it has passed. +pub(super) fn next_time_today( + time: chrono::NaiveTime, + tz: FixedOffset, + now: DateTime, +) -> Option> { + let dt = now + .date_naive() + .and_time(time) + .and_local_timezone(tz) + .single()?; + let candidate = dt.with_timezone(&Utc); + Some(if candidate < now { + now + Duration::days(1) + } else { + candidate + }) +} + /// Decode an epoch integer that may be in milliseconds, seconds, or unparseable. fn date_from_epoch(value: Option) -> Option> { let raw = value?; @@ -247,7 +333,7 @@ fn time_range_string(start: Option>, end: Option>) - let (Some(start), Some(end)) = (start, end) else { return "N/A".to_string(); }; - let tz = FixedOffset::east_opt(8 * 3600).unwrap(); + let tz = CN_TZ; format!( "{}-{}(UTC+8)", start.with_timezone(&tz).format("%H:%M"), @@ -262,7 +348,7 @@ fn weekly_time_range_string( ) -> Option { let start = start?; let end = end?; - let tz = FixedOffset::east_opt(8 * 3600).unwrap(); + let tz = CN_TZ; Some(format!( "{} - {}(UTC+8)", start.with_timezone(&tz).format("%m/%d %H:%M"), @@ -282,25 +368,15 @@ fn reset_description( { let interval = (resets - now).num_seconds(); if interval < 60 { - return format!("Resets in {} seconds", interval); - } else if interval < 3600 { - let minutes = interval / 60; - return format!( - "Resets in {} minute{}", - minutes, - if minutes == 1 { "" } else { "s" } - ); - } else if interval < 86400 { - let hours = interval / 3600; - return format!( - "Resets in {} hour{}", - hours, - if hours == 1 { "" } else { "s" } - ); - } else { - let days = interval / 86400; - return format!("Resets in {} day{}", days, if days == 1 { "" } else { "s" }); + return format!("Resets in {interval} seconds"); } + let (unit_seconds, unit) = [(86_400, "day"), (3_600, "hour")] + .into_iter() + .find(|(unit_seconds, _)| interval >= *unit_seconds) + .unwrap_or((60, "minute")); + let count = interval / unit_seconds; + let plural = if count == 1 { "" } else { "s" }; + return format!("Resets in {count} {unit}{plural}"); } format!("{window_type}: {time_range}") } @@ -308,23 +384,22 @@ fn reset_description( /// Build a `RemainsRow` from the raw interval/weekly fields (upstream /// `makeServiceUsage`). Returns `None` when the row is a placeholder that /// should be skipped. -#[allow( - clippy::too_many_arguments, - reason = "signature mirrors the flat upstream makeServiceUsage payload fields one-to-one" -)] fn make_remains_row( service_type: &str, window_type_override: Option<&str>, - total: Option, - remaining: Option, - remaining_percent: Option, - status: Option, - start: Option, - end: Option, - remains_time: Option, + lane: &LaneFields, now: DateTime, is_weekly: bool, ) -> Option { + let LaneFields { + total, + remaining, + remaining_percent, + status, + start, + end, + remains: remains_time, + } = *lane; if is_unavailable_quota_placeholder( service_type, window_type_override, @@ -434,7 +509,7 @@ fn parse_resets_at_from_time_range( now: DateTime, ) -> Option> { let lower = window_type.trim().to_lowercase(); - let tz = FixedOffset::east_opt(8 * 3600).unwrap(); + let tz = CN_TZ; if lower == "today" { let parts: Vec<&str> = time_range.splitn(2, '-').collect(); @@ -459,13 +534,7 @@ fn parse_resets_at_from_time_range( re.replace_all(end_part, "").trim().to_string() }; let time = chrono::NaiveTime::parse_from_str(&end_clean, "%H:%M").ok()?; - let today = now.date_naive(); - let dt = today.and_time(time).and_local_timezone(tz).single()?; - let mut candidate = dt.with_timezone(&Utc); - if candidate < now { - candidate = now + Duration::days(1); - } - return Some(candidate); + return next_time_today(time, tz, now); } None @@ -506,30 +575,15 @@ fn parse_plan_name_from_data(data: &Value) -> Option { let name = scalar_string(get_field(e, "model_name", "modelName")) .map(|n| n.trim().to_lowercase()) .unwrap_or_default(); + let lane = INTERVAL_LANE.read(e); name == "video" && is_unavailable_quota_placeholder( "Text to Video", None, - value_i64(get_field( - e, - "current_interval_status", - "currentIntervalStatus", - )), - value_i64(get_field( - e, - "current_interval_total_count", - "currentIntervalTotalCount", - )), - value_i64(get_field( - e, - "current_interval_usage_count", - "currentIntervalUsageCount", - )), - json::lenient_f64(get_field( - e, - "current_interval_remaining_percent", - "currentIntervalRemainingPercent", - )), + lane.status, + lane.total, + lane.remaining, + lane.remaining_percent, ) }); if has_text_generation && has_unavailable_video { @@ -599,35 +653,9 @@ fn parse_remains( let service_type = map_model_name_to_service_type(&model_name); // Interval row - if let Some(mut row) = make_remains_row( - &service_type, - None, - value_i64(get_field( - entry, - "current_interval_total_count", - "currentIntervalTotalCount", - )), - value_i64(get_field( - entry, - "current_interval_usage_count", - "currentIntervalUsageCount", - )), - json::lenient_f64(get_field( - entry, - "current_interval_remaining_percent", - "currentIntervalRemainingPercent", - )), - value_i64(get_field( - entry, - "current_interval_status", - "currentIntervalStatus", - )), - value_i64(get_field(entry, "start_time", "startTime")), - value_i64(get_field(entry, "end_time", "endTime")), - value_i64(get_field(entry, "remains_time", "remainsTime")), - now, - false, - ) { + if let Some(mut row) = + make_remains_row(&service_type, None, &INTERVAL_LANE.read(entry), now, false) + { row.model_name = model_name.clone(); rows.push(row); } @@ -637,29 +665,7 @@ fn parse_remains( && let Some(mut row) = make_remains_row( &service_type, Some("Weekly"), - value_i64(get_field( - entry, - "current_weekly_total_count", - "currentWeeklyTotalCount", - )), - value_i64(get_field( - entry, - "current_weekly_usage_count", - "currentWeeklyUsageCount", - )), - json::lenient_f64(get_field( - entry, - "current_weekly_remaining_percent", - "currentWeeklyRemainingPercent", - )), - value_i64(get_field( - entry, - "current_weekly_status", - "currentWeeklyStatus", - )), - value_i64(get_field(entry, "weekly_start_time", "weeklyStartTime")), - value_i64(get_field(entry, "weekly_end_time", "weeklyEndTime")), - value_i64(get_field(entry, "weekly_remains_time", "weeklyRemainsTime")), + &WEEKLY_LANE.read(entry), now, true, ) @@ -733,6 +739,192 @@ mod tests { Utc.with_ymd_and_hms(2026, 8, 3, 12, 0, 0).unwrap() } + #[test] + fn reset_description_buckets_by_largest_whole_unit() { + let cases: [(Option, &str); 14] = [ + (None, "Today: 00:00-24:00"), + (Some(-10), "Today: 00:00-24:00"), + (Some(0), "Today: 00:00-24:00"), + (Some(1), "Resets in 1 seconds"), + (Some(59), "Resets in 59 seconds"), + (Some(60), "Resets in 1 minute"), + (Some(119), "Resets in 1 minute"), + (Some(120), "Resets in 2 minutes"), + (Some(3_599), "Resets in 59 minutes"), + (Some(3_600), "Resets in 1 hour"), + (Some(7_200), "Resets in 2 hours"), + (Some(86_399), "Resets in 23 hours"), + (Some(86_400), "Resets in 1 day"), + (Some(172_800), "Resets in 2 days"), + ]; + for (offset, expected) in cases { + let resets = offset.map(|seconds| now() + Duration::seconds(seconds)); + assert_eq!( + reset_description("Today", "00:00-24:00", now(), resets), + expected, + "offset {offset:?}" + ); + } + } + + #[test] + fn time_range_resets_read_utc_plus_8_end_times() { + let at = |y, mo, d, h, mi| Some(Utc.with_ymd_and_hms(y, mo, d, h, mi, 0).unwrap()); + let tomorrow = Some(now() + Duration::days(1)); + let cases = [ + ( + "2026/08/03 00:00 - 2026/08/04 00:00", + "Today", + at(2026, 8, 3, 16, 0), + ), + ("2026/08/03 00:00", "today", None), + ("garbage - later", "Today", None), + ("10:00-23:00(UTC+8)", "5 hours", at(2026, 8, 3, 15, 0)), + ("10:00 - 23:30", "2h", at(2026, 8, 3, 15, 30)), + ("08:00-19:00(UTC+8)", "5 hours", tomorrow), + ("10:00-23:00", "Month", at(2026, 8, 3, 15, 0)), + ("10:00", "5 hours", None), + ("10:00-25:99", "5 hours", None), + ("10:00-23:00", "Weekly", None), + ]; + for (range, window, expected) in cases { + assert_eq!( + parse_resets_at_from_time_range(range, window, now()), + expected, + "{range} / {window}" + ); + } + } + + #[test] + fn time_range_labels_render_in_utc_plus_8() { + let start = Some(Utc.with_ymd_and_hms(2026, 8, 3, 0, 0, 0).unwrap()); + let end = Some(Utc.with_ymd_and_hms(2026, 8, 10, 5, 30, 0).unwrap()); + assert_eq!(time_range_string(start, end), "08:00-13:30(UTC+8)"); + assert_eq!(time_range_string(start, None), "N/A"); + assert_eq!( + weekly_time_range_string(start, end).as_deref(), + Some("08/03 08:00 - 08/10 13:30(UTC+8)") + ); + assert_eq!(weekly_time_range_string(None, end), None); + } + + fn camel_keys(value: &Value) -> Value { + match value { + Value::Object(map) => Value::Object( + map.iter() + .map(|(key, value)| { + let mut camel = String::new(); + let mut upper = false; + for ch in key.chars() { + if ch == '_' { + upper = true; + } else if upper { + camel.extend(ch.to_uppercase()); + upper = false; + } else { + camel.push(ch); + } + } + (camel, camel_keys(value)) + }) + .collect(), + ), + Value::Array(items) => Value::Array(items.iter().map(camel_keys).collect()), + other => other.clone(), + } + } + + #[test] + fn remains_lanes_read_every_snake_and_camel_key() { + let snake = serde_json::json!({ + "model_remains": [{ + "model_name": "General", + "current_interval_total_count": 200, + "current_interval_usage_count": 50, + "current_interval_status": 0, + "start_time": 1785754800, + "end_time": 1785758340, + "remains_time": 600, + "current_weekly_total_count": 1000, + "current_weekly_usage_count": 100, + "current_weekly_remaining_percent": 40, + "current_weekly_status": 0, + "weekly_start_time": 1785585600, + "weekly_end_time": 1786190400, + "weekly_remains_time": 30 + }] + }); + let camel = camel_keys(&snake); + assert!(camel["modelRemains"][0].get("weeklyRemainsTime").is_some()); + + let rows_for = |json: &Value| { + let MiniMaxCodingPlanSnapshot::Remains { plan_name, rows } = + parse_coding_plan_value(json, now()).unwrap() + else { + panic!("expected Remains"); + }; + assert_eq!(plan_name, None); + rows + }; + let rows = rows_for(&snake); + assert_eq!(format!("{rows:?}"), format!("{:?}", rows_for(&camel))); + assert_eq!(rows.len(), 2); + + let interval = &rows[0]; + assert_eq!(interval.model_name, "General"); + assert_eq!(interval.service_type, "general"); + assert_eq!(interval.window_type, "Custom"); + // usage_count is the remaining count: (200 - 50) / 200. + assert_eq!(interval.percent, 75.0); + assert_eq!(interval.window_minutes, Some(59)); + assert_eq!(interval.resets_at, Some(now() + Duration::seconds(600))); + assert_eq!( + interval.reset_description.as_deref(), + Some("Resets in 10 minutes") + ); + assert!(!interval.is_weekly); + + let weekly = &rows[1]; + assert_eq!(weekly.window_type, "Weekly"); + assert_eq!(weekly.percent, 60.0); + assert_eq!(weekly.window_minutes, Some(10080)); + assert_eq!(weekly.resets_at, Some(now() + Duration::days(5))); + assert_eq!( + weekly.reset_description.as_deref(), + Some("Resets in 5 days") + ); + assert!(weekly.is_weekly); + } + + #[test] + fn plan_name_infers_plus_from_an_unavailable_video_lane() { + let entries = |video_remaining_percent: i64| { + serde_json::json!({ + "model_remains": [ + {"model_name": "MiniMax-M2", "current_interval_total_count": 10, + "current_interval_usage_count": 5}, + {"model_name": " Video ", "current_interval_status": 3, + "current_interval_total_count": 0, "current_interval_usage_count": 0, + "current_interval_remaining_percent": video_remaining_percent} + ] + }) + }; + for (remaining, expected) in [(100, Some("Plus")), (50, None)] { + let snake = entries(remaining); + assert_eq!( + parse_plan_name_from_data(&snake).as_deref(), + expected, + "snake {remaining}" + ); + assert_eq!( + parse_plan_name_from_data(&camel_keys(&snake)).as_deref(), + expected, + "camel {remaining}" + ); + } + } + #[test] fn parses_remains_json_snake_case() { let json = serde_json::json!({ @@ -865,39 +1057,30 @@ mod tests { } #[test] - fn base_resp_1004_is_auth_required() { - let json = serde_json::json!({ - "data": { - "base_resp": { "status_code": 1004 }, - "model_remains": [] + fn base_resp_status_maps_to_errors() { + // `None` expects AuthRequired; `Some` expects Other with that message. + let cases = [ + (serde_json::json!({ "status_code": 1004 }), None), + ( + serde_json::json!({ "status_code": 2000, "status_msg": "please log in" }), + None, + ), + ( + serde_json::json!({ "status_code": 2000, "status_msg": "quota sync failed" }), + Some("quota sync failed"), + ), + ]; + for (base_resp, expected) in cases { + let json = serde_json::json!({ + "data": { "base_resp": base_resp, "model_remains": [] } + }); + let err = parse_coding_plan_value(&json, now()).unwrap_err(); + match (expected, err) { + (None, ProviderError::AuthRequired) => {} + (Some(message), ProviderError::Other(actual)) => assert_eq!(actual, message), + (expected, err) => panic!("{base_resp}: expected {expected:?}, got {err:?}"), } - }); - let err = parse_coding_plan_value(&json, now()).unwrap_err(); - assert!(matches!(err, ProviderError::AuthRequired)); - } - - #[test] - fn base_resp_status_msg_login_is_auth_required() { - let json = serde_json::json!({ - "data": { - "base_resp": { "status_code": 2000, "status_msg": "please log in" }, - "model_remains": [] - } - }); - let err = parse_coding_plan_value(&json, now()).unwrap_err(); - assert!(matches!(err, ProviderError::AuthRequired)); - } - - #[test] - fn base_resp_other_status_is_other_error() { - let json = serde_json::json!({ - "data": { - "base_resp": { "status_code": 2000, "status_msg": "quota sync failed" }, - "model_remains": [] - } - }); - let err = parse_coding_plan_value(&json, now()).unwrap_err(); - assert!(matches!(err, ProviderError::Other(_))); + } } #[test] diff --git a/rust/src/providers/minimax/coding_plan_html.rs b/rust/src/providers/minimax/coding_plan_html.rs index a2fafdb909..7c3df0e90b 100644 --- a/rust/src/providers/minimax/coding_plan_html.rs +++ b/rust/src/providers/minimax/coding_plan_html.rs @@ -9,31 +9,45 @@ use serde_json::Value; use crate::core::{ProviderError, RateWindow, UsageSnapshot}; -use super::coding_plan::{MiniMaxCodingPlanSnapshot, parse_coding_plan_value}; -#[cfg(test)] -use super::coding_plan::{RemainsRow, ServiceRow}; +use super::coding_plan::{ + MiniMaxCodingPlanSnapshot, RemainsRow, ServiceRow, next_time_today, parse_coding_plan_value, +}; + +impl ServiceRow { + fn rate_window(&self) -> RateWindow { + RateWindow::with_details( + self.percent, + None, + self.resets_at, + self.reset_description.clone(), + ) + } +} + +impl RemainsRow { + fn rate_window(&self) -> RateWindow { + RateWindow::with_details( + self.percent, + self.window_minutes, + self.resets_at, + self.reset_description.clone(), + ) + } +} /// Recursively walk the JSON object tree for the first object having /// `model_remains` or `modelRemains` (upstream `findCodingPlanPayload`). fn find_coding_plan_payload(obj: &Value) -> Option<&Value> { - if let Value::Object(map) = obj { - if map.contains_key("model_remains") || map.contains_key("modelRemains") { - return Some(obj); - } - for value in map.values() { - if let Some(matched) = find_coding_plan_payload(value) { - return Some(matched); - } - } - } - if let Value::Array(arr) = obj { - for value in arr { - if let Some(matched) = find_coding_plan_payload(value) { - return Some(matched); - } + match obj { + Value::Object(map) + if map.contains_key("model_remains") || map.contains_key("modelRemains") => + { + Some(obj) } + Value::Object(map) => map.values().find_map(find_coding_plan_payload), + Value::Array(arr) => arr.iter().find_map(find_coding_plan_payload), + _ => None, } - None } /// Extract `__NEXT_DATA__` JSON from HTML (upstream `nextDataJSONData`). @@ -281,13 +295,7 @@ fn date_for_time(time: &str, tz_hint: Option<&str>, now: DateTime) -> Optio None => FixedOffset::east_opt(0).unwrap(), }; let time_only = chrono::NaiveTime::parse_from_str(time.trim(), "%H:%M").ok()?; - let today = now.date_naive(); - let dt = today.and_time(time_only).and_local_timezone(tz).single()?; - let mut candidate = dt.with_timezone(&Utc); - if candidate < now { - candidate = now + Duration::days(1); - } - Some(candidate) + next_time_today(time_only, tz, now) } /// Parse a `UTC±H[:MM]`/`GMT±H[:MM]` hint into a `FixedOffset`. Non-numeric hints @@ -361,37 +369,19 @@ pub(super) fn to_usage_snapshot( if rows.is_empty() { return Err(ProviderError::Parse("Missing coding plan data.".into())); } - let first = &rows[0]; - let primary = RateWindow::with_details( - first.percent, - None, - first.resets_at, - first.reset_description.clone(), - ); - let mut usage = UsageSnapshot::new(primary); + let mut usage = UsageSnapshot::new(rows[0].rate_window()); // Secondary = second row if present - if rows.len() >= 2 { - let second = &rows[1]; - usage = usage.with_secondary(RateWindow::with_details( - second.percent, - None, - second.resets_at, - second.reset_description.clone(), - )); + if let Some(second) = rows.get(1) { + usage = usage.with_secondary(second.rate_window()); } // Extra windows for remaining rows - for row in rows.iter().skip(if rows.len() >= 2 { 2 } else { 1 }) { + for row in rows.iter().skip(2) { usage = usage.with_extra_rate_window( &row.service_type, format!("{} · {}", row.service_type, row.window_type), - RateWindow::with_details( - row.percent, - None, - row.resets_at, - row.reset_description.clone(), - ), + row.rate_window(), ); } @@ -422,22 +412,11 @@ pub(super) fn to_usage_snapshot( .or_else(|| rows.first()) .ok_or_else(|| ProviderError::Parse("Missing coding plan data.".into()))?; - let primary = RateWindow::with_details( - primary_row.percent, - primary_row.window_minutes, - primary_row.resets_at, - primary_row.reset_description.clone(), - ); - let mut usage = UsageSnapshot::new(primary); + let mut usage = UsageSnapshot::new(primary_row.rate_window()); // Secondary = first qualifying weekly row if let Some(weekly) = rows.iter().find(|r| r.is_weekly) { - usage = usage.with_secondary(RateWindow::with_details( - weekly.percent, - weekly.window_minutes, - weekly.resets_at, - weekly.reset_description.clone(), - )); + usage = usage.with_secondary(weekly.rate_window()); } // Every row as an extra window @@ -448,16 +427,7 @@ pub(super) fn to_usage_snapshot( row.service_type.clone() }; let title = format!("{} · {}", row.service_type, row.window_type); - usage = usage.with_extra_rate_window( - id, - title, - RateWindow::with_details( - row.percent, - row.window_minutes, - row.resets_at, - row.reset_description.clone(), - ), - ); + usage = usage.with_extra_rate_window(id, title, row.rate_window()); } if let Some(pn) = plan_name { @@ -510,6 +480,24 @@ mod tests { assert_eq!(reset, now() + Duration::seconds(2 * 86_400)); } + #[test] + fn reset_clock_times_roll_forward_with_the_zone_hint() { + let at = |d, h, mi| Some(Utc.with_ymd_and_hms(2026, 8, d, h, mi, 0).unwrap()); + let tomorrow = Some(now() + Duration::days(1)); + let cases = [ + ("Resets at 15:30", at(3, 15, 30)), + ("Resets at 11:59", tomorrow), + ("Reset at 13:00 (UTC+8)", tomorrow), + ("resets at 23:00 (GMT-5:30)", at(4, 4, 30)), + ("Resets at 20:00 (PST)", at(3, 20, 0)), + ("Resets at 20:00 (UTC+0130)", at(3, 18, 30)), + ("Resets at 25:00", None), + ]; + for (text, expected) in cases { + assert_eq!(parse_resets_at_from_text(text, now()), expected, "{text}"); + } + } + #[test] fn text_without_a_reset_count_is_unmatched() { assert_eq!(parse_resets_at_from_text("no reset info", now()), None); @@ -645,6 +633,184 @@ mod tests { assert_eq!(usage.login_method.as_deref(), Some("Text Generation Pro")); } + type WindowRow = ( + String, + String, + f64, + Option, + Option>, + Option, + ); + + fn window_row(id: &str, title: &str, window: &RateWindow) -> WindowRow { + ( + id.to_string(), + title.to_string(), + window.used_percent, + window.window_minutes, + window.resets_at, + window.reset_description.clone(), + ) + } + + fn extra_rows(usage: &UsageSnapshot) -> Vec { + usage + .extra_rate_windows + .iter() + .map(|named| window_row(&named.id, &named.title, &named.window)) + .collect() + } + + fn service_row(service_type: &str, window_type: &str, percent: f64) -> ServiceRow { + ServiceRow { + service_type: service_type.to_string(), + window_type: window_type.to_string(), + time_range: String::new(), + percent, + resets_at: Some(now() + Duration::hours(1)), + reset_description: Some(format!("{service_type} reset")), + } + } + + #[test] + fn to_usage_snapshot_services_maps_every_row_window() { + let single = MiniMaxCodingPlanSnapshot::Services(vec![service_row("Text", "Today", 5.0)]); + let usage = to_usage_snapshot(&single, now()).unwrap(); + assert!(usage.secondary.is_none()); + assert!(usage.extra_rate_windows.is_empty()); + assert_eq!(usage.login_method, None); + + let rows = vec![ + service_row("Text", "Today", 10.0), + service_row("Speech Max", "Week", 20.0), + service_row("Video", "Month", 30.0), + ]; + let snapshot = MiniMaxCodingPlanSnapshot::Services(rows); + let usage = to_usage_snapshot(&snapshot, now()).unwrap(); + let reset = Some(now() + Duration::hours(1)); + assert_eq!( + window_row("", "", &usage.primary), + window_row( + "", + "", + &RateWindow::with_details(10.0, None, reset, Some("Text reset".into())) + ) + ); + assert_eq!( + window_row("", "", usage.secondary.as_ref().unwrap()), + window_row( + "", + "", + &RateWindow::with_details(20.0, None, reset, Some("Speech Max reset".into())) + ) + ); + assert_eq!( + extra_rows(&usage), + vec![( + "Video".to_string(), + "Video · Month".to_string(), + 30.0, + None, + reset, + Some("Video reset".to_string()), + )] + ); + assert_eq!(usage.login_method.as_deref(), Some("Speech Max")); + } + + #[test] + fn to_usage_snapshot_remains_maps_every_row_window() { + let row = |window_type: &str, percent: f64, minutes: u32, unlimited: bool, weekly: bool| { + RemainsRow { + service_type: "general".to_string(), + model_name: "General".to_string(), + window_type: window_type.to_string(), + percent, + window_minutes: Some(minutes), + resets_at: Some(now() + Duration::minutes(i64::from(minutes))), + reset_description: Some(format!("{window_type} reset")), + is_unlimited: unlimited, + is_weekly: weekly, + } + }; + let snapshot = MiniMaxCodingPlanSnapshot::Remains { + plan_name: None, + rows: vec![ + row("Unlimited", 0.0, 60, true, false), + row("Today", 40.0, 1440, false, false), + row("Weekly", 70.0, 10080, false, true), + ], + }; + let usage = to_usage_snapshot(&snapshot, now()).unwrap(); + let at = |minutes: i64| Some(now() + Duration::minutes(minutes)); + assert_eq!( + window_row("", "", &usage.primary), + ( + "".into(), + "".into(), + 40.0, + Some(1440), + at(1440), + Some("Today reset".into()) + ) + ); + assert_eq!( + window_row("", "", usage.secondary.as_ref().unwrap()), + ( + "".into(), + "".into(), + 70.0, + Some(10080), + at(10080), + Some("Weekly reset".into()) + ) + ); + assert_eq!( + extra_rows(&usage), + vec![ + ( + "general".to_string(), + "general · Unlimited".to_string(), + 0.0, + Some(60), + at(60), + Some("Unlimited reset".to_string()), + ), + ( + "general".to_string(), + "general · Today".to_string(), + 40.0, + Some(1440), + at(1440), + Some("Today reset".to_string()), + ), + ( + "general:weekly".to_string(), + "general · Weekly".to_string(), + 70.0, + Some(10080), + at(10080), + Some("Weekly reset".to_string()), + ), + ] + ); + assert_eq!(usage.login_method, None); + } + + #[test] + fn coding_plan_payload_search_finds_the_first_nested_object() { + let json = serde_json::json!({ + "a": [1, {"b": {"modelRemains": [], "tag": "camel"}}], + "z": {"model_remains": [], "tag": "snake"} + }); + let found = find_coding_plan_payload(&json).unwrap(); + assert_eq!(found["tag"], "camel"); + let direct = serde_json::json!([{"model_remains": [], "tag": "array"}]); + assert_eq!(find_coding_plan_payload(&direct).unwrap()["tag"], "array"); + assert!(find_coding_plan_payload(&serde_json::json!({"a": [1, "x"]})).is_none()); + assert!(find_coding_plan_payload(&serde_json::json!("model_remains")).is_none()); + } + #[test] fn oversized_html_duration_is_omitted() { assert_eq!(minutes_from_duration(f64::MAX, "hours"), None); diff --git a/rust/src/providers/minimax/http_tests.rs b/rust/src/providers/minimax/http_tests.rs new file mode 100644 index 0000000000..785148a74a --- /dev/null +++ b/rust/src/providers/minimax/http_tests.rs @@ -0,0 +1,137 @@ +//! Wire-level pins for the MiniMax console and remains GET requests. + +use super::*; + +pub(super) const JSON_ACCEPT: &str = "application/json, text/plain, */*"; + +/// Console GET with the cookie, JSON accept, XHR and browser headers. +pub(super) fn console_mock(server: &mut mockito::Server, path: &str) -> mockito::Mock { + browser_mock(server, path, JSON_ACCEPT, "XMLHttpRequest".into()) +} + +/// Console GET with the cookie, `accept`, `xhr` marker and browser headers. +fn browser_mock( + server: &mut mockito::Server, + path: &str, + accept: &str, + xhr: mockito::Matcher, +) -> mockito::Mock { + let base = MiniMaxRegion::Global.base_url(); + server + .mock("GET", path) + .match_header("cookie", "session=fixture") + .match_header("accept", accept) + .match_header("x-requested-with", xhr) + .match_header("user-agent", MiniMaxProvider::WEB_USER_AGENT) + .match_header("accept-language", "en-US,en;q=0.9") + .match_header("origin", base) + .match_header( + "referer", + format!("{base}/user-center/payment/coding-plan").as_str(), + ) +} + +#[tokio::test] +async fn cookie_remains_request_maps_statuses() { + let cases: [(usize, Option<&str>); 5] = [ + (401, None), + (403, None), + ( + 404, + Some("Parse:MiniMax remains returned status 404 Not Found"), + ), + ( + 405, + Some("Parse:MiniMax remains returned status 405 Method Not Allowed"), + ), + ( + 500, + Some("Other:MiniMax remains returned status 500 Internal Server Error"), + ), + ]; + for (status, expected) in cases { + let mut server = mockito::Server::new_async().await; + let mock = console_mock(&mut server, "/remains") + .with_status(status) + .expect(1) + .create_async() + .await; + let result = MiniMaxProvider::new() + .fetch_remains_once( + "session=fixture", + &format!("{}/remains", server.url()), + MiniMaxRegion::Global, + Utc::now(), + ) + .await; + mock.assert_async().await; + assert_status_error(result.map(|_| ()), expected, status); + } +} + +#[tokio::test] +async fn cookie_remains_request_parses_a_json_body() { + let mut server = mockito::Server::new_async().await; + let mock = console_mock(&mut server, "/remains") + .with_header("content-type", "application/json") + .with_body(r#"{"base_resp":{"status_code":2062,"status_msg":"no plan"}}"#) + .expect(1) + .create_async() + .await; + let result = MiniMaxProvider::new() + .fetch_remains_once( + "session=fixture", + &format!("{}/remains", server.url()), + MiniMaxRegion::Global, + Utc::now(), + ) + .await; + mock.assert_async().await; + match result { + Err(ProviderError::Other(message)) => assert_eq!(message, "no plan"), + other => panic!("expected Other(\"no plan\"), got {other:?}"), + } +} + +#[tokio::test] +async fn coding_plan_page_request_sends_html_accept_without_the_xhr_marker() { + let mut server = mockito::Server::new_async().await; + let mock = browser_mock( + &mut server, + "/coding-plan", + "text/html,application/xhtml+xml,application/xml;q=0.9,*/*;q=0.8", + mockito::Matcher::Missing, + ) + .expect(1) + .create_async() + .await; + let client = http_client().expect("client"); + let response = console_get( + &client, + &format!("{}/coding-plan", server.url()), + "session=fixture", + MiniMaxRegion::Global, + HTML_ACCEPT, + false, + ) + .send() + .await + .expect("send"); + mock.assert_async().await; + assert_eq!(response.status(), reqwest::StatusCode::OK); +} + +/// `None` expects AuthRequired; otherwise `":"`. +pub(super) fn assert_status_error( + result: Result<(), ProviderError>, + expected: Option<&str>, + status: usize, +) { + let actual = match result { + Err(ProviderError::AuthRequired) => None, + Err(ProviderError::Parse(message)) => Some(format!("Parse:{message}")), + Err(ProviderError::Other(message)) => Some(format!("Other:{message}")), + other => panic!("status {status}: unexpected {other:?}"), + }; + assert_eq!(actual.as_deref(), expected, "status {status}"); +} diff --git a/rust/src/providers/minimax/mod.rs b/rust/src/providers/minimax/mod.rs index 9967e5df9f..668f13f867 100755 --- a/rust/src/providers/minimax/mod.rs +++ b/rust/src/providers/minimax/mod.rs @@ -3,6 +3,7 @@ //! Fetches usage data from MiniMax AI API //! MiniMax stores API keys locally or in environment +mod billing; mod coding_plan; mod coding_plan_html; mod local_storage; @@ -17,65 +18,78 @@ mod token_plan; pub use local_storage::{ImportError, MiniMaxLocalStorageImporter, MiniMaxSession}; use async_trait::async_trait; -use chrono::{DateTime, Duration, TimeZone, Utc}; -use serde::Deserialize; -use std::collections::HashMap; +use chrono::{DateTime, Utc}; use std::path::PathBuf; use crate::core::{ - CostSnapshot, FetchContext, Provider, ProviderError, ProviderFetchResult, ProviderId, - RateWindow, SourceMode, UsageSnapshot, + FetchContext, Provider, ProviderError, ProviderFetchResult, ProviderId, RateWindow, SourceMode, + UsageSnapshot, }; -use crate::providers::json; +use billing::{MiniMaxBillingSummary, attach_billing_summary, parse_billing_summary}; const CODING_PLAN_PATH: &str = "/user-center/payment/coding-plan"; const CODING_PLAN_QUERY: &str = "cycle_type=3"; - -#[derive(Debug, Deserialize)] -struct MiniMaxBillingHistoryPayload { - #[serde(default)] - base_resp: Option, - #[serde(default)] - charge_records: Vec, -} - -#[derive(Debug, Deserialize)] -struct MiniMaxBaseResponse { - status_code: Option, - status_msg: Option, +const JSON_ACCEPT: &str = "application/json, text/plain, */*"; +const HTML_ACCEPT: &str = "text/html,application/xhtml+xml,application/xml;q=0.9,*/*;q=0.8"; + +fn http_client() -> Result { + crate::core::credentialed_http_client_builder() + .timeout(std::time::Duration::from_secs(30)) + .build() + .map_err(|e| ProviderError::Other(e.to_string())) } -#[derive(Debug, Deserialize)] -struct MiniMaxBillingRecord { - consume_token: Option, - consume_input_token: Option, - consume_output_token: Option, - consume_cash: Option, - consume_cash_after_voucher: Option, - created_at: Option, - ymd: Option, - consume_time: Option, - method: Option, - model: Option, - result: Option, - status: Option, +fn is_auth_status(status: reqwest::StatusCode) -> bool { + status == reqwest::StatusCode::UNAUTHORIZED || status == reqwest::StatusCode::FORBIDDEN } -#[derive(Debug, Clone)] -struct MiniMaxBillingSummary { - today_tokens: i64, - last_30_days_tokens: i64, - today_cash: Option, - last_30_days_cash: Option, - top_methods: Vec, - top_models: Vec, +/// 401/403 read as AuthRequired and other failures as +/// "MiniMax {what} returned status {status}". With `not_found_is_parse`, +/// 404/405 read as Parse so the caller tries its next URL. +fn check_status( + status: reqwest::StatusCode, + what: &str, + not_found_is_parse: bool, +) -> Result<(), ProviderError> { + if is_auth_status(status) { + return Err(ProviderError::AuthRequired); + } + if status.is_success() { + return Ok(()); + } + let message = format!("MiniMax {what} returned status {status}"); + if not_found_is_parse + && (status == reqwest::StatusCode::NOT_FOUND + || status == reqwest::StatusCode::METHOD_NOT_ALLOWED) + { + return Err(ProviderError::Parse(message)); + } + Err(ProviderError::Other(message)) } -#[derive(Debug, Clone)] -struct MiniMaxBillingBreakdown { - name: String, - tokens: i64, - cash: Option, +/// Browser-shaped console GET: cookie, `accept`, the optional XHR marker, +/// then the Chrome UA, language, origin and coding-plan referer. +fn console_get( + client: &reqwest::Client, + url: &str, + cookie_header: &str, + region: MiniMaxRegion, + accept: &str, + xhr: bool, +) -> reqwest::RequestBuilder { + let base = region.base_url(); + let mut request = client + .get(url) + .header("Cookie", cookie_header) + .header("Accept", accept); + if xhr { + request = request.header("X-Requested-With", "XMLHttpRequest"); + } + request + .header("User-Agent", MiniMaxProvider::WEB_USER_AGENT) + .header("Accept-Language", "en-US,en;q=0.9") + .header("Origin", base) + .header("Referer", format!("{base}/user-center/payment/coding-plan")) } /// MiniMax API region @@ -299,10 +313,7 @@ impl MiniMaxProvider { api_key: &str, region: MiniMaxRegion, ) -> Result { - let client = crate::core::credentialed_http_client_builder() - .timeout(std::time::Duration::from_secs(30)) - .build() - .map_err(|e| ProviderError::Other(e.to_string()))?; + let client = http_client()?; let base_url = region.api_base_url(); let resp = client @@ -314,19 +325,7 @@ impl MiniMaxProvider { .header("MM-API-Source", "CodexBar") .send() .await?; - - if resp.status() == reqwest::StatusCode::UNAUTHORIZED - || resp.status() == reqwest::StatusCode::FORBIDDEN - { - return Err(ProviderError::AuthRequired); - } - - if !resp.status().is_success() { - return Err(ProviderError::Other(format!( - "MiniMax API returned status {}", - resp.status() - ))); - } + check_status(resp.status(), "API", false)?; let json: serde_json::Value = resp .json() @@ -398,31 +397,17 @@ impl MiniMaxProvider { cookie_header: &str, region: MiniMaxRegion, ) -> Result { - let client = crate::core::credentialed_http_client_builder() - .timeout(std::time::Duration::from_secs(30)) - .build() - .map_err(|e| ProviderError::Other(e.to_string()))?; + let client = http_client()?; let response = client .get(region.billing_history_url()) .query(&[("page", "1"), ("limit", "100"), ("aggregate", "false")]) .header("Cookie", cookie_header) - .header("Accept", "application/json, text/plain, */*") + .header("Accept", JSON_ACCEPT) .header("X-Requested-With", "XMLHttpRequest") .send() .await?; - - if response.status() == reqwest::StatusCode::UNAUTHORIZED - || response.status() == reqwest::StatusCode::FORBIDDEN - { - return Err(ProviderError::AuthRequired); - } - if !response.status().is_success() { - return Err(ProviderError::Other(format!( - "MiniMax billing returned status {}", - response.status() - ))); - } + check_status(response.status(), "billing", false)?; let json: serde_json::Value = response .json() @@ -442,39 +427,19 @@ impl MiniMaxProvider { cookie_header: &str, region: MiniMaxRegion, ) -> Result { - let client = crate::core::credentialed_http_client_builder() - .timeout(std::time::Duration::from_secs(30)) - .build() - .map_err(|e| ProviderError::Other(e.to_string()))?; - - let coding_url = region.coding_plan_url(); - let base = region.base_url(); - let response = client - .get(&coding_url) - .header("Cookie", cookie_header) - .header( - "Accept", - "text/html,application/xhtml+xml,application/xml;q=0.9,*/*;q=0.8", - ) - .header("User-Agent", Self::WEB_USER_AGENT) - .header("Accept-Language", "en-US,en;q=0.9") - .header("Origin", base) - .header( - "Referer", - &format!("{base}/user-center/payment/coding-plan"), - ) - .send() - .await?; - - let status = response.status(); - if status == reqwest::StatusCode::UNAUTHORIZED || status == reqwest::StatusCode::FORBIDDEN { - return Err(ProviderError::AuthRequired); - } - if !status.is_success() { - return Err(ProviderError::Other(format!( - "MiniMax coding plan returned status {status}" - ))); - } + let client = http_client()?; + + let response = console_get( + &client, + ®ion.coding_plan_url(), + cookie_header, + region, + HTML_ACCEPT, + false, + ) + .send() + .await?; + check_status(response.status(), "coding plan", false)?; let content_type = response .headers() @@ -563,41 +528,12 @@ impl MiniMaxProvider { region: MiniMaxRegion, now: DateTime, ) -> Result { - let client = crate::core::credentialed_http_client_builder() - .timeout(std::time::Duration::from_secs(30)) - .build() - .map_err(|e| ProviderError::Other(e.to_string()))?; + let client = http_client()?; - let base = region.base_url(); - let response = client - .get(url) - .header("Cookie", cookie_header) - .header("Accept", "application/json, text/plain, */*") - .header("X-Requested-With", "XMLHttpRequest") - .header("User-Agent", Self::WEB_USER_AGENT) - .header("Accept-Language", "en-US,en;q=0.9") - .header("Origin", base) - .header( - "Referer", - &format!("{base}/user-center/payment/coding-plan"), - ) + let response = console_get(&client, url, cookie_header, region, JSON_ACCEPT, true) .send() .await?; - - let status = response.status(); - if status == reqwest::StatusCode::UNAUTHORIZED || status == reqwest::StatusCode::FORBIDDEN { - return Err(ProviderError::AuthRequired); - } - if !status.is_success() { - let msg = format!("MiniMax remains returned status {status}"); - // 404/405 → try next URL; other → stop - if status == reqwest::StatusCode::NOT_FOUND - || status == reqwest::StatusCode::METHOD_NOT_ALLOWED - { - return Err(ProviderError::Parse(msg)); - } - return Err(ProviderError::Other(msg)); - } + check_status(response.status(), "remains", true)?; let content_type = response .headers() @@ -713,250 +649,6 @@ impl MiniMaxProvider { } } -fn parse_billing_summary(json: &serde_json::Value) -> Result { - let payload: MiniMaxBillingHistoryPayload = serde_json::from_value(json.clone()) - .map_err(|e| ProviderError::Parse(format!("Failed to parse MiniMax billing: {e}")))?; - if let Some(base) = payload.base_resp - && let Some(status) = base.status_code - && status != 0 - { - return Err(ProviderError::Other( - base.status_msg - .unwrap_or_else(|| format!("MiniMax billing status {status}")), - )); - } - if payload.charge_records.is_empty() { - return Err(ProviderError::Parse( - "MiniMax billing records not present".to_string(), - )); - } - Ok(aggregate_billing(&payload.charge_records, Utc::now())) -} - -fn aggregate_billing( - records: &[MiniMaxBillingRecord], - now: DateTime, -) -> MiniMaxBillingSummary { - let today_start = now.date_naive().and_hms_opt(0, 0, 0).unwrap().and_utc(); - let window_start = today_start - Duration::days(29); - let mut today_tokens = 0; - let mut last_30_days_tokens = 0; - let mut today_cash = 0.0; - let mut today_has_cash = false; - let mut last_30_days_cash = 0.0; - let mut last_30_has_cash = false; - let mut method_totals: HashMap = HashMap::new(); - let mut model_totals: HashMap = HashMap::new(); - - for record in records { - if !billing_record_succeeded(record) { - continue; - } - let Some(date) = record_date(record) else { - continue; - }; - if date < window_start || date > now { - continue; - } - let tokens = record_token_count(record); - let cash = record_cash(record); - last_30_days_tokens += tokens; - if let Some(cash) = cash { - last_30_days_cash += cash; - last_30_has_cash = true; - } - if date >= today_start { - today_tokens += tokens; - if let Some(cash) = cash { - today_cash += cash; - today_has_cash = true; - } - } - add_breakdown(&mut method_totals, record.method.as_deref(), tokens, cash); - add_breakdown(&mut model_totals, record.model.as_deref(), tokens, cash); - } - - MiniMaxBillingSummary { - today_tokens, - last_30_days_tokens, - today_cash: today_has_cash.then_some(today_cash), - last_30_days_cash: last_30_has_cash.then_some(last_30_days_cash), - top_methods: top_breakdowns(method_totals), - top_models: top_breakdowns(model_totals), - } -} - -fn billing_record_succeeded(record: &MiniMaxBillingRecord) -> bool { - let status = - scalar_string(record.result.as_ref()).or_else(|| scalar_string(record.status.as_ref())); - match status - .as_deref() - .map(str::trim) - .filter(|value| !value.is_empty()) - { - None => true, - Some(value) => value.eq_ignore_ascii_case("success"), - } -} - -fn result_from_billing_summary( - summary: MiniMaxBillingSummary, - source_label: &str, -) -> ProviderFetchResult { - let usage = UsageSnapshot::new(RateWindow::with_details( - 0.0, - None, - None, - Some(format!( - "{} tokens today", - format_count(summary.today_tokens) - )), - )) - .with_secondary(RateWindow::with_details( - 0.0, - None, - None, - Some(format!( - "{} tokens over last 30 days", - format_count(summary.last_30_days_tokens) - )), - )) - .with_login_method("MiniMax billing"); - attach_billing_summary(ProviderFetchResult::new(usage, source_label), summary) -} - -fn attach_billing_summary( - mut result: ProviderFetchResult, - summary: MiniMaxBillingSummary, -) -> ProviderFetchResult { - result.usage = result.usage.with_extra_rate_window( - "billing-tokens-today", - "Tokens today", - RateWindow::with_details(0.0, None, None, Some(format_count(summary.today_tokens))), - ); - result.usage = result.usage.with_extra_rate_window( - "billing-tokens-30d", - "Tokens (30 days)", - RateWindow::with_details( - 0.0, - None, - None, - Some(format_count(summary.last_30_days_tokens)), - ), - ); - if let Some(cash) = summary.today_cash { - result.usage = result.usage.with_extra_rate_window( - "billing-cash-today", - "Spend today", - RateWindow::with_details(0.0, None, None, Some(format!("${cash:.2}"))), - ); - } - if let Some(cash) = summary.last_30_days_cash { - result.usage = result.usage.with_extra_rate_window( - "billing-cash-30d", - "Spend (30 days)", - RateWindow::with_details(0.0, None, None, Some(format!("${cash:.2}"))), - ); - result.cost = Some(CostSnapshot::new(cash, "USD", "Last 30 days")); - } - for (idx, item) in summary.top_methods.iter().enumerate() { - result.usage = result.usage.with_extra_rate_window( - format!("billing-method-{idx}"), - format!("Method: {}", item.name), - RateWindow::with_details(0.0, None, None, Some(breakdown_description(item))), - ); - } - for (idx, item) in summary.top_models.iter().enumerate() { - result.usage = result.usage.with_extra_rate_window( - format!("billing-model-{idx}"), - format!("Model: {}", item.name), - RateWindow::with_details(0.0, None, None, Some(breakdown_description(item))), - ); - } - result -} - -fn add_breakdown( - totals: &mut HashMap, - raw_name: Option<&str>, - tokens: i64, - cash: Option, -) { - let Some(name) = raw_name.map(str::trim).filter(|s| !s.is_empty()) else { - return; - }; - let total = totals.entry(name.to_string()).or_default(); - total.0 += tokens; - if let Some(cash) = cash { - total.1 += cash; - total.2 = true; - } -} - -fn top_breakdowns(totals: HashMap) -> Vec { - let mut items: Vec<_> = totals - .into_iter() - .map(|(name, (tokens, cash, has_cash))| MiniMaxBillingBreakdown { - name, - tokens, - cash: has_cash.then_some(cash), - }) - .collect(); - items.sort_by(|a, b| b.tokens.cmp(&a.tokens).then_with(|| a.name.cmp(&b.name))); - items.truncate(3); - items -} - -fn breakdown_description(item: &MiniMaxBillingBreakdown) -> String { - match item.cash { - Some(cash) => format!("{} tokens / ${cash:.2}", format_count(item.tokens)), - None => format!("{} tokens", format_count(item.tokens)), - } -} - -fn record_token_count(record: &MiniMaxBillingRecord) -> i64 { - let direct = value_i64(record.consume_token.as_ref()).unwrap_or(0); - if direct > 0 { - return direct; - } - value_i64(record.consume_input_token.as_ref()).unwrap_or(0) - + value_i64(record.consume_output_token.as_ref()).unwrap_or(0) -} - -fn record_cash(record: &MiniMaxBillingRecord) -> Option { - json::lenient_f64(record.consume_cash_after_voucher.as_ref()) - .or_else(|| json::lenient_f64(record.consume_cash.as_ref())) -} - -fn record_date(record: &MiniMaxBillingRecord) -> Option> { - if let Some(created_at) = value_i64(record.created_at.as_ref()) { - let seconds = if created_at > 1_000_000_000_000 { - created_at / 1000 - } else { - created_at - }; - return Utc.timestamp_opt(seconds, 0).single(); - } - if let Some(ymd) = record.ymd.as_deref() { - for format in ["%Y-%m-%d", "%Y%m%d", "%Y/%m/%d"] { - if let Ok(date) = chrono::NaiveDate::parse_from_str(ymd.trim(), format) { - return Some(date.and_hms_opt(0, 0, 0)?.and_utc()); - } - } - } - if let Some(text) = record.consume_time.as_deref() { - for format in ["%Y-%m-%d %H:%M:%S", "%Y/%m/%d %H:%M:%S"] { - if let Ok(date) = chrono::NaiveDateTime::parse_from_str(text.trim(), format) { - return Some(date.and_utc()); - } - } - if let Ok(date) = DateTime::parse_from_rfc3339(text.trim()) { - return Some(date.with_timezone(&Utc)); - } - } - None -} - fn value_i64(value: Option<&serde_json::Value>) -> Option { match value? { serde_json::Value::Number(number) => number.as_i64(), @@ -1040,229 +732,7 @@ impl Provider for MiniMaxProvider { } #[cfg(test)] -mod tests { - use super::*; - - #[test] - fn minimax_region_defaults_to_global_io_urls() { - let region = MiniMaxRegion::from_settings_value(None); - assert_eq!(region, MiniMaxRegion::Global); - assert_eq!(region.settings_value(), "global"); - assert_eq!(region.cookie_domain(), "platform.minimax.io"); - assert_eq!( - region.coding_plan_url(), - "https://platform.minimax.io/user-center/payment/coding-plan?cycle_type=3" - ); - assert_eq!( - MiniMaxProvider::dashboard_url_for_region(None), - "https://platform.minimax.io/user-center/payment/coding-plan?cycle_type=3" - ); - } - - #[test] - fn minimax_region_accepts_legacy_china_value() { - for value in ["cn", "china", "china-mainland", "china_mainland"] { - let region = MiniMaxRegion::from_settings_value(Some(value)); - assert_eq!(region, MiniMaxRegion::ChinaMainland); - assert_eq!(region.settings_value(), "cn"); - assert_eq!(region.cookie_domain(), "platform.minimaxi.com"); - assert_eq!( - region.coding_plan_url(), - "https://platform.minimaxi.com/user-center/payment/coding-plan?cycle_type=3" - ); - } - } - - #[test] - fn minimax_region_cookie_search_domains_and_www_base_url() { - let global = MiniMaxRegion::Global; - assert_eq!(global.www_base_url(), "https://www.minimax.io"); - assert_eq!(global.cookie_search_domains(), ["minimax.io"]); - assert_eq!( - global.coding_plan_remains_url(), - "https://platform.minimax.io/v1/api/openplatform/coding_plan/remains" - ); - assert_eq!( - global.www_remains_url(), - "https://www.minimax.io/v1/api/openplatform/coding_plan/remains" - ); - - let cn = MiniMaxRegion::ChinaMainland; - assert_eq!(cn.www_base_url(), "https://www.minimaxi.com"); - assert_eq!(cn.cookie_search_domains(), ["minimaxi.com"]); - assert_eq!( - cn.coding_plan_remains_url(), - "https://platform.minimaxi.com/v1/api/openplatform/coding_plan/remains" - ); - assert_eq!( - cn.www_remains_url(), - "https://www.minimaxi.com/v1/api/openplatform/coding_plan/remains" - ); - } - - #[test] - fn aggregates_billing_history_records() { - let records = vec![ - MiniMaxBillingRecord { - consume_token: Some(serde_json::json!(1200)), - consume_input_token: None, - consume_output_token: None, - consume_cash: Some(serde_json::json!("0.42")), - consume_cash_after_voucher: None, - created_at: None, - ymd: Some("2026-12-16".to_string()), - consume_time: None, - method: Some("chat".to_string()), - model: Some("abab6.5".to_string()), - result: Some(serde_json::json!("SUCCESS")), - status: None, - }, - MiniMaxBillingRecord { - consume_token: None, - consume_input_token: Some(serde_json::json!(300)), - consume_output_token: Some(serde_json::json!(500)), - consume_cash: None, - consume_cash_after_voucher: Some(serde_json::json!(0.21)), - created_at: None, - ymd: Some("2026-12-15".to_string()), - consume_time: None, - method: Some("completion".to_string()), - model: Some("abab6.5".to_string()), - result: None, - status: None, - }, - ]; - let now = Utc.with_ymd_and_hms(2026, 12, 16, 12, 0, 0).unwrap(); - let summary = aggregate_billing(&records, now); - assert_eq!(summary.last_30_days_tokens, 2000); - assert!((summary.last_30_days_cash.unwrap() - 0.63).abs() < 0.001); - assert_eq!(summary.top_models[0].name, "abab6.5"); - assert_eq!(summary.top_models[0].tokens, 2000); - } +mod http_tests; - #[test] - fn parses_billing_payload_into_result_extras() { - let json = serde_json::json!({ - "base_resp": { "status_code": 0 }, - "charge_records": [{ - "consume_token": "2500", - "consume_cash_after_voucher": "1.25", - "ymd": Utc::now().format("%Y-%m-%d").to_string(), - "method": "chat", - "model": "abab6.5" - }] - }); - let summary = parse_billing_summary(&json).unwrap(); - let result = result_from_billing_summary(summary, "web-billing"); - assert!(result.cost.is_some()); - assert!( - result - .usage - .extra_rate_windows - .iter() - .any(|window| window.id == "billing-tokens-30d") - ); - } - - #[test] - fn parses_plan_title_from_coding_plan_fields() { - let provider = MiniMaxProvider::new(); - for (field, expected) in [ - ("plan_name", "MiniMax Star"), - ("current_plan_title", "Coding Plan Pro"), - ("current_subscribe_title", "Max"), - ("combo_title", "Combo Star"), - ] { - let snapshot = provider - .parse_usage_response(&serde_json::json!({ - "base_resp": { "status_code": 0 }, - field: expected, - "used_amount": 0, - "total_quota": 100 - })) - .unwrap(); - assert_eq!(snapshot.login_method.as_deref(), Some(expected)); - } - - let snapshot = provider - .parse_usage_response(&serde_json::json!({ - "base_resp": { "status_code": 0 }, - "current_combo_card": { "title": "Card Title" }, - "used_amount": 0, - "total_quota": 100 - })) - .unwrap(); - assert_eq!(snapshot.login_method.as_deref(), Some("Card Title")); - } - - #[test] - fn filters_failed_billing_records() { - let records = vec![ - MiniMaxBillingRecord { - consume_token: Some(serde_json::json!(1000)), - consume_input_token: None, - consume_output_token: None, - consume_cash: None, - consume_cash_after_voucher: None, - created_at: None, - ymd: Some("2026-05-17".to_string()), - consume_time: None, - method: Some("chat".to_string()), - model: Some("MiniMax-M1".to_string()), - result: Some(serde_json::json!("SUCCESS")), - status: None, - }, - MiniMaxBillingRecord { - consume_token: Some(serde_json::json!(2000)), - consume_input_token: None, - consume_output_token: None, - consume_cash: None, - consume_cash_after_voucher: None, - created_at: None, - ymd: Some("2026-05-17".to_string()), - consume_time: None, - method: Some("chat".to_string()), - model: Some("MiniMax-M1".to_string()), - result: Some(serde_json::json!("FAILED")), - status: None, - }, - MiniMaxBillingRecord { - consume_token: Some(serde_json::json!(3000)), - consume_input_token: None, - consume_output_token: None, - consume_cash: None, - consume_cash_after_voucher: None, - created_at: None, - ymd: Some("2026-05-17".to_string()), - consume_time: None, - method: Some("audio".to_string()), - model: Some("speech".to_string()), - result: None, - status: None, - }, - MiniMaxBillingRecord { - consume_token: Some(serde_json::json!(4000)), - consume_input_token: None, - consume_output_token: None, - consume_cash: None, - consume_cash_after_voucher: None, - created_at: None, - ymd: Some("2026-05-17".to_string()), - consume_time: None, - method: Some("video".to_string()), - model: Some("video".to_string()), - result: None, - status: Some(serde_json::json!(0)), - }, - ]; - let now = Utc.with_ymd_and_hms(2026, 5, 17, 12, 0, 0).unwrap(); - let summary = aggregate_billing(&records, now); - assert_eq!(summary.today_tokens, 4000); - assert_eq!(summary.last_30_days_tokens, 4000); - assert_eq!(summary.top_methods.len(), 2); - assert_eq!(summary.top_methods[0].name, "audio"); - assert_eq!(summary.top_methods[0].tokens, 3000); - assert_eq!(summary.top_methods[1].name, "chat"); - assert_eq!(summary.top_methods[1].tokens, 1000); - } -} +#[cfg(test)] +mod tests; diff --git a/rust/src/providers/minimax/remains_api.rs b/rust/src/providers/minimax/remains_api.rs index f1608fe940..8ae1944d03 100644 --- a/rust/src/providers/minimax/remains_api.rs +++ b/rust/src/providers/minimax/remains_api.rs @@ -7,7 +7,7 @@ use chrono::Utc; use crate::core::{FetchContext, ProviderError, ProviderFetchResult}; -use super::{MiniMaxRegion, coding_plan, coding_plan_html}; +use super::{JSON_ACCEPT, MiniMaxRegion, check_status, coding_plan, coding_plan_html, http_client}; fn resolve_plain_api_key(explicit: Option<&str>, environment: Option<&str>) -> Option { explicit @@ -46,31 +46,13 @@ async fn fetch_remains_once_via_api_key( api_key: &str, url: &str, ) -> Result { - let client = crate::core::credentialed_http_client_builder() - .timeout(std::time::Duration::from_secs(30)) - .build() - .map_err(|e| ProviderError::Other(e.to_string()))?; - - let response = client + let response = http_client()? .get(url) .header("Authorization", format!("Bearer {api_key}")) - .header("Accept", "application/json, text/plain, */*") + .header("Accept", JSON_ACCEPT) .send() .await?; - - let status = response.status(); - if status == reqwest::StatusCode::UNAUTHORIZED || status == reqwest::StatusCode::FORBIDDEN { - return Err(ProviderError::AuthRequired); - } - if !status.is_success() { - let message = format!("MiniMax remains (api key) returned status {status}"); - if status == reqwest::StatusCode::NOT_FOUND - || status == reqwest::StatusCode::METHOD_NOT_ALLOWED - { - return Err(ProviderError::Parse(message)); - } - return Err(ProviderError::Other(message)); - } + check_status(response.status(), "remains (api key)", true)?; let json: serde_json::Value = response .json() @@ -81,7 +63,40 @@ async fn fetch_remains_once_via_api_key( #[cfg(test)] mod tests { - use super::resolve_plain_api_key; + use super::super::http_tests::{JSON_ACCEPT, assert_status_error}; + use super::{fetch_remains_once_via_api_key, resolve_plain_api_key}; + + #[tokio::test] + async fn api_key_remains_request_maps_statuses() { + let cases: [(usize, Option<&str>); 4] = [ + (401, None), + (403, None), + ( + 404, + Some("Parse:MiniMax remains (api key) returned status 404 Not Found"), + ), + ( + 503, + Some("Other:MiniMax remains (api key) returned status 503 Service Unavailable"), + ), + ]; + for (status, expected) in cases { + let mut server = mockito::Server::new_async().await; + let mock = server + .mock("GET", "/remains") + .match_header("authorization", "Bearer fixture-key") + .match_header("accept", JSON_ACCEPT) + .with_status(status) + .expect(1) + .create_async() + .await; + let result = + fetch_remains_once_via_api_key("fixture-key", &format!("{}/remains", server.url())) + .await; + mock.assert_async().await; + assert_status_error(result.map(|_| ()), expected, status); + } + } #[test] fn plain_api_key_prefers_explicit_then_environment_without_mutating_process_env() { diff --git a/rust/src/providers/minimax/tests.rs b/rust/src/providers/minimax/tests.rs new file mode 100644 index 0000000000..3c74846d67 --- /dev/null +++ b/rust/src/providers/minimax/tests.rs @@ -0,0 +1,270 @@ +use super::billing::*; +use super::*; +use chrono::TimeZone; + +#[test] +fn minimax_region_defaults_to_global_io_urls() { + let region = MiniMaxRegion::from_settings_value(None); + assert_eq!(region, MiniMaxRegion::Global); + assert_eq!(region.settings_value(), "global"); + assert_eq!(region.cookie_domain(), "platform.minimax.io"); + assert_eq!( + region.coding_plan_url(), + "https://platform.minimax.io/user-center/payment/coding-plan?cycle_type=3" + ); + assert_eq!( + MiniMaxProvider::dashboard_url_for_region(None), + "https://platform.minimax.io/user-center/payment/coding-plan?cycle_type=3" + ); +} + +#[test] +fn minimax_region_accepts_legacy_china_value() { + for value in ["cn", "china", "china-mainland", "china_mainland"] { + let region = MiniMaxRegion::from_settings_value(Some(value)); + assert_eq!(region, MiniMaxRegion::ChinaMainland); + assert_eq!(region.settings_value(), "cn"); + assert_eq!(region.cookie_domain(), "platform.minimaxi.com"); + assert_eq!( + region.coding_plan_url(), + "https://platform.minimaxi.com/user-center/payment/coding-plan?cycle_type=3" + ); + } +} + +#[test] +fn minimax_region_cookie_search_domains_and_www_base_url() { + let global = MiniMaxRegion::Global; + assert_eq!(global.www_base_url(), "https://www.minimax.io"); + assert_eq!(global.cookie_search_domains(), ["minimax.io"]); + assert_eq!( + global.coding_plan_remains_url(), + "https://platform.minimax.io/v1/api/openplatform/coding_plan/remains" + ); + assert_eq!( + global.www_remains_url(), + "https://www.minimax.io/v1/api/openplatform/coding_plan/remains" + ); + + let cn = MiniMaxRegion::ChinaMainland; + assert_eq!(cn.www_base_url(), "https://www.minimaxi.com"); + assert_eq!(cn.cookie_search_domains(), ["minimaxi.com"]); + assert_eq!( + cn.coding_plan_remains_url(), + "https://platform.minimaxi.com/v1/api/openplatform/coding_plan/remains" + ); + assert_eq!( + cn.www_remains_url(), + "https://www.minimaxi.com/v1/api/openplatform/coding_plan/remains" + ); +} + +/// A dated billing record with every other field absent. +fn record(ymd: &str, method: &str, model: &str) -> MiniMaxBillingRecord { + MiniMaxBillingRecord { + consume_token: None, + consume_input_token: None, + consume_output_token: None, + consume_cash: None, + consume_cash_after_voucher: None, + created_at: None, + ymd: Some(ymd.to_string()), + consume_time: None, + method: Some(method.to_string()), + model: Some(model.to_string()), + result: None, + status: None, + } +} + +#[test] +fn aggregates_billing_history_records() { + let records = vec![ + MiniMaxBillingRecord { + consume_token: Some(serde_json::json!(1200)), + consume_cash: Some(serde_json::json!("0.42")), + result: Some(serde_json::json!("SUCCESS")), + ..record("2026-12-16", "chat", "abab6.5") + }, + MiniMaxBillingRecord { + consume_input_token: Some(serde_json::json!(300)), + consume_output_token: Some(serde_json::json!(500)), + consume_cash_after_voucher: Some(serde_json::json!(0.21)), + ..record("2026-12-15", "completion", "abab6.5") + }, + ]; + let now = Utc.with_ymd_and_hms(2026, 12, 16, 12, 0, 0).unwrap(); + let summary = aggregate_billing(&records, now); + assert_eq!(summary.last_30_days_tokens, 2000); + assert!((summary.last_30_days_cash.unwrap() - 0.63).abs() < 0.001); + assert_eq!(summary.top_models[0].name, "abab6.5"); + assert_eq!(summary.top_models[0].tokens, 2000); +} + +#[test] +fn parses_billing_payload_into_result_extras() { + let json = serde_json::json!({ + "base_resp": { "status_code": 0 }, + "charge_records": [{ + "consume_token": "2500", + "consume_cash_after_voucher": "1.25", + "ymd": Utc::now().format("%Y-%m-%d").to_string(), + "method": "chat", + "model": "abab6.5" + }] + }); + let summary = parse_billing_summary(&json).unwrap(); + let result = attach_billing_summary( + ProviderFetchResult::new(UsageSnapshot::new(RateWindow::new(0.0)), "web-billing"), + summary, + ); + assert!(result.cost.is_some()); + assert!( + result + .usage + .extra_rate_windows + .iter() + .any(|window| window.id == "billing-tokens-30d") + ); +} + +#[test] +fn billing_summary_windows_follow_a_fixed_order() { + let breakdown = |name: &str, tokens: i64, cash: Option| MiniMaxBillingBreakdown { + name: name.to_string(), + tokens, + cash, + }; + let summary = MiniMaxBillingSummary { + today_tokens: 1234, + last_30_days_tokens: 56789, + today_cash: Some(0.5), + last_30_days_cash: Some(12.345), + top_methods: vec![ + breakdown("chat", 5000, Some(1.0)), + breakdown("audio", 10, None), + ], + top_models: vec![breakdown("abab6.5", 7000, None)], + }; + let base = || ProviderFetchResult::new(UsageSnapshot::new(RateWindow::new(1.0)), "web"); + let windows = |result: &ProviderFetchResult| -> Vec<(String, String, Option)> { + result + .usage + .extra_rate_windows + .iter() + .map(|named| { + assert_eq!(named.window.used_percent, 0.0); + assert_eq!(named.window.window_minutes, None); + assert_eq!(named.window.resets_at, None); + ( + named.id.clone(), + named.title.clone(), + named.window.reset_description.clone(), + ) + }) + .collect() + }; + let row = |id: &str, title: &str, detail: &str| { + (id.to_string(), title.to_string(), Some(detail.to_string())) + }; + + let result = attach_billing_summary(base(), summary.clone()); + assert_eq!( + windows(&result), + vec![ + row("billing-tokens-today", "Tokens today", "1,234"), + row("billing-tokens-30d", "Tokens (30 days)", "56,789"), + row("billing-cash-today", "Spend today", "$0.50"), + row("billing-cash-30d", "Spend (30 days)", "$12.35"), + row("billing-method-0", "Method: chat", "5,000 tokens / $1.00"), + row("billing-method-1", "Method: audio", "10 tokens"), + row("billing-model-0", "Model: abab6.5", "7,000 tokens"), + ] + ); + let cost = result.cost.expect("30-day spend sets cost"); + assert_eq!(cost.used, 12.345); + assert_eq!(cost.currency_code, "USD"); + assert_eq!(result.usage.primary.used_percent, 1.0); + + let no_cash = MiniMaxBillingSummary { + today_cash: None, + last_30_days_cash: None, + top_methods: Vec::new(), + top_models: Vec::new(), + ..summary + }; + let result = attach_billing_summary(base(), no_cash); + assert_eq!( + windows(&result), + vec![ + row("billing-tokens-today", "Tokens today", "1,234"), + row("billing-tokens-30d", "Tokens (30 days)", "56,789"), + ] + ); + assert!(result.cost.is_none()); +} + +#[test] +fn parses_plan_title_from_coding_plan_fields() { + let provider = MiniMaxProvider::new(); + for (field, expected) in [ + ("plan_name", "MiniMax Star"), + ("current_plan_title", "Coding Plan Pro"), + ("current_subscribe_title", "Max"), + ("combo_title", "Combo Star"), + ] { + let snapshot = provider + .parse_usage_response(&serde_json::json!({ + "base_resp": { "status_code": 0 }, + field: expected, + "used_amount": 0, + "total_quota": 100 + })) + .unwrap(); + assert_eq!(snapshot.login_method.as_deref(), Some(expected)); + } + + let snapshot = provider + .parse_usage_response(&serde_json::json!({ + "base_resp": { "status_code": 0 }, + "current_combo_card": { "title": "Card Title" }, + "used_amount": 0, + "total_quota": 100 + })) + .unwrap(); + assert_eq!(snapshot.login_method.as_deref(), Some("Card Title")); +} + +#[test] +fn filters_failed_billing_records() { + let records = vec![ + MiniMaxBillingRecord { + consume_token: Some(serde_json::json!(1000)), + result: Some(serde_json::json!("SUCCESS")), + ..record("2026-05-17", "chat", "MiniMax-M1") + }, + MiniMaxBillingRecord { + consume_token: Some(serde_json::json!(2000)), + result: Some(serde_json::json!("FAILED")), + ..record("2026-05-17", "chat", "MiniMax-M1") + }, + MiniMaxBillingRecord { + consume_token: Some(serde_json::json!(3000)), + ..record("2026-05-17", "audio", "speech") + }, + MiniMaxBillingRecord { + consume_token: Some(serde_json::json!(4000)), + status: Some(serde_json::json!(0)), + ..record("2026-05-17", "video", "video") + }, + ]; + let now = Utc.with_ymd_and_hms(2026, 5, 17, 12, 0, 0).unwrap(); + let summary = aggregate_billing(&records, now); + assert_eq!(summary.today_tokens, 4000); + assert_eq!(summary.last_30_days_tokens, 4000); + assert_eq!(summary.top_methods.len(), 2); + assert_eq!(summary.top_methods[0].name, "audio"); + assert_eq!(summary.top_methods[0].tokens, 3000); + assert_eq!(summary.top_methods[1].name, "chat"); + assert_eq!(summary.top_methods[1].tokens, 1000); +} diff --git a/rust/src/providers/minimax/token_plan.rs b/rust/src/providers/minimax/token_plan.rs index c7d3ba8ea0..113bc3cd7c 100644 --- a/rust/src/providers/minimax/token_plan.rs +++ b/rust/src/providers/minimax/token_plan.rs @@ -10,8 +10,9 @@ use chrono::{DateTime, Utc}; use crate::core::{ProviderError, ProviderFetchResult, RateWindow, UsageSnapshot}; use super::{ - MiniMaxProvider, MiniMaxRegion, attach_billing_summary, coding_plan, coding_plan_html, - format_count, scalar_string, value_i64, + JSON_ACCEPT, MiniMaxProvider, MiniMaxRegion, attach_billing_summary, coding_plan, + coding_plan_html, console_get, format_count, http_client, is_auth_status, scalar_string, + value_i64, }; /// Token Plan console endpoints (www host; issue #254): the charge-API @@ -82,33 +83,23 @@ pub(crate) async fn fetch_token_plan_with_cookie( Ok(result) } -/// Shared GET envelope for the token-plan console endpoints — copied from -/// fetch_remains_once. Network errors propagate; callers map statuses. +/// Shared GET envelope for the token-plan console endpoints (the remains +/// request shape). Network errors propagate; callers map statuses. async fn token_plan_get( cookie_header: &str, url: &str, region: MiniMaxRegion, ) -> Result { - let client = crate::core::credentialed_http_client_builder() - .timeout(std::time::Duration::from_secs(30)) - .build() - .map_err(|e| ProviderError::Other(e.to_string()))?; - - let base = region.base_url(); - Ok(client - .get(url) - .header("Cookie", cookie_header) - .header("Accept", "application/json, text/plain, */*") - .header("X-Requested-With", "XMLHttpRequest") - .header("User-Agent", MiniMaxProvider::WEB_USER_AGENT) - .header("Accept-Language", "en-US,en;q=0.9") - .header("Origin", base) - .header( - "Referer", - &format!("{base}/user-center/payment/coding-plan"), - ) - .send() - .await?) + Ok(console_get( + &http_client()?, + url, + cookie_header, + region, + JSON_ACCEPT, + true, + ) + .send() + .await?) } /// (a) `charge/token_plan/usage`. 401/403 → AuthRequired; any other @@ -123,7 +114,7 @@ async fn fetch_token_plan_usage_once( let response = token_plan_get(cookie_header, &url, region).await?; let status = response.status(); - if status == reqwest::StatusCode::UNAUTHORIZED || status == reqwest::StatusCode::FORBIDDEN { + if is_auth_status(status) { return Err(ProviderError::AuthRequired); } if !status.is_success() { @@ -289,9 +280,41 @@ fn assemble_token_plan_result( #[cfg(test)] mod tests { + use super::super::http_tests::console_mock; use super::*; use chrono::TimeZone; + #[tokio::test] + async fn token_plan_backend_request_sends_console_headers() { + let mut server = mockito::Server::new_async().await; + let ok = console_mock(&mut server, "/summary") + .with_body(r#"{"ok":true}"#) + .expect(1) + .create_async() + .await; + let failed = console_mock(&mut server, "/credit") + .with_status(500) + .expect(1) + .create_async() + .await; + let summary = fetch_token_plan_backend_json( + "session=fixture", + &format!("{}/summary", server.url()), + MiniMaxRegion::Global, + ) + .await; + let credit = fetch_token_plan_backend_json( + "session=fixture", + &format!("{}/credit", server.url()), + MiniMaxRegion::Global, + ) + .await; + ok.assert_async().await; + failed.assert_async().await; + assert_eq!(summary, Some(serde_json::json!({ "ok": true }))); + assert_eq!(credit, None); + } + fn fixed_now() -> DateTime { // 2026-08-03T12:00:00Z = 1,785,758,400 unix seconds. Utc.timestamp_opt(1_785_758_400, 0).unwrap() diff --git a/rust/src/providers/opencode/billing.rs b/rust/src/providers/opencode/billing.rs index 2ac874dcb0..c5a4e4420d 100644 --- a/rust/src/providers/opencode/billing.rs +++ b/rust/src/providers/opencode/billing.rs @@ -6,11 +6,10 @@ //! parsing, the HTTP fallback, and the presentation mapping. use serde_json::Value; -use uuid::Uuid; use crate::core::{ProviderError, RateWindow, UsageSnapshot}; -use super::{BASE_URL, OpenCodeProvider, SERVER_URL}; +use super::OpenCodeProvider; /// Customer/billing server function carrying the monthly spend fields /// pay-as-you-go workspaces bill against (upstream 0.49.5 #2504/#2697). @@ -139,28 +138,12 @@ impl OpenCodeProvider { cookie_header: &str, ) -> Result, ProviderError> { let referer = format!("https://opencode.ai/workspace/{workspace_id}"); - let args = serde_json::json!([workspace_id]); - let encoded_args = Self::url_encode(&args.to_string()); - let url = format!( - "{}?id={}&args={}", - SERVER_URL, BILLING_SERVER_ID, encoded_args - ); - let response = self - .client - .get(&url) - .header("Cookie", cookie_header) - .header("X-Server-Id", BILLING_SERVER_ID) - .header("X-Server-Instance", format!("server-fn:{}", Uuid::new_v4())) - .header( - "User-Agent", - "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36", - ) - .header("Origin", BASE_URL) - .header("Referer", referer) - .header( - "Accept", - "text/javascript, application/json;q=0.9, */*;q=0.8", + .server_fn_get( + BILLING_SERVER_ID, + Some(workspace_id), + cookie_header, + &referer, ) .send() .await; diff --git a/rust/src/providers/opencode/mod.rs b/rust/src/providers/opencode/mod.rs index faac59aa29..b5b1979699 100755 --- a/rust/src/providers/opencode/mod.rs +++ b/rust/src/providers/opencode/mod.rs @@ -4,14 +4,6 @@ //! Uses browser cookies for authentication pub mod billing; -pub mod scraper; - -// Re-exports for advanced scraping -#[allow( - unused_imports, - reason = "imports needed for future OpenCode provider wiring" -)] -pub use scraper::{OpenCodeError, OpenCodeUsageFetcher, OpenCodeUsageSnapshot}; use async_trait::async_trait; use chrono::{DateTime, Utc}; @@ -77,26 +69,42 @@ impl OpenCodeProvider { self.parse_subscription(&subscription) } - /// Fetch workspace ID from server - async fn fetch_workspace_id(&self, cookie_header: &str) -> Result { - let url = format!("{}?id={}", SERVER_URL, WORKSPACES_SERVER_ID); - - let response = self - .client + /// GET a SolidStart server function, optionally with a one-element + /// `[workspace_id]` argument list. + pub(super) fn server_fn_get( + &self, + server_id: &str, + workspace_arg: Option<&str>, + cookie_header: &str, + referer: &str, + ) -> reqwest::RequestBuilder { + let mut url = format!("{}?id={}", SERVER_URL, server_id); + if let Some(workspace_id) = workspace_arg { + let args = serde_json::json!([workspace_id]); + url.push_str("&args="); + url.push_str(&Self::url_encode(&args.to_string())); + } + self.client .get(&url) .header("Cookie", cookie_header) - .header("X-Server-Id", WORKSPACES_SERVER_ID) + .header("X-Server-Id", server_id) .header("X-Server-Instance", format!("server-fn:{}", Uuid::new_v4())) .header( "User-Agent", "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36", ) .header("Origin", BASE_URL) - .header("Referer", BASE_URL) + .header("Referer", referer) .header( "Accept", "text/javascript, application/json;q=0.9, */*;q=0.8", ) + } + + /// Fetch workspace ID from server + async fn fetch_workspace_id(&self, cookie_header: &str) -> Result { + let response = self + .server_fn_get(WORKSPACES_SERVER_ID, None, cookie_header, BASE_URL) .send() .await?; @@ -133,28 +141,12 @@ impl OpenCodeProvider { cookie_header: &str, ) -> Result { let referer = format!("https://opencode.ai/workspace/{}/billing", workspace_id); - let args = serde_json::json!([workspace_id]); - let encoded_args = Self::url_encode(&args.to_string()); - let url = format!( - "{}?id={}&args={}", - SERVER_URL, SUBSCRIPTION_SERVER_ID, encoded_args - ); - let response = self - .client - .get(&url) - .header("Cookie", cookie_header) - .header("X-Server-Id", SUBSCRIPTION_SERVER_ID) - .header("X-Server-Instance", format!("server-fn:{}", Uuid::new_v4())) - .header( - "User-Agent", - "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36", - ) - .header("Origin", BASE_URL) - .header("Referer", referer) - .header( - "Accept", - "text/javascript, application/json;q=0.9, */*;q=0.8", + .server_fn_get( + SUBSCRIPTION_SERVER_ID, + Some(workspace_id), + cookie_header, + &referer, ) .send() .await?; @@ -619,6 +611,73 @@ impl Provider for OpenCodeProvider { mod tests { use super::*; + fn request_headers(request: &reqwest::Request) -> Vec<(String, String)> { + request + .headers() + .iter() + .map(|(name, value)| (name.to_string(), value.to_str().unwrap().to_string())) + .collect() + } + + #[test] + fn server_fn_get_pins_url_and_header_order_for_each_call_site() { + let provider = OpenCodeProvider::new(); + let cases = [ + (WORKSPACES_SERVER_ID, None, BASE_URL, ""), + ( + SUBSCRIPTION_SERVER_ID, + Some("wrk_a b"), + "https://opencode.ai/workspace/wrk_a b/billing", + "&args=%5B%22wrk_a%20b%22%5D", + ), + ( + billing::BILLING_SERVER_ID, + Some("wrk_a b"), + "https://opencode.ai/workspace/wrk_a b", + "&args=%5B%22wrk_a%20b%22%5D", + ), + ]; + for (server_id, workspace_arg, referer, args) in cases { + let request = provider + .server_fn_get(server_id, workspace_arg, "auth=synthetic", referer) + .build() + .unwrap(); + assert_eq!(request.method(), reqwest::Method::GET); + assert_eq!( + request.url().as_str(), + format!("{SERVER_URL}?id={server_id}{args}") + ); + let headers = request_headers(&request); + let names: Vec<_> = headers.iter().map(|(name, _)| name.as_str()).collect(); + assert_eq!( + names, + [ + "cookie", + "x-server-id", + "x-server-instance", + "user-agent", + "origin", + "referer", + "accept" + ] + ); + assert_eq!(headers[0].1, "auth=synthetic"); + assert_eq!(headers[1].1, server_id); + let instance = headers[2].1.strip_prefix("server-fn:").unwrap(); + assert!(Uuid::parse_str(instance).is_ok()); + assert_eq!( + headers[3].1, + "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36" + ); + assert_eq!(headers[4].1, BASE_URL); + assert_eq!(headers[5].1, referer); + assert_eq!( + headers[6].1, + "text/javascript, application/json;q=0.9, */*;q=0.8" + ); + } + } + #[test] fn parses_json_renewal_window() { let provider = OpenCodeProvider::new(); @@ -650,29 +709,91 @@ mod tests { } #[test] - fn parses_weekly_only_json_without_rolling() { - let provider = OpenCodeProvider::new(); - let now = DateTime::::from_timestamp(1_700_000_000, 0).unwrap(); - let payload = serde_json::json!({ - "weeklyUsage": { "usagePercent": 76.0, "resetInSec": 86400 } - }); - - let snap = provider.parse_usage_json(&payload, now).expect("snapshot"); - assert!((snap.primary.used_percent - 76.0).abs() < f64::EPSILON); - assert!(snap.secondary.is_none()); - } - - #[test] - fn parses_rolling_only_json_without_weekly() { + fn parse_usage_json_window_shapes() { let provider = OpenCodeProvider::new(); let now = DateTime::::from_timestamp(1_700_000_000, 0).unwrap(); - let payload = serde_json::json!({ - "rollingUsage": { "usagePercent": 12.5, "resetInSec": 600 } - }); - - let snap = provider.parse_usage_json(&payload, now).expect("snapshot"); - assert!((snap.primary.used_percent - 12.5).abs() < f64::EPSILON); - assert!(snap.secondary.is_none()); + let cases = [ + ( + "weekly only, without rolling", + serde_json::json!({ + "weeklyUsage": { "usagePercent": 76.0, "resetInSec": 86400 } + }), + 76.0, + None, + f64::EPSILON, + ), + ( + "rolling only, without weekly", + serde_json::json!({ + "rollingUsage": { "usagePercent": 12.5, "resetInSec": 600 } + }), + 12.5, + None, + f64::EPSILON, + ), + ( + "nested data.subscription shape", + serde_json::json!({ + "data": { + "subscription": { + "rollingUsage": { "usedPercent": 18.5, "resetInSeconds": 1200 }, + "weeklyUsage": { "percentUsed": 42, "resetsInSec": 86400 } + } + } + }), + 18.5, + Some(42.0), + f64::EPSILON, + ), + // The 0.25 week fraction scales to 25%. + ( + "fiveHour and week aliases", + serde_json::json!({ + "fiveHourUsage": { "pct": 11, "resetInSec": 300 }, + "weekUsage": { "value": 0.25, "resetInSec": 7200 } + }), + 11.0, + Some(25.0), + f64::EPSILON, + ), + // Upstream #2331: used=1, limit=100 is 1%, not 100% (false exhausted). + ( + "sub-one-percent computed used/limit is not rescaled to 100", + serde_json::json!({ + "rollingUsage": { "used": 1, "limit": 100, "resetInSec": 600 }, + "weeklyUsage": { "used": 1, "limit": 200, "resetInSec": 86400 } + }), + 1.0, + Some(0.5), + 0.001, + ), + ( + "direct fractional usagePercent still scales", + serde_json::json!({ + "rollingUsage": { "usagePercent": 0.25, "resetInSec": 600 } + }), + 25.0, + None, + 0.001, + ), + ]; + for (name, payload, primary, secondary, tolerance) in cases { + let snap = provider.parse_usage_json(&payload, now).expect(name); + assert!( + (snap.primary.used_percent - primary).abs() < tolerance, + "{name}: primary {}", + snap.primary.used_percent + ); + match (snap.secondary.as_ref(), secondary) { + (None, None) => {} + (Some(window), Some(expected)) => assert!( + (window.used_percent - expected).abs() < tolerance, + "{name}: secondary {}", + window.used_percent + ), + (actual, expected) => panic!("{name}: secondary {actual:?} vs {expected:?}"), + } + } } #[test] @@ -693,74 +814,6 @@ mod tests { assert!((snap.primary.used_percent - 55.0).abs() < f64::EPSILON); } - #[test] - fn parses_nested_data_subscription_shape() { - let provider = OpenCodeProvider::new(); - let now = DateTime::::from_timestamp(1_700_000_000, 0).unwrap(); - let payload = serde_json::json!({ - "data": { - "subscription": { - "rollingUsage": { "usedPercent": 18.5, "resetInSeconds": 1200 }, - "weeklyUsage": { "percentUsed": 42, "resetsInSec": 86400 } - } - } - }); - let snap = provider - .parse_usage_json(&payload, now) - .expect("nested snapshot"); - assert!((snap.primary.used_percent - 18.5).abs() < f64::EPSILON); - assert!((snap.secondary.as_ref().unwrap().used_percent - 42.0).abs() < f64::EPSILON); - } - - #[test] - fn parses_five_hour_and_week_aliases() { - let provider = OpenCodeProvider::new(); - let now = DateTime::::from_timestamp(1_700_000_000, 0).unwrap(); - let payload = serde_json::json!({ - "fiveHourUsage": { "pct": 11, "resetInSec": 300 }, - "weekUsage": { "value": 0.25, "resetInSec": 7200 } - }); - let snap = provider - .parse_usage_json(&payload, now) - .expect("alias snapshot"); - assert!((snap.primary.used_percent - 11.0).abs() < f64::EPSILON); - // 0.25 fraction scales to 25% - assert!((snap.secondary.as_ref().unwrap().used_percent - 25.0).abs() < f64::EPSILON); - } - - #[test] - fn sub_one_percent_computed_used_limit_is_not_rescaled_to_100() { - // Upstream #2331: used=1, limit=100 → 1%, not 100% (false exhausted). - let provider = OpenCodeProvider::new(); - let now = DateTime::::from_timestamp(1_700_000_000, 0).unwrap(); - let payload = serde_json::json!({ - "rollingUsage": { "used": 1, "limit": 100, "resetInSec": 600 }, - "weeklyUsage": { "used": 1, "limit": 200, "resetInSec": 86400 } - }); - let snap = provider.parse_usage_json(&payload, now).expect("snapshot"); - assert!( - (snap.primary.used_percent - 1.0).abs() < 0.001, - "primary {}", - snap.primary.used_percent - ); - assert!( - (snap.secondary.as_ref().unwrap().used_percent - 0.5).abs() < 0.001, - "secondary {}", - snap.secondary.as_ref().unwrap().used_percent - ); - } - - #[test] - fn direct_fractional_usage_percent_still_scales() { - let provider = OpenCodeProvider::new(); - let now = DateTime::::from_timestamp(1_700_000_000, 0).unwrap(); - let payload = serde_json::json!({ - "rollingUsage": { "usagePercent": 0.25, "resetInSec": 600 } - }); - let snap = provider.parse_usage_json(&payload, now).expect("snapshot"); - assert!((snap.primary.used_percent - 25.0).abs() < 0.001); - } - #[test] fn parse_subscription_accepts_realistic_server_fn_body() { let provider = OpenCodeProvider::new(); diff --git a/rust/src/providers/opencode/scraper.rs b/rust/src/providers/opencode/scraper.rs deleted file mode 100755 index 5260d505fb..0000000000 --- a/rust/src/providers/opencode/scraper.rs +++ /dev/null @@ -1,700 +0,0 @@ -//! OpenCode Advanced Web Scraping -//! -//! Multi-stage fetcher for opencode.ai with: -//! - Workspace ID resolution -//! - Subscription info fetching via custom X-Server-Id headers -//! - Robust JSON parsing with multiple fallback strategies - -use chrono::{DateTime, Duration, Utc}; -use regex_lite::Regex; -use serde::{Deserialize, Serialize}; -use std::sync::OnceLock; - -/// OpenCode server IDs for API endpoints -const WORKSPACES_SERVER_ID: &str = - "def39973159c7f0483d8793a822b8dbb10d067e12c65455fcb4608459ba0234f"; -const SUBSCRIPTION_SERVER_ID: &str = - "7abeebee372f304e050aaaf92be863f4a86490e382f8c79db68fd94040d691b4"; -const BASE_URL: &str = "https://opencode.ai"; -const SERVER_URL: &str = "https://opencode.ai/_server"; - -/// Keys to look for when parsing percent values -const PERCENT_KEYS: &[&str] = &[ - "usagePercent", - "usedPercent", - "percentUsed", - "percent", - "usage_percent", - "used_percent", - "utilization", - "utilizationPercent", - "utilization_percent", - // Note: raw "usage" is often a token count — handled via used/limit only. -]; - -/// Keys to look for when parsing reset time in seconds -const RESET_IN_KEYS: &[&str] = &[ - "resetInSec", - "resetInSeconds", - "resetSeconds", - "reset_sec", - "reset_in_sec", - "resetsInSec", - "resetsInSeconds", - "resetIn", - "resetSec", -]; - -/// Keys to look for when parsing reset timestamp -const RESET_AT_KEYS: &[&str] = &[ - "resetAt", - "resetsAt", - "reset_at", - "resets_at", - "nextReset", - "next_reset", - "renewAt", - "renew_at", -]; - -/// OpenCode usage snapshot -#[derive(Debug, Clone, Serialize, Deserialize)] -pub struct OpenCodeUsageSnapshot { - /// Rolling window usage percent (5-hour window) - pub rolling_usage_percent: f64, - /// Weekly usage percent - pub weekly_usage_percent: f64, - /// Rolling window reset in seconds - pub rolling_reset_in_sec: i64, - /// Weekly reset in seconds - pub weekly_reset_in_sec: i64, - /// Account renewal date when exposed by OpenCode. - pub renews_at: Option>, - /// When this snapshot was captured - pub updated_at: DateTime, -} - -impl OpenCodeUsageSnapshot { - /// Get rolling window reset time - pub fn rolling_resets_at(&self) -> DateTime { - self.updated_at + Duration::seconds(self.rolling_reset_in_sec) - } - - /// Get weekly reset time - pub fn weekly_resets_at(&self) -> DateTime { - self.updated_at + Duration::seconds(self.weekly_reset_in_sec) - } -} - -/// OpenCode usage fetcher errors -#[derive(Debug, Clone)] -pub enum OpenCodeError { - /// Session cookie is invalid or expired - InvalidCredentials, - /// Network error - NetworkError(String), - /// API error - ApiError(String), - /// Parse error - ParseFailed(String), -} - -impl std::fmt::Display for OpenCodeError { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - match self { - OpenCodeError::InvalidCredentials => { - write!(f, "OpenCode session cookie is invalid or expired") - } - OpenCodeError::NetworkError(msg) => write!(f, "OpenCode network error: {}", msg), - OpenCodeError::ApiError(msg) => write!(f, "OpenCode API error: {}", msg), - OpenCodeError::ParseFailed(msg) => write!(f, "OpenCode parse error: {}", msg), - } - } -} - -impl std::error::Error for OpenCodeError {} - -/// Server request configuration -struct ServerRequest { - server_id: String, - args: Option>, - method: String, - referer: String, -} - -/// OpenCode usage fetcher -pub struct OpenCodeUsageFetcher; - -impl OpenCodeUsageFetcher { - /// Fetch usage data from OpenCode - pub async fn fetch_usage( - cookie_header: &str, - timeout_secs: u64, - workspace_id_override: Option<&str>, - ) -> Result { - let now = Utc::now(); - - // Get or resolve workspace ID - let workspace_id = if let Some(override_id) = workspace_id_override { - Self::normalize_workspace_id(override_id) - .ok_or_else(|| OpenCodeError::ParseFailed("Invalid workspace ID".to_string()))? - } else { - Self::fetch_workspace_id(cookie_header, timeout_secs).await? - }; - - // Fetch subscription info - let subscription_text = - Self::fetch_subscription_info(&workspace_id, cookie_header, timeout_secs).await?; - - // Parse subscription - Self::parse_subscription(&subscription_text, now) - } - - /// Fetch workspace ID from OpenCode - async fn fetch_workspace_id( - cookie_header: &str, - timeout_secs: u64, - ) -> Result { - let request = ServerRequest { - server_id: WORKSPACES_SERVER_ID.to_string(), - args: None, - method: "GET".to_string(), - referer: BASE_URL.to_string(), - }; - - let text = Self::fetch_server_text(&request, cookie_header, timeout_secs).await?; - - if Self::looks_signed_out(&text) { - return Err(OpenCodeError::InvalidCredentials); - } - - // Try parsing workspace IDs - let mut ids = Self::parse_workspace_ids(&text); - if ids.is_empty() { - ids = Self::parse_workspace_ids_from_json(&text); - } - - if ids.is_empty() { - // Retry with POST - let retry_request = ServerRequest { - server_id: WORKSPACES_SERVER_ID.to_string(), - args: Some(vec![]), - method: "POST".to_string(), - referer: BASE_URL.to_string(), - }; - - let fallback = - Self::fetch_server_text(&retry_request, cookie_header, timeout_secs).await?; - - if Self::looks_signed_out(&fallback) { - return Err(OpenCodeError::InvalidCredentials); - } - - ids = Self::parse_workspace_ids(&fallback); - if ids.is_empty() { - ids = Self::parse_workspace_ids_from_json(&fallback); - } - - if ids.is_empty() { - return Err(OpenCodeError::ParseFailed( - "Missing workspace id".to_string(), - )); - } - } - - Ok(ids.into_iter().next().unwrap()) - } - - /// Fetch subscription info for a workspace - async fn fetch_subscription_info( - workspace_id: &str, - cookie_header: &str, - timeout_secs: u64, - ) -> Result { - let referer = format!("{}/workspace/{}/billing", BASE_URL, workspace_id); - let request = ServerRequest { - server_id: SUBSCRIPTION_SERVER_ID.to_string(), - args: Some(vec![serde_json::Value::String(workspace_id.to_string())]), - method: "GET".to_string(), - referer, - }; - - let text = Self::fetch_server_text(&request, cookie_header, timeout_secs).await?; - - if Self::looks_signed_out(&text) { - return Err(OpenCodeError::InvalidCredentials); - } - - // Check if we got valid data, otherwise retry with POST - if Self::parse_subscription_json(&text, Utc::now()).is_none() - && Self::extract_double( - r#"rollingUsage[^}]*?usagePercent\s*:\s*([0-9]+(?:\.[0-9]+)?)"#, - &text, - ) - .is_none() - { - let retry_request = ServerRequest { - server_id: SUBSCRIPTION_SERVER_ID.to_string(), - args: Some(vec![serde_json::Value::String(workspace_id.to_string())]), - method: "POST".to_string(), - referer: format!("{}/workspace/{}/billing", BASE_URL, workspace_id), - }; - - let fallback = - Self::fetch_server_text(&retry_request, cookie_header, timeout_secs).await?; - - if Self::looks_signed_out(&fallback) { - return Err(OpenCodeError::InvalidCredentials); - } - - return Ok(fallback); - } - - Ok(text) - } - - /// Normalize a workspace ID from various formats - pub fn normalize_workspace_id(raw: &str) -> Option { - let trimmed = raw.trim(); - - // Already a valid workspace ID - if trimmed.starts_with("wrk_") && trimmed.len() > 4 { - return Some(trimmed.to_string()); - } - - // Try to parse as URL - if let Ok(url) = reqwest::Url::parse(trimmed) { - let parts: Vec<&str> = url.path_segments().map(|s| s.collect()).unwrap_or_default(); - if let Some(idx) = parts.iter().position(|&p| p == "workspace") - && parts.len() > idx + 1 - { - let candidate = parts[idx + 1]; - if candidate.starts_with("wrk_") && candidate.len() > 4 { - return Some(candidate.to_string()); - } - } - } - - // Try regex extraction - static WRK_REGEX: OnceLock = OnceLock::new(); - let regex = WRK_REGEX.get_or_init(|| Regex::new(r"wrk_[A-Za-z0-9]+").unwrap()); - - regex.find(trimmed).map(|m| m.as_str().to_string()) - } - - /// Fetch server text with custom headers - async fn fetch_server_text( - request: &ServerRequest, - cookie_header: &str, - timeout_secs: u64, - ) -> Result { - let url = Self::build_server_url(&request.server_id, &request.args, &request.method); - - let client = crate::core::credentialed_http_client_builder() - .timeout(std::time::Duration::from_secs(timeout_secs)) - .build() - .map_err(|e| OpenCodeError::NetworkError(e.to_string()))?; - - let mut req = if request.method == "GET" { - client.get(&url) - } else { - client.post(&url) - }; - - req = req - .header("Cookie", cookie_header) - .header("X-Server-Id", &request.server_id) - .header("X-Server-Instance", format!("server-fn:{}", uuid::Uuid::new_v4())) - .header("User-Agent", "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 Chrome/120.0.0.0 Safari/537.36") - .header("Origin", BASE_URL) - .header("Referer", &request.referer) - .header("Accept", "text/javascript, application/json;q=0.9, */*;q=0.8"); - - if request.method != "GET" - && let Some(args) = &request.args - { - req = req.header("Content-Type", "application/json").json(args); - } - - let response = req - .send() - .await - .map_err(|e| OpenCodeError::NetworkError(e.to_string()))?; - - let status = response.status(); - let text = response - .text() - .await - .map_err(|e| OpenCodeError::NetworkError(e.to_string()))?; - - if !status.is_success() { - if Self::looks_signed_out(&text) { - return Err(OpenCodeError::InvalidCredentials); - } - if status.as_u16() == 401 || status.as_u16() == 403 { - return Err(OpenCodeError::InvalidCredentials); - } - return Err(OpenCodeError::ApiError(format!("HTTP {}", status))); - } - - Ok(text) - } - - /// Build server request URL - fn build_server_url( - server_id: &str, - args: &Option>, - method: &str, - ) -> String { - if method != "GET" { - return SERVER_URL.to_string(); - } - - let mut url = format!("{}?id={}", SERVER_URL, server_id); - if let Some(args) = args - && !args.is_empty() - && let Ok(json_str) = serde_json::to_string(args) - { - url.push_str(&format!("&args={}", Self::url_encode(&json_str))); - } - url - } - - /// URL encode a string for query parameters - fn url_encode(s: &str) -> String { - let mut result = String::with_capacity(s.len() * 3); - for c in s.chars() { - match c { - 'a'..='z' | 'A'..='Z' | '0'..='9' | '-' | '_' | '.' | '~' => { - result.push(c); - } - _ => { - for b in c.to_string().as_bytes() { - result.push_str(&format!("%{:02X}", b)); - } - } - } - } - result - } - - /// Parse subscription text into snapshot - pub fn parse_subscription( - text: &str, - now: DateTime, - ) -> Result { - // Try JSON parsing first - if let Some(snapshot) = Self::parse_subscription_json(text, now) { - return Ok(snapshot); - } - - // Fall back to regex parsing - let rolling_percent = Self::extract_double( - r#"rollingUsage[^}]*?usagePercent\s*:\s*([0-9]+(?:\.[0-9]+)?)"#, - text, - ); - let rolling_reset = - Self::extract_int(r#"rollingUsage[^}]*?resetInSec\s*:\s*([0-9]+)"#, text); - let weekly_percent = Self::extract_double( - r#"weeklyUsage[^}]*?usagePercent\s*:\s*([0-9]+(?:\.[0-9]+)?)"#, - text, - ); - let weekly_reset = Self::extract_int(r#"weeklyUsage[^}]*?resetInSec\s*:\s*([0-9]+)"#, text); - - match (rolling_percent, rolling_reset, weekly_percent, weekly_reset) { - (Some(rp), Some(rr), Some(wp), Some(wr)) => Ok(OpenCodeUsageSnapshot { - rolling_usage_percent: rp, - weekly_usage_percent: wp, - rolling_reset_in_sec: rr, - weekly_reset_in_sec: wr, - renews_at: crate::providers::extract_renewal(text), - updated_at: now, - }), - _ => Err(OpenCodeError::ParseFailed( - "Missing usage fields".to_string(), - )), - } - } - - /// Parse subscription JSON - fn parse_subscription_json(text: &str, now: DateTime) -> Option { - let json: serde_json::Value = serde_json::from_str(text).ok()?; - let renews_at = Self::find_datetime(&json, &["renewAt", "renew_at"]); - - // Try direct parsing - if let Some(mut snapshot) = Self::parse_usage_dict(&json, now) { - snapshot.renews_at = renews_at; - return Some(snapshot); - } - - // Try nested keys - for key in ["data", "result", "usage", "billing", "payload"] { - if let Some(nested) = json.get(key) - && let Some(mut snapshot) = Self::parse_usage_dict(nested, now) - { - snapshot.renews_at = snapshot.renews_at.or(renews_at); - return Some(snapshot); - } - } - - None - } - - /// Parse usage from a JSON object - fn parse_usage_dict( - json: &serde_json::Value, - now: DateTime, - ) -> Option { - let rolling_keys = ["rollingUsage", "rolling", "rolling_usage"]; - let weekly_keys = ["weeklyUsage", "weekly", "weekly_usage"]; - - let rolling = rolling_keys.iter().find_map(|k| json.get(k)); - let weekly = weekly_keys.iter().find_map(|k| json.get(k)); - - if let (Some(rolling), Some(weekly)) = (rolling, weekly) { - let rolling_window = Self::parse_window(rolling, now)?; - let weekly_window = Self::parse_window(weekly, now)?; - - return Some(OpenCodeUsageSnapshot { - rolling_usage_percent: rolling_window.0, - weekly_usage_percent: weekly_window.0, - rolling_reset_in_sec: rolling_window.1, - weekly_reset_in_sec: weekly_window.1, - renews_at: Self::find_datetime(json, &["renewAt", "renew_at"]), - updated_at: now, - }); - } - - None - } - - /// Parse a window object into (percent, reset_in_sec) - fn parse_window(json: &serde_json::Value, _now: DateTime) -> Option<(f64, i64)> { - // Direct percent field (may be 0..1 fraction or 0..100). - let direct = PERCENT_KEYS - .iter() - .find_map(|k| json.get(k).and_then(|v| v.as_f64())); - let percent_is_direct = direct.is_some(); - - let percent = direct.or_else(|| { - // Computed used/limit is already 0..100 — never apply fraction rescale (#2331). - let used = ["used", "usage", "consumed", "count", "usedTokens"] - .iter() - .find_map(|k| json.get(*k).and_then(|v| v.as_f64())); - let limit = ["limit", "total", "allowance"] - .iter() - .find_map(|k| json.get(*k).and_then(|v| v.as_f64())); - match (used, limit) { - (Some(u), Some(l)) if l > 0.0 => Some((u / l) * 100.0), - _ => None, - } - })?; - - let reset_in = RESET_IN_KEYS - .iter() - .find_map(|k| json.get(k).and_then(|v| v.as_i64())) - .unwrap_or(0); - - let normalized_percent = if percent_is_direct && (0.0..=1.0).contains(&percent) { - percent * 100.0 - } else { - percent.clamp(0.0, 100.0) - }; - Some((normalized_percent.clamp(0.0, 100.0), reset_in.max(0))) - } - - fn find_datetime(json: &serde_json::Value, keys: &[&str]) -> Option> { - for key in keys { - if let Some(value) = json.get(*key) - && let Some(parsed) = Self::date_from_value(value) - { - return Some(parsed); - } - } - - if let Some(obj) = json.as_object() { - for value in obj.values() { - if let Some(parsed) = Self::find_datetime(value, keys) { - return Some(parsed); - } - } - } - None - } - - fn date_from_value(value: &serde_json::Value) -> Option> { - if let Some(number) = value.as_i64() { - return Self::date_from_timestamp(number as f64); - } - if let Some(number) = value.as_f64() { - return Self::date_from_timestamp(number); - } - let text = value.as_str()?.trim(); - if text.is_empty() { - return None; - } - if let Ok(number) = text.parse::() { - return Self::date_from_timestamp(number); - } - DateTime::parse_from_rfc3339(text) - .ok() - .map(|dt| dt.with_timezone(&Utc)) - } - - fn date_from_timestamp(number: f64) -> Option> { - if !number.is_finite() || number <= 0.0 { - return None; - } - let seconds = if number > 10_000_000_000.0 { - number / 1000.0 - } else { - number - }; - #[allow( - clippy::cast_possible_truncation, - reason = "timestamps are normalized to whole seconds (ms input divided by 1000) before the cast" - )] - DateTime::::from_timestamp(seconds as i64, 0) - } - - /// Parse workspace IDs from text using regex - fn parse_workspace_ids(text: &str) -> Vec { - // Try JSON parsing first - let json_ids = Self::parse_workspace_ids_from_json(text); - if !json_ids.is_empty() { - return json_ids; - } - - // Fall back to regex for JavaScript object notation (unquoted keys) - static WRK_ID_REGEX: OnceLock = OnceLock::new(); - let regex = WRK_ID_REGEX.get_or_init(|| { - Regex::new(r#""id"\s*:\s*"(wrk_[^"]+)"|id\s*:\s*"(wrk_[^"]+)""#).unwrap() - }); - - regex - .captures_iter(text) - .filter_map(|cap| { - cap.get(1) - .or_else(|| cap.get(2)) - .map(|m| m.as_str().to_string()) - }) - .collect() - } - - /// Parse workspace IDs from JSON - fn parse_workspace_ids_from_json(text: &str) -> Vec { - let json: serde_json::Value = match serde_json::from_str(text) { - Ok(v) => v, - Err(_) => return vec![], - }; - - let mut results = Vec::new(); - Self::collect_workspace_ids(&json, &mut results); - results - } - - /// Recursively collect workspace IDs from JSON - fn collect_workspace_ids(value: &serde_json::Value, out: &mut Vec) { - match value { - serde_json::Value::String(s) if s.starts_with("wrk_") && !out.contains(s) => { - out.push(s.clone()); - } - serde_json::Value::Array(arr) => { - for item in arr { - Self::collect_workspace_ids(item, out); - } - } - serde_json::Value::Object(obj) => { - for v in obj.values() { - Self::collect_workspace_ids(v, out); - } - } - _ => {} - } - } - - /// Check if response indicates signed out - fn looks_signed_out(text: &str) -> bool { - let lower = text.to_lowercase(); - lower.contains("login") || lower.contains("sign in") || lower.contains("auth/authorize") - } - - /// Extract a double from text using regex - fn extract_double(pattern: &str, text: &str) -> Option { - let regex = Regex::new(pattern).ok()?; - let captures = regex.captures(text)?; - let value_str = captures.get(1)?.as_str(); - value_str.parse().ok() - } - - /// Extract an int from text using regex - fn extract_int(pattern: &str, text: &str) -> Option { - let regex = Regex::new(pattern).ok()?; - let captures = regex.captures(text)?; - let value_str = captures.get(1)?.as_str(); - value_str.parse().ok() - } -} - -#[cfg(test)] -mod tests { - use super::*; - - #[test] - fn test_normalize_workspace_id_direct() { - assert_eq!( - OpenCodeUsageFetcher::normalize_workspace_id("wrk_abc123"), - Some("wrk_abc123".to_string()) - ); - } - - #[test] - fn test_normalize_workspace_id_url() { - assert_eq!( - OpenCodeUsageFetcher::normalize_workspace_id( - "https://opencode.ai/workspace/wrk_xyz789/billing" - ), - Some("wrk_xyz789".to_string()) - ); - } - - #[test] - fn test_normalize_workspace_id_invalid() { - assert_eq!( - OpenCodeUsageFetcher::normalize_workspace_id("invalid"), - None - ); - } - - #[test] - fn test_parse_workspace_ids() { - let text = r#"{"workspaces":[{"id":"wrk_abc123"},{"id":"wrk_def456"}]}"#; - let ids = OpenCodeUsageFetcher::parse_workspace_ids(text); - assert_eq!(ids, vec!["wrk_abc123", "wrk_def456"]); - } - - #[test] - fn test_looks_signed_out() { - assert!(OpenCodeUsageFetcher::looks_signed_out( - "Please login to continue" - )); - assert!(OpenCodeUsageFetcher::looks_signed_out("Sign in required")); - assert!(!OpenCodeUsageFetcher::looks_signed_out( - "Welcome to OpenCode" - )); - } - - #[test] - fn test_parse_subscription_json() { - let json = r#"{ - "rollingUsage": {"usagePercent": 45.5, "resetInSec": 3600}, - "weeklyUsage": {"usagePercent": 20.0, "resetInSec": 604800} - }"#; - - let snapshot = OpenCodeUsageFetcher::parse_subscription(json, Utc::now()).unwrap(); - assert!((snapshot.rolling_usage_percent - 45.5).abs() < 0.01); - assert!((snapshot.weekly_usage_percent - 20.0).abs() < 0.01); - assert_eq!(snapshot.rolling_reset_in_sec, 3600); - assert_eq!(snapshot.weekly_reset_in_sec, 604800); - } -} diff --git a/rust/src/providers/opencodego/console.rs b/rust/src/providers/opencodego/console.rs index 6a1d0daa34..31a28dbc13 100644 --- a/rust/src/providers/opencodego/console.rs +++ b/rust/src/providers/opencodego/console.rs @@ -64,7 +64,7 @@ pub(super) async fn fetch_workspace_id( parse_workspace_ids(&text) .into_iter() .next() - .ok_or_else(|| ProviderError::Parse("Missing OpenCode Console workspace ID".to_string())) + .ok_or_else(|| missing("workspace ID")) } pub(super) async fn fetch_usage( @@ -145,26 +145,21 @@ fn parse_workspace_ids(text: &str) -> Vec { } fn parse_usage(text: &str, _now: DateTime) -> Result { - let root: Value = serde_json::from_str(text) - .map_err(|_| ProviderError::Parse("Invalid OpenCode Console usage payload".to_string()))?; + let root: Value = serde_json::from_str(text).map_err(|_| invalid("usage payload"))?; if root.is_null() || root.get("access").is_some_and(Value::is_null) { return Ok(ConsoleUsage::NoSubscription); } let access = root .get("access") .and_then(Value::as_object) - .ok_or_else(|| { - ProviderError::Parse("Invalid OpenCode Console usage payload".to_string()) - })?; + .ok_or_else(|| invalid("usage payload"))?; let meters = access .get("meters") .and_then(Value::as_object) - .ok_or_else(|| { - ProviderError::Parse("Invalid OpenCode Console usage payload".to_string()) - })?; - let rolling = meters.get("fiveHour").ok_or_else(|| { - ProviderError::Parse("Invalid OpenCode Console usage payload".to_string()) - })?; + .ok_or_else(|| invalid("usage payload"))?; + let rolling = meters + .get("fiveHour") + .ok_or_else(|| invalid("usage payload"))?; let ends_at = access.get("endsAt").and_then(date_value); let primary = meter_window(rolling, Some(300), None)?; let mut snapshot = UsageSnapshot::new(primary).with_login_method("OpenCode Go"); @@ -192,14 +187,12 @@ fn meter_window( window_minutes: Option, reset_override: Option>, ) -> Result { - let object = meter - .as_object() - .ok_or_else(|| ProviderError::Parse("Invalid OpenCode Console usage meter".to_string()))?; - let used = numeric_value(object.get("usedMicroCents")) - .ok_or_else(|| ProviderError::Parse("Missing OpenCode Console usage amount".to_string()))?; + let object = meter.as_object().ok_or_else(|| invalid("usage meter"))?; + let used = + numeric_value(object.get("usedMicroCents")).ok_or_else(|| missing("usage amount"))?; let limit = numeric_value(object.get("limitMicroCents")) .filter(|value| *value > 0.0) - .ok_or_else(|| ProviderError::Parse("Missing OpenCode Console usage limit".to_string()))?; + .ok_or_else(|| missing("usage limit"))?; let resets_at = reset_override.or_else(|| object.get("resetsAt").and_then(date_value)); Ok(RateWindow::with_details( ((used / limit) * 100.0).clamp(0.0, 100.0), @@ -210,47 +203,43 @@ fn meter_window( } fn parse_balance(text: &str) -> Result, ProviderError> { - let root: Value = serde_json::from_str(text).map_err(|_| { - ProviderError::Parse("Invalid OpenCode Console billing payload".to_string()) - })?; + let root: Value = serde_json::from_str(text).map_err(|_| invalid("billing payload"))?; let billing_mode = root .get("billingMode") .and_then(Value::as_str) .filter(|mode| matches!(*mode, "prepaid" | "legacy" | "seat" | "credit")) - .ok_or_else(|| { - ProviderError::Parse("Invalid OpenCode Console billing payload".to_string()) - })?; + .ok_or_else(|| invalid("billing payload"))?; let mode = root .get("mode") .and_then(Value::as_str) .filter(|mode| matches!(*mode, "pay-as-you-go" | "invoiceable")) - .ok_or_else(|| { - ProviderError::Parse("Invalid OpenCode Console billing payload".to_string()) - })?; + .ok_or_else(|| invalid("billing payload"))?; if billing_mode != "prepaid" || mode != "pay-as-you-go" { return Ok(None); } let raw = root .get("balanceMicroCents") .and_then(Value::as_str) - .ok_or_else(|| ProviderError::Parse("Missing OpenCode Console balance".to_string()))?; + .ok_or_else(|| missing("balance"))?; let digits = raw.strip_prefix('-').unwrap_or(raw); if digits.is_empty() || !digits.bytes().all(|byte| byte.is_ascii_digit()) { - return Err(ProviderError::Parse( - "Invalid OpenCode Console balance".to_string(), - )); + return Err(invalid("balance")); } - let balance = raw - .parse::() - .map_err(|_| ProviderError::Parse("Invalid OpenCode Console balance".to_string()))?; + let balance = raw.parse::().map_err(|_| invalid("balance"))?; if !balance.is_finite() { - return Err(ProviderError::Parse( - "Invalid OpenCode Console balance".to_string(), - )); + return Err(invalid("balance")); } Ok(Some(balance / BILLING_SCALE)) } +fn invalid(what: &str) -> ProviderError { + ProviderError::Parse(format!("Invalid OpenCode Console {what}")) +} + +fn missing(what: &str) -> ProviderError { + ProviderError::Parse(format!("Missing OpenCode Console {what}")) +} + fn numeric_value(value: Option<&Value>) -> Option { match value? { Value::Number(number) => number.as_f64().filter(|value| value.is_finite()), @@ -296,6 +285,76 @@ mod tests { ) } + #[test] + fn console_parse_errors_keep_their_messages() { + let usage_cases = [ + ("not json", "Invalid OpenCode Console usage payload"), + (r#"{"access":5}"#, "Invalid OpenCode Console usage payload"), + (r#"{"access":{}}"#, "Invalid OpenCode Console usage payload"), + ( + r#"{"access":{"meters":{}}}"#, + "Invalid OpenCode Console usage payload", + ), + ( + r#"{"access":{"meters":{"fiveHour":5}}}"#, + "Invalid OpenCode Console usage meter", + ), + ( + r#"{"access":{"meters":{"fiveHour":{}}}}"#, + "Missing OpenCode Console usage amount", + ), + ( + r#"{"access":{"meters":{"fiveHour":{"usedMicroCents":"1"}}}}"#, + "Missing OpenCode Console usage limit", + ), + ]; + for (text, expected) in usage_cases { + let error = parse_usage(text, now()).err().expect(text); + assert!( + matches!(&error, ProviderError::Parse(message) if message == expected), + "{text}: {error:?}" + ); + } + + let huge = "9".repeat(400); + let balance_cases = [ + ( + "not json".to_string(), + "Invalid OpenCode Console billing payload", + ), + ( + r#"{"billingMode":"other","mode":"pay-as-you-go"}"#.to_string(), + "Invalid OpenCode Console billing payload", + ), + ( + r#"{"billingMode":"prepaid"}"#.to_string(), + "Invalid OpenCode Console billing payload", + ), + ( + r#"{"billingMode":"prepaid","mode":"pay-as-you-go"}"#.to_string(), + "Missing OpenCode Console balance", + ), + ( + r#"{"billingMode":"prepaid","mode":"pay-as-you-go","balanceMicroCents":"12a"}"# + .to_string(), + "Invalid OpenCode Console balance", + ), + ( + format!( + r#"{{"billingMode":"prepaid","mode":"pay-as-you-go","balanceMicroCents":"{huge}"}}"# + ), + "Invalid OpenCode Console balance", + ), + ]; + for (text, expected) in balance_cases { + let error = parse_balance(&text).unwrap_err(); + assert!( + matches!(&error, ProviderError::Parse(message) if message == expected), + "{text}: {error:?}" + ); + } + } + #[test] fn parses_console_workspace_ids_and_urls() { assert_eq!( diff --git a/rust/src/providers/opencodego/legacy.rs b/rust/src/providers/opencodego/legacy.rs index c876c1b8c5..691db6b600 100644 --- a/rust/src/providers/opencodego/legacy.rs +++ b/rust/src/providers/opencodego/legacy.rs @@ -200,30 +200,8 @@ fn parse_usage_text(text: &str) -> Result { let weekly = extract_window(text, &["weeklyUsage", "weekly_usage", "weekly"]); let monthly = extract_window(text, &["monthlyUsage", "monthly_usage", "monthly"]); - let primary = RateWindow::with_details( - rolling.0, - Some(300), - Some(now + chrono::Duration::seconds(rolling.1)), - None, - ); - let mut snapshot = UsageSnapshot::new(primary).with_login_method("OpenCode Go"); - if let Some((percent, reset)) = weekly { - snapshot = snapshot.with_secondary(RateWindow::with_details( - percent, - Some(10080), - Some(now + chrono::Duration::seconds(reset)), - None, - )); - } - if let Some((percent, reset)) = monthly { - let resets_at = now + chrono::Duration::seconds(reset); - snapshot = snapshot.with_tertiary(RateWindow::with_details( - percent, - RateWindow::monthly_window_minutes(Some(resets_at)).or(Some(43200)), - Some(resets_at), - None, - )); - } + let at = |(percent, reset): (f64, i64)| (percent, now + chrono::Duration::seconds(reset)); + let mut snapshot = super::go_snapshot(at(rolling), weekly.map(at), monthly.map(at)); if let Some(renews_at) = super::super::extract_renewal(text) { snapshot = snapshot.with_extra_rate_window( "renewal", @@ -410,10 +388,27 @@ mod tests { weeklyUsage: { usagePercent: 13, resetInSec: 86400 } monthlyUsage: { usagePercent: 7, resetInSec: 2592000 } "#; + let before = Utc::now(); let snapshot = parse_usage_text(text).unwrap(); + let after = Utc::now(); assert!((snapshot.primary.used_percent - 42.5).abs() < 0.001); - assert!((snapshot.secondary.unwrap().used_percent - 13.0).abs() < 0.001); - assert!((snapshot.tertiary.unwrap().used_percent - 7.0).abs() < 0.001); + assert!((snapshot.secondary.as_ref().unwrap().used_percent - 13.0).abs() < 0.001); + assert!((snapshot.tertiary.as_ref().unwrap().used_percent - 7.0).abs() < 0.001); + assert_eq!(snapshot.login_method.as_deref(), Some("OpenCode Go")); + let windows = [ + (&snapshot.primary, 3600, Some(300)), + (snapshot.secondary.as_ref().unwrap(), 86400, Some(10080)), + (snapshot.tertiary.as_ref().unwrap(), 2592000, None), + ]; + for (window, reset_in, minutes) in windows { + let resets_at = window.resets_at.unwrap(); + assert!(resets_at >= before + chrono::Duration::seconds(reset_in)); + assert!(resets_at <= after + chrono::Duration::seconds(reset_in)); + let expected = minutes.or(RateWindow::monthly_window_minutes(Some(resets_at))); + assert_eq!(window.window_minutes, expected); + assert!(window.reset_description.is_none()); + } + assert!(snapshot.extra_rate_windows.is_empty()); } #[test] diff --git a/rust/src/providers/opencodego/local.rs b/rust/src/providers/opencodego/local.rs index abb36ff955..f846423206 100644 --- a/rust/src/providers/opencodego/local.rs +++ b/rust/src/providers/opencodego/local.rs @@ -9,7 +9,7 @@ use rusqlite::Connection; use std::path::{Path, PathBuf}; use super::tokens::{RowTokens, TokenSums}; -use crate::core::{ProviderError, ProviderFetchResult, RateWindow, UsageSnapshot}; +use crate::core::{ProviderError, ProviderFetchResult}; const FIVE_HOURS_MS: i64 = 5 * 60 * 60 * 1000; const WEEK_MS: i64 = 7 * 24 * 60 * 60 * 1000; @@ -100,26 +100,12 @@ pub struct LocalUsageSnapshot { impl LocalUsageSnapshot { pub fn to_fetch_result(&self) -> ProviderFetchResult { let now = Utc::now(); - let primary = RateWindow::with_details( - self.rolling_usage_percent, - Some(300), - Some(now + Duration::seconds(self.rolling_reset_in_sec)), - None, + let at = |percent: f64, reset: i64| (percent, now + Duration::seconds(reset)); + let snap = super::go_snapshot( + at(self.rolling_usage_percent, self.rolling_reset_in_sec), + Some(at(self.weekly_usage_percent, self.weekly_reset_in_sec)), + Some(at(self.monthly_usage_percent, self.monthly_reset_in_sec)), ); - let mut snap = UsageSnapshot::new(primary).with_login_method("OpenCode Go"); - snap = snap.with_secondary(RateWindow::with_details( - self.weekly_usage_percent, - Some(10080), - Some(now + Duration::seconds(self.weekly_reset_in_sec)), - None, - )); - let monthly_reset = now + Duration::seconds(self.monthly_reset_in_sec); - snap = snap.with_tertiary(RateWindow::with_details( - self.monthly_usage_percent, - RateWindow::monthly_window_minutes(Some(monthly_reset)).or(Some(43200)), - Some(monthly_reset), - None, - )); // Upstream 0.51 (#2982): local SQLite quota reconstruction is useful // but it is not server-confirmed authority. Keep that distinction in // the data contract so CLI/React can present it without guessing. @@ -164,11 +150,7 @@ pub fn fetch_local_usage(now: DateTime) -> Result last_err = Some(e), } } - Err(last_err.unwrap_or_else(|| { - ProviderError::NotInstalled( - "OpenCode Go not detected. Log in with OpenCode Go or use it locally first.".into(), - ) - })) + Err(last_err.unwrap_or_else(not_detected)) } pub fn fetch_from_paths( @@ -183,17 +165,13 @@ pub fn fetch_from_paths( "OpenCode Go local usage history is unavailable: database not found".into(), ) } else { - ProviderError::NotInstalled( - "OpenCode Go not detected. Log in with OpenCode Go or use it locally first.".into(), - ) + not_detected() }); } let rows = read_rows(db_path)?; if !has_auth && rows.is_empty() { - return Err(ProviderError::NotInstalled( - "OpenCode Go not detected. Log in with OpenCode Go or use it locally first.".into(), - )); + return Err(not_detected()); } if rows.is_empty() { return Err(ProviderError::Other( @@ -204,6 +182,16 @@ pub fn fetch_from_paths( Ok(snapshot_from_rows(&rows, now)) } +fn not_detected() -> ProviderError { + ProviderError::NotInstalled( + "OpenCode Go not detected. Log in with OpenCode Go or use it locally first.".into(), + ) +} + +fn sqlite_err(e: rusqlite::Error) -> ProviderError { + ProviderError::Other(format!("SQLite error reading OpenCode Go usage: {e}")) +} + fn has_auth_key(path: &Path) -> bool { let Ok(text) = std::fs::read_to_string(path) else { return false; @@ -221,9 +209,7 @@ fn has_auth_key(path: &Path) -> bool { fn read_rows(db_path: &Path) -> Result, ProviderError> { let conn = open_readonly_connection(db_path)?; conn.busy_timeout(std::time::Duration::from_millis(250)) - .map_err(|e| { - ProviderError::Other(format!("SQLite error reading OpenCode Go usage: {e}")) - })?; + .map_err(sqlite_err)?; let sql = if has_table(&conn, "part") { MESSAGE_AND_PART_USAGE_SQL @@ -231,9 +217,7 @@ fn read_rows(db_path: &Path) -> Result, ProviderError> { MESSAGE_USAGE_SQL }; - let mut stmt = conn.prepare(sql).map_err(|e| { - ProviderError::Other(format!("SQLite error reading OpenCode Go usage: {e}")) - })?; + let mut stmt = conn.prepare(sql).map_err(sqlite_err)?; let rows = stmt .query_map([], |row| { Ok(UsageRow { @@ -253,15 +237,11 @@ fn read_rows(db_path: &Path) -> Result, ProviderError> { .and_then(|json| RowTokens::parse(&json)), }) }) - .map_err(|e| { - ProviderError::Other(format!("SQLite error reading OpenCode Go usage: {e}")) - })?; + .map_err(sqlite_err)?; let mut out = Vec::new(); for row in rows { - let row = row.map_err(|e| { - ProviderError::Other(format!("SQLite error reading OpenCode Go usage: {e}")) - })?; + let row = row.map_err(sqlite_err)?; if row.created_ms > 0 && row.cost.is_finite() && row.cost >= 0.0 { out.push(row); } @@ -273,7 +253,7 @@ fn read_rows(db_path: &Path) -> Result, ProviderError> { /// WAL-mode databases (upstream #2544). fn open_readonly_connection(db_path: &Path) -> Result { crate::core::open_readonly_sqlite_connection(db_path, std::time::Duration::from_millis(250)) - .map_err(|e| ProviderError::Other(format!("SQLite error reading OpenCode Go usage: {e}"))) + .map_err(sqlite_err) } fn has_table(conn: &Connection, name: &str) -> bool { @@ -388,6 +368,27 @@ fn local_today_from_utc(now: DateTime) -> NaiveDate { Local.from_utc_datetime(&now.naive_utc()).date_naive() } +/// Local midnight that opens a `days`-long window ending today. +fn window_since_ms(now: DateTime, days: u32) -> i64 { + let clamped = crate::cost_reporting_period::clamp_window_days(days); + let since = local_today_from_utc(now) - Duration::days(clamped as i64 - 1); + Local + .from_local_datetime(&since.and_hms_opt(0, 0, 0).unwrap_or_default()) + .single() + .map(|dt| dt.timestamp_millis()) + .unwrap_or(0) +} + +/// Trimmed model id, with blanks bucketed as `UNKNOWN_MODEL_NAME`. +fn model_name(raw: &str) -> &str { + let trimmed = raw.trim(); + if trimmed.is_empty() { + UNKNOWN_MODEL_NAME + } else { + trimmed + } +} + /// Group rows into `(day, model)` cost buckets (upstream #2649). /// /// Rows outside the `[since, now]` window are dropped; model ids are trimmed and @@ -398,14 +399,7 @@ pub fn daily_model_costs( now: DateTime, history_days: u32, ) -> Vec { - let clamped = crate::cost_reporting_period::clamp_window_days(history_days); - let today = local_today_from_utc(now); - let since = today - Duration::days(clamped as i64 - 1); - let since_ms = Local - .from_local_datetime(&since.and_hms_opt(0, 0, 0).unwrap_or_default()) - .single() - .map(|dt| dt.timestamp_millis()) - .unwrap_or(0); + let since_ms = window_since_ms(now, history_days); let now_ms = now.timestamp_millis(); let mut by_day_model: std::collections::BTreeMap< @@ -419,12 +413,7 @@ pub fn daily_model_costs( let Some(key) = day_key_local(row.created_ms) else { continue; }; - let trimmed = row.model.trim(); - let model = if trimmed.is_empty() { - UNKNOWN_MODEL_NAME - } else { - trimmed - }; + let model = model_name(&row.model); let entry = by_day_model.entry(key).or_default(); let bucket = entry.entry(model.to_string()).or_default(); bucket.0 += row.cost; @@ -453,14 +442,7 @@ pub fn model_cost_summary_from_rows( now: DateTime, days: u32, ) -> ModelCostSummary { - let clamped = crate::cost_reporting_period::clamp_window_days(days); - let today = local_today_from_utc(now); - let since = today - Duration::days(clamped as i64 - 1); - let since_ms = Local - .from_local_datetime(&since.and_hms_opt(0, 0, 0).unwrap_or_default()) - .single() - .map(|dt| dt.timestamp_millis()) - .unwrap_or(0); + let since_ms = window_since_ms(now, days); let now_ms = now.timestamp_millis(); let mut total = 0.0; @@ -477,12 +459,7 @@ pub fn model_cost_summary_from_rows( } total += row.cost; request_count = request_count.saturating_add(row.request_count); - let trimmed = row.model.trim(); - let model = if trimmed.is_empty() { - UNKNOWN_MODEL_NAME - } else { - trimmed - }; + let model = model_name(&row.model); *by_model.entry(model.to_string()).or_insert(0.0) += row.cost; tokens.add(row.tokens.as_ref()); by_model_tokens @@ -562,15 +539,11 @@ fn month_bounds_ms(now: DateTime, anchor_ms: Option) -> (i64, i64) { .and_hms_opt(0, 0, 0) .unwrap_or_default(); let start_dt = Utc.from_utc_datetime(&start); - let end_dt = if now.month() == 12 { - Utc.with_ymd_and_hms(now.year() + 1, 1, 1, 0, 0, 0) - .single() - .unwrap_or(start_dt) - } else { - Utc.with_ymd_and_hms(now.year(), now.month() + 1, 1, 0, 0, 0) - .single() - .unwrap_or(start_dt) - }; + let (end_year, end_month) = next_month(now.year(), now.month()); + let end_dt = Utc + .with_ymd_and_hms(end_year, end_month, 1, 0, 0, 0) + .single() + .unwrap_or(start_dt); return (start_dt.timestamp_millis(), end_dt.timestamp_millis()); }; @@ -587,13 +560,17 @@ fn month_bounds_ms(now: DateTime, anchor_ms: Option) -> (i64, i64) { } start = anchored_month(year, month, &anchor); } - let (end_year, end_month) = if month == 12 { + let (end_year, end_month) = next_month(year, month); + let end = anchored_month(end_year, end_month, &anchor); + (start.timestamp_millis(), end.timestamp_millis()) +} + +fn next_month(year: i32, month: u32) -> (i32, u32) { + if month == 12 { (year + 1, 1) } else { (year, month + 1) - }; - let end = anchored_month(end_year, end_month, &anchor); - (start.timestamp_millis(), end.timestamp_millis()) + } } fn anchored_month(year: i32, month: u32, anchor: &DateTime) -> DateTime { @@ -612,12 +589,8 @@ fn anchored_month(year: i32, month: u32, anchor: &DateTime) -> DateTime) -> DateTime PathBuf { - let nanos = SystemTime::now() - .duration_since(UNIX_EPOCH) - .map(|d| d.as_nanos()) - .unwrap_or(0); - std::env::temp_dir().join(format!("opencodego-local-{label}-{nanos}.db")) - } - - fn write_message_db(path: &Path, rows: &[(i64, f64)]) { - let conn = Connection::open(path).unwrap(); - conn.execute_batch( - "CREATE TABLE message ( - id TEXT PRIMARY KEY, - data TEXT, - time_created INTEGER - );", - ) - .unwrap(); - for (i, (created_ms, cost)) in rows.iter().enumerate() { - let data = format!( - r#"{{"providerID":"opencode-go","role":"assistant","cost":{cost},"time":{{"created":{created_ms}}}}}"# - ); - conn.execute( - "INSERT INTO message (id, data, time_created) VALUES (?1, ?2, ?3)", - rusqlite::params![format!("m{i}"), data, created_ms], - ) - .unwrap(); - } - } - - /// Build a message-only DB with optional per-row `modelID` (upstream #2649 fixtures). - fn write_message_db_with_model(path: &Path, rows: &[(i64, f64, Option<&str>)]) { - let conn = Connection::open(path).unwrap(); - conn.execute_batch( - "CREATE TABLE message ( - id TEXT PRIMARY KEY, - data TEXT, - time_created INTEGER - );", - ) - .unwrap(); - for (i, (created_ms, cost, model)) in rows.iter().enumerate() { - let model_json = match model { - Some(m) => format!(r#","modelID":"{m}""#), - None => String::new(), - }; - let data = format!( - r#"{{"providerID":"opencode-go","role":"assistant","cost":{cost}{model_json},"time":{{"created":{created_ms}}}}}"# - ); - conn.execute( - "INSERT INTO message (id, data, time_created) VALUES (?1, ?2, ?3)", - rusqlite::params![format!("m{i}"), data, created_ms], - ) - .unwrap(); - } - } - - /// Insert one assistant message and return its id, optionally with a modelID. - fn insert_message( - conn: &Connection, - id: &str, - created_ms: i64, - cost: Option, - model: Option<&str>, - ) { - let cost_json = match cost { - Some(c) => format!(r#","cost":{c}"#), - None => String::new(), - }; - let model_json = match model { - Some(m) => format!(r#","modelID":"{m}""#), - None => String::new(), - }; - let data = format!( - r#"{{"providerID":"opencode-go","role":"assistant"{cost_json}{model_json},"time":{{"created":{created_ms}}}}}"# - ); - conn.execute( - "INSERT INTO message (id, data, time_created) VALUES (?1, ?2, ?3)", - rusqlite::params![id, data, created_ms], - ) - .unwrap(); - } - - /// Insert a step-finish part carrying a cost, attached to `message_id`. - fn insert_step_finish_part( - conn: &Connection, - id: &str, - message_id: &str, - created_ms: i64, - cost: f64, - ) { - let data = - format!(r#"{{"type":"step-finish","cost":{cost},"time":{{"created":{created_ms}}}}}"#); - conn.execute( - "INSERT INTO part (id, message_id, data, time_created) VALUES (?1, ?2, ?3, ?4)", - rusqlite::params![id, message_id, data, created_ms], - ) - .unwrap(); - } - - fn iso_ms(iso: &str) -> i64 { - chrono::DateTime::parse_from_rfc3339(iso) - .unwrap() - .timestamp_millis() - } - - #[test] - fn not_detected_without_db_or_auth() { - let dir = std::env::temp_dir().join(format!( - "opencodego-missing-{}", - SystemTime::now() - .duration_since(UNIX_EPOCH) - .map(|d| d.as_nanos()) - .unwrap_or(0) - )); - // Best-effort fixture setup; a failure surfaces as the fetch error below. - let _created = std::fs::create_dir_all(&dir); - let auth = dir.join("auth.json"); - let db = dir.join("opencode.db"); - let err = fetch_from_paths(&auth, &db, Utc::now()).unwrap_err(); - assert!(matches!(err, ProviderError::NotInstalled(_))); - // Best-effort teardown; temp dir may already be gone. - let _removed_dir = std::fs::remove_dir_all(&dir); - } - - #[test] - fn sums_session_weekly_monthly_costs() { - let db = temp_db_path("sums"); - let now = Utc.with_ymd_and_hms(2026, 3, 18, 12, 0, 0).unwrap(); // Wednesday - let now_ms = now.timestamp_millis(); - // $6 in the rolling 5h window → 50% of $12 - // $15 in ISO week → 50% of $30 - // $30 in anchored month → 50% of $60 - let session_ms = now_ms - 60_000; - let week_ms = start_of_utc_iso_week_ms(now) + 3_600_000; - let month_anchor_ms = now_ms - 10 * 24 * 60 * 60 * 1000; - write_message_db( - &db, - &[ - (session_ms, 6.0), - (week_ms, 9.0), // plus session = 15 in week if session also in week - (month_anchor_ms, 15.0), - ], - ); - - // auth present so empty-rows path is not used; auth not required when rows exist - let auth = db.with_extension("auth.json"); - // Fixture write; failure would make fetch_from_paths return an error. - let _written = std::fs::write(&auth, r#"{"opencode-go":{"key":"test-key"}}"#); - - let snap = fetch_from_paths(&auth, &db, now).unwrap(); - assert!((snap.rolling_usage_percent - 50.0).abs() < 0.05, "{snap:?}"); - // session 6 + week-only 9 = 15 → 50% - assert!((snap.weekly_usage_percent - 50.0).abs() < 0.05, "{snap:?}"); - // session 6 + week 9 + month 15 = 30 → 50% - assert!((snap.monthly_usage_percent - 50.0).abs() < 0.05, "{snap:?}"); - - // Best-effort teardown; leftover temp files are harmless. - let _removed_db = std::fs::remove_file(&db); - let _removed_auth = std::fs::remove_file(&auth); - } - - #[test] - fn prefers_step_finish_parts_when_present() { - let db = temp_db_path("parts"); - let conn = Connection::open(&db).unwrap(); - conn.execute_batch( - "CREATE TABLE message (id TEXT PRIMARY KEY, data TEXT, time_created INTEGER); - CREATE TABLE part (id TEXT PRIMARY KEY, message_id TEXT, data TEXT, time_created INTEGER);", - ) - .unwrap(); - let now = Utc.with_ymd_and_hms(2026, 3, 18, 12, 0, 0).unwrap(); - let created = now.timestamp_millis() - 1_000; - // Message cost would be $12 (100%), but step-finish parts sum to $3 (25%). - conn.execute( - "INSERT INTO message (id, data, time_created) VALUES (?1, ?2, ?3)", - rusqlite::params![ - "m1", - format!( - r#"{{"providerID":"opencode-go","role":"assistant","cost":12,"time":{{"created":{created}}}}}"# - ), - created - ], - ) - .unwrap(); - conn.execute( - "INSERT INTO part (id, message_id, data, time_created) VALUES (?1, ?2, ?3, ?4)", - rusqlite::params![ - "p1", - "m1", - format!(r#"{{"type":"step-finish","cost":3,"time":{{"created":{created}}}}}"#), - created - ], - ) - .unwrap(); - drop(conn); - - let auth = db.with_extension("auth.json"); - // Fixture write; a failure would surface in fetch_from_paths. - let _written_auth = std::fs::write(&auth, r#"{"opencode-go":{"key":"k"}}"#); - let snap = fetch_from_paths(&auth, &db, now).unwrap(); - assert!( - (snap.rolling_usage_percent - 25.0).abs() < 0.05, - "expected step-finish cost only, got {snap:?}" - ); - // Best-effort teardown; leftover temp files are harmless. - let _removed_db = std::fs::remove_file(&db); - let _removed_auth = std::fs::remove_file(&auth); - } - - #[test] - fn percent_rounds_to_one_decimal() { - assert!((percent(1.0, 12.0) - 8.3).abs() < 0.05); - assert_eq!(percent(0.0, 12.0), 0.0); - assert_eq!(percent(f64::NAN, 12.0), 0.0); - } - - #[test] - fn iso_week_starts_monday_utc() { - // 2026-03-18 is a Wednesday; week start should be 2026-03-16 00:00 UTC. - let wed = Utc.with_ymd_and_hms(2026, 3, 18, 15, 0, 0).unwrap(); - let start = start_of_utc_iso_week_ms(wed); - let expected = Utc - .with_ymd_and_hms(2026, 3, 16, 0, 0, 0) - .unwrap() - .timestamp_millis(); - assert_eq!(start, expected); - assert_eq!(wed.weekday(), Weekday::Wed); - } - - #[test] - fn idle_wal_mode_read_creates_no_sidecars() { - let dir = tempfile::tempdir().unwrap(); - let db = dir.path().join("opencode.db"); - let auth = dir.path().join("auth.json"); - std::fs::write(&auth, r#"{"opencode-go":{"key":"k"}}"#).unwrap(); - - // Build a WAL-mode DB, insert a row, leave journal_mode=WAL, then drop - // any writer-created sidecars so the main file is an idle WAL header. - { - let conn = Connection::open(&db).unwrap(); - conn.execute_batch("PRAGMA journal_mode=WAL;").unwrap(); - conn.execute_batch( - "CREATE TABLE message ( - id TEXT PRIMARY KEY, - data TEXT, - time_created INTEGER - );", - ) - .unwrap(); - let now = Utc::now(); - let created = now.timestamp_millis() - 1_000; - let data = format!( - r#"{{"providerID":"opencode-go","role":"assistant","cost":3,"time":{{"created":{created}}}}}"# - ); - conn.execute( - "INSERT INTO message (id, data, time_created) VALUES ('m1', ?1, ?2)", - rusqlite::params![data, created], - ) - .unwrap(); - // Truncate empties WAL content before close; journal_mode stays WAL. - // Best-effort checkpoint; a truncated WAL is just tidier. - let _checkpointed = conn.query_row("PRAGMA wal_checkpoint(TRUNCATE)", [], |_| Ok(())); - drop(conn); - } - - // Ensure idle-WAL case: header says WAL, no live sidecars. - let wal = crate::core::sqlite_sidecar_path(&db, "-wal"); - let shm = crate::core::sqlite_sidecar_path(&db, "-shm"); - // Prefer rename-away over delete if OS still holds handles. - for side in [&wal, &shm] { - if side.exists() { - let parked = side.with_extension("parked"); - // Best-effort parking; a locked sidecar just stays put. - let _parked = std::fs::rename(side, parked); - } - } - assert!(!wal.exists(), "precondition: no -wal"); - assert!(!shm.exists(), "precondition: no -shm"); - - let snap = fetch_from_paths(&auth, &db, Utc::now()).expect("read idle WAL db"); - assert!(snap.rolling_usage_percent > 0.0, "{snap:?}"); - - assert!( - !wal.exists() && !shm.exists(), - "reader must not create -wal/-shm sidecars" - ); - } - - // ---- A14: per-model daily cost breakdown (upstream #2649) ------------- - - fn a14_now() -> DateTime { - Utc.timestamp_opt(1_772_798_400, 0).unwrap() - } - - fn a14_now_afternoon() -> DateTime { - Utc.timestamp_opt(1_772_798_400 + 4 * 3600, 0).unwrap() - } - - #[test] - fn daily_entries_group_cost_by_model_within_a_day() { - let dir = tempfile::tempdir().unwrap(); - let db = dir.path().join("opencode.db"); - let now = a14_now_afternoon(); - write_message_db_with_model( - &db, - &[ - ( - iso_ms("2026-03-06T11:00:00.000Z"), - 3.0, - Some("claude-sonnet-4-5"), - ), - ( - iso_ms("2026-03-06T12:00:00.000Z"), - 2.0, - Some("gpt-5.1-codex"), - ), - ( - iso_ms("2026-03-06T13:00:00.000Z"), - 1.0, - Some("claude-sonnet-4-5"), - ), - ], - ); - let rows = read_rows(&db).unwrap(); - let buckets = daily_model_costs(&rows, now, 30); - - // Day key is local-calendar; assert on tz-independent model aggregation. - let total: f64 = buckets.iter().map(|b| b.cost).sum(); - assert!((total - 6.0).abs() < 1e-6, "total {total}"); - assert_eq!(buckets.iter().map(|b| b.request_count).sum::(), 3); - let by_model: std::collections::HashMap<&str, f64> = - buckets.iter().map(|b| (b.model.as_str(), b.cost)).collect(); - assert!((by_model["claude-sonnet-4-5"] - 4.0).abs() < 1e-6); - assert!((by_model["gpt-5.1-codex"] - 2.0).abs() < 1e-6); - } - - #[test] - fn step_finish_parts_inherit_their_model_from_the_parent_message() { - let dir = tempfile::tempdir().unwrap(); - let db = dir.path().join("opencode.db"); - let conn = Connection::open(&db).unwrap(); - conn.execute_batch( - "CREATE TABLE message (id TEXT PRIMARY KEY, data TEXT, time_created INTEGER); - CREATE TABLE part (id TEXT PRIMARY KEY, message_id TEXT, data TEXT, time_created INTEGER);", - ) - .unwrap(); - let created = iso_ms("2026-03-06T11:00:00.000Z"); - insert_message(&conn, "m1", created, None, Some("grok-code-fast-1")); - insert_step_finish_part(&conn, "p1", "m1", created, 3.0); - drop(conn); - - let rows = read_rows(&db).unwrap(); - let buckets = daily_model_costs(&rows, a14_now(), 30); - assert_eq!(buckets.len(), 1, "{buckets:?}"); - assert_eq!(buckets[0].model, "grok-code-fast-1"); - assert!((buckets[0].cost - 3.0).abs() < 1e-6); - } - - #[test] - fn messages_without_a_model_fall_back_to_the_unknown_bucket() { - let dir = tempfile::tempdir().unwrap(); - let db = dir.path().join("opencode.db"); - write_message_db_with_model(&db, &[(iso_ms("2026-03-06T11:00:00.000Z"), 4.0, None)]); - let rows = read_rows(&db).unwrap(); - let buckets = daily_model_costs(&rows, a14_now(), 30); - assert_eq!(buckets.len(), 1, "{buckets:?}"); - assert_eq!(buckets[0].model, UNKNOWN_MODEL_NAME); - assert!((buckets[0].cost - 4.0).abs() < 1e-6); - } - - #[test] - fn whitespace_only_model_ids_fall_back_to_the_unknown_bucket() { - let dir = tempfile::tempdir().unwrap(); - let db = dir.path().join("opencode.db"); - write_message_db_with_model( - &db, - &[(iso_ms("2026-03-06T11:00:00.000Z"), 5.0, Some(" "))], - ); - let rows = read_rows(&db).unwrap(); - let buckets = daily_model_costs(&rows, a14_now(), 30); - assert_eq!(buckets.len(), 1, "{buckets:?}"); - assert_eq!(buckets[0].model, UNKNOWN_MODEL_NAME); - assert!((buckets[0].cost - 5.0).abs() < 1e-6); - } - - #[test] - fn model_ids_with_incidental_whitespace_merge_with_the_trimmed_bucket() { - let dir = tempfile::tempdir().unwrap(); - let db = dir.path().join("opencode.db"); - write_message_db_with_model( - &db, - &[ - ( - iso_ms("2026-03-06T11:00:00.000Z"), - 2.0, - Some("claude-sonnet-4-5"), - ), - ( - iso_ms("2026-03-06T12:00:00.000Z"), - 3.0, - Some(" claude-sonnet-4-5 "), - ), - ], - ); - let rows = read_rows(&db).unwrap(); - let buckets = daily_model_costs(&rows, a14_now_afternoon(), 30); - assert_eq!(buckets.len(), 1, "{buckets:?}"); - assert_eq!(buckets[0].model, "claude-sonnet-4-5"); - assert!((buckets[0].cost - 5.0).abs() < 1e-6); - assert_eq!(buckets[0].request_count, 2); - } - - #[test] - fn multiple_days_bucket_separately_and_sort_deterministically() { - let dir = tempfile::tempdir().unwrap(); - let db = dir.path().join("opencode.db"); - write_message_db_with_model( - &db, - &[ - (iso_ms("2026-03-05T11:00:00.000Z"), 1.0, Some("a")), - (iso_ms("2026-03-06T11:00:00.000Z"), 2.0, Some("b")), - (iso_ms("2026-03-07T11:00:00.000Z"), 3.0, Some("a")), - ], - ); - let rows = read_rows(&db).unwrap(); - let buckets = daily_model_costs(&rows, a14_now_afternoon(), 30); - assert!( - buckets - .windows(2) - .all(|w| (w[0].day_key.as_str(), w[0].model.as_str()) - <= (w[1].day_key.as_str(), w[1].model.as_str())), - "not sorted: {buckets:?}" - ); - assert!( - buckets - .iter() - .map(|b| b.day_key.as_str()) - .collect::>() - .len() - >= 2 - ); - } - - #[test] - fn zero_cost_rows_are_kept_and_aggregated() { - let dir = tempfile::tempdir().unwrap(); - let db = dir.path().join("opencode.db"); - write_message_db_with_model( - &db, - &[ - (iso_ms("2026-03-06T11:00:00.000Z"), 0.0, Some("a")), - (iso_ms("2026-03-06T12:00:00.000Z"), 4.0, Some("a")), - ], - ); - let rows = read_rows(&db).unwrap(); - let buckets = daily_model_costs(&rows, a14_now_afternoon(), 30); - assert_eq!(buckets.len(), 1, "{buckets:?}"); - assert!((buckets[0].cost - 4.0).abs() < 1e-6); - assert_eq!(buckets[0].request_count, 2); - } - - #[test] - fn malformed_rows_are_dropped() { - let dir = tempfile::tempdir().unwrap(); - let db = dir.path().join("opencode.db"); - let conn = Connection::open(&db).unwrap(); - conn.execute_batch( - "CREATE TABLE message (id TEXT PRIMARY KEY, data TEXT, time_created INTEGER);", - ) - .unwrap(); - conn.execute( - "INSERT INTO message (id, data, time_created) VALUES (?1, ?2, ?3)", - rusqlite::params![ - "m1", - r#"{"providerID":"opencode-go","role":"user","cost":9,"time":{"created":1772798400000}}"#, - 1772798400000i64 - ], - ) - .unwrap(); - conn.execute( - "INSERT INTO message (id, data, time_created) VALUES (?1, ?2, ?3)", - rusqlite::params![ - "m2", - r#"{"providerID":"opencode-go","role":"assistant","cost":null,"modelID":"x","time":{"created":1772798400000}}"#, - 1772798400000i64 - ], - ) - .unwrap(); - conn.execute( - "INSERT INTO message (id, data, time_created) VALUES (?1, ?2, ?3)", - rusqlite::params![ - "m3", - r#"{"providerID":"opencode-go","role":"assistant","cost":7,"modelID":"good","time":{"created":1772798400000}}"#, - 1772798400000i64 - ], - ) - .unwrap(); - drop(conn); - - let rows = read_rows(&db).unwrap(); - assert_eq!(rows.len(), 1, "only the valid assistant+cost row survives"); - assert_eq!(rows[0].model, "good"); - } - - #[test] - fn rows_outside_history_window_are_dropped() { - let dir = tempfile::tempdir().unwrap(); - let db = dir.path().join("opencode.db"); - let far_past = iso_ms("2025-01-01T00:00:00.000Z"); - let recent = a14_now().timestamp_millis(); - write_message_db_with_model( - &db, - &[(far_past, 1.0, Some("old")), (recent, 2.0, Some("new"))], - ); - let rows = read_rows(&db).unwrap(); - let buckets = daily_model_costs(&rows, a14_now(), 1); - let models: Vec<&str> = buckets.iter().map(|b| b.model.as_str()).collect(); - assert!( - !models.contains(&"old"), - "old row should be outside the 1-day window: {buckets:?}" - ); - } - - #[test] - fn day_boundary_keys_by_local_calendar_day() { - let dir = tempfile::tempdir().unwrap(); - let db = dir.path().join("opencode.db"); - let just_after_utc_midnight = iso_ms("2026-03-06T00:30:00.000Z"); - write_message_db_with_model(&db, &[(just_after_utc_midnight, 1.5, Some("edge"))]); - let rows = read_rows(&db).unwrap(); - let buckets = daily_model_costs(&rows, a14_now_afternoon(), 30); - assert_eq!(buckets.len(), 1, "{buckets:?}"); - assert!( - NaiveDate::parse_from_str(&buckets[0].day_key, "%Y-%m-%d").is_ok(), - "day_key not yyyy-MM-dd: {}", - buckets[0].day_key - ); - } - - #[test] - fn model_cost_summary_aggregates_total_and_by_model() { - let dir = tempfile::tempdir().unwrap(); - let db = dir.path().join("opencode.db"); - write_message_db_with_model( - &db, - &[ - (iso_ms("2026-03-06T11:00:00.000Z"), 3.0, Some("a")), - (iso_ms("2026-03-06T12:00:00.000Z"), 1.0, Some("b")), - (iso_ms("2026-03-06T13:00:00.000Z"), 2.0, None), - ], - ); - let rows = read_rows(&db).unwrap(); - let summary = model_cost_summary_from_rows(&rows, a14_now_afternoon(), 30); - assert!((summary.total_cost_usd - 6.0).abs() < 1e-6, "{summary:?}"); - assert_eq!(summary.request_count, 3); - assert!((summary.by_model["a"] - 3.0).abs() < 1e-6); - assert!((summary.by_model["b"] - 1.0).abs() < 1e-6); - assert!((summary.by_model[UNKNOWN_MODEL_NAME] - 2.0).abs() < 1e-6); - assert!(summary.period_start.is_some() && summary.period_end.is_some()); - } - - #[test] - fn daily_series_sums_models_per_day_via_pure_aggregation() { - let dir = tempfile::tempdir().unwrap(); - let db = dir.path().join("opencode.db"); - write_message_db_with_model( - &db, - &[ - (iso_ms("2026-03-06T11:00:00.000Z"), 3.0, Some("a")), - (iso_ms("2026-03-06T12:00:00.000Z"), 2.0, Some("b")), - ], - ); - let rows = read_rows(&db).unwrap(); - let buckets = daily_model_costs(&rows, a14_now_afternoon(), 30); - let mut by_day: std::collections::BTreeMap = std::collections::BTreeMap::new(); - for b in &buckets { - *by_day.entry(b.day_key.clone()).or_insert(0.0) += b.cost; - } - let series: Vec<(String, f64)> = by_day.into_iter().collect(); - assert_eq!(series.len(), 1, "{series:?}"); - assert!((series[0].1 - 5.0).abs() < 1e-6); - } - - #[test] - fn daily_aggregation_is_independent_of_zen_wait() { - let dir = tempfile::tempdir().unwrap(); - let db = dir.path().join("opencode.db"); - write_message_db_with_model(&db, &[(iso_ms("2026-03-06T11:00:00.000Z"), 2.0, Some("a"))]); - let rows = read_rows(&db).unwrap(); - let now = a14_now(); - let b1 = daily_model_costs(&rows, now, 30); - let b2 = daily_model_costs(&rows, now, 30); - assert_eq!(b1, b2, "pure aggregation must be deterministic"); - assert_eq!(b1.len(), 1); - } -} +mod tests; diff --git a/rust/src/providers/opencodego/local/test_db.rs b/rust/src/providers/opencodego/local/test_db.rs new file mode 100644 index 0000000000..ee0c2f6143 --- /dev/null +++ b/rust/src/providers/opencodego/local/test_db.rs @@ -0,0 +1,83 @@ +//! Synthetic OpenCode SQLite fixtures shared by the local reader tests. + +use rusqlite::Connection; +use std::path::Path; + +pub(super) fn iso_ms(iso: &str) -> i64 { + chrono::DateTime::parse_from_rfc3339(iso) + .unwrap() + .timestamp_millis() +} + +pub(super) fn open_db(path: &Path, with_part_table: bool) -> Connection { + let conn = Connection::open(path).unwrap(); + conn.execute_batch( + "CREATE TABLE message (id TEXT PRIMARY KEY, data TEXT, time_created INTEGER);", + ) + .unwrap(); + if with_part_table { + conn.execute_batch( + "CREATE TABLE part (id TEXT PRIMARY KEY, message_id TEXT, data TEXT, time_created INTEGER);", + ) + .unwrap(); + } + conn +} + +/// Insert one assistant message. `tokens` is spliced in as raw JSON so +/// malformed shapes can be exercised. +pub(super) fn insert_message( + conn: &Connection, + id: &str, + created_ms: i64, + cost: Option, + model: Option<&str>, + tokens: Option<&str>, +) { + let cost_json = cost.map_or_else(String::new, |c| format!(r#","cost":{c}"#)); + let model_json = model.map_or_else(String::new, |m| format!(r#","modelID":"{m}""#)); + let tokens_json = tokens.map_or_else(String::new, |t| format!(r#","tokens":{t}"#)); + let data = format!( + r#"{{"providerID":"opencode-go","role":"assistant"{cost_json}{model_json}{tokens_json},"time":{{"created":{created_ms}}}}}"# + ); + conn.execute( + "INSERT INTO message (id, data, time_created) VALUES (?1, ?2, ?3)", + rusqlite::params![id, data, created_ms], + ) + .unwrap(); +} + +/// Insert a step-finish part carrying a cost, attached to `message_id`. +pub(super) fn insert_step_finish( + conn: &Connection, + id: &str, + message_id: &str, + created_ms: i64, + cost: f64, + tokens: Option<&str>, +) { + let tokens_json = tokens.map_or_else(String::new, |t| format!(r#","tokens":{t}"#)); + let data = format!( + r#"{{"type":"step-finish","cost":{cost}{tokens_json},"time":{{"created":{created_ms}}}}}"# + ); + conn.execute( + "INSERT INTO part (id, message_id, data, time_created) VALUES (?1, ?2, ?3, ?4)", + rusqlite::params![id, message_id, data, created_ms], + ) + .unwrap(); +} + +/// Message-only DB with one assistant row per `(created_ms, cost, modelID)`. +pub(super) fn write_message_db(path: &Path, rows: &[(i64, f64, Option<&str>)]) { + let conn = open_db(path, false); + for (i, (created_ms, cost, model)) in rows.iter().enumerate() { + insert_message( + &conn, + &format!("m{i}"), + *created_ms, + Some(*cost), + *model, + None, + ); + } +} diff --git a/rust/src/providers/opencodego/local/tests.rs b/rust/src/providers/opencodego/local/tests.rs new file mode 100644 index 0000000000..b606312ff3 --- /dev/null +++ b/rust/src/providers/opencodego/local/tests.rs @@ -0,0 +1,636 @@ +use super::test_db::{insert_message, insert_step_finish, iso_ms, open_db, write_message_db}; +use super::*; +use crate::core::RateWindow; +use chrono::Weekday; +use std::time::{SystemTime, UNIX_EPOCH}; + +#[test] +fn local_fetch_result_keeps_estimate_source_and_reset_windows() { + let result = LocalUsageSnapshot { + rolling_usage_percent: 12.0, + weekly_usage_percent: 23.0, + monthly_usage_percent: 34.0, + rolling_reset_in_sec: 300, + weekly_reset_in_sec: 1_000, + monthly_reset_in_sec: 2_000, + } + .to_fetch_result(); + + assert_eq!( + result.source_label, + super::super::LOCAL_ESTIMATE_SOURCE_LABEL + ); + assert_eq!(result.usage.primary.used_percent, 12.0); + assert_eq!(result.usage.secondary.as_ref().unwrap().used_percent, 23.0); + assert_eq!(result.usage.tertiary.as_ref().unwrap().used_percent, 34.0); + assert!(result.usage.primary.resets_at.is_some()); + assert!(result.usage.secondary.as_ref().unwrap().resets_at.is_some()); + assert!(result.usage.tertiary.as_ref().unwrap().resets_at.is_some()); +} + +#[test] +fn local_fetch_result_pins_window_minutes_and_reset_offsets() { + let before = Utc::now(); + let result = LocalUsageSnapshot { + rolling_usage_percent: 12.0, + weekly_usage_percent: 23.0, + monthly_usage_percent: 34.0, + rolling_reset_in_sec: 300, + weekly_reset_in_sec: 1_000, + monthly_reset_in_sec: 2_000, + } + .to_fetch_result(); + let after = Utc::now(); + let usage = &result.usage; + assert_eq!(usage.login_method.as_deref(), Some("OpenCode Go")); + let windows = [ + (&usage.primary, 300, Some(300)), + (usage.secondary.as_ref().unwrap(), 1_000, Some(10080)), + (usage.tertiary.as_ref().unwrap(), 2_000, None), + ]; + for (window, reset_in, minutes) in windows { + let resets_at = window.resets_at.unwrap(); + assert!(resets_at >= before + Duration::seconds(reset_in)); + assert!(resets_at <= after + Duration::seconds(reset_in)); + let expected = minutes.or(RateWindow::monthly_window_minutes(Some(resets_at))); + assert_eq!(window.window_minutes, expected); + assert!(window.reset_description.is_none()); + } + assert!(usage.extra_rate_windows.is_empty()); +} + +fn temp_db_path(label: &str) -> PathBuf { + let nanos = SystemTime::now() + .duration_since(UNIX_EPOCH) + .map(|d| d.as_nanos()) + .unwrap_or(0); + std::env::temp_dir().join(format!("opencodego-local-{label}-{nanos}.db")) +} + +#[test] +fn unreadable_database_reports_the_sqlite_error_prefix() { + let db = temp_db_path("garbage"); + std::fs::write(&db, b"not a sqlite database, just synthetic bytes").unwrap(); + let err = read_rows(&db).unwrap_err(); + // Best-effort teardown; the temp file may already be gone. + let _removed = std::fs::remove_file(&db); + assert!( + matches!(&err, ProviderError::Other(message) + if message.starts_with("SQLite error reading OpenCode Go usage: ")), + "{err:?}" + ); +} + +#[test] +fn not_detected_without_db_or_auth() { + let dir = std::env::temp_dir().join(format!( + "opencodego-missing-{}", + SystemTime::now() + .duration_since(UNIX_EPOCH) + .map(|d| d.as_nanos()) + .unwrap_or(0) + )); + // Best-effort fixture setup; a failure surfaces as the fetch error below. + let _created = std::fs::create_dir_all(&dir); + let auth = dir.join("auth.json"); + let db = dir.join("opencode.db"); + let err = fetch_from_paths(&auth, &db, Utc::now()).unwrap_err(); + assert!(matches!( + &err, + ProviderError::NotInstalled(message) + if message == "OpenCode Go not detected. Log in with OpenCode Go or use it locally first." + )); + // Best-effort teardown; temp dir may already be gone. + let _removed_dir = std::fs::remove_dir_all(&dir); +} + +#[test] +fn sums_session_weekly_monthly_costs() { + let db = temp_db_path("sums"); + let now = Utc.with_ymd_and_hms(2026, 3, 18, 12, 0, 0).unwrap(); // Wednesday + let now_ms = now.timestamp_millis(); + // $6 in the rolling 5h window → 50% of $12 + // $15 in ISO week → 50% of $30 + // $30 in anchored month → 50% of $60 + let session_ms = now_ms - 60_000; + let week_ms = start_of_utc_iso_week_ms(now) + 3_600_000; + let month_anchor_ms = now_ms - 10 * 24 * 60 * 60 * 1000; + write_message_db( + &db, + &[ + (session_ms, 6.0, None), + (week_ms, 9.0, None), // plus session = 15 in week if session also in week + (month_anchor_ms, 15.0, None), + ], + ); + + // auth present so empty-rows path is not used; auth not required when rows exist + let auth = db.with_extension("auth.json"); + // Fixture write; failure would make fetch_from_paths return an error. + let _written = std::fs::write(&auth, r#"{"opencode-go":{"key":"test-key"}}"#); + + let snap = fetch_from_paths(&auth, &db, now).unwrap(); + assert!((snap.rolling_usage_percent - 50.0).abs() < 0.05, "{snap:?}"); + // session 6 + week-only 9 = 15 → 50% + assert!((snap.weekly_usage_percent - 50.0).abs() < 0.05, "{snap:?}"); + // session 6 + week 9 + month 15 = 30 → 50% + assert!((snap.monthly_usage_percent - 50.0).abs() < 0.05, "{snap:?}"); + + // Best-effort teardown; leftover temp files are harmless. + let _removed_db = std::fs::remove_file(&db); + let _removed_auth = std::fs::remove_file(&auth); +} + +#[test] +fn prefers_step_finish_parts_when_present() { + let db = temp_db_path("parts"); + let conn = open_db(&db, true); + let now = Utc.with_ymd_and_hms(2026, 3, 18, 12, 0, 0).unwrap(); + let created = now.timestamp_millis() - 1_000; + // Message cost would be $12 (100%), but step-finish parts sum to $3 (25%). + insert_message(&conn, "m1", created, Some(12.0), None, None); + insert_step_finish(&conn, "p1", "m1", created, 3.0, None); + drop(conn); + + let auth = db.with_extension("auth.json"); + // Fixture write; a failure would surface in fetch_from_paths. + let _written_auth = std::fs::write(&auth, r#"{"opencode-go":{"key":"k"}}"#); + let snap = fetch_from_paths(&auth, &db, now).unwrap(); + assert!( + (snap.rolling_usage_percent - 25.0).abs() < 0.05, + "expected step-finish cost only, got {snap:?}" + ); + // Best-effort teardown; leftover temp files are harmless. + let _removed_db = std::fs::remove_file(&db); + let _removed_auth = std::fs::remove_file(&auth); +} + +#[test] +fn percent_rounds_to_one_decimal() { + assert!((percent(1.0, 12.0) - 8.3).abs() < 0.05); + assert_eq!(percent(0.0, 12.0), 0.0); + assert_eq!(percent(f64::NAN, 12.0), 0.0); +} + +#[test] +fn month_bounds_follow_calendar_or_anchor_day() { + let at = |iso: &str| { + DateTime::parse_from_rfc3339(iso) + .unwrap() + .with_timezone(&Utc) + }; + let cases = [ + // (now, anchor, expected start, expected end) + ( + "2026-12-15T12:00:00Z", + None, + "2026-12-01T00:00:00Z", + "2027-01-01T00:00:00Z", + ), + ( + "2026-03-18T12:00:00Z", + None, + "2026-03-01T00:00:00Z", + "2026-04-01T00:00:00Z", + ), + // Day 31 clamps to Feb 28, which is after now, so the window steps back. + ( + "2026-02-20T10:00:00Z", + Some("2026-01-31T08:30:00Z"), + "2026-01-31T08:30:00Z", + "2026-02-28T08:30:00Z", + ), + ( + "2026-12-20T00:00:00Z", + Some("2026-03-15T00:00:00Z"), + "2026-12-15T00:00:00Z", + "2027-01-15T00:00:00Z", + ), + ( + "2026-01-10T00:00:00Z", + Some("2025-11-15T06:00:00Z"), + "2025-12-15T06:00:00Z", + "2026-01-15T06:00:00Z", + ), + ( + "2026-11-10T12:00:00Z", + Some("2026-08-31T00:00:00Z"), + "2026-10-31T00:00:00Z", + "2026-11-30T00:00:00Z", + ), + ]; + for (now, anchor, start, end) in cases { + let bounds = month_bounds_ms(at(now), anchor.map(|a| at(a).timestamp_millis())); + assert_eq!( + bounds, + (at(start).timestamp_millis(), at(end).timestamp_millis()), + "{now} {anchor:?}" + ); + } +} + +#[test] +fn daily_and_summary_windows_share_the_local_midnight_cutoff() { + let now = a14_now(); + let since = local_today_from_utc(now) - Duration::days(1); + let since_ms = Local + .from_local_datetime(&since.and_hms_opt(0, 0, 0).unwrap()) + .single() + .unwrap() + .timestamp_millis(); + let row = |created_ms: i64, model: &str| UsageRow { + created_ms, + cost: 1.0, + request_count: 1, + model: model.to_string(), + tokens: None, + }; + let rows = [ + row(since_ms - 1, "dropped"), + row(since_ms, " kept "), + row(now.timestamp_millis(), ""), + row(now.timestamp_millis() + 1, "future"), + ]; + let daily = daily_model_costs(&rows, now, 2); + let models: Vec<_> = daily.iter().map(|b| b.model.as_str()).collect(); + let mut sorted = models.clone(); + sorted.sort_unstable(); + assert_eq!(sorted, ["kept", UNKNOWN_MODEL_NAME]); + let summary = model_cost_summary_from_rows(&rows, now, 2); + assert_eq!(summary.request_count, 2); + assert!((summary.total_cost_usd - 2.0).abs() < 1e-9); + let mut keys: Vec<_> = summary.by_model.keys().map(String::as_str).collect(); + keys.sort_unstable(); + assert_eq!(keys, ["kept", UNKNOWN_MODEL_NAME]); +} + +#[test] +fn iso_week_starts_monday_utc() { + // 2026-03-18 is a Wednesday; week start should be 2026-03-16 00:00 UTC. + let wed = Utc.with_ymd_and_hms(2026, 3, 18, 15, 0, 0).unwrap(); + let start = start_of_utc_iso_week_ms(wed); + let expected = Utc + .with_ymd_and_hms(2026, 3, 16, 0, 0, 0) + .unwrap() + .timestamp_millis(); + assert_eq!(start, expected); + assert_eq!(wed.weekday(), Weekday::Wed); +} + +#[test] +fn idle_wal_mode_read_creates_no_sidecars() { + let dir = tempfile::tempdir().unwrap(); + let db = dir.path().join("opencode.db"); + let auth = dir.path().join("auth.json"); + std::fs::write(&auth, r#"{"opencode-go":{"key":"k"}}"#).unwrap(); + + // Build a WAL-mode DB, insert a row, leave journal_mode=WAL, then drop + // any writer-created sidecars so the main file is an idle WAL header. + { + let conn = Connection::open(&db).unwrap(); + conn.execute_batch("PRAGMA journal_mode=WAL;").unwrap(); + conn.execute_batch( + "CREATE TABLE message ( + id TEXT PRIMARY KEY, + data TEXT, + time_created INTEGER + );", + ) + .unwrap(); + let now = Utc::now(); + let created = now.timestamp_millis() - 1_000; + let data = format!( + r#"{{"providerID":"opencode-go","role":"assistant","cost":3,"time":{{"created":{created}}}}}"# + ); + conn.execute( + "INSERT INTO message (id, data, time_created) VALUES ('m1', ?1, ?2)", + rusqlite::params![data, created], + ) + .unwrap(); + // Truncate empties WAL content before close; journal_mode stays WAL. + // Best-effort checkpoint; a truncated WAL is just tidier. + let _checkpointed = conn.query_row("PRAGMA wal_checkpoint(TRUNCATE)", [], |_| Ok(())); + drop(conn); + } + + // Ensure idle-WAL case: header says WAL, no live sidecars. + let wal = crate::core::sqlite_sidecar_path(&db, "-wal"); + let shm = crate::core::sqlite_sidecar_path(&db, "-shm"); + // Prefer rename-away over delete if OS still holds handles. + for side in [&wal, &shm] { + if side.exists() { + let parked = side.with_extension("parked"); + // Best-effort parking; a locked sidecar just stays put. + let _parked = std::fs::rename(side, parked); + } + } + assert!(!wal.exists(), "precondition: no -wal"); + assert!(!shm.exists(), "precondition: no -shm"); + + let snap = fetch_from_paths(&auth, &db, Utc::now()).expect("read idle WAL db"); + assert!(snap.rolling_usage_percent > 0.0, "{snap:?}"); + + assert!( + !wal.exists() && !shm.exists(), + "reader must not create -wal/-shm sidecars" + ); +} + +// ---- A14: per-model daily cost breakdown (upstream #2649) ------------- + +fn a14_now() -> DateTime { + Utc.timestamp_opt(1_772_798_400, 0).unwrap() +} + +fn a14_now_afternoon() -> DateTime { + Utc.timestamp_opt(1_772_798_400 + 4 * 3600, 0).unwrap() +} + +#[test] +fn daily_entries_group_cost_by_model_within_a_day() { + let dir = tempfile::tempdir().unwrap(); + let db = dir.path().join("opencode.db"); + let now = a14_now_afternoon(); + write_message_db( + &db, + &[ + ( + iso_ms("2026-03-06T11:00:00.000Z"), + 3.0, + Some("claude-sonnet-4-5"), + ), + ( + iso_ms("2026-03-06T12:00:00.000Z"), + 2.0, + Some("gpt-5.1-codex"), + ), + ( + iso_ms("2026-03-06T13:00:00.000Z"), + 1.0, + Some("claude-sonnet-4-5"), + ), + ], + ); + let rows = read_rows(&db).unwrap(); + let buckets = daily_model_costs(&rows, now, 30); + + // Day key is local-calendar; assert on tz-independent model aggregation. + let total: f64 = buckets.iter().map(|b| b.cost).sum(); + assert!((total - 6.0).abs() < 1e-6, "total {total}"); + assert_eq!(buckets.iter().map(|b| b.request_count).sum::(), 3); + let by_model: std::collections::HashMap<&str, f64> = + buckets.iter().map(|b| (b.model.as_str(), b.cost)).collect(); + assert!((by_model["claude-sonnet-4-5"] - 4.0).abs() < 1e-6); + assert!((by_model["gpt-5.1-codex"] - 2.0).abs() < 1e-6); +} + +#[test] +fn step_finish_parts_inherit_their_model_from_the_parent_message() { + let dir = tempfile::tempdir().unwrap(); + let db = dir.path().join("opencode.db"); + let conn = open_db(&db, true); + let created = iso_ms("2026-03-06T11:00:00.000Z"); + insert_message(&conn, "m1", created, None, Some("grok-code-fast-1"), None); + insert_step_finish(&conn, "p1", "m1", created, 3.0, None); + drop(conn); + + let rows = read_rows(&db).unwrap(); + let buckets = daily_model_costs(&rows, a14_now(), 30); + assert_eq!(buckets.len(), 1, "{buckets:?}"); + assert_eq!(buckets[0].model, "grok-code-fast-1"); + assert!((buckets[0].cost - 3.0).abs() < 1e-6); +} + +#[test] +fn messages_without_a_model_fall_back_to_the_unknown_bucket() { + let dir = tempfile::tempdir().unwrap(); + let db = dir.path().join("opencode.db"); + write_message_db(&db, &[(iso_ms("2026-03-06T11:00:00.000Z"), 4.0, None)]); + let rows = read_rows(&db).unwrap(); + let buckets = daily_model_costs(&rows, a14_now(), 30); + assert_eq!(buckets.len(), 1, "{buckets:?}"); + assert_eq!(buckets[0].model, UNKNOWN_MODEL_NAME); + assert!((buckets[0].cost - 4.0).abs() < 1e-6); +} + +#[test] +fn whitespace_only_model_ids_fall_back_to_the_unknown_bucket() { + let dir = tempfile::tempdir().unwrap(); + let db = dir.path().join("opencode.db"); + write_message_db( + &db, + &[(iso_ms("2026-03-06T11:00:00.000Z"), 5.0, Some(" "))], + ); + let rows = read_rows(&db).unwrap(); + let buckets = daily_model_costs(&rows, a14_now(), 30); + assert_eq!(buckets.len(), 1, "{buckets:?}"); + assert_eq!(buckets[0].model, UNKNOWN_MODEL_NAME); + assert!((buckets[0].cost - 5.0).abs() < 1e-6); +} + +#[test] +fn model_ids_with_incidental_whitespace_merge_with_the_trimmed_bucket() { + let dir = tempfile::tempdir().unwrap(); + let db = dir.path().join("opencode.db"); + write_message_db( + &db, + &[ + ( + iso_ms("2026-03-06T11:00:00.000Z"), + 2.0, + Some("claude-sonnet-4-5"), + ), + ( + iso_ms("2026-03-06T12:00:00.000Z"), + 3.0, + Some(" claude-sonnet-4-5 "), + ), + ], + ); + let rows = read_rows(&db).unwrap(); + let buckets = daily_model_costs(&rows, a14_now_afternoon(), 30); + assert_eq!(buckets.len(), 1, "{buckets:?}"); + assert_eq!(buckets[0].model, "claude-sonnet-4-5"); + assert!((buckets[0].cost - 5.0).abs() < 1e-6); + assert_eq!(buckets[0].request_count, 2); +} + +#[test] +fn multiple_days_bucket_separately_and_sort_deterministically() { + let dir = tempfile::tempdir().unwrap(); + let db = dir.path().join("opencode.db"); + write_message_db( + &db, + &[ + (iso_ms("2026-03-05T11:00:00.000Z"), 1.0, Some("a")), + (iso_ms("2026-03-06T11:00:00.000Z"), 2.0, Some("b")), + (iso_ms("2026-03-07T11:00:00.000Z"), 3.0, Some("a")), + ], + ); + let rows = read_rows(&db).unwrap(); + let buckets = daily_model_costs(&rows, a14_now_afternoon(), 30); + assert!( + buckets + .windows(2) + .all(|w| (w[0].day_key.as_str(), w[0].model.as_str()) + <= (w[1].day_key.as_str(), w[1].model.as_str())), + "not sorted: {buckets:?}" + ); + assert!( + buckets + .iter() + .map(|b| b.day_key.as_str()) + .collect::>() + .len() + >= 2 + ); +} + +#[test] +fn zero_cost_rows_are_kept_and_aggregated() { + let dir = tempfile::tempdir().unwrap(); + let db = dir.path().join("opencode.db"); + write_message_db( + &db, + &[ + (iso_ms("2026-03-06T11:00:00.000Z"), 0.0, Some("a")), + (iso_ms("2026-03-06T12:00:00.000Z"), 4.0, Some("a")), + ], + ); + let rows = read_rows(&db).unwrap(); + let buckets = daily_model_costs(&rows, a14_now_afternoon(), 30); + assert_eq!(buckets.len(), 1, "{buckets:?}"); + assert!((buckets[0].cost - 4.0).abs() < 1e-6); + assert_eq!(buckets[0].request_count, 2); +} + +#[test] +fn malformed_rows_are_dropped() { + let dir = tempfile::tempdir().unwrap(); + let db = dir.path().join("opencode.db"); + let conn = open_db(&db, false); + conn.execute( + "INSERT INTO message (id, data, time_created) VALUES (?1, ?2, ?3)", + rusqlite::params![ + "m1", + r#"{"providerID":"opencode-go","role":"user","cost":9,"time":{"created":1772798400000}}"#, + 1772798400000i64 + ], + ) + .unwrap(); + conn.execute( + "INSERT INTO message (id, data, time_created) VALUES (?1, ?2, ?3)", + rusqlite::params![ + "m2", + r#"{"providerID":"opencode-go","role":"assistant","cost":null,"modelID":"x","time":{"created":1772798400000}}"#, + 1772798400000i64 + ], + ) + .unwrap(); + conn.execute( + "INSERT INTO message (id, data, time_created) VALUES (?1, ?2, ?3)", + rusqlite::params![ + "m3", + r#"{"providerID":"opencode-go","role":"assistant","cost":7,"modelID":"good","time":{"created":1772798400000}}"#, + 1772798400000i64 + ], + ) + .unwrap(); + drop(conn); + + let rows = read_rows(&db).unwrap(); + assert_eq!(rows.len(), 1, "only the valid assistant+cost row survives"); + assert_eq!(rows[0].model, "good"); +} + +#[test] +fn rows_outside_history_window_are_dropped() { + let dir = tempfile::tempdir().unwrap(); + let db = dir.path().join("opencode.db"); + let far_past = iso_ms("2025-01-01T00:00:00.000Z"); + let recent = a14_now().timestamp_millis(); + write_message_db( + &db, + &[(far_past, 1.0, Some("old")), (recent, 2.0, Some("new"))], + ); + let rows = read_rows(&db).unwrap(); + let buckets = daily_model_costs(&rows, a14_now(), 1); + let models: Vec<&str> = buckets.iter().map(|b| b.model.as_str()).collect(); + assert!( + !models.contains(&"old"), + "old row should be outside the 1-day window: {buckets:?}" + ); +} + +#[test] +fn day_boundary_keys_by_local_calendar_day() { + let dir = tempfile::tempdir().unwrap(); + let db = dir.path().join("opencode.db"); + let just_after_utc_midnight = iso_ms("2026-03-06T00:30:00.000Z"); + write_message_db(&db, &[(just_after_utc_midnight, 1.5, Some("edge"))]); + let rows = read_rows(&db).unwrap(); + let buckets = daily_model_costs(&rows, a14_now_afternoon(), 30); + assert_eq!(buckets.len(), 1, "{buckets:?}"); + assert!( + NaiveDate::parse_from_str(&buckets[0].day_key, "%Y-%m-%d").is_ok(), + "day_key not yyyy-MM-dd: {}", + buckets[0].day_key + ); +} + +#[test] +fn model_cost_summary_aggregates_total_and_by_model() { + let dir = tempfile::tempdir().unwrap(); + let db = dir.path().join("opencode.db"); + write_message_db( + &db, + &[ + (iso_ms("2026-03-06T11:00:00.000Z"), 3.0, Some("a")), + (iso_ms("2026-03-06T12:00:00.000Z"), 1.0, Some("b")), + (iso_ms("2026-03-06T13:00:00.000Z"), 2.0, None), + ], + ); + let rows = read_rows(&db).unwrap(); + let summary = model_cost_summary_from_rows(&rows, a14_now_afternoon(), 30); + assert!((summary.total_cost_usd - 6.0).abs() < 1e-6, "{summary:?}"); + assert_eq!(summary.request_count, 3); + assert!((summary.by_model["a"] - 3.0).abs() < 1e-6); + assert!((summary.by_model["b"] - 1.0).abs() < 1e-6); + assert!((summary.by_model[UNKNOWN_MODEL_NAME] - 2.0).abs() < 1e-6); + assert!(summary.period_start.is_some() && summary.period_end.is_some()); +} + +#[test] +fn daily_series_sums_models_per_day_via_pure_aggregation() { + let dir = tempfile::tempdir().unwrap(); + let db = dir.path().join("opencode.db"); + write_message_db( + &db, + &[ + (iso_ms("2026-03-06T11:00:00.000Z"), 3.0, Some("a")), + (iso_ms("2026-03-06T12:00:00.000Z"), 2.0, Some("b")), + ], + ); + let rows = read_rows(&db).unwrap(); + let buckets = daily_model_costs(&rows, a14_now_afternoon(), 30); + let mut by_day: std::collections::BTreeMap = std::collections::BTreeMap::new(); + for b in &buckets { + *by_day.entry(b.day_key.clone()).or_insert(0.0) += b.cost; + } + let series: Vec<(String, f64)> = by_day.into_iter().collect(); + assert_eq!(series.len(), 1, "{series:?}"); + assert!((series[0].1 - 5.0).abs() < 1e-6); +} + +#[test] +fn daily_aggregation_is_independent_of_zen_wait() { + let dir = tempfile::tempdir().unwrap(); + let db = dir.path().join("opencode.db"); + write_message_db(&db, &[(iso_ms("2026-03-06T11:00:00.000Z"), 2.0, Some("a"))]); + let rows = read_rows(&db).unwrap(); + let now = a14_now(); + let b1 = daily_model_costs(&rows, now, 30); + let b2 = daily_model_costs(&rows, now, 30); + assert_eq!(b1, b2, "pure aggregation must be deterministic"); + assert_eq!(b1.len(), 1); +} diff --git a/rust/src/providers/opencodego/local/tokens_tests.rs b/rust/src/providers/opencodego/local/tokens_tests.rs index c64f449ae2..b789acfd93 100644 --- a/rust/src/providers/opencodego/local/tokens_tests.rs +++ b/rust/src/providers/opencodego/local/tokens_tests.rs @@ -1,5 +1,6 @@ //! Recorded-token coverage for the OpenCode Go local reader (upstream 0.67.0). +use super::test_db::{insert_message, insert_step_finish, iso_ms, open_db}; use super::*; const FULL: &str = @@ -7,55 +8,10 @@ const FULL: &str = /// Older OpenCode rows record no `total`; the five components sum to 20. const NO_TOTAL: &str = r#"{"input":10,"output":5,"reasoning":1,"cache":{"read":4,"write":0}}"#; -fn iso_ms(iso: &str) -> i64 { - chrono::DateTime::parse_from_rfc3339(iso) - .unwrap() - .timestamp_millis() -} - fn now() -> DateTime { Utc.with_ymd_and_hms(2026, 3, 8, 0, 0, 0).unwrap() } -fn open_db(path: &Path, with_part_table: bool) -> Connection { - let conn = Connection::open(path).unwrap(); - conn.execute_batch( - "CREATE TABLE message (id TEXT PRIMARY KEY, data TEXT, time_created INTEGER);", - ) - .unwrap(); - if with_part_table { - conn.execute_batch( - "CREATE TABLE part (id TEXT PRIMARY KEY, message_id TEXT, data TEXT, time_created INTEGER);", - ) - .unwrap(); - } - conn -} - -/// `tokens` is spliced in as raw JSON so malformed shapes can be exercised. -fn insert_message(conn: &Connection, id: &str, created_ms: i64, model: &str, tokens: Option<&str>) { - let tokens_json = tokens.map_or_else(String::new, |t| format!(r#","tokens":{t}"#)); - let data = format!( - r#"{{"providerID":"opencode-go","role":"assistant","cost":1.5,"modelID":"{model}"{tokens_json},"time":{{"created":{created_ms}}}}}"# - ); - conn.execute( - "INSERT INTO message (id, data, time_created) VALUES (?1, ?2, ?3)", - rusqlite::params![id, data, created_ms], - ) - .unwrap(); -} - -fn insert_step_finish(conn: &Connection, id: &str, message_id: &str, ms: i64, tokens: &str) { - let data = format!( - r#"{{"type":"step-finish","cost":0.5,"tokens":{tokens},"time":{{"created":{ms}}}}}"# - ); - conn.execute( - "INSERT INTO part (id, message_id, data, time_created) VALUES (?1, ?2, ?3, ?4)", - rusqlite::params![id, message_id, data, ms], - ) - .unwrap(); -} - fn parse(json: &str) -> Option { RowTokens::parse(json) } @@ -147,8 +103,15 @@ fn daily_entries_carry_message_token_counts_per_day_and_model() { iso_ms("2026-03-06T11:00:00.000Z"), iso_ms("2026-03-06T11:30:00.000Z"), ); - insert_message(&conn, "m1", first, "kimi-k2", Some(FULL)); - insert_message(&conn, "m2", second, "kimi-k2", Some(NO_TOTAL)); + insert_message(&conn, "m1", first, Some(1.5), Some("kimi-k2"), Some(FULL)); + insert_message( + &conn, + "m2", + second, + Some(1.5), + Some("kimi-k2"), + Some(NO_TOTAL), + ); drop(conn); let rows = read_rows(&db).unwrap(); @@ -173,14 +136,21 @@ fn step_finish_tokens_replace_the_parent_message_tokens() { let conn = open_db(&db, true); let created = iso_ms("2026-03-06T11:00:00.000Z"); // The parent message records the session total; the steps record each call. - insert_message(&conn, "m1", created, "kimi-k2", Some(FULL)); + insert_message(&conn, "m1", created, Some(1.5), Some("kimi-k2"), Some(FULL)); let step_a = r#"{"total":30,"input":10,"output":5,"reasoning":5,"cache":{"read":10,"write":0}}"#; let step_b = r#"{"total":20,"input":8,"output":2,"reasoning":0,"cache":{"read":10,"write":0}}"#; - insert_step_finish(&conn, "p1", "m1", created, step_a); - insert_step_finish(&conn, "p2", "m1", created, step_b); + insert_step_finish(&conn, "p1", "m1", created, 0.5, Some(step_a)); + insert_step_finish(&conn, "p2", "m1", created, 0.5, Some(step_b)); // A message with no step-finish parts keeps using its own tokens. - insert_message(&conn, "m2", created, "kimi-k2", Some(NO_TOTAL)); + insert_message( + &conn, + "m2", + created, + Some(1.5), + Some("kimi-k2"), + Some(NO_TOTAL), + ); drop(conn); let rows = read_rows(&db).unwrap(); @@ -198,10 +168,17 @@ fn rows_without_usable_tokens_leave_the_day_incomplete_not_zero() { let db = dir.path().join("opencode.db"); let conn = open_db(&db, false); let created = iso_ms("2026-03-06T11:00:00.000Z"); - insert_message(&conn, "m1", created, "kimi-k2", Some(FULL)); - insert_message(&conn, "m2", created, "kimi-k2", None); - insert_message(&conn, "m3", created, "kimi-k2", Some("5")); - insert_message(&conn, "m4", created, "kimi-k2", Some(r#"{"total":-4}"#)); + insert_message(&conn, "m1", created, Some(1.5), Some("kimi-k2"), Some(FULL)); + insert_message(&conn, "m2", created, Some(1.5), Some("kimi-k2"), None); + insert_message(&conn, "m3", created, Some(1.5), Some("kimi-k2"), Some("5")); + insert_message( + &conn, + "m4", + created, + Some(1.5), + Some("kimi-k2"), + Some(r#"{"total":-4}"#), + ); drop(conn); let rows = read_rows(&db).unwrap(); @@ -223,14 +200,22 @@ fn model_summary_carries_per_model_tokens_and_marks_partial_models() { let db = dir.path().join("opencode.db"); let conn = open_db(&db, false); let created = iso_ms("2026-03-06T11:00:00.000Z"); - insert_message(&conn, "m1", created, "kimi-k2", Some(FULL)); - insert_message(&conn, "m2", created, "kimi-k2", Some(NO_TOTAL)); - insert_message(&conn, "m3", created, "glm-5", None); + insert_message(&conn, "m1", created, Some(1.5), Some("kimi-k2"), Some(FULL)); + insert_message( + &conn, + "m2", + created, + Some(1.5), + Some("kimi-k2"), + Some(NO_TOTAL), + ); + insert_message(&conn, "m3", created, Some(1.5), Some("glm-5"), None); insert_message( &conn, "m4", created, - "mimo", + Some(1.5), + Some("mimo"), Some(r#"{"total":9,"input":4}"#), ); drop(conn); diff --git a/rust/src/providers/opencodego/mod.rs b/rust/src/providers/opencodego/mod.rs index 531d64c02d..e03fd8dda0 100644 --- a/rust/src/providers/opencodego/mod.rs +++ b/rust/src/providers/opencodego/mod.rs @@ -328,20 +328,8 @@ impl OpenCodeGoProvider { } Err(console_error) if cookies.can_recover_with_legacy(&console_error) => { zen_task.abort(); - let session = match transport.resolve_legacy_session(&cookies).await { - Ok(session) => session, - Err(error) => { - return cookies.select_legacy_result(console_error, Err(error)); - } - }; - return Self::fetch_legacy_with_session( - ctx, - &cookies, - console_error, - transport, - session, - ) - .await; + return Self::fetch_legacy_with_cookies(ctx, &cookies, console_error, transport) + .await; } Err(error) => { zen_task.abort(); @@ -445,6 +433,30 @@ impl OpenCodeGoProvider { } } +/// `(used percent, reset time)` of one OpenCode Go window. +type GoWindow = (f64, chrono::DateTime); + +/// Rolling 5h, weekly and calendar-month windows in the shared Go snapshot shape. +fn go_snapshot( + rolling: GoWindow, + weekly: Option, + monthly: Option, +) -> UsageSnapshot { + let window = |(percent, resets_at): GoWindow, minutes| { + RateWindow::with_details(percent, minutes, Some(resets_at), None) + }; + let mut snapshot = + UsageSnapshot::new(window(rolling, Some(300))).with_login_method("OpenCode Go"); + if let Some(weekly) = weekly { + snapshot = snapshot.with_secondary(window(weekly, Some(10080))); + } + if let Some(monthly) = monthly { + let minutes = RateWindow::monthly_window_minutes(Some(monthly.1)).or(Some(43200)); + snapshot = snapshot.with_tertiary(window(monthly, minutes)); + } + snapshot +} + impl Default for OpenCodeGoProvider { fn default() -> Self { Self::new() diff --git a/rust/src/providers/opencodego/usage_api.rs b/rust/src/providers/opencodego/usage_api.rs index 13427919d9..283c1ee6dd 100644 --- a/rust/src/providers/opencodego/usage_api.rs +++ b/rust/src/providers/opencodego/usage_api.rs @@ -1,7 +1,7 @@ use chrono::{DateTime, Utc}; use reqwest::Client; -use crate::core::{FetchContext, ProviderError, ProviderFetchResult, RateWindow, UsageSnapshot}; +use crate::core::{FetchContext, ProviderError, ProviderFetchResult, UsageSnapshot}; const USAGE_API_URL: &str = "https://opencode.ai/zen/go/v1/usage"; @@ -71,30 +71,11 @@ fn parse_usage_text(text: &str, now: DateTime) -> Result = windows + .iter() + .map(|w| (w.window_minutes, w.resets_at.map(|r| r.to_rfc3339()))) + .collect(); + assert_eq!( + shape, + [ + (Some(300), Some("2026-08-12T02:00:00+00:00".to_string())), + (Some(10080), Some("2026-08-18T00:00:00+00:00".to_string())), + // Aug 1 to Sep 1 is a 31-day calendar month. + (Some(44640), Some("2026-09-01T00:00:00+00:00".to_string())), + ] + ); + assert!(windows.iter().all(|w| w.reset_description.is_none())); + assert!(snapshot.extra_rate_windows.is_empty()); } #[test] diff --git a/rust/src/providers/openrouter/activity.rs b/rust/src/providers/openrouter/activity.rs index d164320cf8..3d98bf93e6 100644 --- a/rust/src/providers/openrouter/activity.rs +++ b/rust/src/providers/openrouter/activity.rs @@ -71,35 +71,21 @@ pub(super) fn parse_activity_cost( )); } for (index, row) in rows.iter().enumerate() { - let object = row.as_object().ok_or_else(|| { - ProviderError::Parse(format!( - "OpenRouter activity.data[{index}] must be an object" - )) - })?; + let object = row + .as_object() + .ok_or_else(|| row_err(index, " must be an object"))?; let raw_day = object .get("date") .and_then(Value::as_str) .map(str::trim) .filter(|value| !value.is_empty()) - .ok_or_else(|| { - ProviderError::Parse(format!( - "OpenRouter activity.data[{index}].date is missing" - )) - })?; - let day = normalize_activity_day(raw_day).ok_or_else(|| { - ProviderError::Parse(format!( - "OpenRouter activity.data[{index}].date must be YYYY-MM-DD or YYYY-MM-DD HH:MM:SS" - )) - })?; - let parsed_day = chrono::NaiveDate::parse_from_str(day, "%Y-%m-%d").map_err(|_| { - ProviderError::Parse(format!( - "OpenRouter activity.data[{index}].date must be a real calendar date" - )) - })?; + .ok_or_else(|| row_err(index, ".date is missing"))?; + let day = normalize_activity_day(raw_day) + .ok_or_else(|| row_err(index, ".date must be YYYY-MM-DD or YYYY-MM-DD HH:MM:SS"))?; + let parsed_day = chrono::NaiveDate::parse_from_str(day, "%Y-%m-%d") + .map_err(|_| row_err(index, ".date must be a real calendar date"))?; if parsed_day > latest_completed { - return Err(ProviderError::Parse(format!( - "OpenRouter activity.data[{index}].date must be a completed UTC day" - ))); + return Err(row_err(index, ".date must be a completed UTC day")); } if parsed_day < cutoff { continue; @@ -111,9 +97,7 @@ pub(super) fn parse_activity_cost( .map(str::trim) .unwrap_or(""); if model.len() > 64 { - return Err(ProviderError::Parse(format!( - "OpenRouter activity.data[{index}].model exceeds 64 characters" - ))); + return Err(row_err(index, ".model exceeds 64 characters")); } let prompt = nonnegative_integer(object.get("prompt_tokens"), index, "prompt_tokens")?; let completion = @@ -126,9 +110,7 @@ pub(super) fn parse_activity_cost( .checked_add(completion) .is_none_or(|total| total > MAX_SAFE_INTEGER) { - return Err(ProviderError::Parse(format!( - "OpenRouter activity.data[{index}] token total overflowed" - ))); + return Err(row_err(index, " token total overflowed")); } let requests = nonnegative_integer(object.get("requests"), index, "requests")?; let metered = nonnegative_number(object.get("usage"), index, "usage")?; @@ -249,25 +231,24 @@ fn normalize_activity_day(raw: &str) -> Option<&str> { Some(&raw[..10]) } +fn row_err(index: usize, detail: &str) -> ProviderError { + ProviderError::Parse(format!("OpenRouter activity.data[{index}]{detail}")) +} + fn nonnegative_integer( value: Option<&Value>, index: usize, field: &str, ) -> Result { - let value = value.ok_or_else(|| { - ProviderError::Parse(format!( - "OpenRouter activity.data[{index}].{field} is missing" - )) - })?; - let value = value.as_u64().ok_or_else(|| { - ProviderError::Parse(format!( - "OpenRouter activity.data[{index}].{field} must be a nonnegative integer" - )) - })?; + let value = value.ok_or_else(|| row_err(index, &format!(".{field} is missing")))?; + let value = value + .as_u64() + .ok_or_else(|| row_err(index, &format!(".{field} must be a nonnegative integer")))?; if value > MAX_SAFE_INTEGER { - return Err(ProviderError::Parse(format!( - "OpenRouter activity.data[{index}].{field} must be a nonnegative safe integer" - ))); + return Err(row_err( + index, + &format!(".{field} must be a nonnegative safe integer"), + )); } Ok(value) } @@ -280,11 +261,7 @@ fn nonnegative_number( value .and_then(Value::as_f64) .filter(|value| value.is_finite() && *value >= 0.0) - .ok_or_else(|| { - ProviderError::Parse(format!( - "OpenRouter activity.data[{index}].{field} must be finite and nonnegative" - )) - }) + .ok_or_else(|| row_err(index, &format!(".{field} must be finite and nonnegative"))) } #[cfg(test)] @@ -297,6 +274,64 @@ mod tests { .with_timezone(&Utc) } + #[test] + fn activity_row_errors_name_the_row_index_and_field() { + let long_model = "m".repeat(65); + let cases = [ + (serde_json::json!(5), " must be an object"), + (serde_json::json!({}), ".date is missing"), + ( + serde_json::json!({"date":"2026/08/21"}), + ".date must be YYYY-MM-DD or YYYY-MM-DD HH:MM:SS", + ), + ( + serde_json::json!({"date":"2026-02-31"}), + ".date must be a real calendar date", + ), + ( + serde_json::json!({"date":"2026-08-22"}), + ".date must be a completed UTC day", + ), + ( + serde_json::json!({"date":"2026-08-21","model":long_model}), + ".model exceeds 64 characters", + ), + ( + serde_json::json!({"date":"2026-08-21"}), + ".prompt_tokens is missing", + ), + ( + serde_json::json!({"date":"2026-08-21","prompt_tokens":-1}), + ".prompt_tokens must be a nonnegative integer", + ), + ( + serde_json::json!({"date":"2026-08-21","prompt_tokens":MAX_SAFE_INTEGER + 1}), + ".prompt_tokens must be a nonnegative safe integer", + ), + ( + serde_json::json!({"date":"2026-08-21","prompt_tokens":MAX_SAFE_INTEGER, + "completion_tokens":1}), + " token total overflowed", + ), + ( + serde_json::json!({"date":"2026-08-21","prompt_tokens":1,"completion_tokens":1, + "requests":1,"usage":-1.0}), + ".usage must be finite and nonnegative", + ), + ]; + for (row, detail) in cases { + let valid = serde_json::json!({"date":"2026-08-21","model":"m","prompt_tokens":1, + "completion_tokens":1,"requests":1,"usage":0.1}); + let payload = serde_json::json!({ "data": [valid, row] }); + let error = parse_activity_cost(&[payload], now()).unwrap_err(); + assert!( + matches!(&error, ProviderError::Parse(message) + if *message == format!("OpenRouter activity.data[1]{detail}")), + "{detail}: {error:?}" + ); + } + } + #[test] fn aggregates_metered_and_byok_spend_without_double_counting_latest_completed_day() { let history = serde_json::json!({"data":[ diff --git a/rust/src/providers/openrouter/mod.rs b/rust/src/providers/openrouter/mod.rs index 8f0bf132bd..6bf7621dd2 100755 --- a/rust/src/providers/openrouter/mod.rs +++ b/rust/src/providers/openrouter/mod.rs @@ -296,26 +296,48 @@ impl OpenRouterProvider { management_key: &str, date: Option<&str>, ) -> Result { - let mut request = client - .get(OPENROUTER_ACTIVITY_URL) - .header("Authorization", format!("Bearer {management_key}")) - .header("Accept", "application/json"); + let mut request = client.get(OPENROUTER_ACTIVITY_URL); if let Some(date) = date { request = request.query(&[("date", date)]); } - let response = request.send().await?; + Self::get_json( + request, + management_key, + "Activity", + AUTH_REJECTED, + Some(diagnostics::ACTIVITY_KEY_REQUIRED), + ) + .await + } + + /// Bearer-authenticated JSON GET. Statuses in `auth_statuses` are typed as + /// auth failures; a 403 shows `forbidden_reason` when one is given. + async fn get_json( + request: reqwest::RequestBuilder, + api_key: &str, + label: &str, + auth_statuses: &[reqwest::StatusCode], + forbidden_reason: Option<&str>, + ) -> Result { + let response = request + .header("Authorization", format!("Bearer {api_key}")) + .header("Accept", "application/json") + .send() + .await?; let status = response.status(); - if status == reqwest::StatusCode::FORBIDDEN { - return Err(Degraded::http("Activity", status, AUTH_REJECTED) - .with_reason(diagnostics::ACTIVITY_KEY_REQUIRED)); - } if !status.is_success() { - return Err(Degraded::http("Activity", status, AUTH_REJECTED)); + let degraded = Degraded::http(label, status, auth_statuses); + return Err(match forbidden_reason { + Some(reason) if status == reqwest::StatusCode::FORBIDDEN => { + degraded.with_reason(reason) + } + _ => degraded, + }); } response - .json::() + .json::() .await - .map_err(|error| Degraded::body("Activity", error)) + .map_err(|error| Degraded::body(label, error)) } fn build_client() -> Result { @@ -330,25 +352,14 @@ impl OpenRouterProvider { api_key: &str, ) -> Result { let credits_url = format!("{}/credits", OPENROUTER_API_BASE); - let resp = client - .get(&credits_url) - .header("Authorization", format!("Bearer {}", api_key)) - .header("Accept", "application/json") - .send() - .await?; - - if !resp.status().is_success() { - return Err(Degraded::http( - "credits", - resp.status(), - &[reqwest::StatusCode::UNAUTHORIZED], - )); - } - - let response = resp - .json::() - .await - .map_err(|error| Degraded::body("credits", error))?; + let response: CreditsResponse = Self::get_json( + client.get(&credits_url), + api_key, + "credits", + &[reqwest::StatusCode::UNAUTHORIZED], + None, + ) + .await?; response.data.validate().map_err(Degraded::invalid)?; Ok(response) } @@ -430,21 +441,8 @@ impl OpenRouterProvider { async fn fetch_key_data(client: &reqwest::Client, api_key: &str) -> Result { let key_url = format!("{}/key", OPENROUTER_API_BASE); - let key_resp = client - .get(&key_url) - .header("Authorization", format!("Bearer {}", api_key)) - .header("Accept", "application/json") - .send() - .await?; - - if !key_resp.status().is_success() { - return Err(Degraded::http("key", key_resp.status(), AUTH_REJECTED)); - } - - let response = key_resp - .json::() - .await - .map_err(|error| Degraded::body("key", error))?; + let response: KeyResponse = + Self::get_json(client.get(&key_url), api_key, "key", AUTH_REJECTED, None).await?; response.data.validate().map_err(Degraded::invalid)?; Ok(response.data) } @@ -455,27 +453,23 @@ impl OpenRouterProvider { } fn add_spend_windows(usage: &mut UsageSnapshot, key_data: &KeyData) { - Self::add_spend_window( - usage, - key_data.usage_daily, - "daily-spend", - "Daily spend", - "today", - ); - Self::add_spend_window( - usage, - key_data.usage_weekly, - "weekly-spend", - "Weekly spend", - "this week", - ); - Self::add_spend_window( - usage, - key_data.usage_monthly, - "monthly-spend", - "Monthly spend", - "this month", - ); + for (value, id, label, period) in [ + (key_data.usage_daily, "daily-spend", "Daily spend", "today"), + ( + key_data.usage_weekly, + "weekly-spend", + "Weekly spend", + "this week", + ), + ( + key_data.usage_monthly, + "monthly-spend", + "Monthly spend", + "this month", + ), + ] { + Self::add_spend_window(usage, value, id, label, period); + } } fn key_quota_metrics(key_data: &KeyData) -> Option<(f64, f64, f64)> { @@ -506,10 +500,6 @@ impl OpenRouterProvider { /// reads 0%. Without it, fall back to the period usage matching the /// declared reset window, then cumulative usage; with no usable source the /// meter stays hidden. - fn add_key_quota(usage: &mut UsageSnapshot, key_data: &KeyData) { - Self::add_key_quota_with_suffix(usage, key_data, "Spending cap, not balance"); - } - fn add_key_quota_with_suffix(usage: &mut UsageSnapshot, key_data: &KeyData, suffix: &str) { let Some(key_window) = Self::key_quota_window(key_data, suffix) else { return; diff --git a/rust/src/providers/openrouter/tests.rs b/rust/src/providers/openrouter/tests.rs index 4c41d4dffd..e264a2d508 100644 --- a/rust/src/providers/openrouter/tests.rs +++ b/rust/src/providers/openrouter/tests.rs @@ -73,9 +73,11 @@ fn key_data( } } +const KEY_CAP_SUFFIX: &str = "Spending cap, not balance"; + fn key_quota_percent(key_data: KeyData) -> Option { let mut usage = UsageSnapshot::new(RateWindow::new(0.0)); - OpenRouterProvider::add_key_quota(&mut usage, &key_data); + OpenRouterProvider::add_key_quota_with_suffix(&mut usage, &key_data, KEY_CAP_SUFFIX); usage.secondary.map(|window| window.used_percent) } @@ -89,7 +91,7 @@ fn key_limit_copy_stays_distinct_from_account_balance() { total_usage: 3.1, }; let mut usage = OpenRouterProvider::build_credits_usage(&credits); - OpenRouterProvider::add_key_quota( + OpenRouterProvider::add_key_quota_with_suffix( &mut usage, &key_data( Some(30.0), @@ -100,6 +102,7 @@ fn key_limit_copy_stays_distinct_from_account_balance() { None, Some(0.0), ), + KEY_CAP_SUFFIX, ); assert_eq!(usage.login_method.as_deref(), Some("$1.90 balance")); let key = usage.secondary.expect("key spending cap"); @@ -267,114 +270,82 @@ fn activity_cost_wins_while_uncapped_key_spend_windows_remain() { } #[test] -fn server_remaining_replaces_lifetime_usage_for_meter() { - // limit 50, server says 12.50 left this period → 75% used, even though - // cumulative lifetime usage would imply a different ratio. - let pct = key_quota_percent(key_data( - Some(50.0), - Some(12.5), - None, - Some(40.0), - None, - None, - None, - )); - assert_eq!(pct, Some(75.0)); -} - -#[test] -fn negative_server_remaining_reads_exhausted() { - // Upstream: "treat negative remaining as exhausted quota". - let pct = key_quota_percent(key_data( - Some(50.0), - Some(-3.0), - None, - Some(10.0), - None, - None, - None, - )); - assert_eq!(pct, Some(100.0)); -} - -#[test] -fn above_limit_server_remaining_reads_zero() { - // Inclusive [0, keyLimit] clamp: a server remaining above the - // configured limit renders 0% used, not a suppressed meter. - let pct = key_quota_percent(key_data( - Some(50.0), - Some(75.0), - None, - Some(10.0), - None, - None, - None, - )); - assert_eq!(pct, Some(0.0)); -} - -#[test] -fn reset_window_usage_is_the_preferred_fallback() { - // No remaining: `limit_reset: "monthly"` picks usage_monthly (25/50). - let pct = key_quota_percent(key_data( - Some(50.0), - None, - Some("monthly"), - Some(40.0), - Some(1.0), - Some(2.0), - Some(25.0), - )); - assert_eq!(pct, Some(50.0)); - // Case-insensitive reset label. - let pct = key_quota_percent(key_data( - Some(50.0), - None, - Some("WEEKLY"), - Some(40.0), - Some(1.0), - Some(2.0), - Some(25.0), - )); - assert_eq!(pct, Some(4.0)); -} - -#[test] -fn cumulative_usage_is_the_last_fallback() { - let pct = key_quota_percent(key_data( - Some(50.0), - None, - None, - Some(20.0), - Some(1.0), - None, - None, - )); - assert_eq!(pct, Some(40.0)); -} - -#[test] -fn no_usable_quota_source_hides_the_meter() { - assert_eq!( - key_quota_percent(key_data(Some(50.0), None, None, None, None, None, None)), - None +fn key_meter_prefers_server_remaining_then_reset_window_then_cumulative_usage() { + type Case = ( + &'static str, + (Option, Option, Option<&'static str>, Option), + (Option, Option, Option), + Option, ); - assert_eq!( - key_quota_percent(key_data( + let cases: [Case; 9] = [ + // limit 50, server says 12.50 left this period: 75% used, even though + // cumulative lifetime usage would imply a different ratio. + ( + "server remaining replaces lifetime usage", + (Some(50.0), Some(12.5), None, Some(40.0)), + (None, None, None), + Some(75.0), + ), + // Upstream: "treat negative remaining as exhausted quota". + ( + "negative server remaining reads exhausted", + (Some(50.0), Some(-3.0), None, Some(10.0)), + (None, None, None), + Some(100.0), + ), + // Inclusive [0, keyLimit] clamp: a server remaining above the + // configured limit renders 0% used, not a suppressed meter. + ( + "above-limit server remaining reads zero", + (Some(50.0), Some(75.0), None, Some(10.0)), + (None, None, None), Some(0.0), - Some(5.0), + ), + // No remaining: `limit_reset: "monthly"` picks usage_monthly (25/50). + ( + "reset window usage is the preferred fallback", + (Some(50.0), None, Some("monthly"), Some(40.0)), + (Some(1.0), Some(2.0), Some(25.0)), + Some(50.0), + ), + // Case-insensitive reset label. + ( + "reset window label is case-insensitive", + (Some(50.0), None, Some("WEEKLY"), Some(40.0)), + (Some(1.0), Some(2.0), Some(25.0)), + Some(4.0), + ), + ( + "cumulative usage is the last fallback", + (Some(50.0), None, None, Some(20.0)), + (Some(1.0), None, None), + Some(40.0), + ), + ( + "no usable quota source hides the meter", + (Some(50.0), None, None, None), + (None, None, None), None, - Some(1.0), + ), + ( + "zero limit hides the meter", + (Some(0.0), Some(5.0), None, Some(1.0)), + (None, None, None), None, + ), + ( + "missing limit hides the meter", + (None, Some(5.0), None, Some(1.0)), + (None, None, None), None, - None - )), - None - ); - assert_eq!( - key_quota_percent(key_data(None, Some(5.0), None, Some(1.0), None, None, None)), - None - ); + ), + ]; + for (name, (limit, remaining, reset, usage), (daily, weekly, monthly), expected) in cases { + let pct = key_quota_percent(key_data( + limit, remaining, reset, usage, daily, weekly, monthly, + )); + assert_eq!(pct, expected, "{name}"); + } } #[test] @@ -560,3 +531,116 @@ fn missing_key_message_explains_primary_and_management_fields() { "Enter a regular API key or a Management API key in the API key field, or set OPENROUTER_API_KEY. In Settings, the optional Management API key field does not replace it." ); } + +// ── Shared JSON GET: status typing and safe reasons per call site ── + +#[tokio::test] +async fn get_json_keeps_each_call_sites_auth_set_and_reason() { + use reqwest::StatusCode; + + let mut server = mockito::Server::new_async().await; + let client = reqwest::Client::new(); + let credits_auth = &[StatusCode::UNAUTHORIZED][..]; + // (label, auth statuses, forbidden reason, HTTP status, body, + // expected error, expected reason) + type Case<'a> = ( + &'a str, + &'a [StatusCode], + Option<&'a str>, + usize, + &'a str, + &'a str, + &'a str, + ); + let cases: [Case<'_>; 7] = [ + ("credits", credits_auth, None, 200, r#"{"ok":1}"#, "", ""), + ( + "credits", + credits_auth, + None, + 401, + "", + "auth", + "Request returned HTTP 401", + ), + ( + "credits", + credits_auth, + None, + 403, + "", + "OpenRouter credits request returned HTTP 403 Forbidden", + "Request returned HTTP 403", + ), + ( + "key", + AUTH_REJECTED, + None, + 403, + "", + "auth", + "Request returned HTTP 403", + ), + ( + "key", + AUTH_REJECTED, + None, + 200, + "not json", + "OpenRouter key response was invalid:", + "Response was invalid", + ), + ( + "Activity", + AUTH_REJECTED, + Some(ACTIVITY_KEY_REQUIRED), + 403, + "", + "auth", + ACTIVITY_KEY_REQUIRED, + ), + ( + "Activity", + AUTH_REJECTED, + Some(ACTIVITY_KEY_REQUIRED), + 500, + "", + "OpenRouter Activity request returned HTTP 500 Internal Server Error", + "Request returned HTTP 500", + ), + ]; + for (index, (label, auth, forbidden, status, body, error, reason)) in + cases.into_iter().enumerate() + { + let path = format!("/case-{index}"); + let mock = server + .mock("GET", path.as_str()) + .match_header("authorization", "Bearer sk-or-synthetic") + .match_header("accept", "application/json") + .with_status(status) + .with_body(body) + .create_async() + .await; + let request = client.get(format!("{}{path}", server.url())); + let result: Result = + OpenRouterProvider::get_json(request, "sk-or-synthetic", label, auth, forbidden).await; + mock.assert_async().await; + match result { + Ok(value) => { + assert_eq!(error, "", "{label} {status}"); + assert_eq!(value, serde_json::json!({ "ok": 1 })); + } + Err(degraded) => { + assert_eq!(degraded.reason, reason, "{label} {status}"); + match degraded.error { + ProviderError::AuthRequired => assert_eq!(error, "auth", "{label} {status}"), + ProviderError::Other(message) => assert_eq!(message, error), + ProviderError::Parse(message) => { + assert!(message.starts_with(error), "{message}") + } + other => panic!("{label} {status}: unexpected {other:?}"), + } + } + } + } +}