diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index a86b331..455c094 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -17,7 +17,7 @@ jobs: - uses: actions/checkout@v6 with: repository: Nanako0129/TokenBar - ref: 1e00e7b769a1b0d20b2de077b0a3aab52a24c088 + ref: 4dfed5ffab26e2707a8cd82ee99864520b75892b path: tokenbar-mac - uses: actions/checkout@v6 with: diff --git a/Cargo.lock b/Cargo.lock index e9ed667..8590bc5 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1552,15 +1552,18 @@ dependencies = [ "dirs", "fs2", "hmac", + "hyper-util", "parking_lot", "rayon", "reqwest", + "rustls", "security-framework", "serde", "serde_json", "sha2", "tokio", "tokscale-core", + "tower-service", "windows-sys 0.61.2", ] diff --git a/Fixtures/CrossCheck/provider-quota-pace-v3.json b/Fixtures/CrossCheck/provider-quota-pace-v3.json new file mode 100644 index 0000000..7270a12 --- /dev/null +++ b/Fixtures/CrossCheck/provider-quota-pace-v3.json @@ -0,0 +1,201 @@ +{ + "schemaVersion": 3, + "payload": { + "generatedAt": "2026-07-10T12:00:00.000Z", + "agents": [ + { + "clientId": "provider-fixture.invalid", + "source": "fixture.invalid", + "updatedAt": "2026-07-10T12:00:00.000Z", + "identity": null, + "windows": [ + { + "cardId": "ahead.invalid", + "label": "Ahead quota", + "usedPercent": 72.0, + "remainingPercent": 28.0, + "resetsAt": "2026-07-10T15:00:00Z", + "windowMinutes": 300, + "paceStatus": { + "state": "available", + "windowKey": "quota.ahead.invalid", + "durationSeconds": 18000, + "durationSource": "provider", + "completeCycles": 5 + }, + "historicalPace": { + "expectedUsedPercent": 32.0, + "etaSeconds": 3600.0, + "willLastToReset": false, + "runOutProbability": 0.75 + } + }, + { + "cardId": "behind.invalid", + "label": "Behind quota", + "usedPercent": 28.0, + "remainingPercent": 72.0, + "resetsAt": "2026-07-15T12:00:00Z", + "windowMinutes": 10080, + "paceStatus": { + "state": "available", + "windowKey": "quota.behind.invalid", + "durationSeconds": 604800, + "durationSource": "contract", + "completeCycles": 7 + }, + "historicalPace": { + "expectedUsedPercent": 56.0, + "willLastToReset": true, + "runOutProbability": 0.2 + } + }, + { + "cardId": "learning-history.invalid", + "label": "Learning history", + "usedPercent": 40.0, + "remainingPercent": 60.0, + "resetsAt": "2026-07-10T15:00:00Z", + "windowMinutes": 300, + "paceStatus": { + "state": "learningHistory", + "windowKey": "quota.learning-history.invalid", + "durationSeconds": 18000, + "durationSource": "provider", + "completeCycles": 2 + } + }, + { + "cardId": "learning-duration.invalid", + "label": "Learning duration", + "usedPercent": 40.0, + "remainingPercent": 60.0, + "resetsAt": "2026-07-10T15:00:00Z", + "paceStatus": { + "state": "learningDuration", + "windowKey": "quota.learning-duration.invalid", + "durationSource": "observed", + "completeCycles": 0 + } + }, + { + "cardId": "missing-reset.invalid", + "label": "Missing reset", + "usedPercent": 50.0, + "remainingPercent": 50.0, + "paceStatus": { + "state": "unavailable", + "windowKey": "quota.missing-reset.invalid", + "completeCycles": 0, + "reason": "missingReset" + } + }, + { + "cardId": "shared-first.invalid", + "label": "Shared label", + "usedPercent": 10.0, + "remainingPercent": 90.0, + "resetsAt": "2026-07-10T15:00:00Z", + "windowMinutes": 300, + "paceStatus": { + "state": "learningHistory", + "windowKey": "quota.shared-first.invalid", + "durationSeconds": 18000, + "durationSource": "provider", + "completeCycles": 2 + } + }, + { + "cardId": "shared-second.invalid", + "label": "Shared label", + "usedPercent": 20.0, + "remainingPercent": 80.0, + "resetsAt": "2026-07-10T15:00:00Z", + "windowMinutes": 300, + "paceStatus": { + "state": "learningHistory", + "windowKey": "quota.shared-second.invalid", + "durationSeconds": 18000, + "durationSource": "provider", + "completeCycles": 2 + } + } + ], + "credits": null, + "error": null + } + ] + }, + "cases": [ + { + "name": "available-ahead-historical", + "kind": "pace", + "clientId": "provider-fixture.invalid", + "cardId": "ahead.invalid", + "mode": "historical", + "now": "2026-07-10T12:00:00Z" + }, + { + "name": "available-behind-historical", + "kind": "pace", + "clientId": "provider-fixture.invalid", + "cardId": "behind.invalid", + "mode": "historical", + "now": "2026-07-10T12:00:00Z" + }, + { + "name": "learning-history-historical-uses-linear", + "kind": "pace", + "clientId": "provider-fixture.invalid", + "cardId": "learning-history.invalid", + "mode": "historical", + "now": "2026-07-10T12:00:00Z" + }, + { + "name": "selection-exact-card-id", + "kind": "selection", + "selection": "provider-fixture.invalid|shared-second.invalid" + }, + { + "name": "selection-unique-legacy-label", + "kind": "selection", + "selection": "provider-fixture.invalid|Ahead quota" + }, + { + "name": "selection-ambiguous-legacy-label-preserved", + "kind": "selection", + "selection": "provider-fixture.invalid|Shared label" + }, + { + "name": "legacy-missing-pace-status", + "kind": "legacy", + "now": "2026-07-10T12:00:00Z", + "rawWindow": "{\"cardId\":\"legacy.invalid\",\"label\":\"Legacy\",\"usedPercent\":20,\"remainingPercent\":80,\"resetsAt\":\"2026-07-10T15:00:00Z\",\"windowMinutes\":300}" + }, + { + "name": "malformed-pace-status-null", + "kind": "malformed", + "rawWindow": "{\"cardId\":\"malformed-null.invalid\",\"label\":\"Malformed null\",\"usedPercent\":20,\"remainingPercent\":80,\"resetsAt\":\"2026-07-10T15:00:00Z\",\"windowMinutes\":300,\"paceStatus\":null}" + }, + { + "name": "malformed-unknown-state", + "kind": "malformed", + "rawWindow": "{\"cardId\":\"malformed-state.invalid\",\"label\":\"Malformed state\",\"usedPercent\":20,\"remainingPercent\":80,\"resetsAt\":\"2026-07-10T15:00:00Z\",\"windowMinutes\":300,\"paceStatus\":{\"state\":\"unknownState\",\"windowKey\":\"quota.malformed-state.invalid\",\"durationSeconds\":18000,\"durationSource\":\"provider\",\"completeCycles\":0}}" + }, + { + "name": "malformed-available-without-historical", + "kind": "malformed", + "rawWindow": "{\"cardId\":\"malformed-available.invalid\",\"label\":\"Malformed available\",\"usedPercent\":20,\"remainingPercent\":80,\"resetsAt\":\"2026-07-10T15:00:00Z\",\"windowMinutes\":300,\"paceStatus\":{\"state\":\"available\",\"windowKey\":\"quota.malformed-available.invalid\",\"durationSeconds\":18000,\"durationSource\":\"provider\",\"completeCycles\":5}}" + }, + { + "name": "malformed-duration-window-minutes-mismatch", + "kind": "malformed", + "rawWindow": "{\"cardId\":\"malformed-duration.invalid\",\"label\":\"Malformed duration\",\"usedPercent\":20,\"remainingPercent\":80,\"resetsAt\":\"2026-07-10T15:00:00Z\",\"windowMinutes\":301,\"paceStatus\":{\"state\":\"learningHistory\",\"windowKey\":\"quota.malformed-duration.invalid\",\"durationSeconds\":18000,\"durationSource\":\"provider\",\"completeCycles\":0}}" + }, + { + "name": "malformed-percentages-not-complementary", + "kind": "malformed", + "rawWindow": "{\"cardId\":\"malformed-percent.invalid\",\"label\":\"Malformed percentages\",\"usedPercent\":80,\"remainingPercent\":80,\"resetsAt\":\"2026-07-10T15:00:00Z\",\"windowMinutes\":300}" + } + ] +} diff --git a/crates/tb_core_ffi/Cargo.toml b/crates/tb_core_ffi/Cargo.toml index 5a56717..dd21ad8 100644 --- a/crates/tb_core_ffi/Cargo.toml +++ b/crates/tb_core_ffi/Cargo.toml @@ -19,6 +19,10 @@ serde_json = { version = "1.0", features = ["raw_value"] } # Version specs mirror the Tauri app (TokenBar-tokcat/src-tauri/Cargo.toml) so # the ported report/agent modules compile against the same dependency surface. reqwest = { version = "0.13", default-features = false, features = ["json", "rustls"] } +# Public resolver and typed TLS error surfaces used by provider diagnostics. +hyper-util = { version = "0.1", default-features = false, features = ["client-legacy"] } +tower-service = "0.3" +rustls = { version = "0.23", default-features = false } base64 = "0.22" hmac = "0.12" sha2 = "0.10" diff --git a/crates/tb_core_ffi/src/agent_antigravity.rs b/crates/tb_core_ffi/src/agent_antigravity.rs index 7a25d54..5cd9ef4 100644 --- a/crates/tb_core_ffi/src/agent_antigravity.rs +++ b/crates/tb_core_ffi/src/agent_antigravity.rs @@ -9,10 +9,10 @@ //! listening ports with platform-native process tools, and call the local //! Connect-RPC `GetUserStatus` over loopback TLS. Live, no token refresh, //! no disk writes. -//! 2. **OAuth remote (`oauth`)** — otherwise read the shared Google creds at -//! `~/.gemini/oauth_creds.json`, refresh against Google (client id/secret -//! scanned from the installed Antigravity.app binary), and hit the -//! `cloudcode-pa.googleapis.com` Code Assist quota endpoints. +//! 2. **OAuth remote (`oauth`)** — otherwise read the shared Google creds under +//! `GEMINI_CLI_HOME` (falling back to `~/.gemini`), refresh against Google +//! (client id/secret scanned from the installed Antigravity.app binary), and +//! hit the `cloudcode-pa.googleapis.com` Code Assist quota endpoints. //! //! Both yield per-model "remaining fraction + reset" which map to `UsageWindow`s. @@ -20,7 +20,11 @@ use crate::agent_account_scope::{ self, AccountScope, AccountScopeError, AuthoritativeIdKind, RefreshCheckpoint, RefreshScopeTransaction, }; -use crate::agent_usage::{clean_plan, parse_datetime, percent_encode, AgentIdentity, UsageWindow}; +use crate::agent_usage::{ + clean_plan, parse_datetime, percent_encode, provider_http_client_builder, read_response_body, + request_after_verified_binding, AgentIdentity, ProviderCacheBinding, ProviderFetchFailure, + ResponseReadFailure, TransportErrorFacts, TransportPhase, UsageWindow, +}; #[cfg(any(windows, test))] use base64::Engine; use chrono::{DateTime, Utc}; @@ -36,32 +40,92 @@ use std::sync::OnceLock; const LANG_SERVICE: &str = "/exa.language_server_pb.LanguageServerService/GetUserStatus"; const CODE_ASSIST_BASE: &str = "https://cloudcode-pa.googleapis.com/v1internal"; const GOOGLE_TOKEN_URL: &str = "https://oauth2.googleapis.com/token"; -const GOOGLE_USERINFO_URL: &str = "https://www.googleapis.com/oauth2/v2/userinfo"; const REFRESH_SAFETY_SECS: i64 = 60; +#[derive(Debug)] pub(crate) struct Fetched { pub source: String, pub identity: Option, pub account_scope: Result, + pub cache_binding: Option, pub windows: Vec, } -/// Auto: prefer the live Local IDE API; fall back to the OAuth remote API. -pub(crate) async fn fetch(now: DateTime) -> Result { - match fetch_local_ide(now).await { - Ok(local) if !local.windows.is_empty() => Ok(local), - local_result => { - let local_err = local_result.err(); - match fetch_oauth_remote(now).await { - Ok(remote) => Ok(remote), - Err(remote_err) => Err(local_err - .map(|le| format!("{remote_err} (local IDE: {le})")) - .unwrap_or(remote_err)), +#[derive(Debug)] +enum LocalAttempt { + Success(Fetched), + RouteMiss, +} + +#[derive(Debug)] +enum PrimaryAttempt { + Success(Fetched), + Forbidden(Context), + SchemaContradiction { + context: Context, + failure: ProviderFetchFailure, + }, + Transient(Context), + FinalFailure(ProviderFetchFailure), +} + +async fn fetch_with< + Local, + LocalFuture, + Primary, + PrimaryFuture, + Secondary, + SecondaryFuture, + Context, +>( + local_attempt: Local, + primary_remote_attempt: Primary, + secondary_remote_attempt: Secondary, +) -> Result +where + Local: FnOnce() -> LocalFuture, + LocalFuture: std::future::Future, + Primary: FnOnce() -> PrimaryFuture, + PrimaryFuture: std::future::Future>, + Secondary: FnOnce(Context) -> SecondaryFuture, + SecondaryFuture: std::future::Future>, +{ + if let LocalAttempt::Success(fetched) = local_attempt().await { + return Ok(fetched); + } + + match primary_remote_attempt().await { + PrimaryAttempt::Success(fetched) => Ok(fetched), + PrimaryAttempt::Forbidden(context) => secondary_remote_attempt(context).await, + PrimaryAttempt::SchemaContradiction { context, failure } => { + match secondary_remote_attempt(context).await { + Ok(fetched) => Ok(fetched), + Err(_) => Err(failure), } } + PrimaryAttempt::Transient(context) => match secondary_remote_attempt(context).await { + Ok(fetched) => Ok(fetched), + Err(failure) => Err(failure), + }, + PrimaryAttempt::FinalFailure(failure) => Err(failure), } } +/// Auto: prefer the live Local IDE API; fall back to the OAuth remote API. +pub(crate) async fn fetch(now: DateTime) -> Result { + fetch_with( + || async move { + match fetch_local_ide(now).await { + Ok(local) if !local.windows.is_empty() => LocalAttempt::Success(local), + Ok(_) | Err(_) => LocalAttempt::RouteMiss, + } + }, + || fetch_oauth_primary(now), + |context| fetch_oauth_secondary(context, now), + ) + .await +} + // ── Local IDE API ─────────────────────────────────────────────────────────── struct ProcInfo { @@ -114,8 +178,14 @@ async fn fetch_local_ide(now: DateTime) -> Result { continue; }; match parse_user_status(&text, now) { - Ok(mut fetched) if local_result_is_acceptable(&fetched) => { + Ok(mut fetched) if !fetched.windows.is_empty() => { fetched.account_scope = resolve_local_account_scope(fetched.identity.as_ref()); + fetched.cache_binding = fetched + .account_scope + .as_ref() + .ok() + .cloned() + .map(ProviderCacheBinding::primary); return Ok(fetched); } Ok(_) => last_err = "local API returned no model quotas".to_string(), @@ -125,10 +195,6 @@ async fn fetch_local_ide(now: DateTime) -> Result { Err(last_err) } -fn local_result_is_acceptable(fetched: &Fetched) -> bool { - !fetched.windows.is_empty() -} - fn local_api_candidates(processes: Vec<(ProcInfo, Vec)>) -> Vec<(u16, String)> { let mut candidates = Vec::new(); for (proc, ports) in processes { @@ -184,7 +250,8 @@ fn detect_process() -> Result { return Ok(ProcInfo { pid, csrf_token: csrf, - extension_port: extract_flag(cmd, "--extension_server_port").and_then(|s| s.parse().ok()), + extension_port: extract_flag(cmd, "--extension_server_port") + .and_then(|s| s.parse().ok()), extension_csrf: extract_flag(cmd, "--extension_server_csrf_token"), }); } @@ -246,12 +313,12 @@ const WINDOWS_DISCOVERY_PREFIX: &str = "ANTIGRAVITY_V1"; #[cfg(windows)] const CREATE_NO_WINDOW: u32 = 0x0800_0000; -#[cfg(windows)] +#[cfg(any(windows, test))] const WINDOWS_DISCOVERY_SCRIPT: &str = r#" $ErrorActionPreference = 'Stop' try { [Console]::OutputEncoding = [System.Text.UTF8Encoding]::new($false) - $processes = @(Get-CimInstance -ClassName Win32_Process -Filter "Name = 'language_server.exe'") + $processes = @(Get-CimInstance -ClassName Win32_Process -Filter "Name LIKE 'language_server%.exe'") if ($processes.Count -eq 0) { [Console]::Out.WriteLine("ANTIGRAVITY_V1`tN") exit 0 @@ -489,7 +556,6 @@ struct ModelCandidate { model_id: Option, fraction: f64, reset: Option>, - reset_was_supplied: bool, source_index: usize, label: String, } @@ -502,16 +568,12 @@ fn quota_window( label: String, fraction: f64, reset: Option>, - reset_was_supplied: bool, now: DateTime, card_id: String, window_key: Option, ) -> Option { - valid_remaining_fraction(fraction).then(|| { - UsageWindow::from_fraction(label, fraction, reset, now) - .with_identity(card_id, window_key) - .with_observed_duration_evidence(now, reset_was_supplied) - }) + UsageWindow::try_from_provider_fraction(label, fraction, reset, now) + .map(|window| window.with_identity(card_id, window_key, None, None)) } fn parse_user_status(body: &str, now: DateTime) -> Result { @@ -527,7 +589,7 @@ fn parse_user_status(body: &str, now: DateTime) -> Result .unwrap_or_default(); let mut selected: BTreeMap = BTreeMap::new(); let mut missing_model = Vec::new(); - for (source_index, config) in configs.into_iter().enumerate() { + for (index, config) in configs.into_iter().enumerate() { let Ok(config) = serde_json::from_str::(config.get()) else { continue; }; @@ -537,9 +599,6 @@ fn parse_user_status(body: &str, now: DateTime) -> Result let Some(fraction) = quota.get("remainingFraction").and_then(Value::as_f64) else { continue; }; - let reset_was_supplied = quota - .get("resetTime") - .is_some_and(|value| !value.is_null()); let reset = quota .get("resetTime") .and_then(Value::as_str) @@ -561,8 +620,7 @@ fn parse_user_status(body: &str, now: DateTime) -> Result model_id: model_id.clone(), fraction, reset, - reset_was_supplied, - source_index, + source_index: index, label, }; let Some(model_id) = model_id else { @@ -596,13 +654,15 @@ fn parse_user_status(body: &str, now: DateTime) -> Result let key = format!("model.{model_id}.v1"); (key.clone(), Some(key)) } - None => (format!("row.cli.config.{}.v1", candidate.source_index), None), + None => ( + format!("row.cli.config.{}.v1", candidate.source_index), + None, + ), }; quota_window( candidate.label, candidate.fraction, candidate.reset, - candidate.reset_was_supplied, now, card_id, window_key, @@ -614,17 +674,21 @@ fn parse_user_status(body: &str, now: DateTime) -> Result .user_tier .and_then(|t| t.name) .filter(|s| !s.trim().is_empty()) - .or_else(|| status.plan_status.and_then(|p| p.plan_info).and_then(local_plan_name)); + .or_else(|| { + status + .plan_status + .and_then(|p| p.plan_info) + .and_then(local_plan_name) + }); + let email = status.email.filter(|value| !value.trim().is_empty()); Ok(Fetched { source: "cli".to_string(), - identity: Some(AgentIdentity { - email: status.email.filter(|s| !s.trim().is_empty()), - plan, - }), - // Parsing stays pure. The authenticated loopback fetch resolves this only - // after accepting the response row. + identity: Some(AgentIdentity { email, plan }), + // Parsing remains pure and hermetic. fetch_local_ide resolves this only + // after the authenticated loopback response has been accepted. account_scope: Err(AccountScopeError::NoTrustedEvidence), + cache_binding: None, windows, }) } @@ -632,22 +696,10 @@ fn parse_user_status(body: &str, now: DateTime) -> Result fn resolve_local_account_scope( identity: Option<&AgentIdentity>, ) -> Result { - resolve_local_account_scope_with(identity, resolve_email_account_scope) -} - -fn resolve_email_account_scope(email: &str) -> Result { - agent_account_scope::resolve_authoritative("antigravity", AuthoritativeIdKind::Email, email) -} - -fn resolve_local_account_scope_with( - identity: Option<&AgentIdentity>, - resolve: impl FnOnce(&str) -> Result, -) -> Result { let email = identity .and_then(|identity| identity.email.as_deref()) - .filter(|email| !email.trim().is_empty()) .ok_or(AccountScopeError::NoTrustedEvidence)?; - resolve(email) + agent_account_scope::resolve_authoritative("antigravity", AuthoritativeIdKind::Email, email) } fn local_plan_name(info: LocalPlanInfo) -> Option { @@ -665,141 +717,143 @@ fn local_plan_name(info: LocalPlanInfo) -> Option { // ── OAuth remote (Google Code Assist) ───────────────────────────────────────── -#[derive(Deserialize)] -struct GoogleUserInfo { - email: Option, - #[serde(default)] - verified_email: bool, -} - -struct RemoteAuth { +struct RemoteContext { + client: reqwest::Client, access_token: String, - account_scope: Result, + project: Option, + plan: Option, + account_scope: AccountScope, + cache_binding: Option, +} + +impl RemoteContext { + fn finish(self, windows: Vec) -> Fetched { + Fetched { + source: "oauth".to_string(), + // google_accounts.active is unrelated local state, not authenticated + // by the credential that fetched these quotas. + identity: Some(remote_identity(self.plan)), + account_scope: Ok(self.account_scope), + cache_binding: self.cache_binding, + windows, + } + } } -async fn fetch_oauth_remote(now: DateTime) -> Result { - let creds_path = gemini_home() - .map(|home| home.join("oauth_creds.json")) - .ok_or_else(|| "Could not resolve ~/.gemini".to_string())?; - let creds = load_remote_credentials(&creds_path)?; - let client = reqwest::Client::builder() - .timeout(std::time::Duration::from_secs(30)) - .build() - .map_err(|e| format!("build Antigravity client: {e}"))?; - let auth = prepare_remote_auth_with( - &creds, - now, - || refresh_access_token(&creds_path, now), - |access_token| fetch_remote_account_scope(&client, access_token), - ) - .await?; - - let code_assist_body = code_assist_post( - &client, - "loadCodeAssist", - &json!({ - "metadata": { "ideType": "ANTIGRAVITY", "platform": "PLATFORM_UNSPECIFIED", "pluginType": "GEMINI" } - }), - &auth.access_token, - ) - .await?; - let code_assist: Value = serde_json::from_str(&code_assist_body) - .map_err(|e| format!("decode Antigravity loadCodeAssist: {e}"))?; - let project = project_id(&code_assist); - let plan = resolve_remote_plan(&code_assist); - let windows = fetch_model_quotas(&client, &auth.access_token, project.as_deref(), now).await?; +enum PrimaryQuotaAttempt { + Success(Vec), + Forbidden, + SchemaContradiction(ProviderFetchFailure), + Transient, + Terminal(ProviderFetchFailure), +} - Ok(remote_fetched(plan, auth.account_scope, windows)) +async fn fetch_oauth_primary(now: DateTime) -> PrimaryAttempt { + let context = match prepare_remote_context(now).await { + Ok(context) => context, + Err(failure) => return PrimaryAttempt::FinalFailure(failure), + }; + match fetch_available_models(&context, now).await { + PrimaryQuotaAttempt::Success(windows) => PrimaryAttempt::Success(context.finish(windows)), + PrimaryQuotaAttempt::Forbidden => PrimaryAttempt::Forbidden(context), + PrimaryQuotaAttempt::SchemaContradiction(failure) => { + PrimaryAttempt::SchemaContradiction { context, failure } + } + PrimaryQuotaAttempt::Transient => PrimaryAttempt::Transient(context), + PrimaryQuotaAttempt::Terminal(failure) => PrimaryAttempt::FinalFailure(failure), + } } -async fn prepare_remote_auth_with( - creds: &Value, +async fn fetch_oauth_secondary( + context: RemoteContext, now: DateTime, - refresh: Refresh, - resolve_account_scope: Resolve, -) -> Result, String> -where - Refresh: FnOnce() -> RefreshFuture, - RefreshFuture: std::future::Future< - Output = Result<(Value, String, Result), String>, - >, - Resolve: FnOnce(String) -> ResolveFuture, - ResolveFuture: std::future::Future>, -{ - let mut access_token = remote_access_token(creds)?; - if remote_credentials_need_refresh(creds, now) { - // Credential lineage protects refresh rotation only; it is never history identity. - let (_, final_access_token, _credential_lineage) = refresh().await?; - access_token = final_access_token; - } - let account_scope = resolve_account_scope(access_token.clone()).await; - Ok(RemoteAuth { - access_token, - account_scope, - }) +) -> Result { + let windows = fetch_user_quota(&context, now).await?; + Ok(context.finish(windows)) } -async fn fetch_remote_account_scope( - client: &reqwest::Client, - access_token: String, -) -> Result { - let response = client - .get(GOOGLE_USERINFO_URL) - .bearer_auth(&access_token) - .send() - .await - .map_err(|_| AccountScopeError::NoTrustedEvidence)?; - let status = response.status(); - require_google_user_info_success(status)?; - let body = response - .text() - .await - .map_err(|_| AccountScopeError::NoTrustedEvidence)?; - resolve_google_user_info_response_with(status, &body, resolve_email_account_scope) -} - -fn require_google_user_info_success( - status: reqwest::StatusCode, -) -> Result<(), AccountScopeError> { - status - .is_success() - .then_some(()) - .ok_or(AccountScopeError::NoTrustedEvidence) -} - -fn resolve_google_user_info_response_with( - status: reqwest::StatusCode, - body: &str, - resolve: impl FnOnce(&str) -> Result, -) -> Result { - require_google_user_info_success(status)?; - let user_info: GoogleUserInfo = - serde_json::from_str(body).map_err(|_| AccountScopeError::NoTrustedEvidence)?; - if !user_info.verified_email { - return Err(AccountScopeError::NoTrustedEvidence); - } - let email = user_info - .email - .filter(|email| !email.trim().is_empty()) - .ok_or(AccountScopeError::NoTrustedEvidence)?; - resolve(&email) -} +async fn prepare_remote_context(now: DateTime) -> Result { + let creds_path = gemini_home() + .map(|home| home.join("oauth_creds.json")) + .ok_or_else(|| { + ProviderFetchFailure::terminal("Antigravity credential location could not be resolved.") + })?; + let creds = load_remote_credentials(&creds_path).map_err(|_| { + ProviderFetchFailure::terminal("Antigravity is not logged in. Re-login in Antigravity.") + })?; + let verified = if remote_credentials_need_refresh(&creds, now) { + refresh_access_token(&creds_path, now).await.map( + |(_, access_token, account_scope, cache_binding)| { + (access_token, account_scope, cache_binding) + }, + ) + } else { + remote_access_token(&creds) + .map_err(|_| { + ProviderFetchFailure::terminal("Antigravity credentials have no access token.") + }) + .and_then(|access_token| { + resolve_remote_account_scope(&creds_path, &creds) + .map(|account_scope| { + let cache_binding = ProviderCacheBinding::primary(account_scope.clone()); + (access_token, account_scope, Some(cache_binding)) + }) + .map_err(|_| { + ProviderFetchFailure::terminal( + "Antigravity account identity could not be verified.", + ) + }) + }) + }; -fn remote_fetched( - plan: Option, - account_scope: Result, - windows: Vec, -) -> Fetched { - Fetched { - source: "oauth".to_string(), - identity: Some(remote_identity(plan)), - account_scope, - windows, - } + request_after_verified_binding( + verified, + |(access_token, account_scope, cache_binding)| async move { + let client = provider_http_client_builder() + .timeout(std::time::Duration::from_secs(30)) + .build() + .map_err(|_| { + ProviderFetchFailure::terminal( + "Antigravity usage client could not be created.", + ) + })?; + let code_assist_body = code_assist_post( + &client, + "loadCodeAssist", + &json!({ + "metadata": { "ideType": "ANTIGRAVITY", "platform": "PLATFORM_UNSPECIFIED", "pluginType": "GEMINI" } + }), + &access_token, + cache_binding.clone(), + false, + ) + .await + .map_err(|failure| match failure { + CodeAssistPostFailure::Forbidden => ProviderFetchFailure::terminal( + "Antigravity loadCodeAssist permission was denied.", + ), + CodeAssistPostFailure::Failure(failure) => failure, + })?; + let code_assist: Value = serde_json::from_str(&code_assist_body).map_err(|_| { + ProviderFetchFailure::terminal( + "Antigravity loadCodeAssist response could not be decoded.", + ) + })?; + + Ok(RemoteContext { + client, + access_token, + project: project_id(&code_assist), + plan: resolve_remote_plan(&code_assist), + account_scope, + cache_binding, + }) + }, + ) + .await } fn remote_identity(plan: Option) -> AgentIdentity { - // UserInfo email is scope evidence only and must not enter the wire identity. AgentIdentity { email: None, plan } } @@ -838,12 +892,26 @@ fn remote_scope_location(path: &Path) -> Result { agent_account_scope::canonical_file_location(path, Some("refresh_token")) } +fn resolve_remote_account_scope( + path: &Path, + creds: &Value, +) -> Result { + let marker = remote_refresh_marker(creds).ok_or(AccountScopeError::NoTrustedEvidence)?; + agent_account_scope::resolve_credential( + "antigravity", + "google-oauth-creds", + &remote_scope_location(path)?, + marker, + ) +} + async fn refresh_access_token( creds_path: &Path, now: DateTime, -) -> Result<(Value, String, Result), String> { - let refresh = agent_account_scope::begin_refresh("antigravity") - .map_err(|_| "Antigravity credential refresh lock is unavailable.".to_string())?; +) -> Result<(Value, String, AccountScope, Option), ProviderFetchFailure> { + let refresh = agent_account_scope::begin_refresh("antigravity").map_err(|_| { + ProviderFetchFailure::terminal("Antigravity credential refresh lock is unavailable.") + })?; refresh_access_token_with( creds_path, now, @@ -855,13 +923,21 @@ async fn refresh_access_token( .await } -async fn request_access_token(refresh_token: String) -> Result { - let client = resolve_oauth_client() - .ok_or_else(|| "Antigravity OAuth client not found. Install Antigravity.app or set ANTIGRAVITY_OAUTH_CLIENT_ID/SECRET.".to_string())?; - let http = reqwest::Client::builder() +async fn request_access_token( + refresh_token: String, + attempt_binding: ProviderCacheBinding, +) -> Result { + let client = resolve_oauth_client().ok_or_else(|| { + ProviderFetchFailure::terminal( + "Antigravity OAuth client was not found. Install Antigravity.app or configure its OAuth client.", + ) + })?; + let http = provider_http_client_builder() .timeout(std::time::Duration::from_secs(30)) .build() - .map_err(|e| format!("build refresh client: {e}"))?; + .map_err(|_| { + ProviderFetchFailure::terminal("Antigravity refresh client could not be created.") + })?; let form = format!( "client_id={}&client_secret={}&refresh_token={}&grant_type=refresh_token", percent_encode(&client.0), @@ -877,14 +953,33 @@ async fn request_access_token(refresh_token: String) -> Result { .body(form) .send() .await - .map_err(|e| format!("Antigravity token refresh failed: {e}"))?; - if !response.status().is_success() { - return Err("Antigravity token refresh rejected. Re-login in Antigravity.".to_string()); - } - response - .json() - .await - .map_err(|e| format!("decode refresh response: {e}")) + .map_err(|error| { + ProviderFetchFailure::from_send_error( + "Antigravity token refresh failed. Retrying automatically.", + Some(attempt_binding.clone()), + &error, + ) + })?; + let status = response.status().as_u16(); + let body = read_response_body(status, false, || async { + response.text().await.map_err(|error| { + TransportErrorFacts::from_reqwest(&error, TransportPhase::ResponseBody) + }) + }) + .await + .map_err(|failure| match failure { + ResponseReadFailure::Transient(diagnostic) => ProviderFetchFailure::transient( + "Antigravity token refresh failed. Retrying automatically.", + Some(attempt_binding), + diagnostic, + ), + ResponseReadFailure::Terminal(_) => ProviderFetchFailure::terminal( + "Antigravity token refresh was rejected. Re-login in Antigravity.", + ), + })?; + serde_json::from_str(&body).map_err(|_| { + ProviderFetchFailure::terminal("Antigravity token refresh response could not be decoded.") + }) } async fn refresh_access_token_with( @@ -894,77 +989,128 @@ async fn refresh_access_token_with( request: Request, save: Save, mut checkpoint: Checkpoint, -) -> Result<(Value, String, Result), String> +) -> Result<(Value, String, AccountScope, Option), ProviderFetchFailure> where R: RefreshScopeTransaction + ?Sized, - Request: FnOnce(String) -> RequestFuture, - RequestFuture: std::future::Future>, + Request: FnOnce(String, ProviderCacheBinding) -> RequestFuture, + RequestFuture: std::future::Future>, Save: FnOnce(&Value) -> std::io::Result<()>, - Checkpoint: FnMut(RefreshCheckpoint) -> Result<(), String>, + Checkpoint: FnMut(RefreshCheckpoint) -> Result<(), ProviderFetchFailure>, { - let mut creds = load_remote_credentials(creds_path)?; + let creds = load_remote_credentials(creds_path).map_err(|_| { + ProviderFetchFailure::terminal("Antigravity credentials could not be reloaded.") + })?; checkpoint(RefreshCheckpoint::Reloaded)?; - let location = remote_scope_location(creds_path) - .map_err(|_| "Antigravity auth location cannot be scoped safely.".to_string())?; - if !remote_credentials_need_refresh(&creds, Utc::now()) { - let access_token = remote_access_token(&creds)?; - let scope = match remote_refresh_marker(&creds) { - Some(marker) => refresh.resolve_current("google-oauth-creds", &location, marker), - None => Err(AccountScopeError::NoTrustedEvidence), - }; - return Ok((creds, access_token, scope)); - } - + let location = remote_scope_location(creds_path).map_err(|_| { + ProviderFetchFailure::terminal("Antigravity auth location could not be verified.") + })?; let old_marker = remote_refresh_marker(&creds) - .ok_or_else(|| "Antigravity access token expired and no refresh token".to_string())? + .ok_or_else(|| { + ProviderFetchFailure::terminal("Antigravity credential has no trusted refresh marker.") + })? .to_vec(); + let pre_scope = refresh + .resolve_current("google-oauth-creds", &location, &old_marker) + .map_err(|_| { + ProviderFetchFailure::terminal("Antigravity account identity could not be verified.") + })?; + let pre_binding = ProviderCacheBinding::primary(pre_scope.clone()); + if !remote_credentials_need_refresh(&creds, now) { + let access_token = remote_access_token(&creds).map_err(|_| { + ProviderFetchFailure::terminal("Antigravity credentials have no access token.") + })?; + return Ok((creds, access_token, pre_scope, Some(pre_binding))); + } + let refresh_token = std::str::from_utf8(&old_marker) - .map_err(|_| "Antigravity refresh credential is not valid text.".to_string())? + .map_err(|_| ProviderFetchFailure::terminal("Antigravity refresh credential is invalid."))? .to_string(); - let json = request(refresh_token).await?; + let json = request(refresh_token, pre_binding).await?; checkpoint(RefreshCheckpoint::NetworkReturned)?; let access_token = json .get("access_token") .and_then(Value::as_str) - .ok_or_else(|| "refresh response missing access_token".to_string())? + .map(str::trim) + .filter(|token| !token.is_empty()) + .ok_or_else(|| { + ProviderFetchFailure::terminal( + "Antigravity token refresh response had no access token.", + ) + })? .to_string(); - if let Some(obj) = creds.as_object_mut() { - obj.insert("access_token".into(), Value::String(access_token.clone())); - if let Some(expires_in) = json.get("expires_in").and_then(Value::as_f64) { - let expiry = now.timestamp_millis() as f64 + expires_in * 1000.0; - obj.insert("expiry_date".into(), json!(expiry)); - } - if let Some(id_token) = json.get("id_token").and_then(Value::as_str) { - obj.insert("id_token".into(), Value::String(id_token.to_string())); - } - if let Some(replacement) = json - .get("refresh_token") - .and_then(Value::as_str) - .map(str::trim) - .filter(|token| !token.is_empty()) - { - obj.insert( - "refresh_token".into(), - Value::String(replacement.to_string()), - ); - } + // The provider refresh lock serializes TokenBar writers, but the credential + // file has no cross-process compare-and-swap. Re-reading closes the network + // wait race; an external writer can still race this check and atomic rename. + let mut current_creds = load_remote_credentials(creds_path).map_err(|_| { + ProviderFetchFailure::terminal( + "Antigravity credentials changed during refresh; refusing stale write-back.", + ) + })?; + let current_marker = remote_refresh_marker(¤t_creds).ok_or_else(|| { + ProviderFetchFailure::terminal( + "Antigravity credentials changed during refresh; refusing stale write-back.", + ) + })?; + if current_marker != old_marker.as_slice() { + return Err(ProviderFetchFailure::terminal( + "Antigravity credentials changed during refresh; refusing stale write-back.", + )); } - let new_marker = remote_refresh_marker(&creds); - let marker_rotated = new_marker.is_some_and(|marker| marker != old_marker.as_slice()); - let scope = match new_marker { - Some(new_marker) => { - refresh.transfer("google-oauth-creds", &location, &old_marker, new_marker) - } - None => Err(AccountScopeError::NoTrustedEvidence), - }; - checkpoint(RefreshCheckpoint::MetadataHandled)?; - if marker_rotated && scope.is_err() { - return Ok((creds, access_token, scope)); + + let obj = current_creds.as_object_mut().ok_or_else(|| { + ProviderFetchFailure::terminal( + "Antigravity credentials changed during refresh; refusing stale write-back.", + ) + })?; + obj.insert("access_token".into(), Value::String(access_token.clone())); + if let Some(expires_in) = json.get("expires_in").and_then(Value::as_f64) { + let expiry = now.timestamp_millis() as f64 + expires_in * 1000.0; + obj.insert("expiry_date".into(), json!(expiry)); + } + if let Some(id_token) = json.get("id_token").and_then(Value::as_str) { + obj.insert("id_token".into(), Value::String(id_token.to_string())); + } + if let Some(replacement) = json + .get("refresh_token") + .and_then(Value::as_str) + .map(str::trim) + .filter(|token| !token.is_empty()) + { + obj.insert( + "refresh_token".into(), + Value::String(replacement.to_string()), + ); } - let _ = save(&creds); + let new_marker = remote_refresh_marker(¤t_creds) + .ok_or_else(|| { + ProviderFetchFailure::terminal( + "Antigravity refreshed credential has no trusted marker.", + ) + })? + .to_vec(); + let account_scope = refresh + .transfer("google-oauth-creds", &location, &old_marker, &new_marker) + .map_err(|_| { + ProviderFetchFailure::terminal("Antigravity credential lineage could not be preserved.") + })?; + checkpoint(RefreshCheckpoint::MetadataHandled)?; + let persisted = save(¤t_creds).is_ok(); checkpoint(RefreshCheckpoint::CredentialsPersisted)?; - Ok((creds, access_token, scope)) + let cache_binding = if persisted { + Some(ProviderCacheBinding::primary( + refresh + .resolve_current("google-oauth-creds", &location, &new_marker) + .map_err(|_| { + ProviderFetchFailure::terminal( + "Antigravity account identity could not be verified after refresh.", + ) + })?, + )) + } else { + None + }; + Ok((current_creds, access_token, account_scope, cache_binding)) } fn write_creds_atomic(path: &Path, creds: &Value) -> std::io::Result<()> { @@ -1000,22 +1146,31 @@ fn write_creds_atomic(path: &Path, creds: &Value) -> std::io::Result<()> { let _ = std::fs::remove_file(&tmp); return Err(error); } - if let Err(error) = std::fs::rename(&tmp, path) { + if let Err(error) = tokscale_core::fs_atomic::replace_file(&tmp, path) { let _ = std::fs::remove_file(&tmp); return Err(error); } #[cfg(unix)] - std::fs::File::open(directory)?.sync_all()?; + { + let _ = std::fs::File::open(directory).and_then(|dir| dir.sync_all()); + } Ok(()) } +enum CodeAssistPostFailure { + Forbidden, + Failure(ProviderFetchFailure), +} + async fn code_assist_post( client: &reqwest::Client, - method: &str, + method: &'static str, body: &Value, access_token: &str, -) -> Result { - let resp = client + attempt_binding: Option, + forbidden_is_route_miss: bool, +) -> Result { + let response = client .post(format!("{CODE_ASSIST_BASE}:{method}")) .bearer_auth(access_token) .header(reqwest::header::USER_AGENT, "antigravity") @@ -1023,44 +1178,113 @@ async fn code_assist_post( .json(body) .send() .await - .map_err(|e| format!("Antigravity {method} request failed: {e}"))?; - let status = resp.status(); - if status == reqwest::StatusCode::UNAUTHORIZED { - return Err("Antigravity Google auth expired. Re-login in Antigravity.".to_string()); - } - if status == reqwest::StatusCode::FORBIDDEN { - return Err(format!("Antigravity {method} permission denied")); - } - if !status.is_success() { - return Err(format!("Antigravity {method} returned {}", status.as_u16())); + .map_err(|error| { + CodeAssistPostFailure::Failure(ProviderFetchFailure::from_send_error( + format!("Antigravity {method} request failed. Retrying automatically."), + attempt_binding.clone(), + &error, + )) + })?; + let status = response.status().as_u16(); + if status == 403 && forbidden_is_route_miss { + return Err(CodeAssistPostFailure::Forbidden); } - resp.text() - .await - .map_err(|e| format!("read Antigravity {method}: {e}")) + read_response_body(status, false, || async { + response.text().await.map_err(|error| { + TransportErrorFacts::from_reqwest(&error, TransportPhase::ResponseBody) + }) + }) + .await + .map_err(|failure| { + CodeAssistPostFailure::Failure(match failure { + ResponseReadFailure::Transient(diagnostic) => ProviderFetchFailure::transient( + format!("Antigravity {method} request failed. Retrying automatically."), + attempt_binding, + diagnostic, + ), + ResponseReadFailure::Terminal(401) => ProviderFetchFailure::terminal( + "Antigravity Google auth expired. Re-login in Antigravity.", + ), + ResponseReadFailure::Terminal(403) => ProviderFetchFailure::terminal(format!( + "Antigravity {method} permission was denied." + )), + ResponseReadFailure::Terminal(status) => ProviderFetchFailure::terminal(format!( + "Antigravity {method} rejected the request (status {status})." + )), + }) + }) } -async fn fetch_model_quotas( - client: &reqwest::Client, - access_token: &str, - project: Option<&str>, +async fn fetch_available_models( + context: &RemoteContext, now: DateTime, -) -> Result, String> { - let body = match project { - Some(p) => json!({ "project": p }), +) -> PrimaryQuotaAttempt { + let body = match context.project.as_deref() { + Some(project) => json!({ "project": project }), None => json!({}), }; - // Primary: fetchAvailableModels (per-model quotaInfo). Fall back to - // retrieveUserQuota buckets if the catalog endpoint is denied. - match code_assist_post(client, "fetchAvailableModels", &body, access_token) - .await - .and_then(|body| models_from_available(&body, now)) + let response = match code_assist_post( + &context.client, + "fetchAvailableModels", + &body, + &context.access_token, + context.cache_binding.clone(), + true, + ) + .await { - Ok(windows) if !windows.is_empty() => Ok(windows), - _ => { - let quota = code_assist_post(client, "retrieveUserQuota", &body, access_token).await?; - buckets_from_quota("a, now) + Ok(response) => response, + Err(CodeAssistPostFailure::Forbidden) => return PrimaryQuotaAttempt::Forbidden, + Err(CodeAssistPostFailure::Failure(failure @ ProviderFetchFailure::Transient { .. })) => { + let _ = failure; + return PrimaryQuotaAttempt::Transient; + } + Err(CodeAssistPostFailure::Failure(failure)) => { + return PrimaryQuotaAttempt::Terminal(failure); + } + }; + match models_from_available(&response, now) { + Ok(windows) if !windows.is_empty() => PrimaryQuotaAttempt::Success(windows), + Ok(_) | Err(_) => PrimaryQuotaAttempt::SchemaContradiction(ProviderFetchFailure::terminal( + "Antigravity fetchAvailableModels returned no usable quota windows.", + )), + } +} + +async fn fetch_user_quota( + context: &RemoteContext, + now: DateTime, +) -> Result, ProviderFetchFailure> { + let body = match context.project.as_deref() { + Some(project) => json!({ "project": project }), + None => json!({}), + }; + let response = code_assist_post( + &context.client, + "retrieveUserQuota", + &body, + &context.access_token, + context.cache_binding.clone(), + false, + ) + .await + .map_err(|failure| match failure { + CodeAssistPostFailure::Forbidden => { + ProviderFetchFailure::terminal("Antigravity retrieveUserQuota permission was denied.") } + CodeAssistPostFailure::Failure(failure) => failure, + })?; + let windows = buckets_from_quota(&response, now).map_err(|_| { + ProviderFetchFailure::terminal( + "Antigravity retrieveUserQuota response could not be decoded.", + ) + })?; + if windows.is_empty() { + return Err(ProviderFetchFailure::terminal( + "Antigravity retrieveUserQuota returned no usable quota windows.", + )); } + Ok(windows) } fn project_id(code_assist: &Value) -> Option { @@ -1117,9 +1341,6 @@ fn models_from_available(body: &str, now: DateTime) -> Result) -> Result) -> Result, fraction: f64, reset: Option>, - reset_was_supplied: bool, source_index: usize, } @@ -1214,9 +1432,6 @@ fn buckets_from_quota(body: &str, now: DateTime) -> Result let Some(fraction) = bucket.get("remainingFraction").and_then(Value::as_f64) else { continue; }; - let reset_was_supplied = bucket - .get("resetTime") - .is_some_and(|value| !value.is_null()); let reset = bucket .get("resetTime") .and_then(Value::as_str) @@ -1231,7 +1446,6 @@ fn buckets_from_quota(body: &str, now: DateTime) -> Result model_id: model_id.clone(), fraction, reset, - reset_was_supplied, source_index, }; let Some(model_id) = model_id else { @@ -1270,7 +1484,6 @@ fn buckets_from_quota(body: &str, now: DateTime) -> Result label, candidate.fraction, candidate.reset, - candidate.reset_was_supplied, now, card_id, window_key, @@ -1312,11 +1525,10 @@ fn binding_candidate_is_better( (false, true) => return false, _ => {} } - if candidate_fraction < current_fraction { - return true; - } - if candidate_fraction > current_fraction { - return false; + match candidate_fraction.total_cmp(¤t_fraction) { + std::cmp::Ordering::Less => return true, + std::cmp::Ordering::Greater => return false, + std::cmp::Ordering::Equal => {} } let candidate_reset = candidate_reset.filter(|reset| *reset > now); let current_reset = current_reset.filter(|reset| *reset > now); @@ -1359,7 +1571,6 @@ fn discover_client_from_app() -> Option<(String, String)> { None } -#[cfg(target_os = "macos")] fn client_artifact_candidates() -> Vec { let relative = [ "Contents/Resources/bin/language_server", @@ -1379,39 +1590,6 @@ fn client_artifact_candidates() -> Vec { .collect() } -#[cfg(any(windows, test))] -fn windows_client_artifact_candidates( - local_app_data: Option, - home: Option, -) -> Vec { - const RELATIVE: &str = "Programs/antigravity/resources/bin/language_server.exe"; - - let mut candidates = Vec::new(); - if let Some(root) = local_app_data.filter(|path| !path.as_os_str().is_empty()) { - candidates.push(root.join(RELATIVE)); - } - if let Some(home) = home.filter(|path| !path.as_os_str().is_empty()) { - let fallback = home.join("AppData/Local").join(RELATIVE); - if candidates.last() != Some(&fallback) { - candidates.push(fallback); - } - } - candidates -} - -#[cfg(windows)] -fn client_artifact_candidates() -> Vec { - windows_client_artifact_candidates( - std::env::var_os("LOCALAPPDATA").map(PathBuf::from), - crate::user_home_dir(), - ) -} - -#[cfg(not(any(target_os = "macos", windows)))] -fn client_artifact_candidates() -> Vec { - Vec::new() -} - fn is_token_byte(b: u8) -> bool { b.is_ascii_alphanumeric() || b == b'-' || b == b'_' } @@ -1428,26 +1606,17 @@ fn scan_client_ids(data: &[u8]) -> Vec { let mut out: Vec = Vec::new(); let mut i = 0; while let Some(pos) = find_sub(&data[i..], suffix) { - let suffix_start = i + pos; - let end = suffix_start + suffix.len(); - let mut span_start = suffix_start; - while span_start > 0 && is_token_byte(data[span_start - 1]) { - span_start -= 1; + let end = i + pos + suffix.len(); + let mut start = i + pos; + while start > 0 && is_token_byte(data[start - 1]) { + start -= 1; } - for start in span_start..suffix_start { - if !data[start].is_ascii_digit() || (start > 0 && data[start - 1].is_ascii_digit()) { - continue; - } - if let Ok(candidate) = std::str::from_utf8(&data[start..end]) { - if valid_client_id(candidate) { - if !out.contains(&candidate.to_string()) { - out.push(candidate.to_string()); - } - break; - } + if let Ok(candidate) = std::str::from_utf8(&data[start..end]) { + if valid_client_id(candidate) && !out.contains(&candidate.to_string()) { + out.push(candidate.to_string()); } } - i = end; + i = i + pos + suffix.len(); } out } @@ -1500,27 +1669,66 @@ fn preferred_client(ids: &[String], secrets: &[String]) -> Option<(String, Strin // ── shared ──────────────────────────────────────────────────────────────────── fn gemini_home() -> Option { - crate::user_home_dir().map(|home| home.join(".gemini")) + gemini_home_from( + std::env::var("GEMINI_CLI_HOME"), + crate::user_home_dir().as_deref(), + ) +} + +fn gemini_home_from( + gemini_cli_home: Result, + user_home: Option<&Path>, +) -> Option { + let root = match gemini_cli_home { + Ok(root) if !root.trim().is_empty() => root, + Ok(_) | Err(_) => format!("{}/.gemini", user_home?.to_string_lossy()), + }; + Some(PathBuf::from(root)) } #[cfg(test)] mod tests { use super::*; use crate::agent_account_scope::test_support::TestRefreshScope; - use sha2::{Digest as _, Sha256}; - fn fixture_authoritative_scope(email: &str) -> Result<[u8; 32], AccountScopeError> { - let mut digest = Sha256::new(); - digest.update(b"antigravity\0email\0"); - digest.update(email.trim().to_ascii_lowercase().as_bytes()); - Ok(digest.finalize().into()) + #[test] + fn gemini_home_uses_nonempty_configured_root_unchanged() { + let configured = " /tmp/gemini-cli-home "; + assert_eq!( + gemini_home_from(Ok(configured.to_string()), None), + Some(PathBuf::from(configured)) + ); + } + + #[test] + fn gemini_home_falls_back_on_environment_errors() { + let home = PathBuf::from("resolved-home"); + let fallback = Some(PathBuf::from("resolved-home/.gemini")); + for error in [ + std::env::VarError::NotPresent, + std::env::VarError::NotUnicode(std::ffi::OsString::new()), + ] { + assert_eq!(gemini_home_from(Err(error), Some(&home)), fallback); + } + } + + #[test] + fn gemini_home_falls_back_for_trim_empty_root() { + let home = PathBuf::from("resolved-home"); + assert_eq!( + gemini_home_from(Ok(" \t\n ".to_string()), Some(&home)), + Some(PathBuf::from("resolved-home/.gemini")) + ); } #[test] fn extracts_flags_both_forms() { let cmd = "/x/language_server --app_data_dir /Users/me/.gemini/antigravity --csrf_token=ABC123 --extension_server_port 4567"; assert_eq!(extract_flag(cmd, "--csrf_token").as_deref(), Some("ABC123")); - assert_eq!(extract_flag(cmd, "--extension_server_port").as_deref(), Some("4567")); + assert_eq!( + extract_flag(cmd, "--extension_server_port").as_deref(), + Some("4567") + ); assert!(is_language_server(&cmd.to_lowercase())); assert!(is_antigravity(&cmd.to_lowercase())); } @@ -1537,42 +1745,36 @@ mod tests { format!("{WINDOWS_DISCOVERY_PREFIX}\tP\t{pid}\t{encoded}\t{ports}") } - #[test] - fn parses_windows_discovery_fixture_and_ignores_malformed_rows() { - let csrf = "fixture-primary-secret"; - let extension_csrf = "fixture-extension-secret"; - let command = format!( - r#""C:\Program Files\Antigravity\language_server.exe" --app_data_dir antigravity --csrf_token={csrf} --extension_server_port 4567 --extension_server_csrf_token={extension_csrf}"# - ); - let valid = windows_process_fixture(4242, &command, "61234,54321,61234"); - let fixture = format!( - "garbage\n{WINDOWS_DISCOVERY_PREFIX}\tP\tnot-a-pid\tbad-base64\t80\n{valid}\r\n" - ); + fn windows_discovery_error(stdout: &str) -> WindowsDiscoveryError { + match parse_windows_discovery(stdout) { + Ok(_) => panic!("Windows discovery fixture unexpectedly succeeded"), + Err(error) => error, + } + } - let processes = match parse_windows_discovery(&fixture) { - Ok(processes) => processes, - Err(_) => panic!("valid Windows discovery fixture was rejected"), - }; - assert_eq!(processes.len(), 1); - let (proc, ports) = &processes[0]; - assert_eq!(proc.pid, 4242); - assert!(proc.csrf_token == csrf); - assert_eq!(proc.extension_port, Some(4567)); - assert!(proc.extension_csrf.as_deref() == Some(extension_csrf)); - assert_eq!(ports.as_slice(), &[54321, 61234]); + #[test] + fn windows_discovery_queries_language_server_executable_family() { + assert!(WINDOWS_DISCOVERY_SCRIPT.contains( + r#"Get-CimInstance -ClassName Win32_Process -Filter "Name LIKE 'language_server%.exe'""# + )); } #[test] - fn retains_all_windows_processes_in_probe_order() { - let first_csrf = "first-fixture-secret"; - let first_extension_csrf = "first-extension-fixture-secret"; - let second_csrf = "second-fixture-secret"; + fn windows_discovery_preserves_per_process_candidate_binding() { + let first_csrf = "first-primary-secret"; + let first_extension_csrf = "first-extension-secret"; + let second_csrf = "second-primary-secret"; let first = windows_process_fixture( 1001, &format!( - r#"C:\Antigravity\language_server.exe --app_data_dir antigravity --csrf_token={first_csrf} --extension_server_port=41999 --extension_server_csrf_token={first_extension_csrf}"# + r#""C:\Program Files\Antigravity\language_server_windows_x64.exe" --app_data_dir="C:\Users\me\AppData\Roaming\Antigravity" --csrf_token={first_csrf} --extension_server_port=41999 --extension_server_csrf_token={first_extension_csrf}"# ), - "41002,41001", + "41002,41001,41002", + ); + let ignored = windows_process_fixture( + 1009, + r#"C:\Other\language_server.exe --csrf_token=decoy-secret"#, + "49999", ); let second = windows_process_fixture( 1002, @@ -1581,11 +1783,13 @@ mod tests { ), "42000", ); - let fixture = format!("{first}\n{second}\n"); + let fixture = format!( + "garbage\n{first}\n{ignored}\n{WINDOWS_DISCOVERY_PREFIX}\tP\tnot-a-pid\tnot-base64\t80\n{second}\r\n" + ); let processes = match parse_windows_discovery(&fixture) { Ok(processes) => processes, - Err(_) => panic!("multi-process Windows discovery fixture was rejected"), + Err(_) => panic!("valid Windows discovery fixture was rejected"), }; assert_eq!(processes.len(), 2); assert_eq!(processes[0].0.pid, 1001); @@ -1593,81 +1797,69 @@ mod tests { assert_eq!(processes[1].0.pid, 1002); assert_eq!(processes[1].1.as_slice(), &[42000]); - let candidates = local_api_candidates(processes); - let candidate_ports: Vec = candidates.iter().map(|candidate| candidate.0).collect(); - assert_eq!(candidate_ports, vec![41001, 41002, 41999, 41999, 42000]); - assert!(candidates[0].1.as_str() == first_csrf); - assert!(candidates[2].1.as_str() == first_extension_csrf); - assert!(candidates[3].1.as_str() == first_csrf); - assert!(candidates[4].1.as_str() == second_csrf); - } - - #[test] - fn classifies_windows_discovery_failures_without_exposing_token() { - assert!(matches!( - parse_windows_discovery(&format!("{WINDOWS_DISCOVERY_PREFIX}\tN\n")), - Err(WindowsDiscoveryError::ProcessNotFound) - )); - - let no_token = windows_process_fixture( - 4242, - r#"C:\Antigravity\language_server.exe --app_data_dir antigravity"#, - "54321", - ); - assert!(matches!( - parse_windows_discovery(&no_token), - Err(WindowsDiscoveryError::TokenMissing) - )); - - let csrf = "fixture-secret-that-must-not-leak"; - let no_ports = windows_process_fixture( - 4242, - &format!( - r#"C:\Antigravity\language_server.exe --app_data_dir antigravity --csrf_token={csrf}"# - ), - "", + assert_eq!( + local_api_candidates(processes), + vec![ + (41001, first_csrf.to_string()), + (41002, first_csrf.to_string()), + (41999, first_extension_csrf.to_string()), + (41999, first_csrf.to_string()), + (42000, second_csrf.to_string()), + ] ); - let error = match parse_windows_discovery(&no_ports) { - Err(error) => error, - Ok(_) => panic!("missing ports should fail discovery"), - }; - assert!(matches!(error, WindowsDiscoveryError::PortsMissing)); - assert!(!error.message().contains(csrf)); - - assert!(matches!( - parse_windows_discovery(&format!( - "{WINDOWS_DISCOVERY_PREFIX}\tP\t4242\tnot-base64\t54321\n" - )), - Err(WindowsDiscoveryError::MalformedOutput) - )); - assert!(!WindowsDiscoveryError::PowerShellFailed - .message() - .contains(csrf)); } #[test] - fn maps_windows_oauth_client_artifact_roots_in_order_without_duplicates() { - const RELATIVE: &str = "Programs/antigravity/resources/bin/language_server.exe"; - let home = PathBuf::from("/Users/example"); - let fallback = home.join("AppData/Local").join(RELATIVE); - let local_app_data = PathBuf::from("/redirected/local"); + fn windows_discovery_rejects_malformed_rows_without_exposing_secrets() { + let sentinel = "sentinel-secret-that-must-not-leak"; + let command = format!( + r#"C:\Antigravity\language_server.exe --app_data_dir antigravity --csrf_token={sentinel}"# + ); + let encoded = base64::engine::general_purpose::STANDARD.encode(&command); + for malformed in [ + "garbage".to_string(), + format!("{WINDOWS_DISCOVERY_PREFIX}\tP\tnot-a-pid\t{encoded}\t54321"), + format!("{WINDOWS_DISCOVERY_PREFIX}\tP\t4242\tnot-base64\t54321"), + windows_process_fixture(4242, &command, "not-a-port"), + windows_process_fixture(4242, &command, "0"), + windows_process_fixture(4242, &command, "65536"), + ] { + assert_eq!( + windows_discovery_error(&malformed), + WindowsDiscoveryError::MalformedOutput + ); + } assert_eq!( - windows_client_artifact_candidates(Some(local_app_data.clone()), Some(home.clone())), - vec![local_app_data.join(RELATIVE), fallback.clone()] + windows_discovery_error(&format!("{WINDOWS_DISCOVERY_PREFIX}\tN\n")), + WindowsDiscoveryError::ProcessNotFound ); assert_eq!( - windows_client_artifact_candidates( - Some(home.join("AppData/Local")), - Some(home.clone()), - ), - vec![fallback.clone()] + windows_discovery_error(&windows_process_fixture( + 4242, + r#"C:\Other\language_server.exe --csrf_token=decoy-secret"#, + "54321", + )), + WindowsDiscoveryError::ProcessNotFound ); assert_eq!( - windows_client_artifact_candidates(None, Some(home)), - vec![fallback] + windows_discovery_error(&windows_process_fixture( + 4242, + r#"C:\Antigravity\language_server.exe --app_data_dir antigravity"#, + "54321", + )), + WindowsDiscoveryError::TokenMissing ); - assert!(windows_client_artifact_candidates(None, None).is_empty()); + + let error = windows_discovery_error(&windows_process_fixture(4242, &command, "")); + assert_eq!(error, WindowsDiscoveryError::PortsMissing); + let display = error.message(); + let debug = format!("{error:?}"); + assert!(!display.contains(sentinel)); + assert!(!debug.contains(sentinel)); + assert!(!WindowsDiscoveryError::PowerShellFailed + .message() + .contains(sentinel)); } #[test] @@ -1675,41 +1867,27 @@ mod tests { let blob = b"junk\x00123-abcDEF_g.apps.googleusercontent.com\x00\x00GOCSPX-abcdefghijklmnopqrstuvwxyz12\x00tail"; let ids = scan_client_ids(blob); let secrets = scan_client_secrets(blob); - assert_eq!(ids, vec!["123-abcDEF_g.apps.googleusercontent.com".to_string()]); + assert_eq!( + ids, + vec!["123-abcDEF_g.apps.googleusercontent.com".to_string()] + ); assert_eq!(secrets.len(), 1); let client = preferred_client(&ids, &secrets).unwrap(); assert_eq!(client.0, "123-abcDEF_g.apps.googleusercontent.com"); assert!(client.1.starts_with("GOCSPX-")); } - #[test] - fn scans_client_id_attached_to_alphanumeric_prefix() { - let blob = b"prefix7x123456789012-attached_ID.apps.googleusercontent.com\x00tail"; - - assert_eq!( - scan_client_ids(blob), - vec!["123456789012-attached_ID.apps.googleusercontent.com".to_string()] - ); - } - - #[test] - fn scans_multiple_client_ids_in_order_without_duplicates() { - let blob = b"123456789012-first.apps.googleusercontent.com\x001234567890123-second.apps.googleusercontent.com\x00123456789012-first.apps.googleusercontent.com"; - - assert_eq!( - scan_client_ids(blob), - vec![ - "123456789012-first.apps.googleusercontent.com".to_string(), - "1234567890123-second.apps.googleusercontent.com".to_string(), - ] - ); - } - #[test] fn prefers_last_id_when_single_secret() { - let ids = vec!["1-a.apps.googleusercontent.com".into(), "2-b.apps.googleusercontent.com".into()]; + let ids = vec![ + "1-a.apps.googleusercontent.com".into(), + "2-b.apps.googleusercontent.com".into(), + ]; let secrets = vec!["GOCSPX-only".into()]; - assert_eq!(preferred_client(&ids, &secrets).unwrap().0, "2-b.apps.googleusercontent.com"); + assert_eq!( + preferred_client(&ids, &secrets).unwrap().0, + "2-b.apps.googleusercontent.com" + ); } #[test] @@ -1730,100 +1908,17 @@ mod tests { }"#; let fetched = parse_user_status(body, now).unwrap(); assert_eq!(fetched.source, "cli"); - assert_eq!(fetched.identity.as_ref().unwrap().email.as_deref(), Some("me@gmail.com")); - assert_eq!(fetched.identity.as_ref().unwrap().plan.as_deref(), Some("Pro")); assert_eq!( - fetched.account_scope, - Err(AccountScopeError::NoTrustedEvidence) + fetched.identity.as_ref().unwrap().email.as_deref(), + Some("me@gmail.com") ); - assert_eq!(fetched.windows.len(), 1); - assert_eq!(fetched.windows[0].label_for_test(), "Gemini 3 Pro"); - assert_eq!(fetched.windows[0].card_id_for_test(), "model.gemini-3-pro.v1"); - assert_eq!( - fetched.windows[0].pace_window_key_for_test(), - Some("model.gemini-3-pro.v1") - ); - assert!((fetched.windows[0].remaining_for_test() - 42.0).abs() < 0.01); - } - - #[test] - fn identity_only_local_response_is_not_accepted_for_scope_resolution() { - let fetched = parse_user_status( - r#"{ - "userStatus": { - "email": "user@example.com", - "cascadeModelConfigData": { "clientModelConfigs": [] } - } - }"#, - Utc::now(), - ) - .unwrap(); - - assert!(fetched.identity.is_some()); - assert!(fetched.windows.is_empty()); - assert!(!local_result_is_acceptable(&fetched)); - } - - #[test] - fn local_scope_resolver_uses_only_exact_non_empty_email() { - let identity = AgentIdentity { - email: Some("Exact.Email+tag@Example.COM".to_string()), - plan: None, - }; assert_eq!( - resolve_local_account_scope_with(Some(&identity), |email| { - assert_eq!(email, "Exact.Email+tag@Example.COM"); - Ok("resolved") - }), - Ok("resolved") + fetched.identity.as_ref().unwrap().plan.as_deref(), + Some("Pro") ); - - let calls = std::cell::Cell::new(0); - for email in [None, Some(" \t ")] { - let identity = AgentIdentity { - email: email.map(str::to_string), - plan: None, - }; - assert_eq!( - resolve_local_account_scope_with(Some(&identity), |_| { - calls.set(calls.get() + 1); - Ok(()) - }), - Err(AccountScopeError::NoTrustedEvidence) - ); - } - assert_eq!(calls.get(), 0); - - let identity = AgentIdentity { - email: Some("failure@example.com".to_string()), - plan: None, - }; - assert_eq!( - resolve_local_account_scope_with(Some(&identity), |email| { - assert_eq!(email, "failure@example.com"); - Err::<(), _>(AccountScopeError::MetadataWrite) - }), - Err(AccountScopeError::MetadataWrite) - ); - } - - #[test] - fn local_config_without_model_has_row_identity() { - let now = Utc::now(); - let body = r#"{ - "userStatus": { - "cascadeModelConfigData": { - "clientModelConfigs": [ - { "label": "Unnamed", "quotaInfo": { "remainingFraction": 0.5 } } - ] - } - } - }"#; - let fetched = parse_user_status(body, now).unwrap(); assert_eq!(fetched.windows.len(), 1); - assert_eq!(fetched.windows[0].card_id_for_test(), "row.cli.config.0.v1"); - assert!(fetched.windows[0].pace_window_key_for_test().is_none()); - assert_eq!(fetched.windows[0].pace_reason_for_test(), Some("windowIdentity")); + assert_eq!(fetched.windows[0].label_for_test(), "Gemini 3 Pro"); + assert!((fetched.windows[0].remaining_for_test() - 42.0).abs() < 0.01); } #[test] @@ -1836,7 +1931,6 @@ mod tests { }); let w = models_from_available(&models.to_string(), now).unwrap(); assert_eq!(w.len(), 1); - assert_eq!(w[0].card_id_for_test(), "model.gemini-3-pro.v1"); let quota = json!({ "buckets": [ @@ -1850,7 +1944,7 @@ mod tests { } #[test] - fn stage3c3a_identity_and_duplicate_rules_are_deterministic() { + fn stage4_antigravity_identity_and_duplicate_rules_are_deterministic() { let now = DateTime::parse_from_rfc3339("2026-07-10T00:00:00Z") .unwrap() .with_timezone(&Utc); @@ -1976,49 +2070,43 @@ mod tests { }); let selected = buckets_from_quota(&duplicate.to_string(), now).unwrap(); assert_eq!(selected.len(), 1); - let selected_wire = serde_json::to_value(&selected[0]).unwrap(); assert_eq!( - selected_wire["resetsAt"], - "2026-07-11T00:00:00.000Z", + selected[0].resets_at_for_test(), + Some("2026-07-11T00:00:00.000Z"), "future reset beats past reset, then earliest future reset wins" ); let same_reset = parse_datetime("2026-07-11T00:00:00Z"); assert!(!binding_candidate_is_better( 0.25, same_reset, 1, 0.25, same_reset, 0, now )); - - let signed_zero = r#"{ - "buckets": [ - { - "modelId": "signed-zero", - "remainingFraction": 0.0, - "resetTime": "2026-07-11T00:00:00Z" - }, - { - "modelId": "signed-zero", - "remainingFraction": -0.0 - } - ] - }"#; - let selected = buckets_from_quota(signed_zero, now).unwrap(); - assert_eq!(selected.len(), 1); - let selected_wire = serde_json::to_value(&selected[0]).unwrap(); - assert_eq!(selected_wire["resetsAt"], "2026-07-11T00:00:00.000Z"); - assert_eq!(selected_wire["paceStatus"]["state"], "learningDuration"); - assert_eq!(selected_wire["paceStatus"]["durationSource"], "observed"); } #[test] - fn stage3c3a_rejects_invalid_fractions_and_isolates_raw_rows() { + fn stage4_antigravity_rejects_invalid_fractions_at_every_source() { + assert!(!valid_remaining_fraction(f64::NAN)); + assert!(!valid_remaining_fraction(f64::INFINITY)); + assert!(!valid_remaining_fraction(-0.01)); + assert!(!valid_remaining_fraction(1.01)); + assert!(valid_remaining_fraction(0.0)); + assert!(valid_remaining_fraction(1.0)); + let now = DateTime::parse_from_rfc3339("2026-07-10T00:00:00Z") .unwrap() .with_timezone(&Utc); - let assert_valid_only = |windows: Vec| { + let assert_rows = |windows: Vec| { assert_eq!(windows.len(), 1); assert_eq!( windows[0].pace_window_key_for_test(), Some("model.valid.v1") ); + for key in ["model.negative.v1", "model.over.v1"] { + assert!( + windows + .iter() + .all(|window| window.pace_window_key_for_test() != Some(key)), + "invalid quota row must not be published: {key}" + ); + } }; let cli = r#"{ @@ -2045,8 +2133,16 @@ mod tests { } } }"#; - assert_valid_only(parse_user_status(cli, now).unwrap().windows); + assert_rows(parse_user_status(cli, now).unwrap().windows); + let assert_overflowing_duplicate = |windows: Vec| { + assert_eq!(windows.len(), 1); + assert_eq!( + windows[0].pace_window_key_for_test(), + Some("model.same-model.v1") + ); + assert!((windows[0].remaining_for_test() - 50.0).abs() < 0.01); + }; let overflowing_cli = r#"{ "userStatus": { "cascadeModelConfigData": { @@ -2063,12 +2159,9 @@ mod tests { } } }"#; - let windows = parse_user_status(overflowing_cli, now).unwrap().windows; - assert_eq!(windows.len(), 1); - assert_eq!(windows[0].pace_window_key_for_test(), Some("model.same-model.v1")); - assert!((windows[0].remaining_for_test() - 50.0).abs() < 0.01); + assert_overflowing_duplicate(parse_user_status(overflowing_cli, now).unwrap().windows); - assert_valid_only( + assert_rows( models_from_available( &json!({ "models": { @@ -2094,12 +2187,9 @@ mod tests { } } }"#; - let windows = models_from_available(overflowing_models, now).unwrap(); - assert_eq!(windows.len(), 1); - assert_eq!(windows[0].pace_window_key_for_test(), Some("model.same-model.v1")); - assert!((windows[0].remaining_for_test() - 50.0).abs() < 0.01); + assert_overflowing_duplicate(models_from_available(overflowing_models, now).unwrap()); - assert_valid_only( + assert_rows( buckets_from_quota( &json!({ "buckets": [ @@ -2121,10 +2211,7 @@ mod tests { { "modelId": "same-model", "remainingFraction": 0.5 } ] }"#; - let windows = buckets_from_quota(overflowing_buckets, now).unwrap(); - assert_eq!(windows.len(), 1); - assert_eq!(windows[0].pace_window_key_for_test(), Some("model.same-model.v1")); - assert!((windows[0].remaining_for_test() - 50.0).abs() < 0.01); + assert_overflowing_duplicate(buckets_from_quota(overflowing_buckets, now).unwrap()); let malformed_cli_row = r#"{ "userStatus": { @@ -2143,9 +2230,7 @@ mod tests { } } }"#; - let windows = parse_user_status(malformed_cli_row, now).unwrap().windows; - assert_eq!(windows.len(), 1); - assert!((windows[0].remaining_for_test() - 50.0).abs() < 0.01); + assert_overflowing_duplicate(parse_user_status(malformed_cli_row, now).unwrap().windows); let malformed_model_row = r#"{ "models": { @@ -2158,9 +2243,7 @@ mod tests { } } }"#; - let windows = models_from_available(malformed_model_row, now).unwrap(); - assert_eq!(windows.len(), 1); - assert!((windows[0].remaining_for_test() - 50.0).abs() < 0.01); + assert_overflowing_duplicate(models_from_available(malformed_model_row, now).unwrap()); let malformed_bucket_fields = r#"{ "buckets": [ @@ -2168,347 +2251,262 @@ mod tests { { "modelId": "valid", "remainingFraction": 0.5 } ] }"#; - let windows = buckets_from_quota(malformed_bucket_fields, now).unwrap(); - assert_eq!(windows.len(), 2); - assert!(windows + let malformed_bucket_windows = buckets_from_quota(malformed_bucket_fields, now).unwrap(); + assert_eq!(malformed_bucket_windows.len(), 2); + assert!(malformed_bucket_windows .iter() .any(|window| window.pace_window_key_for_test() == Some("model.valid.v1"))); - let wire = serde_json::to_value( - windows - .iter() - .find(|window| window.pace_window_key_for_test().is_none()) - .unwrap(), - ) - .unwrap(); - assert_eq!(wire["paceStatus"]["reason"], "windowIdentity"); + let malformed_bucket_wire = serde_json::to_value(&malformed_bucket_windows[0]).unwrap(); + assert_eq!( + malformed_bucket_wire["paceStatus"]["reason"], + "windowIdentity" + ); - assert!(!valid_remaining_fraction(f64::NAN)); - assert!(!valid_remaining_fraction(f64::INFINITY)); - assert!(!valid_remaining_fraction(-0.01)); - assert!(!valid_remaining_fraction(1.01)); - assert!(valid_remaining_fraction(0.0)); - assert!(valid_remaining_fraction(1.0)); + assert!(quota_window( + "Non-finite".to_string(), + f64::NAN, + None, + now, + "model.non-finite.v1".to_string(), + Some("model.non-finite.v1".to_string()), + ) + .is_none()); + assert!(!binding_candidate_is_better( + -0.1, None, 1, 0.5, None, 0, now + )); } #[test] - fn stage3c3a_reset_presence_maps_to_typed_pace_reasons() { - let now = DateTime::parse_from_rfc3339("2026-07-10T00:00:00Z") - .unwrap() - .with_timezone(&Utc); - let body = r#"{ - "userStatus": { - "cascadeModelConfigData": { - "clientModelConfigs": [ - { "modelOrAlias": { "model": "absent" }, "quotaInfo": { "remainingFraction": 0.1 } }, - { "modelOrAlias": { "model": "null" }, "quotaInfo": { "remainingFraction": 0.2, "resetTime": null } }, - { "modelOrAlias": { "model": "malformed" }, "quotaInfo": { "remainingFraction": 0.3, "resetTime": "bogus" } }, - { "modelOrAlias": { "model": "typed" }, "quotaInfo": { "remainingFraction": 0.4, "resetTime": 42 } }, - { "modelOrAlias": { "model": "past" }, "quotaInfo": { "remainingFraction": 0.5, "resetTime": "2026-07-09T00:00:00Z" } }, - { "modelOrAlias": { "model": "future" }, "quotaInfo": { "remainingFraction": 0.6, "resetTime": "2026-07-11T00:00:00Z" } }, - { "label": "Missing identity", "quotaInfo": { "remainingFraction": 0.7, "resetTime": "bogus" } } - ] - } - } - }"#; - let windows = parse_user_status(body, now).unwrap().windows; - assert_eq!(windows.len(), 7); - for (index, expected) in [ - (0, Some("missingReset")), - (1, Some("missingReset")), - (2, Some("invalidEvidence")), - (3, Some("invalidEvidence")), - (4, Some("invalidEvidence")), - (6, Some("windowIdentity")), - ] { - assert_eq!(windows[index].pace_reason_for_test(), expected); - } - assert_eq!(windows[5].pace_reason_for_test(), None); - assert_eq!(windows[5].pace_window_key_for_test(), Some("model.future.v1")); - let future_wire = serde_json::to_value(&windows[5]).unwrap(); - assert_eq!(future_wire["paceStatus"]["state"], "learningDuration"); - assert_eq!(future_wire["paceStatus"]["durationSource"], "observed"); + fn remote_scope_and_presentation_ignore_unbound_active_email() { + let stale_active_email = "stale-other-account@example.com"; + let credentials = json!({ + "access_token": "short-lived-access", + "refresh_token": "bound-google-refresh" + }); + assert_eq!( + remote_refresh_marker(&credentials), + Some(b"bound-google-refresh".as_slice()) + ); + assert_ne!( + remote_refresh_marker(&credentials), + Some(stale_active_email.as_bytes()) + ); + let identity = remote_identity(Some("Paid".to_string())); + assert_eq!(identity.email, None); + assert_eq!(identity.plan.as_deref(), Some("Paid")); - let subsecond_now = DateTime::parse_from_rfc3339("2026-07-10T00:00:00.000500Z") - .unwrap() - .with_timezone(&Utc); - let subsecond = parse_user_status( - r#"{ - "userStatus": { - "cascadeModelConfigData": { - "clientModelConfigs": [{ - "modelOrAlias": { "model": "subsecond" }, - "quotaInfo": { - "remainingFraction": 0.5, - "resetTime": "2026-07-10T00:00:00.000900Z" - } - }] - } - } - }"#, - subsecond_now, - ) - .unwrap(); - let wire = serde_json::to_value(&subsecond.windows[0]).unwrap(); - assert_eq!(wire["paceStatus"]["state"], "learningDuration"); - assert_eq!(wire["paceStatus"]["durationSource"], "observed"); + let access_only = json!({ "access_token": "access-is-not-the-frozen-marker" }); + assert_eq!(remote_refresh_marker(&access_only), None); } #[test] - fn verified_remote_user_info_shares_authoritative_scope_with_local_email() { - let local_identity = AgentIdentity { - email: Some(" Same.User@Example.COM ".to_string()), - plan: None, - }; - let local_scope = - resolve_local_account_scope_with(Some(&local_identity), fixture_authoritative_scope) - .unwrap(); - let remote_scope = resolve_google_user_info_response_with( - reqwest::StatusCode::OK, - r#"{"email":"same.user@example.com","verified_email":true}"#, - fixture_authoritative_scope, - ) - .unwrap(); - let other_remote_scope = resolve_google_user_info_response_with( - reqwest::StatusCode::OK, - r#"{"email":"other.user@example.com","verified_email":true}"#, - fixture_authoritative_scope, + fn local_route_fails_closed_without_authenticated_email() { + let fetched = parse_user_status( + r#"{"userStatus":{"cascadeModelConfigData":{"clientModelConfigs":[]}}}"#, + Utc::now(), ) .unwrap(); + assert_eq!( + fetched.account_scope, + Err(AccountScopeError::NoTrustedEvidence) + ); + } - assert_eq!(remote_scope, local_scope); - assert_ne!(other_remote_scope, local_scope); + #[test] + fn resolves_remote_plan_from_tier() { + assert_eq!( + resolve_remote_plan(&json!({"currentTier":{"id":"free-tier"}})).as_deref(), + Some("Free") + ); + assert_eq!( + resolve_remote_plan(&json!({"planInfo":{"planType":"standard"}})).as_deref(), + Some("Standard") + ); + } + + fn orchestration_fetched(source: &str) -> Fetched { + Fetched { + source: source.to_string(), + identity: None, + account_scope: Err(AccountScopeError::NoTrustedEvidence), + cache_binding: None, + windows: Vec::new(), + } + } + + fn orchestration_transient(display: &str) -> ProviderFetchFailure { + ProviderFetchFailure::transient( + display, + None, + crate::agent_usage::SafeTransportDiagnostic::from_facts( + TransportErrorFacts::synthetic(true, false, TransportPhase::Request, None), + ), + ) } #[tokio::test] - async fn remote_user_info_uses_refreshed_final_token_and_ignores_cached_identity() { - let scope_store = TestRefreshScope::new("antigravity", "antigravity-userinfo-token"); - let lineage_scope = scope_store - .resolve_current( - "google-oauth-creds", - "fixture-location", - b"fixture-refresh-marker", - ) - .unwrap(); - let now = DateTime::parse_from_rfc3339("2026-07-17T00:00:00Z") - .unwrap() - .with_timezone(&Utc); - let credentials = json!({ - "access_token": "fixture-stale-access", - "refresh_token": "fixture-refresh-marker", - "expiry_date": 0, - "id_token": "stale-id-token-must-not-authorize-scope", - "active_email": "stale-active-email@example.com" - }); - let final_token_seen = std::cell::Cell::new(false); - let verified_email_seen = std::cell::Cell::new(false); - let auth = prepare_remote_auth_with( - &credentials, - now, - || { - std::future::ready(Ok(( - json!({"access_token": "fixture-final-access"}), - "fixture-final-access".to_string(), - Ok(lineage_scope), - ))) + async fn orchestration_local_success_and_primary_terminal_do_not_call_later_routes() { + let primary_calls = std::cell::Cell::new(0); + let secondary_calls = std::cell::Cell::new(0); + let local = fetch_with( + || async { LocalAttempt::Success(orchestration_fetched("cli")) }, + || async { + primary_calls.set(primary_calls.get() + 1); + PrimaryAttempt::<()>::FinalFailure(ProviderFetchFailure::terminal("unexpected")) }, - |access_token| { - final_token_seen.set(access_token == "fixture-final-access"); - let scope = resolve_google_user_info_response_with( - reqwest::StatusCode::OK, - r#"{"email":"fresh.user@example.com","verified_email":true}"#, - |email| { - verified_email_seen.set(email == "fresh.user@example.com"); - fixture_authoritative_scope(email) - }, - ); - std::future::ready(scope) + |_: ()| async { + secondary_calls.set(secondary_calls.get() + 1); + Ok(orchestration_fetched("secondary")) }, ) .await .unwrap(); - - assert!(final_token_seen.get()); - assert!(verified_email_seen.get()); - assert!(auth.access_token == "fixture-final-access"); - assert!(auth.account_scope.is_ok()); - scope_store.cleanup(); + assert_eq!(local.source, "cli"); + assert_eq!(primary_calls.get(), 0); + assert_eq!(secondary_calls.get(), 0); + + let failure = fetch_with( + || async { LocalAttempt::RouteMiss }, + || async { + PrimaryAttempt::<()>::FinalFailure(ProviderFetchFailure::terminal("primary 401")) + }, + |_: ()| async { + secondary_calls.set(secondary_calls.get() + 1); + Ok(orchestration_fetched("secondary")) + }, + ) + .await + .unwrap_err(); + assert!(matches!( + failure, + ProviderFetchFailure::Terminal { ref display } if display == "primary 401" + )); + assert_eq!(secondary_calls.get(), 0); } #[tokio::test] - async fn untrusted_user_info_keeps_quota_and_never_falls_back_to_lineage() { - let scope_store = TestRefreshScope::new("antigravity", "antigravity-userinfo-fail"); - let lineage_scope = scope_store - .resolve_current( - "google-oauth-creds", - "fixture-location", - b"fixture-refresh-marker", - ) - .unwrap(); - let now = DateTime::parse_from_rfc3339("2026-07-17T00:00:00Z") - .unwrap() - .with_timezone(&Utc); - let fixtures = [ - ( - reqwest::StatusCode::OK, - r#"{"verified_email":true,"raw":"raw-response-must-not-leak"}"#, - ), - ( - reqwest::StatusCode::OK, - r#"{"email":" \t ","verified_email":true,"raw":"raw-response-must-not-leak"}"#, - ), - ( - reqwest::StatusCode::OK, - r#"{"email":"raw-email-must-not-leak@example.com","verified_email":false}"#, - ), - ( - reqwest::StatusCode::OK, - r#"{"email":"raw-email-must-not-leak@example.com"}"#, - ), - ( - reqwest::StatusCode::OK, - r#"{"email":"raw-email-must-not-leak@example.com","verified_email":"true","raw":"raw-response-must-not-leak"}"#, - ), - ( - reqwest::StatusCode::UNAUTHORIZED, - r#"{"email":"raw-email-must-not-leak@example.com","verified_email":true}"#, - ), - ( - reqwest::StatusCode::FORBIDDEN, - r#"{"email":"raw-email-must-not-leak@example.com","verified_email":true}"#, - ), - ( - reqwest::StatusCode::INTERNAL_SERVER_ERROR, - "raw-response-must-not-leak", - ), - ]; - - for (status, body) in fixtures { - let credentials = json!({ - "access_token": "fixture-stale-access", - "refresh_token": "fixture-refresh-marker", - "expiry_date": 0, - "id_token": "stale-id-token-must-not-authorize-scope" - }); - let resolver_calls = std::cell::Cell::new(0); - let auth: RemoteAuth = prepare_remote_auth_with( - &credentials, - now, - || { - std::future::ready(Ok(( - json!({"access_token": "fixture-final-access"}), - "fixture-final-access".to_string(), - Ok(lineage_scope.clone()), - ))) - }, - |_| { - let scope = resolve_google_user_info_response_with(status, body, |_| { - resolver_calls.set(resolver_calls.get() + 1); - Err(AccountScopeError::MetadataWrite) - }); - std::future::ready(scope) - }, - ) - .await - .unwrap(); + async fn orchestration_primary_forbidden_uses_secondary_result() { + let success = fetch_with( + || async { LocalAttempt::RouteMiss }, + || async { PrimaryAttempt::Forbidden(()) }, + |()| async { Ok(orchestration_fetched("secondary")) }, + ) + .await + .unwrap(); + assert_eq!(success.source, "secondary"); - assert_eq!(resolver_calls.get(), 0); - assert_eq!( - auth.account_scope, - Err(AccountScopeError::NoTrustedEvidence) - ); - let error = auth.account_scope.as_ref().unwrap_err(); - let rendered = format!("{error} {error:?}"); - for sensitive in [ - "fixture-final-access", - "fixture-refresh-marker", - "stale-id-token-must-not-authorize-scope", - "raw-email-must-not-leak@example.com", - "raw-response-must-not-leak", - ] { - assert!(!rendered.contains(sensitive)); - } + let transient = fetch_with( + || async { LocalAttempt::RouteMiss }, + || async { PrimaryAttempt::Forbidden(()) }, + |()| async { Err(orchestration_transient("secondary transient")) }, + ) + .await + .unwrap_err(); + assert!(matches!( + transient, + ProviderFetchFailure::Transient { ref display, .. } + if display == "secondary transient" + )); - let window = quota_window( - "Gemini fixture".to_string(), - 0.5, - None, - false, - now, - "model.fixture.v1".to_string(), - Some("model.fixture.v1".to_string()), - ) - .unwrap(); - let fetched = remote_fetched( - Some("Paid".to_string()), - auth.account_scope, - vec![window], - ); - assert_eq!(fetched.source, "oauth"); - assert_eq!(fetched.windows.len(), 1); - assert!((fetched.windows[0].remaining_for_test() - 50.0).abs() < 0.01); - assert_eq!( - fetched.account_scope, - Err(AccountScopeError::NoTrustedEvidence) - ); - assert_eq!( - fetched - .identity - .as_ref() - .and_then(|identity| identity.email.as_deref()), - None - ); - } - scope_store.cleanup(); + let terminal = fetch_with( + || async { LocalAttempt::RouteMiss }, + || async { PrimaryAttempt::Forbidden(()) }, + |()| async { Err(ProviderFetchFailure::terminal("secondary terminal")) }, + ) + .await + .unwrap_err(); + assert!(matches!( + terminal, + ProviderFetchFailure::Terminal { ref display } + if display == "secondary terminal" + )); } - #[test] - fn atomic_credential_write_replaces_existing_without_orphan_temp() { - let scope = TestRefreshScope::new("antigravity", "antigravity-credential-write"); - let path = scope.root().join("antigravity/oauth_creds.json"); - let directory = path.parent().unwrap(); - std::fs::create_dir_all(directory).unwrap(); - std::fs::write(&path, b"old credentials").unwrap(); + #[tokio::test] + async fn orchestration_schema_and_transient_precedence_is_fail_closed() { + let schema_recovered = fetch_with( + || async { LocalAttempt::RouteMiss }, + || async { + PrimaryAttempt::SchemaContradiction { + context: (), + failure: ProviderFetchFailure::terminal("primary schema"), + } + }, + |()| async { Ok(orchestration_fetched("secondary")) }, + ) + .await + .unwrap(); + assert_eq!(schema_recovered.source, "secondary"); + + let schema_stays_terminal = fetch_with( + || async { LocalAttempt::RouteMiss }, + || async { + PrimaryAttempt::SchemaContradiction { + context: (), + failure: ProviderFetchFailure::terminal("primary schema"), + } + }, + |()| async { Err(orchestration_transient("secondary transient")) }, + ) + .await + .unwrap_err(); + assert!(matches!( + schema_stays_terminal, + ProviderFetchFailure::Terminal { ref display } if display == "primary schema" + )); - let credentials = json!({ - "access_token": "replacement-access", - "refresh_token": "replacement-refresh" - }); - let expected = serde_json::to_vec_pretty(&credentials).unwrap(); - write_creds_atomic(&path, &credentials).unwrap(); - - assert_eq!(std::fs::read(&path).unwrap(), expected); - assert!(!std::fs::read_dir(directory).unwrap().any(|entry| { - entry - .ok() - .is_some_and(|entry| { - entry - .file_name() - .to_string_lossy() - .starts_with(".oauth_creds.json.tokenbar.") - }) - })); - scope.cleanup(); - } + let transient_recovered = fetch_with( + || async { LocalAttempt::RouteMiss }, + || async { PrimaryAttempt::Transient(()) }, + |()| async { Ok(orchestration_fetched("secondary")) }, + ) + .await + .unwrap(); + assert_eq!(transient_recovered.source, "secondary"); - #[test] - fn resolves_remote_plan_from_tier() { - assert_eq!(resolve_remote_plan(&json!({"currentTier":{"id":"free-tier"}})).as_deref(), Some("Free")); - assert_eq!(resolve_remote_plan(&json!({"planInfo":{"planType":"standard"}})).as_deref(), Some("Standard")); + let transient_then_terminal = fetch_with( + || async { LocalAttempt::RouteMiss }, + || async { PrimaryAttempt::Transient(()) }, + |()| async { Err(ProviderFetchFailure::terminal("secondary terminal")) }, + ) + .await + .unwrap_err(); + assert!(matches!( + transient_then_terminal, + ProviderFetchFailure::Terminal { ref display } + if display == "secondary terminal" + )); + + let both_transient = fetch_with( + || async { LocalAttempt::RouteMiss }, + || async { PrimaryAttempt::Transient(()) }, + |()| async { Err(orchestration_transient("secondary transient")) }, + ) + .await + .unwrap_err(); + assert!(matches!( + both_transient, + ProviderFetchFailure::Transient { ref display, .. } + if display == "secondary transient" + )); } fn checkpoint_at( target: Option, - ) -> impl FnMut(RefreshCheckpoint) -> Result<(), String> { + ) -> impl FnMut(RefreshCheckpoint) -> Result<(), ProviderFetchFailure> { move |checkpoint| { if Some(checkpoint) == target { - Err("injected crash".to_string()) + Err(ProviderFetchFailure::terminal("injected crash")) } else { Ok(()) } } } - async fn test_refresh_response(refresh_token: String) -> Result { + async fn test_refresh_response( + refresh_token: String, + _attempt_binding: ProviderCacheBinding, + ) -> Result { assert_eq!(refresh_token, "antigravity-old-refresh"); Ok(json!({ "access_token": "antigravity-new-access", @@ -2543,7 +2541,8 @@ mod tests { scope: &TestRefreshScope, path: &Path, crash: Option, - ) -> Result<(Value, String, Result), String> { + ) -> Result<(Value, String, AccountScope, Option), ProviderFetchFailure> + { let now = DateTime::parse_from_rfc3339("2026-07-17T00:00:00Z") .unwrap() .with_timezone(&Utc); @@ -2565,6 +2564,192 @@ mod tests { .to_string() } + #[tokio::test] + async fn refresh_rejects_concurrent_account_switch_without_touching_b() { + const B_BYTES: &[u8] = br#"{ + "access_token": "account-b-access", + "refresh_token": "account-b-refresh", + "id_token": "account-b-id", + "expiry_date": 4102444800000, + "sibling": {"writer": "b", "revision": 2} +} +"#; + let (scope, path, _, before, _) = setup_refresh("antigravity-target-switch"); + let request_path = path.clone(); + let now = DateTime::parse_from_rfc3339("2026-07-17T00:00:00Z") + .unwrap() + .with_timezone(&Utc); + + let failure = refresh_access_token_with( + &path, + now, + &scope, + move |refresh_token, _attempt_binding| async move { + assert_eq!(refresh_token, "antigravity-old-refresh"); + std::fs::write(&request_path, B_BYTES).unwrap(); + Ok(json!({ + "access_token": "antigravity-new-access", + "refresh_token": "antigravity-new-refresh" + })) + }, + |_| -> std::io::Result<()> { + panic!("target mismatch must not reach credential persistence") + }, + checkpoint_at(None), + ) + .await + .unwrap_err(); + + assert!(matches!(failure, ProviderFetchFailure::Terminal { .. })); + let stored_bytes = std::fs::read(&path).unwrap(); + assert_eq!(stored_bytes, B_BYTES); + assert!(!String::from_utf8_lossy(&stored_bytes).contains("antigravity-new")); + let stored = load_remote_credentials(&path).unwrap(); + assert_eq!(stored["access_token"], "account-b-access"); + assert_eq!(stored["refresh_token"], "account-b-refresh"); + assert_eq!(stored["id_token"], "account-b-id"); + assert_eq!(stored["sibling"]["writer"], "b"); + assert_eq!(scope.metadata_bytes(), before); + scope.cleanup(); + } + + #[tokio::test] + async fn refresh_patches_unchanged_target_and_preserves_current_root_siblings() { + let (scope, path, old_scope, _, location) = setup_refresh("antigravity-target-unchanged"); + std::fs::write( + &path, + serde_json::to_vec_pretty(&json!({ + "access_token": "antigravity-old-access", + "refresh_token": "antigravity-old-refresh", + "id_token": "antigravity-old-id", + "expiry_date": 0, + "stale_only": "must-not-return" + })) + .unwrap(), + ) + .unwrap(); + let current = json!({ + "access_token": "concurrent-access", + "refresh_token": "antigravity-old-refresh", + "id_token": "concurrent-id", + "expiry_date": 0, + "token_type": "current-writer", + "sibling": {"writer": "antigravity-cli", "revision": 2}, + "unrelated": [1, 2, 3] + }); + let current_bytes = serde_json::to_vec_pretty(¤t).unwrap(); + let request_path = path.clone(); + let now = DateTime::parse_from_rfc3339("2026-07-17T00:00:00Z") + .unwrap() + .with_timezone(&Utc); + + let (refreshed, access_token, scope_outcome, cache_binding) = refresh_access_token_with( + &path, + now, + &scope, + move |refresh_token, _attempt_binding| async move { + assert_eq!(refresh_token, "antigravity-old-refresh"); + std::fs::write(&request_path, current_bytes).unwrap(); + Ok(json!({ + "access_token": "antigravity-new-access", + "refresh_token": "antigravity-new-refresh", + "id_token": "antigravity-new-id", + "expires_in": 3600, + "token_type": "provider-response" + })) + }, + |credentials| write_creds_atomic(&path, credentials), + checkpoint_at(None), + ) + .await + .unwrap(); + + assert_eq!(access_token, "antigravity-new-access"); + assert_eq!(scope_outcome, old_scope); + assert_eq!( + cache_binding, + Some(ProviderCacheBinding::primary(old_scope.clone())) + ); + let stored = load_remote_credentials(&path).unwrap(); + assert_eq!(refreshed, stored); + assert_eq!(stored["access_token"], "antigravity-new-access"); + assert_eq!(stored["refresh_token"], "antigravity-new-refresh"); + assert_eq!(stored["id_token"], "antigravity-new-id"); + assert_eq!( + stored["expiry_date"].as_f64(), + Some(now.timestamp_millis() as f64 + 3_600_000.0) + ); + assert_eq!(stored["token_type"], "current-writer"); + assert_eq!(stored["sibling"]["writer"], "antigravity-cli"); + assert_eq!(stored["sibling"]["revision"], 2); + assert_eq!(stored["unrelated"], json!([1, 2, 3])); + assert!(stored.get("stale_only").is_none()); + assert_eq!( + scope + .resolve_current("google-oauth-creds", &location, b"antigravity-new-refresh",) + .unwrap(), + old_scope + ); + scope.cleanup(); + } + + #[tokio::test] + async fn refresh_rejects_concurrent_logout_removal_and_malformed_root_without_restoring_a() { + const LOGGED_OUT_BYTES: &[u8] = br#"{ + "sibling": {"writer": "logout", "revision": 2} +} +"#; + const MALFORMED_BYTES: &[u8] = b"{not-json"; + let cases: [(&str, Option<&[u8]>); 3] = [ + ("logout", Some(LOGGED_OUT_BYTES)), + ("removed", None), + ("malformed", Some(MALFORMED_BYTES)), + ]; + + for (case, current_bytes) in cases { + let (scope, path, _, before, _) = setup_refresh(&format!("antigravity-target-{case}")); + let request_path = path.clone(); + let now = DateTime::parse_from_rfc3339("2026-07-17T00:00:00Z") + .unwrap() + .with_timezone(&Utc); + + let failure = refresh_access_token_with( + &path, + now, + &scope, + move |refresh_token, _attempt_binding| async move { + assert_eq!(refresh_token, "antigravity-old-refresh"); + if let Some(bytes) = current_bytes { + std::fs::write(&request_path, bytes).unwrap(); + } else { + std::fs::remove_file(&request_path).unwrap(); + } + Ok(json!({ + "access_token": "antigravity-new-access", + "refresh_token": "antigravity-new-refresh" + })) + }, + |_| -> std::io::Result<()> { + panic!("missing or malformed target must not reach credential persistence") + }, + checkpoint_at(None), + ) + .await + .unwrap_err(); + + assert!(matches!(failure, ProviderFetchFailure::Terminal { .. })); + assert_eq!(scope.metadata_bytes(), before); + if let Some(expected) = current_bytes { + let stored = std::fs::read(&path).unwrap(); + assert_eq!(stored, expected); + assert!(!String::from_utf8_lossy(&stored).contains("antigravity-new")); + } else { + assert!(!path.exists()); + } + scope.cleanup(); + } + } + #[tokio::test] async fn refresh_crash_boundaries_and_scope_gate_use_production_sequence() { for boundary in [ @@ -2574,12 +2759,13 @@ mod tests { RefreshCheckpoint::CredentialsPersisted, ] { let (scope, path, old_scope, before, location) = setup_refresh("antigravity-crash"); - assert_eq!( - run_refresh(&scope, &path, Some(boundary)) - .await - .unwrap_err(), - "injected crash" - ); + let failure = run_refresh(&scope, &path, Some(boundary)) + .await + .unwrap_err(); + assert!(matches!( + failure, + ProviderFetchFailure::Terminal { ref display } if display == "injected crash" + )); assert_eq!( stored_refresh_token(&path), if boundary == RefreshCheckpoint::CredentialsPersisted { @@ -2621,11 +2807,8 @@ mod tests { let (scope, path, old_scope, before, location) = setup_refresh("antigravity-metadata-fail"); scope.fail_metadata_save(); - let (refreshed, access_token, scope_outcome) = - run_refresh(&scope, &path, None).await.unwrap(); - assert_eq!(access_token, "antigravity-new-access"); - assert_eq!(remote_access_token(&refreshed).unwrap(), access_token); - assert_eq!(scope_outcome, Err(AccountScopeError::MetadataWrite)); + let failure = run_refresh(&scope, &path, None).await.unwrap_err(); + assert!(matches!(failure, ProviderFetchFailure::Terminal { .. })); assert_eq!(scope.metadata_bytes(), before); assert_eq!(stored_refresh_token(&path), "antigravity-old-refresh"); assert_eq!( @@ -2636,41 +2819,40 @@ mod tests { ); scope.cleanup(); - let (scope, path, _old_scope, before, _) = - setup_refresh("antigravity-metadata-fail-unchanged"); - scope.fail_metadata_save(); + let (scope, path, old_scope, _, location) = setup_refresh("antigravity-save-fail"); let now = DateTime::parse_from_rfc3339("2026-07-17T00:00:00Z") .unwrap() .with_timezone(&Utc); - let save_path = path.clone(); - let (refreshed, access_token, scope_outcome) = refresh_access_token_with( + let (refreshed, access_token, scope_outcome, cache_binding) = refresh_access_token_with( &path, now, &scope, - |refresh_token| async move { - assert_eq!(refresh_token, "antigravity-old-refresh"); - Ok(json!({ - "access_token": "antigravity-new-access", - "expires_in": 3600 - })) - }, - move |credentials| write_creds_atomic(&save_path, credentials), + test_refresh_response, + |_| Err(std::io::Error::other("injected save failure")), checkpoint_at(None), ) .await .unwrap(); - assert_eq!(scope_outcome, Err(AccountScopeError::MetadataWrite)); - assert_eq!(scope.metadata_bytes(), before); assert_eq!(access_token, "antigravity-new-access"); assert_eq!(remote_access_token(&refreshed).unwrap(), access_token); - let persisted = load_remote_credentials(&path).unwrap(); - assert_eq!(remote_access_token(&persisted).unwrap(), access_token); + assert_eq!(scope_outcome, old_scope); + assert_eq!(cache_binding, None); assert_eq!(stored_refresh_token(&path), "antigravity-old-refresh"); + assert_eq!( + scope + .resolve_current("google-oauth-creds", &location, b"antigravity-new-refresh") + .unwrap(), + old_scope + ); scope.cleanup(); let (scope, path, old_scope, _, location) = setup_refresh("antigravity-success"); - let (_, _, scope_outcome) = run_refresh(&scope, &path, None).await.unwrap(); - assert_eq!(scope_outcome.unwrap(), old_scope); + let (_, _, scope_outcome, cache_binding) = run_refresh(&scope, &path, None).await.unwrap(); + assert_eq!(scope_outcome, old_scope); + assert_eq!( + cache_binding, + Some(ProviderCacheBinding::primary(old_scope.clone())) + ); assert_eq!( scope .resolve_current("google-oauth-creds", &location, b"antigravity-new-refresh") @@ -2679,4 +2861,47 @@ mod tests { ); scope.cleanup(); } + + #[tokio::test] + async fn refresh_transient_uses_lock_reloaded_binding_not_outer_binding() { + let (scope, path, inner_scope, _, location) = setup_refresh("antigravity-lock-binding"); + let outer_scope = scope + .resolve_current("google-oauth-creds", &location, b"outer-refresh-a") + .unwrap(); + assert_ne!(outer_scope, inner_scope); + let expected = ProviderCacheBinding::primary(inner_scope); + let request_expected = expected.clone(); + let now = DateTime::parse_from_rfc3339("2026-07-17T00:00:00Z") + .unwrap() + .with_timezone(&Utc); + + let failure = refresh_access_token_with( + &path, + now, + &scope, + move |refresh_token, attempt_binding| async move { + assert_eq!(refresh_token, "antigravity-old-refresh"); + assert_eq!(attempt_binding, request_expected); + Err(ProviderFetchFailure::transient( + "Antigravity token refresh failed. Retrying automatically.", + Some(attempt_binding), + crate::agent_usage::SafeTransportDiagnostic::from_facts( + TransportErrorFacts::synthetic(true, false, TransportPhase::Request, None), + ), + )) + }, + |credentials| write_creds_atomic(&path, credentials), + checkpoint_at(None), + ) + .await + .unwrap_err(); + + match failure { + ProviderFetchFailure::Transient { + attempt_binding, .. + } => assert_eq!(attempt_binding, Some(expected)), + ProviderFetchFailure::Terminal { .. } => panic!("timeout must remain transient"), + } + scope.cleanup(); + } } diff --git a/crates/tb_core_ffi/src/agent_copilot.rs b/crates/tb_core_ffi/src/agent_copilot.rs index 8c65929..5d71d21 100644 --- a/crates/tb_core_ffi/src/agent_copilot.rs +++ b/crates/tb_core_ffi/src/agent_copilot.rs @@ -8,17 +8,22 @@ use crate::agent_account_scope::{self, AccountScope, AccountScopeError}; use crate::agent_quota_duration::{copilot_calendar_duration, DurationEvidence}; -use crate::agent_usage::{clean_plan, AgentIdentity, UsageWindow}; +use crate::agent_usage::{ + clean_plan, provider_http_client_builder, read_response_body, request_after_verified_binding, + AgentIdentity, ProviderCacheBinding, ProviderFetchFailure, ResponseReadFailure, + TransportErrorFacts, TransportPhase, UsageWindow, +}; use crate::opencode_integrations::GitHubCopilotCredential; use chrono::{DateTime, NaiveDate, TimeZone, Utc}; -use serde::Deserialize; -use serde_json::value::RawValue; +use serde::{Deserialize, Deserializer}; +use serde_json::{value::RawValue, Value}; const COPILOT_USAGE_URL: &str = "https://api.github.com/copilot_internal/user"; pub(crate) struct CopilotData { pub identity: Option, pub account_scope: Result, + pub cache_binding: ProviderCacheBinding, pub windows: Vec, } @@ -26,26 +31,38 @@ pub(crate) struct CopilotData { struct CopilotUser { #[serde(default)] copilot_plan: Option, - #[serde(default, deserialize_with = "deserialize_optional_raw")] - quota_reset_date: Option>, + #[serde(default, deserialize_with = "deserialize_optional_string")] + quota_reset_date: Option, #[serde(default)] quota_snapshots: Option, } -#[derive(Debug, Clone, Deserialize)] +fn deserialize_optional_string<'de, D>(deserializer: D) -> Result, D::Error> +where + D: Deserializer<'de>, +{ + Ok( + Option::::deserialize(deserializer)?.and_then(|value| match value { + Value::String(value) => Some(value), + _ => None, + }), + ) +} + +#[derive(Debug, Deserialize)] struct QuotaSnapshots { #[serde(default)] - premium_interactions: Option, + premium_interactions: Option>, #[serde(default)] - chat: Option, + chat: Option>, } -#[derive(Debug, Clone, Deserialize)] +#[derive(Debug, Deserialize)] struct QuotaSnapshot { #[serde(default)] - entitlement: f64, + entitlement: Option, #[serde(default)] - remaining: f64, + remaining: Option, #[serde(default)] percent_remaining: Option, } @@ -53,248 +70,247 @@ struct QuotaSnapshot { pub(crate) async fn fetch( now: DateTime, credential: GitHubCopilotCredential, -) -> Result { - fetch_with( - now, - credential, - request_usage, - |semantic_source, canonical_location, marker| { - agent_account_scope::resolve_credential( - "copilot", - semantic_source, - canonical_location, - marker, - ) - }, +) -> Result { + let verified = agent_account_scope::resolve_credential( + "copilot", + credential.semantic_source, + &credential.canonical_location, + &credential.marker, ) + .map(|account_scope| { + let cache_binding = ProviderCacheBinding::primary(account_scope.clone()); + (account_scope, cache_binding) + }) + .map_err(|_| { + ProviderFetchFailure::terminal("GitHub Copilot account identity could not be verified.") + }); + let (account_scope, cache_binding, response) = + request_after_verified_binding(verified, |(account_scope, cache_binding)| async move { + let client = provider_http_client_builder() + .timeout(std::time::Duration::from_secs(30)) + .build() + .map_err(|_| { + ProviderFetchFailure::terminal("Copilot usage client could not be created.") + })?; + let response = client + .get(COPILOT_USAGE_URL) + .header( + reqwest::header::AUTHORIZATION, + format!("token {}", credential.request_token), + ) + .header(reqwest::header::ACCEPT, "application/json") + .header(reqwest::header::USER_AGENT, "GitHubCopilotChat/0.26.7") + .header("Editor-Version", "vscode/1.96.2") + .header("Editor-Plugin-Version", "copilot-chat/0.26.7") + .header("X-Github-Api-Version", "2025-04-01") + .send() + .await + .map_err(|error| { + ProviderFetchFailure::from_send_error( + "Copilot usage request failed. Retrying automatically.", + Some(cache_binding.clone()), + &error, + ) + })?; + Ok((account_scope, cache_binding, response)) + }) + .await?; + let status = response.status().as_u16(); + let body = read_response_body(status, false, || async { + response.text().await.map_err(|error| { + TransportErrorFacts::from_reqwest(&error, TransportPhase::ResponseBody) + }) + }) .await + .map_err(|failure| match failure { + ResponseReadFailure::Transient(diagnostic) => ProviderFetchFailure::transient( + "Copilot usage request failed. Retrying automatically.", + Some(cache_binding.clone()), + diagnostic, + ), + ResponseReadFailure::Terminal(401 | 403) => { + ProviderFetchFailure::terminal("GitHub Copilot token expired or lacks access.") + } + ResponseReadFailure::Terminal(status) => ProviderFetchFailure::terminal(format!( + "Copilot usage API rejected the request (status {status})." + )), + })?; + let (plan, windows) = decode_usage_response(&body, now)?; + Ok(CopilotData { + identity: Some(AgentIdentity { email: None, plan }), + account_scope: Ok(account_scope), + cache_binding, + windows, + }) } -async fn request_usage(request_token: String) -> Result { - let client = reqwest::Client::builder() - .timeout(std::time::Duration::from_secs(30)) - .build() - .map_err(|e| format!("build Copilot client: {e}"))?; - let response = client - .get(COPILOT_USAGE_URL) - .header(reqwest::header::AUTHORIZATION, format!("token {request_token}")) - .header(reqwest::header::ACCEPT, "application/json") - .header(reqwest::header::USER_AGENT, "GitHubCopilotChat/0.26.7") - .header("Editor-Version", "vscode/1.96.2") - .header("Editor-Plugin-Version", "copilot-chat/0.26.7") - .header("X-Github-Api-Version", "2025-04-01") - .send() - .await - .map_err(|e| format!("Copilot usage request failed: {e}"))?; - let status = response.status(); - if status == reqwest::StatusCode::UNAUTHORIZED || status == reqwest::StatusCode::FORBIDDEN { - return Err("GitHub Copilot token expired or lacks access.".to_string()); - } - if !status.is_success() { - return Err(format!("Copilot usage API returned {}.", status.as_u16())); - } - let body = response - .text() - .await - .map_err(|e| format!("read Copilot response: {e}"))?; - serde_json::from_str(&body).map_err(|e| format!("decode Copilot usage: {e}")) +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +enum CopilotMapping { + Usable, + PlaceholderOnly, + Invalid, } -async fn fetch_with( +#[derive(Debug)] +enum CopilotRow { + Usable(Box), + Placeholder, + Invalid, + Absent, +} + +pub(crate) fn decode_usage_response( + body: &str, now: DateTime, - credential: GitHubCopilotCredential, - request: Request, - resolve_scope: ResolveScope, -) -> Result -where - Request: FnOnce(String) -> RequestFuture, - RequestFuture: std::future::Future>, - ResolveScope: FnOnce(&str, &str, &[u8]) -> Result, -{ - let GitHubCopilotCredential { - request_token, - marker, - semantic_source, - canonical_location, - } = credential; - let usage = request(request_token).await?; - let windows = snapshot_windows(&usage, now); - let account_scope = resolve_scope(semantic_source, &canonical_location, &marker); - Ok(CopilotData { - identity: Some(AgentIdentity { - email: None, - plan: usage.copilot_plan.filter(|s| !s.trim().is_empty()).map(clean_plan), - }), - account_scope, - windows, - }) +) -> Result<(Option, Vec), ProviderFetchFailure> { + let usage: CopilotUser = serde_json::from_str(body).map_err(|_| { + ProviderFetchFailure::terminal("Copilot usage response could not be decoded.") + })?; + let (plan, windows, mapping) = map_user(usage, now); + if windows.is_empty() && mapping != CopilotMapping::PlaceholderOnly { + return Err(ProviderFetchFailure::terminal( + "Copilot usage API returned no usable quota windows.", + )); + } + Ok((plan, windows)) } -fn snapshot_windows(usage: &CopilotUser, now: DateTime) -> Vec { - let resets_at = usage.quota_reset_date.as_deref().and_then(parse_reset_raw); - let reset_was_supplied = usage.quota_reset_date.is_some(); - let Some(snapshots) = usage.quota_snapshots.as_ref() else { - return Vec::new(); +fn map_user( + usage: CopilotUser, + now: DateTime, +) -> (Option, Vec, CopilotMapping) { + let resets_at = usage.quota_reset_date.as_deref().and_then(parse_reset_date); + let mut windows = Vec::new(); + let mut saw_placeholder = false; + let mut saw_invalid = false; + if let Some(snapshots) = usage.quota_snapshots { + for row in [ + snapshot_window_with_identity( + "Premium", + "premium_interactions.v1", + snapshots.premium_interactions.as_deref(), + resets_at, + now, + ), + snapshot_window_with_identity( + "Chat", + "chat.v1", + snapshots.chat.as_deref(), + resets_at, + now, + ), + ] { + match row { + CopilotRow::Usable(window) => windows.push(*window), + CopilotRow::Placeholder => saw_placeholder = true, + CopilotRow::Invalid => saw_invalid = true, + CopilotRow::Absent => {} + } + } + } + let plan = usage + .copilot_plan + .filter(|plan| !plan.trim().is_empty()) + .map(clean_plan); + let mapping = if !windows.is_empty() { + CopilotMapping::Usable + } else if saw_placeholder && !saw_invalid { + CopilotMapping::PlaceholderOnly + } else { + CopilotMapping::Invalid }; - [ - snapshot_window_with_identity( - "Premium", - "premium_interactions.v1", - Some("premium_interactions.v1".to_string()), - snapshots.premium_interactions.clone(), - resets_at, - reset_was_supplied, - now, - ), - snapshot_window_with_identity( - "Chat", - "chat.v1", - Some("chat.v1".to_string()), - snapshots.chat.clone(), - resets_at, - reset_was_supplied, - now, - ), - ] - .into_iter() - .flatten() - .collect() + (plan, windows, mapping) } fn snapshot_window_with_identity( label: &str, - card_id: &str, - window_key: Option, - snapshot: Option, + window_key: &str, + raw: Option<&RawValue>, resets_at: Option>, - reset_was_supplied: bool, now: DateTime, -) -> Option { - let snapshot = snapshot?; - // Skip explicit zero-entitlement placeholders (no usable quota signal). - if snapshot.entitlement == 0.0 - && snapshot.remaining == 0.0 - && snapshot.percent_remaining.is_none() +) -> CopilotRow { + let Some(raw) = raw else { + return CopilotRow::Absent; + }; + let Ok(snapshot) = serde_json::from_str::(raw.get()) else { + return CopilotRow::Invalid; + }; + let (Some(entitlement), Some(remaining)) = (snapshot.entitlement, snapshot.remaining) else { + return CopilotRow::Invalid; + }; + if !entitlement.is_finite() + || !remaining.is_finite() + || entitlement < 0.0 + || remaining < 0.0 + || remaining > entitlement { - return None; + return CopilotRow::Invalid; } - let percent_remaining = snapshot.percent_remaining.or_else(|| { - (snapshot.entitlement > 0.0).then(|| (snapshot.remaining / snapshot.entitlement) * 100.0) - })?; - if !percent_remaining.is_finite() || !(0.0..=100.0).contains(&percent_remaining) { - return None; + if entitlement == 0.0 { + return CopilotRow::Placeholder; } - - let window = UsageWindow::from_fraction( - label.to_string(), - percent_remaining / 100.0, - resets_at, - now, - ) - .with_identity(card_id, window_key); - let contract_duration = resets_at.and_then(|reset| copilot_calendar_duration(reset.timestamp())); - Some(match contract_duration { - Some(duration) => window.with_contract_duration_evidence( + let derived_percent = (remaining / entitlement) * 100.0; + let percent_remaining = match snapshot.percent_remaining { + Some(percent) + if percent.is_finite() + && (0.0..=100.0).contains(&percent) + // Provider payloads may round or truncate to a whole percent. + && (percent - derived_percent).abs() <= 1.0 => + { + percent + } + Some(_) => return CopilotRow::Invalid, + None => derived_percent, + }; + let contract_duration = resets_at + .and_then(|reset| copilot_calendar_duration(reset.timestamp())) + .map(DurationEvidence::contract); + CopilotRow::Usable(Box::new( + UsageWindow::from_provider_used_percent( + label.to_string(), + 100.0 - percent_remaining, + resets_at, now, - reset_was_supplied, - DurationEvidence::contract(duration), + ) + .with_identity( + window_key, + Some(window_key.to_string()), + None, + contract_duration, ), - None => window.with_observed_duration_evidence(now, reset_was_supplied), - }) + )) } #[cfg(test)] fn snapshot_window( label: &str, - snapshot: Option, + raw: Option<&RawValue>, resets_at: Option>, now: DateTime, ) -> Option { - let (card_id, window_key) = match label { - "Premium" => ("premium_interactions.v1", Some("premium_interactions.v1".to_string())), - "Chat" => ("chat.v1", Some("chat.v1".to_string())), - _ => ("row.copilot.unknown.v1", None), + let window_key = match label { + "Premium" => "premium_interactions.v1", + "Chat" => "chat.v1", + _ => "row.copilot.unknown.v1", }; - snapshot_window_with_identity( - label, - card_id, - window_key, - snapshot, - resets_at, - resets_at.is_some(), - now, - ) -} - -fn deserialize_optional_raw<'de, D>(deserializer: D) -> Result>, D::Error> -where - D: serde::Deserializer<'de>, -{ - Box::::deserialize(deserializer).map(Some) -} - -fn parse_reset_raw(raw: &RawValue) -> Option> { - let value = serde_json::from_str::(raw.get()).ok()?; - parse_reset_date(&value) + match snapshot_window_with_identity(label, window_key, raw, resets_at, now) { + CopilotRow::Usable(window) => Some(*window), + CopilotRow::Placeholder | CopilotRow::Invalid | CopilotRow::Absent => None, + } } /// Copilot reports `quota_reset_date` as a bare `YYYY-MM-DD`; treat it as UTC midnight. fn parse_reset_date(value: &str) -> Option> { - let value = value.trim(); - let bytes = value.as_bytes(); - if bytes.len() != 10 - || bytes[4] != b'-' - || bytes[7] != b'-' - || !bytes[..4].iter().all(u8::is_ascii_digit) - || !bytes[5..7].iter().all(u8::is_ascii_digit) - || !bytes[8..].iter().all(u8::is_ascii_digit) - { - return None; - } - let date = NaiveDate::parse_from_str(value, "%Y-%m-%d").ok()?; - Some(Utc.from_utc_datetime(&date.and_hms_opt(0, 0, 0)?)) + let date = NaiveDate::parse_from_str(value.trim(), "%Y-%m-%d").ok()?; + Utc.from_utc_datetime(&date.and_hms_opt(0, 0, 0)?).into() } #[cfg(test)] mod tests { use super::*; - use crate::agent_account_scope::test_support::TestRefreshScope; - use crate::agent_account_scope::RefreshScopeTransaction; - use serde_json::json; - use std::sync::atomic::{AtomicU64, Ordering}; - - static TEMP_COUNTER: AtomicU64 = AtomicU64::new(0); - - fn temp_auth_path(tag: &str) -> std::path::PathBuf { - let root = std::env::temp_dir().join(format!( - "tb-copilot-{tag}-{}-{}", - std::process::id(), - TEMP_COUNTER.fetch_add(1, Ordering::Relaxed) - )); - std::fs::create_dir_all(&root).unwrap(); - root.join("auth.json") - } - fn fixture_credential(path: &std::path::Path, marker: &str) -> GitHubCopilotCredential { - let json = json!({"github-copilot": { - "type": "oauth", "refresh": marker, "access": "fake-access-token" - }}); - std::fs::write(path, serde_json::to_vec(&json).unwrap()).unwrap(); - crate::opencode_integrations::github_copilot_credential_from(path, &json).unwrap() - } - - fn fixture_user() -> CopilotUser { - serde_json::from_value(json!({ - "copilot_plan": "individual", - "quota_reset_date": "2026-08-01", - "quota_snapshots": { - "premium_interactions": { - "entitlement": 300, "remaining": 90, "percent_remaining": 30 - }, - "chat": {"entitlement": 100, "remaining": 75, "percent_remaining": 75} - } - })) - .unwrap() + fn raw(value: &str) -> Box { + RawValue::from_string(value.to_string()).unwrap() } #[test] @@ -310,302 +326,344 @@ mod tests { }"#; let usage: CopilotUser = serde_json::from_str(body).unwrap(); let snaps = usage.quota_snapshots.unwrap(); - let premium = snapshot_window("Premium", snaps.premium_interactions, None, now).unwrap(); - assert_eq!(premium.card_id_for_test(), "premium_interactions.v1"); - assert_eq!(premium.pace_window_key_for_test(), Some("premium_interactions.v1")); + let premium = + snapshot_window("Premium", snaps.premium_interactions.as_deref(), None, now).unwrap(); assert!((premium.remaining_for_test() - 30.0).abs() < 0.01); // chat is a zero-entitlement placeholder → skipped - assert!(snapshot_window("Chat", snaps.chat, None, now).is_none()); + assert!(snapshot_window("Chat", snaps.chat.as_deref(), None, now).is_none()); } #[test] - fn maps_exact_calendar_duration_for_each_month_length_on_both_cards() { - let cases = [ - ("2023-03-01", "2023-02-15T00:00:00Z", 28 * 86_400), - ("2024-03-01", "2024-02-15T00:00:00Z", 29 * 86_400), - ("2023-05-01", "2023-04-15T00:00:00Z", 30 * 86_400), - ("2023-08-01", "2023-07-15T00:00:00Z", 31 * 86_400), - ]; - for (reset, now_text, expected_seconds) in cases { - let now = now_text.parse::>().unwrap(); - let usage: CopilotUser = serde_json::from_value(json!({ - "quota_reset_date": reset, + fn quota_reset_date_is_lossy_without_poisoning_valid_snapshots() { + let valid: CopilotUser = serde_json::from_str( + r#"{ + "quota_reset_date": "2026-08-01", "quota_snapshots": { "premium_interactions": { - "entitlement": 300, "remaining": 90, "percent_remaining": 30 - }, - "chat": {"entitlement": 100, "remaining": 75, "percent_remaining": 75} + "entitlement": 100, + "remaining": 60, + "percent_remaining": 60 + } } - })) - .unwrap(); - assert!(usage.quota_reset_date.is_some()); - let windows = snapshot_windows(&usage, now); - assert_eq!(windows.len(), 2, "{reset}"); - for (window, card_id) in windows.iter().zip([ - "premium_interactions.v1", - "chat.v1", - ]) { - let wire = serde_json::to_value(window).unwrap(); - assert_eq!(wire["cardId"], card_id, "{reset}"); - assert_eq!(wire["paceStatus"]["state"], "learningHistory", "{reset}"); - assert_eq!(wire["paceStatus"]["durationSource"], "contract", "{reset}"); - assert_eq!( - wire["paceStatus"]["durationSeconds"], - expected_seconds, - "{reset}" - ); - assert_eq!(wire["windowMinutes"], expected_seconds / 60, "{reset}"); - assert!(wire.get("historicalPace").is_none(), "{reset}"); - } + }"#, + ) + .unwrap(); + assert_eq!(valid.quota_reset_date.as_deref(), Some("2026-08-01")); + + let now = Utc.timestamp_opt(1_751_328_000, 0).single().unwrap(); + for malformed_reset in [ + "null", + "42", + r#"{"credential":"token-secret"}"#, + r#"["token-secret"]"#, + "true", + ] { + let body = format!( + r#"{{ + "quota_reset_date": {malformed_reset}, + "quota_snapshots": {{ + "premium_interactions": {{ + "entitlement": 100, + "remaining": 60, + "percent_remaining": 60 + }} + }} + }}"# + ); + let usage: CopilotUser = serde_json::from_str(&body).unwrap(); + assert_eq!(usage.quota_reset_date, None, "{malformed_reset}"); + let (_, windows, mapping) = map_user(usage, now); + assert_eq!(mapping, CopilotMapping::Usable, "{malformed_reset}"); + assert_eq!(windows.len(), 1, "{malformed_reset}"); + + let wire = serde_json::to_value(&windows[0]).unwrap(); + assert_eq!(wire["usedPercent"], 40.0, "{malformed_reset}"); + assert_eq!(wire["remainingPercent"], 60.0, "{malformed_reset}"); + assert_eq!(wire["cardId"], "premium_interactions.v1"); + assert!(wire.get("resetsAt").is_none(), "{malformed_reset}"); + assert!(wire.get("resetText").is_none(), "{malformed_reset}"); + assert!(wire.get("windowMinutes").is_none(), "{malformed_reset}"); + assert_eq!(wire["paceStatus"]["state"], "unavailable"); + assert_eq!(wire["paceStatus"]["reason"], "missingReset"); + assert!(wire["paceStatus"].get("durationSeconds").is_none()); + assert!(wire["paceStatus"].get("durationSource").is_none()); + assert!(wire.get("historicalPace").is_none(), "{malformed_reset}"); + assert!(!wire.to_string().contains("token-secret")); } } #[test] - fn valid_future_non_month_start_uses_observed_duration_learning() { - let now = "2023-08-01T00:00:00Z".parse::>().unwrap(); - let usage: CopilotUser = serde_json::from_value(json!({ - "quota_reset_date": "2023-08-15", + fn lossy_reset_date_does_not_weaken_quota_row_validation() { + let body = r#"{ + "quota_reset_date": {"ignored":"token-secret"}, "quota_snapshots": { - "premium_interactions": {"entitlement": 300, "remaining": 90}, - "chat": {"entitlement": 100, "remaining": 75} + "premium_interactions": { + "entitlement": 100, + "remaining": 101, + "percent_remaining": 100 + } } - })) - .unwrap(); - let windows = snapshot_windows(&usage, now); - assert_eq!(windows.len(), 2); - for window in windows { - let wire = serde_json::to_value(window).unwrap(); - assert_eq!(wire["paceStatus"]["state"], "learningDuration"); - assert_eq!(wire["paceStatus"]["durationSource"], "observed"); - assert!(wire["paceStatus"].get("durationSeconds").is_none()); - assert!(wire.get("windowMinutes").is_none()); - assert!(wire.get("historicalPace").is_none()); - } + }"#; + assert!(matches!( + decode_usage_response(body, Utc::now()), + Err(ProviderFetchFailure::Terminal { .. }) + )); } #[test] - fn reset_presence_and_validity_fail_closed_without_observed_fallback() { - let now = "2023-08-15T00:00:00Z".parse::>().unwrap(); - let cases = [ - ("missing", json!({}), false, "missingReset"), - ( - "null", - json!({"quota_reset_date": null}), - true, - "invalidEvidence", - ), - ( - "object", - json!({"quota_reset_date": {"reset": "2023-08-01"}}), - true, - "invalidEvidence", - ), - ( - "number", - json!({"quota_reset_date": 42}), - true, - "invalidEvidence", - ), - ( - "bool", - json!({"quota_reset_date": true}), - true, - "invalidEvidence", - ), - ( - "blank", - json!({"quota_reset_date": ""}), - true, - "invalidEvidence", - ), - ( - "malformed", - json!({"quota_reset_date": "not-a-date"}), - true, - "invalidEvidence", - ), - ( - "past", - json!({"quota_reset_date": "2023-08-01"}), - true, - "invalidEvidence", - ), - ]; - for (case, reset, supplied, reason) in cases { - let mut body = json!({ + fn zero_entitlement_and_missing_fields_do_not_create_usable_windows() { + let now = Utc::now(); + let placeholder: CopilotUser = serde_json::from_str( + r#"{ "quota_snapshots": { - "premium_interactions": {"entitlement": 300, "remaining": 90}, - "chat": {"entitlement": 100, "remaining": 75} + "premium_interactions": { + "entitlement": 0, + "remaining": 0, + "percent_remaining": 0 + } } - }); - if let Some(reset) = reset.get("quota_reset_date") { - body["quota_reset_date"] = reset.clone(); - } - let usage: CopilotUser = serde_json::from_value(body).unwrap(); - assert_eq!(usage.quota_reset_date.is_some(), supplied, "{case}"); - let windows = snapshot_windows(&usage, now); - assert_eq!(windows.len(), 2, "{case}"); - for window in windows { - let wire = serde_json::to_value(window).unwrap(); - assert_eq!(wire["paceStatus"]["state"], "unavailable", "{case}"); - assert_eq!(wire["paceStatus"]["reason"], reason, "{case}"); - assert!(wire["paceStatus"].get("durationSeconds").is_none(), "{case}"); - assert!(wire.get("windowMinutes").is_none(), "{case}"); - assert!(wire.get("historicalPace").is_none(), "{case}"); - } + }"#, + ) + .unwrap(); + let (_, windows, mapping) = map_user(placeholder, now); + assert!(windows.is_empty()); + assert_eq!(mapping, CopilotMapping::PlaceholderOnly); + + for malformed in [r#"{}"#, r#"{"entitlement": 10}"#, r#"{"remaining": 10}"#] { + let row = raw(malformed); + assert!(matches!( + snapshot_window_with_identity( + "Premium", + "premium_interactions.v1", + Some(row.as_ref()), + None, + now, + ), + CopilotRow::Invalid + )); } - let early_now = "2023-06-15T00:00:00Z".parse::>().unwrap(); - let future_calendar: CopilotUser = serde_json::from_value(json!({ - "quota_reset_date": "2023-08-01", - "quota_snapshots": { - "premium_interactions": {"entitlement": 300, "remaining": 90}, - "chat": {"entitlement": 100, "remaining": 75} - } - })) + let valid_with_malformed_sibling: CopilotUser = serde_json::from_str( + r#"{ + "quota_snapshots": { + "premium_interactions": { + "entitlement": 100, + "remaining": 60, + "percent_remaining": 60 + }, + "chat": {} + } + }"#, + ) .unwrap(); - for window in snapshot_windows(&future_calendar, early_now) { - let wire = serde_json::to_value(window).unwrap(); - assert_eq!(wire["paceStatus"]["state"], "unavailable"); - assert_eq!(wire["paceStatus"]["reason"], "invalidEvidence"); + let (_, windows, mapping) = map_user(valid_with_malformed_sibling, now); + assert_eq!(mapping, CopilotMapping::Usable); + assert_eq!(windows.len(), 1); + + for malformed in [ + r#"{"entitlement":0,"remaining":1,"percent_remaining":100}"#, + r#"{"entitlement":100,"remaining":101,"percent_remaining":100}"#, + r#"{"entitlement":-1,"remaining":0,"percent_remaining":0}"#, + r#"{"entitlement":100,"remaining":-1,"percent_remaining":0}"#, + ] { + let row = raw(malformed); + assert!(matches!( + snapshot_window_with_identity( + "Premium", + "premium_interactions.v1", + Some(row.as_ref()), + None, + now, + ), + CopilotRow::Invalid + )); } + + let invalid_with_valid_sibling: CopilotUser = serde_json::from_str( + r#"{ + "quota_snapshots": { + "premium_interactions": { + "entitlement": 0, + "remaining": 1, + "percent_remaining": 100 + }, + "chat": { + "entitlement": 100, + "remaining": 75, + "percent_remaining": 75 + } + } + }"#, + ) + .unwrap(); + let (_, windows, mapping) = map_user(invalid_with_valid_sibling, now); + assert_eq!(mapping, CopilotMapping::Usable); + assert_eq!(windows.len(), 1); + assert_eq!(windows[0].label_for_test(), "Chat"); + + let invalid_only: CopilotUser = serde_json::from_str( + r#"{ + "quota_snapshots": { + "premium_interactions": { + "entitlement": 0, + "remaining": 1, + "percent_remaining": 100 + } + } + }"#, + ) + .unwrap(); + let (_, windows, mapping) = map_user(invalid_only, now); + assert!(windows.is_empty()); + assert_eq!(mapping, CopilotMapping::Invalid); } - #[tokio::test] - async fn request_and_lineage_use_the_same_normalized_marker() { - let path = temp_auth_path("lineage"); - let store = TestRefreshScope::new("copilot", "copilot-lineage"); + #[test] + fn stage4_copilot_maps_shared_reset_to_both_quota_cards() { let now = Utc.timestamp_opt(1_751_328_000, 0).single().unwrap(); - let events = std::cell::RefCell::new(Vec::new()); - let first = fetch_with( - now, - fixture_credential(&path, " stable-marker "), - |token| { - events.borrow_mut().push("request"); - assert_eq!(token, "stable-marker"); - std::future::ready(Ok(fixture_user())) - }, - |source, location, marker| { - events.borrow_mut().push("scope"); - assert_eq!(source, "opencode-auth-json"); - assert!(location.ends_with("\0github-copilot")); - assert_eq!(marker, b"stable-marker"); - store.resolve_current(source, location, marker) - }, + let usage: CopilotUser = serde_json::from_str( + r#"{ + "copilot_plan": "individual", + "quota_reset_date": "2026-08-01", + "quota_snapshots": { + "premium_interactions": { + "entitlement": 300, + "remaining": 90, + "percent_remaining": 30 + }, + "chat": { + "entitlement": 100, + "remaining": 75 + } + } + }"#, ) - .await .unwrap(); - assert_eq!(&*events.borrow(), &["request", "scope"]); - assert_eq!(first.windows.len(), 2); + let (plan, windows, mapping) = map_user(usage, now); + assert_eq!(mapping, CopilotMapping::Usable); + assert_eq!(plan.as_deref(), Some("Individual")); + assert_eq!(windows.len(), 2); + let premium = &windows[0]; + let chat = &windows[1]; + + assert_eq!(premium.label_for_test(), "Premium"); + assert_eq!(chat.label_for_test(), "Chat"); assert_eq!( - first.windows[0].pace_window_key_for_test(), + premium.resets_at_for_test(), + Some("2026-08-01T00:00:00.000Z") + ); + assert_eq!(chat.resets_at_for_test(), premium.resets_at_for_test()); + assert_eq!( + premium.window_minutes_for_test(), + Some(44_640), + "first-of-month reset uses the exact preceding calendar month" + ); + assert_eq!(chat.window_minutes_for_test(), Some(44_640)); + assert_eq!( + premium.pace_window_key_for_test(), Some("premium_interactions.v1") ); - assert_eq!(first.windows[1].pace_window_key_for_test(), Some("chat.v1")); - let first_scope = first.account_scope.unwrap(); - let normalized_scope = fetch_with( - now, - fixture_credential(&path, "stable-marker"), - |_| std::future::ready(Ok(fixture_user())), - |source, location, marker| store.resolve_current(source, location, marker), - ) - .await - .unwrap() - .account_scope - .unwrap(); - let different_scope = fetch_with( + assert_eq!(chat.pace_window_key_for_test(), Some("chat.v1")); + for window in &windows { + let wire = serde_json::to_value(window).unwrap(); + assert_eq!(wire["paceStatus"]["durationSource"], "contract"); + assert_eq!(wire["paceStatus"]["durationSeconds"], 2_678_400); + } + + let non_calendar_reset = parse_reset_date("2026-08-15").unwrap(); + let observed_raw = + raw(r#"{ "entitlement": 300, "remaining": 90, "percent_remaining": 30 }"#); + let observed = snapshot_window( + "Premium", + Some(observed_raw.as_ref()), + Some(non_calendar_reset), now, - fixture_credential(&path, "different-marker"), - |_| std::future::ready(Ok(fixture_user())), - |source, location, marker| store.resolve_current(source, location, marker), ) - .await - .unwrap() - .account_scope .unwrap(); - assert_eq!(first_scope, normalized_scope); - assert_ne!(first_scope, different_scope); - store.cleanup(); - let _ = std::fs::remove_dir_all(path.parent().unwrap()); + assert_eq!( + observed.window_minutes_for_test(), + None, + "copilot.premium.observed-fallback" + ); } - #[tokio::test] - async fn api_and_scope_errors_fail_closed_without_losing_successful_gauges() { - let path = temp_auth_path("errors"); - let scope_calls = std::cell::Cell::new(0); - let result = fetch_with( - Utc::now(), - fixture_credential(&path, " api-error-secret "), - |token| { - assert_eq!(token, "api-error-secret"); - std::future::ready(Err("Copilot usage API returned 503.".to_string())) - }, - |_, _, _| { - scope_calls.set(scope_calls.get() + 1); - Err(AccountScopeError::MetadataWrite) - }, - ) - .await; - let error = match result { - Ok(_) => panic!("API failure unexpectedly succeeded"), - Err(error) => error, - }; - assert_eq!(scope_calls.get(), 0); - assert_eq!(error, "Copilot usage API returned 503."); - assert!(!error.contains("api-error-secret")); - assert!(!error.contains(path.to_string_lossy().as_ref())); - - let data = fetch_with( - Utc::now(), - fixture_credential(&path, " scope-error-secret "), - |_| std::future::ready(Ok(fixture_user())), - |source, location, marker| { - assert_eq!(source, "opencode-auth-json"); - assert!(location.ends_with("\0github-copilot")); - assert_eq!(marker, b"scope-error-secret"); - Err(AccountScopeError::MetadataWrite) - }, - ) - .await - .unwrap(); - assert_eq!(data.account_scope, Err(AccountScopeError::MetadataWrite)); - assert_eq!(data.windows.len(), 2); - let _ = std::fs::remove_dir_all(path.parent().unwrap()); + #[test] + fn rejects_invalid_remaining_percentages_before_wire() { + let now = Utc::now(); + let out_of_range = + raw(r#"{ "entitlement": 300, "remaining": 90, "percent_remaining": 101 }"#); + assert!(snapshot_window("Premium", Some(out_of_range.as_ref()), None, now,).is_none()); + let non_finite = + raw(r#"{ "entitlement": 100, "remaining": "NaN", "percent_remaining": null }"#); + assert!(snapshot_window("Chat", Some(non_finite.as_ref()), None, now).is_none()); } #[test] - fn unknown_test_window_has_no_history_identity() { - let window = snapshot_window( - "Other", - Some(QuotaSnapshot { - entitlement: 10.0, - remaining: 5.0, - percent_remaining: None, - }), - None, - Utc::now(), - ) - .unwrap(); - assert_eq!(window.card_id_for_test(), "row.copilot.unknown.v1"); - assert!(window.pace_window_key_for_test().is_none()); - assert_eq!(window.pace_reason_for_test(), Some("windowIdentity")); + fn rejects_percentages_that_contradict_absolute_quota() { + let now = Utc::now(); + let contradictory = + raw(r#"{ "entitlement": 100, "remaining": 0, "percent_remaining": 100 }"#); + assert!(matches!( + snapshot_window_with_identity( + "Premium", + "premium_interactions.v1", + Some(contradictory.as_ref()), + None, + now, + ), + CopilotRow::Invalid + )); + + for rounded in [66, 67] { + let payload = raw(&format!( + r#"{{ "entitlement": 3, "remaining": 2, "percent_remaining": {rounded} }}"# + )); + assert!(matches!( + snapshot_window_with_identity( + "Premium", + "premium_interactions.v1", + Some(payload.as_ref()), + None, + now, + ), + CopilotRow::Usable(_) + )); + } } #[test] - fn parses_only_trimmed_exact_utc_calendar_dates() { - assert_eq!( - parse_reset_date(" 2026-07-01 \t"), - Some("2026-07-01T00:00:00Z".parse::>().unwrap()) - ); - for invalid in [ - "not-a-date", - "2026-7-01", - "2026-07-1", - "2026-07-01T00:00:00Z", - "2026/07/01", - "2026-02-29", - ] { - assert!(parse_reset_date(invalid).is_none(), "{invalid}"); + fn malformed_snapshot_percentage_does_not_poison_valid_sibling() { + let now = Utc.timestamp_opt(1_751_328_000, 0).single().unwrap(); + for invalid in ["1e400", r#""NaN""#, "100"] { + let usage: CopilotUser = serde_json::from_str(&format!( + r#"{{ + "quota_reset_date": "2026-08-01", + "quota_snapshots": {{ + "premium_interactions": {{ + "entitlement": 300, + "remaining": 90, + "percent_remaining": {invalid} + }}, + "chat": {{ + "entitlement": 100, + "remaining": 75, + "percent_remaining": 75 + }} + }} + }}"# + )) + .unwrap(); + let (_, windows, mapping) = map_user(usage, now); + assert_eq!(mapping, CopilotMapping::Usable); + assert_eq!(windows.len(), 1); + assert_eq!(windows[0].label_for_test(), "Chat"); + assert!((windows[0].remaining_for_test() - 75.0).abs() < 0.01); } } + + #[test] + fn parses_reset_date() { + assert!(parse_reset_date("2026-07-01").is_some()); + assert!(parse_reset_date("not-a-date").is_none()); + } } diff --git a/crates/tb_core_ffi/src/agent_grok.rs b/crates/tb_core_ffi/src/agent_grok.rs index 639b3ea..24b2b24 100644 --- a/crates/tb_core_ffi/src/agent_grok.rs +++ b/crates/tb_core_ffi/src/agent_grok.rs @@ -1,21 +1,31 @@ -//! Grok Build subscription quota (weekly SuperGrok credits). +//! Grok Build subscription quota — two meters, shown as two windows. //! //! Grok Build stores OIDC credentials at `$GROK_HOME/auth.json` (default //! `~/.grok/auth.json`). TokenBar refreshes the access token against -//! `auth.x.ai` and reads weekly credit usage from the same private billing +//! `auth.x.ai` and reads usage from two views of the same private billing //! endpoint the CLI uses: //! -//! GET https://cli-chat-proxy.grok.com/v1/billing?format=credits +//! Weekly: GET /v1/billing?format=credits -> creditUsagePercent / GrokBuild +//! Monthly: GET /v1/billing -> used / monthlyLimit //! -//! Prefer the `GrokBuild` product percent when present; fall back to overall -//! `creditUsagePercent`. Omit the card entirely when no Grok auth is on disk -//! (same stance as Copilot). +//! The weekly view is the SuperGrok weekly credit meter (the primary "will I +//! run out this week" number). Right after a weekly reset with no usage yet, it +//! OMITS the percent fields while still reporting the period — that is a genuine +//! 0%, not an error, and was the original cause of the card erroring. The +//! monthly view is the included-allowance meter (percent = used / monthlyLimit) +//! over a monthly period. The monthly call is best-effort with a short timeout: +//! a failure or hang there never sinks the card when the weekly meter succeeded. +//! Omit the card entirely when no Grok auth is on disk (same stance as Copilot). use crate::agent_account_scope::{ self, AccountScope, AccountScopeError, RefreshCheckpoint, RefreshScopeTransaction, }; use crate::agent_quota_duration::DurationEvidence; -use crate::agent_usage::{AgentIdentity, UsageWindow}; +use crate::agent_usage::{ + provider_http_client_builder, read_response_body, request_after_verified_binding, + AgentIdentity, ProviderCacheBinding, ProviderFetchFailure, ResponseReadFailure, + TransportErrorFacts, TransportPhase, UsageWindow, +}; use chrono::{DateTime, SecondsFormat, Utc}; use serde::Deserialize; use serde_json::{value::RawValue, Value}; @@ -23,13 +33,30 @@ use std::fs; use std::path::{Path, PathBuf}; const GROK_TOKEN_URL: &str = "https://auth.x.ai/oauth2/token"; -const GROK_BILLING_URL: &str = "https://cli-chat-proxy.grok.com/v1/billing?format=credits"; +/// Weekly SuperGrok credits meter. Returns `creditUsagePercent` / a `GrokBuild` +/// product percent over a weekly `currentPeriod`. Right after a weekly reset, +/// with no usage recorded yet, xAI OMITS those percent fields — the period is +/// still reported, so that state is a genuine 0%, not an error. +const GROK_CREDITS_URL: &str = "https://cli-chat-proxy.grok.com/v1/billing?format=credits"; +/// Monthly included-allowance meter. The default view reports `monthlyLimit` + +/// `used` over a monthly billing period (percent = used / monthlyLimit). +const GROK_MONTHLY_URL: &str = "https://cli-chat-proxy.grok.com/v1/billing"; /// Refresh a few minutes early so a clock-skewed expiry doesn't 401 the billing call. const ACCESS_SKEW_SECS: i64 = 120; +/// Best-effort monthly GET budget. Grok is joined in `agent_usage::run`, so a +/// full 30s hang on the additive meter would stall the whole quota payload after +/// the weekly meter is already ready. +const MONTHLY_TIMEOUT_SECS: u64 = 5; +/// The exact `currentPeriod.type` xAI stamps on the weekly credits meter. Matched +/// exactly (not by substring) so `..._BIWEEKLY` / `..._NOT_WEEKLY` can't pass. +const WEEKLY_PERIOD_TYPE: &str = "USAGE_PERIOD_TYPE_WEEKLY"; +const WEEKLY_WINDOW_KEY: &str = "billing.weekly.v1"; +const MONTHLY_WINDOW_KEY: &str = "billing.monthly.v1"; pub(crate) struct GrokData { pub identity: Option, pub account_scope: Result, + pub cache_binding: Option, pub windows: Vec, } @@ -47,9 +74,9 @@ struct GrokCredentials { } impl GrokCredentials { - fn scope_marker(&self) -> Option<&[u8]> { + fn scope_marker(&self) -> Option<&str> { let marker = self.refresh_token.trim(); - (!marker.is_empty()).then_some(marker.as_bytes()) + (!marker.is_empty()).then_some(marker) } fn scope_location(&self) -> Result { @@ -64,7 +91,7 @@ impl GrokCredentials { "grok", "grok-auth-json", &self.scope_location()?, - marker, + marker.as_bytes(), ) } } @@ -83,32 +110,49 @@ struct BillingResponse { struct BillingConfig { #[serde(default)] current_period: Option, - #[serde(default, deserialize_with = "deserialize_optional_raw")] + /// Unified-billing allowance and consumption (`{ "val": n }`). Percent is + /// `used / monthly_limit`. Present in the default billing view. + #[serde(default)] + monthly_limit: Option, + #[serde(default)] + used: Option, + /// RPC-shaped consumption (`usage.totalUsed`); accepted defensively in case + /// an account nests the consumed amount under `usage` instead of `used`. + #[serde(default)] + usage: Option, + /// Credits-view percent fields (`?format=credits`). + #[serde(default)] credit_usage_percent: Option>, #[serde(default)] product_usage: Option>, - #[serde(default, deserialize_with = "deserialize_optional_raw")] - billing_period_start: Option>, - #[serde(default, deserialize_with = "deserialize_optional_raw")] - billing_period_end: Option>, - #[serde(default, deserialize_with = "deserialize_optional_raw")] - used: Option>, - #[serde(default, deserialize_with = "deserialize_optional_raw")] - monthly_limit: Option>, - #[serde(default, deserialize_with = "deserialize_optional_raw")] - on_demand_used: Option>, - #[serde(default, deserialize_with = "deserialize_optional_raw")] - on_demand_cap: Option>, + #[serde(default)] + billing_period_start: Option, + #[serde(default)] + billing_period_end: Option, +} + +/// xAI wraps billing amounts as `{ "val": n }` on the wire. +#[derive(Debug, Deserialize)] +struct CentVal { + #[serde(default)] + val: Option, +} + +#[derive(Debug, Deserialize)] +#[serde(rename_all = "camelCase")] +struct UnifiedUsage { + #[serde(default)] + total_used: Option, } #[derive(Debug, Deserialize)] struct UsagePeriod { #[serde(default, rename = "type")] period_type: Option, - #[serde(default, deserialize_with = "deserialize_optional_raw")] - start: Option>, - #[serde(default, deserialize_with = "deserialize_optional_raw")] - end: Option>, + #[serde(default)] + start: Option, + #[serde(default)] + end: Option, } #[derive(Debug, Deserialize)] @@ -116,25 +160,10 @@ struct UsagePeriod { struct ProductUsage { #[serde(default)] product: Option, - #[serde(default, deserialize_with = "deserialize_optional_raw")] + #[serde(default)] usage_percent: Option>, } -enum LegacyPercentageEvidence { - Absent, - Invalid, - GrokBuild(f64), - Credit(f64), - CreditZero(f64), -} - -enum PercentageEvidence { - Absent, - Disabled, - Invalid, - Valid(f64), -} - #[derive(Debug, Deserialize)] struct TokenResponse { access_token: String, @@ -144,467 +173,392 @@ struct TokenResponse { expires_in: Option, } -/// Fetch Grok quota when local auth exists. Returns `None` when the user has -/// never signed into Grok Build (no card). Returns `Err` when auth exists but -/// the fetch fails so the card can show an error state. -pub(crate) async fn fetch(now: DateTime) -> Option> { +/// Fetch Grok quota when local auth exists. A genuinely absent credential +/// omits the optional card; every present-credential failure stays typed. +pub(crate) async fn fetch(now: DateTime) -> Result, ProviderFetchFailure> { let credentials = match load_credentials() { - Ok(Some(c)) => c, - Ok(None) => return None, - Err(e) => return Some(Err(e)), + Ok(Some(credentials)) => credentials, + Ok(None) => return Ok(None), + Err(_) => { + return Err(ProviderFetchFailure::terminal( + "Grok credentials could not be loaded.", + )); + } }; - Some(fetch_with_credentials(credentials, now).await) + fetch_with_credentials(credentials, now).await.map(Some) } async fn fetch_with_credentials( - mut credentials: GrokCredentials, + credentials: GrokCredentials, now: DateTime, -) -> Result { - let mut refreshed_scope = None; - if credentials_needs_refresh(&credentials, now) { - let refreshed = - refresh_credentials(&credentials.auth_path, &credentials.entry_key, false).await?; - credentials = refreshed.0; - refreshed_scope = merge_refreshed_scope(refreshed_scope, refreshed.1); +) -> Result { + let verified = if credentials_needs_refresh(&credentials, now) { + refresh_credentials(&credentials.auth_path, &credentials.entry_key, false).await + } else { + credentials + .resolve_account_scope() + .map(|account_scope| { + let cache_binding = ProviderCacheBinding::primary(account_scope.clone()); + (credentials, account_scope, Some(cache_binding)) + }) + .map_err(|_| { + ProviderFetchFailure::terminal("Grok account identity could not be verified.") + }) + }; + + let (mut credentials, mut account_scope, mut cache_binding, client, response) = + request_after_verified_binding( + verified, + |(credentials, account_scope, cache_binding)| async move { + let client = provider_http_client_builder() + .timeout(std::time::Duration::from_secs(30)) + .build() + .map_err(|_| { + ProviderFetchFailure::terminal("Grok billing client could not be created.") + })?; + let response = client + .get(GROK_CREDITS_URL) + .bearer_auth(&credentials.access_token) + .header(reqwest::header::ACCEPT, "application/json") + .header(reqwest::header::USER_AGENT, "TokenBar") + .send() + .await + .map_err(|error| { + ProviderFetchFailure::from_send_error( + "Grok weekly billing request failed. Retrying automatically.", + cache_binding.clone(), + &error, + ) + })?; + Ok((credentials, account_scope, cache_binding, client, response)) + }, + ) + .await?; + let status = response.status().as_u16(); + + let weekly_body = if matches!(status, 401 | 403) { + if credentials.scope_marker().is_none() { + return Err(ProviderFetchFailure::terminal( + "Grok OAuth token expired or invalid. Run `grok` to log in again.", + )); + } + let (refreshed, scope, binding) = + refresh_credentials(&credentials.auth_path, &credentials.entry_key, true).await?; + credentials = refreshed; + account_scope = scope; + cache_binding = binding; + let retry = client + .get(GROK_CREDITS_URL) + .bearer_auth(&credentials.access_token) + .header(reqwest::header::ACCEPT, "application/json") + .header(reqwest::header::USER_AGENT, "TokenBar") + .send() + .await + .map_err(|error| { + ProviderFetchFailure::from_send_error( + "Grok weekly billing retry failed. Retrying automatically.", + cache_binding.clone(), + &error, + ) + })?; + let retry_status = retry.status().as_u16(); + read_response_body(retry_status, false, || async { + retry.text().await.map_err(|error| { + TransportErrorFacts::from_reqwest(&error, TransportPhase::ResponseBody) + }) + }) + .await + .map_err(|failure| grok_response_failure(failure, cache_binding.clone()))? + } else { + read_response_body(status, false, || async { + response.text().await.map_err(|error| { + TransportErrorFacts::from_reqwest(&error, TransportPhase::ResponseBody) + }) + }) + .await + .map_err(|failure| grok_response_failure(failure, cache_binding.clone()))? + }; + + let monthly_body = fetch_monthly_best_effort(&credentials).await; + build_grok_data( + &weekly_body, + monthly_body.as_deref(), + &credentials, + now, + Ok(account_scope), + cache_binding, + ) +} + +fn grok_response_failure( + failure: ResponseReadFailure, + attempt_binding: Option, +) -> ProviderFetchFailure { + match failure { + ResponseReadFailure::Transient(diagnostic) => ProviderFetchFailure::transient( + "Grok weekly billing request failed. Retrying automatically.", + attempt_binding, + diagnostic, + ), + ResponseReadFailure::Terminal(401 | 403) => ProviderFetchFailure::terminal( + "Grok OAuth token expired or invalid. Run `grok` to log in again.", + ), + ResponseReadFailure::Terminal(status) => ProviderFetchFailure::terminal(format!( + "Grok billing API rejected the request (status {status})." + )), } +} +/// GET the monthly default billing view. Never refreshes credentials on 4xx — +/// the weekly call already owns token repair — and uses a short timeout so a +/// hung additive request cannot stall the joined quota payload. +async fn fetch_monthly_best_effort(credentials: &GrokCredentials) -> Option { let client = reqwest::Client::builder() - .timeout(std::time::Duration::from_secs(30)) + .timeout(std::time::Duration::from_secs(MONTHLY_TIMEOUT_SECS)) .build() - .map_err(|e| format!("build Grok billing client: {e}"))?; - + .ok()?; let response = client - .get(GROK_BILLING_URL) + .get(GROK_MONTHLY_URL) .bearer_auth(&credentials.access_token) .header(reqwest::header::ACCEPT, "application/json") .header(reqwest::header::USER_AGENT, "TokenBar") .send() .await - .map_err(|e| format!("Grok billing request failed: {e}"))?; - - let status = response.status(); - let body = response - .text() - .await - .map_err(|e| format!("read Grok billing response: {e}"))?; - - if status == reqwest::StatusCode::UNAUTHORIZED || status == reqwest::StatusCode::FORBIDDEN { - // One retry after a forced refresh in case the access token was revoked - // mid-window while the refresh token still works. - if !credentials.refresh_token.is_empty() { - let refreshed = - refresh_credentials(&credentials.auth_path, &credentials.entry_key, true).await?; - credentials = refreshed.0; - refreshed_scope = merge_refreshed_scope(refreshed_scope, refreshed.1); - let retry = client - .get(GROK_BILLING_URL) - .bearer_auth(&credentials.access_token) - .header(reqwest::header::ACCEPT, "application/json") - .header(reqwest::header::USER_AGENT, "TokenBar") - .send() - .await - .map_err(|e| format!("Grok billing retry failed: {e}"))?; - let retry_status = retry.status(); - let retry_body = retry - .text() - .await - .map_err(|e| format!("read Grok billing retry: {e}"))?; - if !retry_status.is_success() { - return Err(format!( - "Grok billing API returned {}.", - retry_status.as_u16() - )); - } - let account_scope = - refreshed_scope.unwrap_or_else(|| credentials.resolve_account_scope()); - return map_billing(&retry_body, &credentials, now, account_scope); - } - return Err("Grok OAuth token expired or invalid. Run `grok` to log in again.".to_string()); - } - if !status.is_success() { - return Err(format!("Grok billing API returned {}.", status.as_u16())); + .ok()?; + if !response.status().is_success() { + return None; } - - let account_scope = refreshed_scope.unwrap_or_else(|| credentials.resolve_account_scope()); - map_billing(&body, &credentials, now, account_scope) -} - -fn merge_refreshed_scope( - current: Option>, - next: Result, -) -> Option> { - Some(match current { - None => next, - Some(Err(first_error)) => Err(first_error), - Some(Ok(current_scope)) => match next { - Err(error) => Err(error), - Ok(next_scope) if next_scope == current_scope => Ok(current_scope), - Ok(_) => Err(AccountScopeError::MetadataConflict), - }, - }) + response.text().await.ok() } -fn map_billing( - body: &str, +/// Assemble the card's windows from the weekly credits view (required) and the +/// monthly unified view (additive only). Success requires a weekly window — +/// monthly-only is not enough to show the card, so a missing/unusable weekly +/// meter errors even when monthly parsed cleanly. +fn build_grok_data( + credits_body: &str, + monthly_body: Option<&str>, credentials: &GrokCredentials, now: DateTime, account_scope: Result, -) -> Result { - let payload: BillingResponse = - serde_json::from_str(body).map_err(|e| format!("decode Grok billing response: {e}"))?; - let config = payload - .config - .ok_or_else(|| "Grok billing response missing config.".to_string())?; - - let (used_percent, included_monthly_fallback, included_invalid, legacy_invalid) = - match legacy_percentage_evidence(&config) { - LegacyPercentageEvidence::GrokBuild(percent) - | LegacyPercentageEvidence::Credit(percent) => (Some(percent), false, false, false), - LegacyPercentageEvidence::CreditZero(percent) => { - match included_percentage_evidence(&config) { - PercentageEvidence::Valid(included) => (Some(included), true, false, false), - PercentageEvidence::Disabled => (None, false, false, false), - PercentageEvidence::Absent | PercentageEvidence::Invalid => { - (Some(percent), false, false, false) - } - } - } - LegacyPercentageEvidence::Invalid => (None, false, false, true), - LegacyPercentageEvidence::Absent => match included_percentage_evidence(&config) { - PercentageEvidence::Valid(percent) => (Some(percent), true, false, false), - PercentageEvidence::Disabled | PercentageEvidence::Absent => { - (None, false, false, false) - } - PercentageEvidence::Invalid => (None, false, true, false), - }, - }; - let (on_demand_recognized, extra_window) = extra_usage_window(&config, now); - if legacy_invalid || (included_invalid && extra_window.is_none()) { - return Err( - "Grok billing response has no creditUsagePercent or GrokBuild usage.".to_string(), - ); - } - if used_percent.is_none() && !on_demand_recognized { - return Err( - "Grok billing response has no creditUsagePercent or GrokBuild usage.".to_string(), - ); - } + cache_binding: Option, +) -> Result { + let credits: BillingResponse = serde_json::from_str(credits_body).map_err(|_| { + ProviderFetchFailure::terminal("Grok billing response could not be decoded.") + })?; - let mut windows = Vec::with_capacity(2); - if let Some(used_percent) = used_percent { - let period = period_details(&config); - let kind = period.kind.or( - included_monthly_fallback.then_some(("Monthly", "billing.monthly.v1")), - ); - let window = match kind { - Some((label, window_key)) => { - let mut window = UsageWindow::from_used_percent( - label.to_string(), - used_percent, - period.end, - now, - ) - .with_identity(window_key, Some(window_key.to_string())); - - if period.invalid_evidence { - let invalid_provider = period - .end - .map(|end| DurationEvidence::provider(provider_reset_at(end, now), 0)); - window = window.with_provider_duration_evidence(now, true, invalid_provider); - } else if let (Some(end), Some(duration_seconds)) = - (period.end, period.duration_seconds) - { - window = window.with_provider_duration_evidence( - now, - true, - Some(DurationEvidence::provider( - provider_reset_at(end, now), - duration_seconds, - )), - ); - } else if period.end_was_supplied { - window = window.with_observed_duration_evidence(now, true); + let weekly = credits + .config + .as_ref() + .and_then(|config| weekly_window(config, now)) + .ok_or_else(|| { + ProviderFetchFailure::terminal("Grok billing response had no usable weekly usage.") + })?; + + let mut windows = vec![weekly]; + if let Some(body) = monthly_body { + if let Ok(monthly) = serde_json::from_str::(body) { + if let Some(config) = monthly.config.as_ref() { + if let Some(window) = monthly_window(config, now) { + windows.push(window); } - window - } - None => { - UsageWindow::from_used_percent("Unknown".to_string(), used_percent, period.end, now) - .with_identity("row.billing.unknown.v1", None) } - }; - windows.push(window); - } - if let Some(window) = extra_window { - windows.push(window); + } } Ok(GrokData { identity: Some(AgentIdentity { email: credentials.email.clone(), - plan: payload + plan: credits .subscription_tiers .filter(|s| !s.trim().is_empty()) .map(|s| s.trim().to_string()), }), account_scope, + cache_binding, windows, }) } -fn legacy_percentage_evidence(config: &BillingConfig) -> LegacyPercentageEvidence { - if let Some(products) = config.product_usage.as_ref() { - let mut grok_build_percent = None; - for product in products { - if !product - .product - .as_deref() - .is_some_and(|name| name.eq_ignore_ascii_case("GrokBuild")) - { - continue; - } - if let Some(raw) = product.usage_percent.as_deref() { - let Some(percent) = valid_percentage(raw) else { - return LegacyPercentageEvidence::Invalid; - }; - grok_build_percent.get_or_insert(percent); - } - } - if let Some(percent) = grok_build_percent { - return LegacyPercentageEvidence::GrokBuild(percent); - } - } - match config.credit_usage_percent.as_deref() { - Some(raw) => match valid_percentage(raw) { - Some(percent) if decimal_mantissa_is_zero(raw.get()) => { - LegacyPercentageEvidence::CreditZero(percent) - } - Some(percent) => LegacyPercentageEvidence::Credit(percent), - None => LegacyPercentageEvidence::Invalid, - }, - None => LegacyPercentageEvidence::Absent, - } -} - -fn included_percentage_evidence(config: &BillingConfig) -> PercentageEvidence { - let (used, monthly_limit) = match (config.used.as_deref(), config.monthly_limit.as_deref()) { - (None, None) => return PercentageEvidence::Absent, - (Some(used), Some(monthly_limit)) => (used, monthly_limit), - _ => return PercentageEvidence::Invalid, - }; - let (Some(used), Some(monthly_limit)) = - (raw_amount_value(used), raw_amount_value(monthly_limit)) - else { - return PercentageEvidence::Invalid; +/// Weekly SuperGrok credits window. When the credits view reports a valid +/// percent, use it. Empty-week 0% is allowed only when percent fields are truly +/// *absent* and a self-contained weekly `currentPeriod` proves the meter exists +/// (fresh reset, no usage yet). Present-but-unparsable or out-of-range percent +/// is a failed reading — never synthesize 0% for bad data. A malformed, partial, +/// or non-weekly period with no usable percent is "unknown", not zero. +fn weekly_window(config: &BillingConfig, now: DateTime) -> Option { + let used_percent = match weekly_used_percent(config) { + WeeklyPercent::Value(pct) => pct, + // Absent percent is a genuine 0% ONLY when a self-contained weekly + // `currentPeriod` proves the meter exists: the weekly type AND a + // positive start→end window must come from that same object, not a mix + // of a weekly-typed period and unrelated flat billing dates. + WeeklyPercent::Absent if self_contained_weekly_period(config).is_some() => 0.0, + WeeklyPercent::Absent | WeeklyPercent::Invalid => return None, }; - if used == 0.0 && monthly_limit == 0.0 { - return PercentageEvidence::Disabled; - } - if used < 0.0 || monthly_limit <= 0.0 { - return PercentageEvidence::Invalid; + let period = period_bounds(config); + let mut window = UsageWindow::from_provider_used_percent( + "Weekly".to_string(), + used_percent, + period.end, + now, + ) + .with_identity( + WEEKLY_WINDOW_KEY, + Some(WEEKLY_WINDOW_KEY.to_string()), + period.duration, + None, + ); + if period.invalid_evidence { + window.unavailable("invalidEvidence"); } - PercentageEvidence::Valid(((used / monthly_limit) * 100.0).min(100.0)) + Some(window) } -fn extra_usage_window(config: &BillingConfig, now: DateTime) -> (bool, Option) { - let Some(cap) = config - .on_demand_cap - .as_deref() - .and_then(raw_amount_value) - else { - return (false, None); - }; - if cap < 0.0 { - return (false, None); - } - if cap == 0.0 { - return (true, None); - } - - let Some(used) = config - .on_demand_used - .as_deref() - .and_then(raw_amount_value) - else { - return (false, None); - }; - if used < 0.0 { - return (false, None); +/// The positive window length of a self-contained weekly `currentPeriod` — the +/// exact weekly type plus a `start < end` pair parsed from that same object — +/// or `None`. This is the "the weekly meter exists" signal that lets an absent +/// percent be read as 0%; it deliberately does not consult the flat +/// `billingPeriod*` fields, so a partial or non-weekly period cannot borrow an +/// unrelated window to masquerade as a weekly meter. +fn self_contained_weekly_period(config: &BillingConfig) -> Option<(DateTime, DateTime)> { + let period = config.current_period.as_ref()?; + if period.period_type.as_deref() != Some(WEEKLY_PERIOD_TYPE) { + return None; } - - let used_percent = ((used / cap) * 100.0).min(100.0); - let reset = config - .billing_period_end - .as_deref() - .and_then(parse_timestamp_raw); - let window = - UsageWindow::from_used_percent("Extra usage".to_string(), used_percent, reset, now) - .with_identity("extra_usage.v1", Some("extra_usage.v1".to_string())) - .with_non_recurring(); - (true, Some(window)) + let start = period.start.as_deref().and_then(parse_timestamp)?; + let end = period.end.as_deref().and_then(parse_timestamp)?; + (end > start).then_some((start, end)) } -fn raw_amount_value(raw: &RawValue) -> Option { - let raw_value = raw.get(); - strict_json_number(raw_value).or_else(|| { - let fields: std::collections::BTreeMap> = - serde_json::from_str(raw_value).ok()?; - strict_json_number(fields.get("val")?.get()) - }) +/// Result of reading weekly percent fields. Distinguishes true omission (empty +/// week may be 0%) from a present-but-bad value (must not become 0%). +#[derive(Debug, Clone, Copy, PartialEq)] +enum WeeklyPercent { + Value(f64), + Absent, + Invalid, } -fn strict_json_number(raw_value: &str) -> Option { - let value = serde_json::from_str::(raw_value).ok()?; - // Reject any syntactically nonzero number that underflows to signed zero. - if value == 0.0 && !decimal_mantissa_is_zero(raw_value) { - return None; +/// Prefer `GrokBuild.usagePercent`; only when that field is absent (product +/// missing or key omitted) fall through to `creditUsagePercent`. A present but +/// unparsable/out-of-range value on the chosen field is `Invalid`, not `Absent`. +fn weekly_used_percent(config: &BillingConfig) -> WeeklyPercent { + if let Some(products) = config.product_usage.as_ref() { + for product in products { + let name = product.product.as_deref().unwrap_or(""); + if name.eq_ignore_ascii_case("GrokBuild") { + if let Some(usage_percent) = product.usage_percent.as_deref() { + return match valid_percentage(usage_percent) { + Some(pct) => WeeklyPercent::Value(pct), + None => WeeklyPercent::Invalid, + }; + } + // GrokBuild row present but usagePercent key omitted — try overall. + break; + } + } + } + match config.credit_usage_percent.as_deref() { + Some(raw) => match valid_percentage(raw) { + Some(pct) => WeeklyPercent::Value(pct), + None => WeeklyPercent::Invalid, + }, + None => WeeklyPercent::Absent, } - value.is_finite().then_some(value) -} - -fn decimal_mantissa_is_zero(value: &str) -> bool { - let value = value.trim(); - let unsigned = value.strip_prefix('-').unwrap_or(value); - let mantissa = unsigned.split(['e', 'E']).next().unwrap_or(unsigned); - !mantissa - .bytes() - .any(|digit| digit.is_ascii_digit() && digit != b'0') } -fn deserialize_optional_raw<'de, D>(deserializer: D) -> Result>, D::Error> -where - D: serde::Deserializer<'de>, -{ - Box::::deserialize(deserializer).map(Some) +/// Monthly included-allowance window: percent = used / monthlyLimit. Only +/// emitted when the consumed amount is *explicitly present* and non-negative — +/// xAI reports a genuine zero as `{ "val": 0 }` (the `history` cycles do +/// exactly this), so an absent `used`/`usage.totalUsed` means "unknown", not +/// zero. A negative consumption is invalid meter data and must not clamp into +/// a healthy 0%-used window. +fn monthly_window(config: &BillingConfig, now: DateTime) -> Option { + let limit = config + .monthly_limit + .as_ref() + .and_then(|c| c.val) + .filter(|v| *v > 0)?; + let used = config + .used + .as_ref() + .and_then(|c| c.val) + .or_else(|| { + config + .usage + .as_ref() + .and_then(|u| u.total_used.as_ref()) + .and_then(|c| c.val) + }) + .filter(|v| *v >= 0)?; + let used_percent = (used as f64 / limit as f64 * 100.0).clamp(0.0, 100.0); + let period = period_bounds(config); + let mut window = UsageWindow::from_provider_used_percent( + "Monthly".to_string(), + used_percent, + period.end, + now, + ) + .with_identity( + MONTHLY_WINDOW_KEY, + Some(MONTHLY_WINDOW_KEY.to_string()), + period.duration, + None, + ); + if period.invalid_evidence { + window.unavailable("invalidEvidence"); + } + Some(window) } fn valid_percentage(raw: &RawValue) -> Option { - let value = raw.get(); - serde_json::from_str::(value) + serde_json::from_str::(raw.get()) .ok() - .filter(|pct| pct.is_finite() && decimal_percentage_in_range(value)) -} - -fn decimal_percentage_in_range(value: &str) -> bool { - let (negative, unsigned) = match value.strip_prefix('-') { - Some(unsigned) => (true, unsigned), - None => (false, value), - }; - let exponent_index = unsigned.find(|character| matches!(character, 'e' | 'E')); - let (mantissa, exponent_text) = match exponent_index { - Some(index) => (&unsigned[..index], Some(&unsigned[index + 1..])), - None => (unsigned, None), - }; - let fractional_digits = mantissa - .split_once('.') - .map_or(0, |(_, fraction)| fraction.len()); - let digits = mantissa - .bytes() - .filter(|byte| byte.is_ascii_digit()) - .collect::>(); - let Some(first_nonzero) = digits.iter().position(|digit| *digit != b'0') else { - return true; - }; - if negative { - return false; - } - let exponent = match exponent_text { - Some(text) => match text.parse::() { - Ok(exponent) => exponent, - Err(_) => return text.starts_with('-'), - }, - None => 0, - }; - - let significant = &digits[first_nonzero..]; - let integer_digits = exponent - .saturating_add(significant.len() as i128) - .saturating_sub(fractional_digits as i128); - match integer_digits.cmp(&3) { - std::cmp::Ordering::Less => true, - std::cmp::Ordering::Greater => false, - std::cmp::Ordering::Equal => { - significant[0] == b'1' && significant[1..].iter().all(|digit| *digit == b'0') - } - } + .filter(|pct| pct.is_finite() && (0.0..=100.0).contains(pct)) } -struct PeriodDetails { - kind: Option<(&'static str, &'static str)>, +struct PeriodMeta { end: Option>, - end_was_supplied: bool, - duration_seconds: Option, + duration: Option, invalid_evidence: bool, } -fn period_details(config: &BillingConfig) -> PeriodDetails { +/// Reset instant and duration evidence from a config's period. Prefers the +/// explicit `currentPeriod`, falling back to the flat `billingPeriodStart/End`. +/// Identity (label / window_key) is fixed per meter and is not derived here. +fn period_bounds(config: &BillingConfig) -> PeriodMeta { let period = config.current_period.as_ref(); - let period_type = period - .and_then(|period| period.period_type.as_deref()) - .unwrap_or("") - .trim() - .to_ascii_uppercase(); - let kind = match period_type.as_str() { - "WEEKLY" | "USAGE_PERIOD_TYPE_WEEKLY" => Some(("Weekly", "billing.weekly.v1")), - "MONTHLY" | "USAGE_PERIOD_TYPE_MONTHLY" => Some(("Monthly", "billing.monthly.v1")), - _ => None, - }; - - // Select each primary field independently. A present but malformed primary - // value on a recognized period must not be hidden by a valid billing-level - // fallback. Unrecognized nested periods are unrelated to the synthesized - // monthly included card, so only billing-level dates apply to them. - let (start_raw, end_raw) = if kind.is_some() { - ( - period - .and_then(|period| period.start.as_deref()) - .or(config.billing_period_start.as_deref()), - period - .and_then(|period| period.end.as_deref()) - .or(config.billing_period_end.as_deref()), - ) - } else { - ( - config.billing_period_start.as_deref(), - config.billing_period_end.as_deref(), - ) - }; - let start = start_raw.and_then(parse_timestamp_raw); - let end = end_raw.and_then(parse_timestamp_raw); - let (duration_seconds, invalid_evidence) = match (start_raw, end_raw, start, end) { - (Some(_), Some(_), Some(start), Some(end)) if end > start => { - (Some((end - start).num_seconds()), false) - } + let start_raw = period + .and_then(|period| period.start.as_deref()) + .or(config.billing_period_start.as_deref()); + let end_raw = period + .and_then(|period| period.end.as_deref()) + .or(config.billing_period_end.as_deref()); + let start = start_raw.and_then(parse_timestamp); + let end = end_raw.and_then(parse_timestamp); + let (duration, invalid_evidence) = match (start_raw, end_raw, start, end) { + (Some(_), Some(_), Some(start), Some(end)) if end > start => ( + Some(DurationEvidence::provider( + end.timestamp(), + (end - start).num_seconds(), + )), + false, + ), (Some(_), Some(_), Some(_), Some(_)) => (None, true), (Some(_), Some(_), _, _) => (None, true), - (Some(_), None, Some(_), _) => (None, false), - (Some(_), None, None, _) => (None, true), + (Some(_), None, _, _) => (None, false), (None, Some(_), _, Some(_)) => (None, false), (None, Some(_), _, None) => (None, true), _ => (None, false), }; - PeriodDetails { - kind, + PeriodMeta { end, - end_was_supplied: end_raw.is_some(), - duration_seconds, + duration, invalid_evidence, } } -fn provider_reset_at(reset: DateTime, now: DateTime) -> i64 { - if reset > now { - reset.timestamp().max(now.timestamp().saturating_add(1)) - } else { - reset.timestamp() - } -} - -fn parse_timestamp_raw(value: &RawValue) -> Option> { - serde_json::from_str::(value.get()) - .ok() - .and_then(|value| parse_timestamp(&value)) -} - fn parse_timestamp(value: &str) -> Option> { DateTime::parse_from_rfc3339(value.trim()) .ok() @@ -631,9 +585,10 @@ async fn refresh_credentials( auth_path: &Path, entry_key: &str, force: bool, -) -> Result<(GrokCredentials, Result), String> { - let refresh = agent_account_scope::begin_refresh("grok") - .map_err(|_| "Grok credential refresh lock is unavailable.".to_string())?; +) -> Result<(GrokCredentials, AccountScope, Option), ProviderFetchFailure> { + let refresh = agent_account_scope::begin_refresh("grok").map_err(|_| { + ProviderFetchFailure::terminal("Grok credential refresh lock is unavailable.") + })?; refresh_credentials_with( auth_path, entry_key, @@ -649,11 +604,12 @@ async fn refresh_credentials( async fn request_refresh( refresh_token: String, client_id: String, -) -> Result { - let client = reqwest::Client::builder() + attempt_binding: ProviderCacheBinding, +) -> Result { + let client = provider_http_client_builder() .timeout(std::time::Duration::from_secs(30)) .build() - .map_err(|e| format!("build Grok token client: {e}"))?; + .map_err(|_| ProviderFetchFailure::terminal("Grok refresh client could not be created."))?; let form = [ ("grant_type", "refresh_token"), ("refresh_token", refresh_token.as_str()), @@ -681,21 +637,35 @@ async fn request_refresh( .body(form) .send() .await - .map_err(|e| format!("Grok token refresh failed: {e}"))?; + .map_err(|error| { + ProviderFetchFailure::from_send_error( + "Grok token refresh failed. Retrying automatically.", + Some(attempt_binding.clone()), + &error, + ) + })?; - let status = response.status(); - let body = response - .text() - .await - .map_err(|e| format!("read Grok token refresh response: {e}"))?; - if !status.is_success() { - return Err(format!( - "Grok token refresh returned {}. Run `grok` to log in again.", - status.as_u16() - )); - } + let status = response.status().as_u16(); + let body = read_response_body(status, false, || async { + response.text().await.map_err(|error| { + TransportErrorFacts::from_reqwest(&error, TransportPhase::ResponseBody) + }) + }) + .await + .map_err(|failure| match failure { + ResponseReadFailure::Transient(diagnostic) => ProviderFetchFailure::transient( + "Grok token refresh failed. Retrying automatically.", + Some(attempt_binding), + diagnostic, + ), + ResponseReadFailure::Terminal(_) => ProviderFetchFailure::terminal( + "Grok token refresh was rejected. Run `grok` to log in again.", + ), + })?; - serde_json::from_str(&body).map_err(|e| format!("decode Grok token refresh: {e}")) + serde_json::from_str(&body).map_err(|_| { + ProviderFetchFailure::terminal("Grok token refresh response could not be decoded.") + }) } async fn refresh_credentials_with( @@ -706,82 +676,107 @@ async fn refresh_credentials_with( request: Request, save: Save, mut checkpoint: Checkpoint, -) -> Result<(GrokCredentials, Result), String> +) -> Result<(GrokCredentials, AccountScope, Option), ProviderFetchFailure> where R: RefreshScopeTransaction + ?Sized, - Request: FnOnce(String, String) -> RequestFuture, - RequestFuture: std::future::Future>, + Request: FnOnce(String, String, ProviderCacheBinding) -> RequestFuture, + RequestFuture: std::future::Future>, Save: FnOnce(&GrokCredentials) -> Result<(), String>, - Checkpoint: FnMut(RefreshCheckpoint) -> Result<(), String>, + Checkpoint: FnMut(RefreshCheckpoint) -> Result<(), ProviderFetchFailure>, { - // Another TokenBar process may have refreshed while this caller waited. - // Reload the exact request-bearing entry only after the provider lock is held. - let mut credentials = load_credentials_entry_from(auth_path, Some(entry_key))? - .ok_or_else(|| "Grok auth entry disappeared during refresh.".to_string())?; + let mut credentials = load_credentials_entry_from(auth_path, Some(entry_key)) + .map_err(|_| ProviderFetchFailure::terminal("Grok credentials could not be reloaded."))? + .ok_or_else(|| { + ProviderFetchFailure::terminal("Grok auth entry disappeared during refresh.") + })?; checkpoint(RefreshCheckpoint::Reloaded)?; - if !force && !credentials_needs_refresh(&credentials, Utc::now()) { - let scope = match credentials.scope_marker() { - Some(marker) => refresh.resolve_current( - "grok-auth-json", - &credentials - .scope_location() - .map_err(|_| "Grok auth location cannot be scoped safely.".to_string())?, - marker, - ), - None => Err(AccountScopeError::NoTrustedEvidence), - }; - return Ok((credentials, scope)); - } - if credentials.refresh_token.trim().is_empty() { - return Err( - "Grok OAuth token needs refresh but auth.json has no refresh token.".to_string(), - ); - } - if credentials.client_id.trim().is_empty() { - return Err("Grok auth.json is missing oidc_client_id.".to_string()); - } - + let needs_refresh = force || credentials_needs_refresh(&credentials, Utc::now()); let old_marker = credentials .scope_marker() - .ok_or_else(|| "Grok OAuth token needs refresh but auth.json has no refresh token.".to_string())? - .to_vec(); - let refresh_token = credentials.refresh_token.trim().to_string(); - let tokens = request(refresh_token, credentials.client_id.clone()).await?; + .ok_or_else(|| { + ProviderFetchFailure::terminal("Grok OAuth credential has no trusted refresh token.") + })? + .to_string(); + if needs_refresh && credentials.client_id.trim().is_empty() { + return Err(ProviderFetchFailure::terminal( + "Grok auth.json is missing oidc_client_id.", + )); + } + let location = credentials + .scope_location() + .map_err(|_| ProviderFetchFailure::terminal("Grok auth location could not be verified."))?; + let pre_scope = refresh + .resolve_current("grok-auth-json", &location, old_marker.as_bytes()) + .map_err(|_| { + ProviderFetchFailure::terminal("Grok account identity could not be verified.") + })?; + let pre_binding = ProviderCacheBinding::primary(pre_scope.clone()); + if !needs_refresh { + return Ok((credentials, pre_scope, Some(pre_binding))); + } + + let tokens = request( + old_marker.clone(), + credentials.client_id.clone(), + pre_binding, + ) + .await?; checkpoint(RefreshCheckpoint::NetworkReturned)?; - credentials.access_token = tokens.access_token; + let access_token = tokens.access_token.trim(); + if access_token.is_empty() { + return Err(ProviderFetchFailure::terminal( + "Grok token refresh response had no access token.", + )); + } + credentials.access_token = access_token.to_string(); if let Some(refresh_token) = tokens .refresh_token - .map(|token| token.trim().to_string()) + .as_deref() + .map(str::trim) .filter(|token| !token.is_empty()) { - credentials.refresh_token = refresh_token; + credentials.refresh_token = refresh_token.to_string(); } if let Some(expires_in) = tokens.expires_in { credentials.expires_at = Some(Utc::now() + chrono::Duration::seconds(expires_in.max(0))); } let new_marker = credentials .scope_marker() - .ok_or_else(|| "Grok refreshed credential has no usable refresh token.".to_string())?; - let refresh_token_rotated = new_marker != old_marker.as_slice(); - let location = credentials - .scope_location() - .map_err(|_| "Grok auth location cannot be scoped safely.".to_string())?; - let scope = refresh.transfer("grok-auth-json", &location, &old_marker, new_marker); + .ok_or_else(|| { + ProviderFetchFailure::terminal( + "Grok refreshed credential has no trusted refresh token.", + ) + })? + .to_string(); + let scope = refresh + .transfer( + "grok-auth-json", + &location, + old_marker.as_bytes(), + new_marker.as_bytes(), + ) + .map_err(|_| { + ProviderFetchFailure::terminal("Grok credential lineage could not be preserved.") + })?; checkpoint(RefreshCheckpoint::MetadataHandled)?; - // A rotated marker may reach disk only after its lineage transfer is durable. - // The refreshed access token remains usable in memory for this poll. - if refresh_token_rotated && scope.is_err() { - return Ok((credentials, scope)); - } - - // If write-back fails, the still-stored old marker resolves the same scope. - if let Err(error) = save(&credentials) { - eprintln!("tb_core_ffi: failed to persist refreshed Grok credentials: {error}"); - } + let persisted = save(&credentials).is_ok(); checkpoint(RefreshCheckpoint::CredentialsPersisted)?; + let cache_binding = if persisted { + Some(ProviderCacheBinding::primary( + refresh + .resolve_current("grok-auth-json", &location, new_marker.as_bytes()) + .map_err(|_| { + ProviderFetchFailure::terminal( + "Grok account identity could not be verified after refresh.", + ) + })?, + )) + } else { + None + }; - Ok((credentials, scope)) + Ok((credentials, scope, cache_binding)) } fn load_credentials() -> Result, String> { @@ -792,6 +787,20 @@ fn load_credentials_from(auth_path: &Path) -> Result, St load_credentials_entry_from(auth_path, None) } +fn credential_token( + entry: &serde_json::Map, + field: &str, +) -> Result, String> { + match entry.get(field) { + None => Ok(None), + Some(Value::String(token)) => { + let token = token.trim(); + Ok((!token.is_empty()).then(|| token.to_string())) + } + Some(_) => Err(format!("Grok auth entry {field} is not a string.")), + } +} + fn load_credentials_entry_from( auth_path: &Path, expected_entry_key: Option<&str>, @@ -830,17 +839,9 @@ fn load_credentials_entry_from( .as_object() .ok_or_else(|| "Grok auth entry is not an object.".to_string())?; - let access_token = obj - .get("key") - .and_then(|v| v.as_str()) - .unwrap_or("") - .to_string(); - let refresh_token = obj - .get("refresh_token") - .and_then(|v| v.as_str()) - .unwrap_or("") - .to_string(); - if access_token.is_empty() && refresh_token.is_empty() { + let access_token = credential_token(obj, "key")?; + let refresh_token = credential_token(obj, "refresh_token")?; + if access_token.is_none() && refresh_token.is_none() { return Ok(None); } @@ -864,8 +865,8 @@ fn load_credentials_entry_from( Ok(Some(GrokCredentials { auth_path: auth_path.to_path_buf(), entry_key, - access_token, - refresh_token, + access_token: access_token.unwrap_or_default(), + refresh_token: refresh_token.unwrap_or_default(), client_id, expires_at, email, @@ -882,13 +883,7 @@ fn load_credentials_entry_from( /// key with no `::` separator has no client-id segment and is not the shape /// Grok writes, so it is rejected too (fail-closed). fn is_grok_auth_entry_key(key: &str) -> bool { - matches!( - key.split_once("::"), - Some((issuer, client_id)) - if issuer == "https://auth.x.ai" - && !client_id.trim().is_empty() - && !client_id.contains("::") - ) + matches!(key.split_once("::"), Some((issuer, _)) if issuer == "https://auth.x.ai") } fn client_id_from_entry_key(key: &str) -> Option { @@ -900,12 +895,27 @@ fn client_id_from_entry_key(key: &str) -> Option { } fn save_credentials(credentials: &GrokCredentials) -> Result<(), String> { - let mut raw = credentials.raw_json.clone(); + let expected_entry = credentials + .raw_json + .as_object() + .and_then(|map| map.get(&credentials.entry_key)) + .ok_or_else(|| "Grok auth entry missing from the loaded credentials.".to_string())?; + let data = fs::read(&credentials.auth_path) + .map_err(|error| format!("reload Grok auth.json before saving: {error}"))?; + let mut raw: Value = serde_json::from_slice(&data) + .map_err(|error| format!("parse Grok auth.json before saving: {error}"))?; + let current_entry = raw + .as_object() + .and_then(|map| map.get(&credentials.entry_key)) + .ok_or_else(|| "Grok auth entry disappeared before saving.".to_string())?; + if current_entry != expected_entry { + return Err("Grok auth entry changed during refresh.".to_string()); + } let entry = raw .as_object_mut() - .and_then(|m| m.get_mut(&credentials.entry_key)) - .and_then(|v| v.as_object_mut()) - .ok_or_else(|| "Grok auth entry missing while saving.".to_string())?; + .and_then(|map| map.get_mut(&credentials.entry_key)) + .and_then(Value::as_object_mut) + .ok_or_else(|| "Grok auth entry is not an object while saving.".to_string())?; entry.insert( "key".to_string(), @@ -983,752 +993,357 @@ fn grok_home() -> PathBuf { mod tests { use super::*; use crate::agent_account_scope::test_support::TestRefreshScope; + use crate::agent_usage::SafeTransportDiagnostic; + use std::sync::{ + atomic::{AtomicUsize, Ordering}, + Arc, + }; - #[test] - fn prefers_grok_build_product_percent() { - let config: BillingConfig = serde_json::from_str( - r#"{ - "creditUsagePercent": 50.0, - "productUsage": [ - { "product": "GrokChat", "usagePercent": 10.0 }, - { "product": "GrokBuild", "usagePercent": 4.0 } - ] - }"#, - ) - .unwrap(); - assert!(matches!( - legacy_percentage_evidence(&config), - LegacyPercentageEvidence::GrokBuild(percent) if (percent - 4.0).abs() < 0.01 - )); + /// The two real payloads captured from the live endpoint, side by side. + const WEEKLY_CREDITS_BODY: &str = r#"{ + "config": { + "currentPeriod": { + "type": "USAGE_PERIOD_TYPE_WEEKLY", + "start": "2026-07-15T00:00:00+00:00", + "end": "2026-07-22T00:00:00+00:00" + }, + "creditUsagePercent": 4.0, + "productUsage": [ { "product": "GrokBuild", "usagePercent": 4.0 } ], + "isUnifiedBillingUser": true, + "billingPeriodStart": "2026-07-15T00:00:00+00:00", + "billingPeriodEnd": "2026-07-22T00:00:00+00:00" + }, + "subscriptionTiers": "X Premium+" + }"#; + const MONTHLY_BODY: &str = r#"{ + "config": { + "monthlyLimit": { "val": 15000 }, + "used": { "val": 216 }, + "billingPeriodStart": "2026-07-01T00:00:00+00:00", + "billingPeriodEnd": "2026-08-01T00:00:00+00:00" + } + }"#; + /// The empty-week credits payload: period present, percent fields omitted. + const EMPTY_WEEK_CREDITS_BODY: &str = r#"{ + "config": { + "currentPeriod": { + "type": "USAGE_PERIOD_TYPE_WEEKLY", + "start": "2026-07-15T00:00:00+00:00", + "end": "2026-07-22T00:00:00+00:00" + }, + "onDemandCap": { "val": 0 }, + "isUnifiedBillingUser": true, + "billingPeriodStart": "2026-07-15T00:00:00+00:00", + "billingPeriodEnd": "2026-07-22T00:00:00+00:00" + } + }"#; + + fn test_credentials() -> GrokCredentials { + GrokCredentials { + auth_path: PathBuf::from("/tmp/unused"), + entry_key: "k".into(), + access_token: "t".into(), + refresh_token: "r".into(), + client_id: "c".into(), + expires_at: None, + email: Some("user@example.com".into()), + raw_json: Value::Object(Default::default()), + } + } + + fn now() -> DateTime { + DateTime::parse_from_rfc3339("2026-07-21T00:00:00Z") + .unwrap() + .with_timezone(&Utc) + } + + fn scope_none() -> Result { + Err(AccountScopeError::NoTrustedEvidence) } #[test] - fn falls_back_to_overall_credit_percent() { - let config: BillingConfig = serde_json::from_str( - r#"{ - "creditUsagePercent": 12.5, - "productUsage": [ - { "product": "GrokChat" }, - { "product": "GrokBuild" } - ] - }"#, + fn builds_both_weekly_and_monthly_windows() { + let data = build_grok_data( + WEEKLY_CREDITS_BODY, + Some(MONTHLY_BODY), + &test_credentials(), + now(), + scope_none(), + None, ) .unwrap(); - assert!(matches!( - legacy_percentage_evidence(&config), - LegacyPercentageEvidence::Credit(percent) if (percent - 12.5).abs() < 0.01 - )); + assert_eq!(data.windows.len(), 2); + // [0] Weekly: 4% used -> 96% remaining, resets 2026-07-22. + assert_eq!(data.windows[0].label_for_test(), "Weekly"); + assert_eq!( + data.windows[0].pace_window_key_for_test(), + Some("billing.weekly.v1") + ); + assert!((data.windows[0].remaining_for_test() - 96.0).abs() < 0.01); + // [1] Monthly: 216/15000 = 1.44% used -> 98.56% remaining, resets 2026-08-01. + assert_eq!(data.windows[1].label_for_test(), "Monthly"); + assert_eq!( + data.windows[1].pace_window_key_for_test(), + Some("billing.monthly.v1") + ); + assert!((data.windows[1].remaining_for_test() - 98.56).abs() < 0.01); + assert_eq!( + data.identity.as_ref().and_then(|i| i.email.as_deref()), + Some("user@example.com") + ); + assert_eq!( + data.identity.as_ref().and_then(|i| i.plan.as_deref()), + Some("X Premium+") + ); } #[test] - fn rejects_invalid_percentages_without_falling_back() { - for invalid in [ - r#"{ "creditUsagePercent": 12.5, "productUsage": [{ "product": "GrokBuild", "usagePercent": 150.0 }] }"#, - r#"{ "creditUsagePercent": 12.5, "productUsage": [{ "product": "GrokBuild", "usagePercent": -1.0 }] }"#, - r#"{ "creditUsagePercent": 12.5, "productUsage": [{ "product": "GrokBuild", "usagePercent": -1e-400 }] }"#, - r#"{ "creditUsagePercent": 12.5, "productUsage": [{ "product": "GrokBuild", "usagePercent": -1e-9999999999999999999999999999999999999999 }] }"#, - r#"{ "creditUsagePercent": 12.5, "productUsage": [{ "product": "GrokBuild", "usagePercent": 100.0000000000000000001 }] }"#, - r#"{ "creditUsagePercent": 12.5, "productUsage": [{ "product": "GrokBuild", "usagePercent": 1e400 }] }"#, - r#"{ "creditUsagePercent": 12.5, "productUsage": [{ "product": "GrokBuild", "usagePercent": "NaN" }] }"#, - r#"{ "creditUsagePercent": 12.5, "productUsage": [{ "product": "GrokBuild", "usagePercent": null }] }"#, - ] { - let config: BillingConfig = serde_json::from_str(invalid).unwrap(); - assert!( - matches!( - legacy_percentage_evidence(&config), - LegacyPercentageEvidence::Invalid - ), - "invalid GrokBuild usage must not fall back to overall credit usage" - ); - } - - for valid in [ - "-0.0", - "99.9999999999999999999", - "1e2", - "0.100e3", - "1e-9999999999999999999999999999999999999999", - "0.01e-170141183460469231731687303715884105728", - ] { - let config: BillingConfig = serde_json::from_str(&format!( - r#"{{ "creditUsagePercent": {valid} }}"# - )) - .unwrap(); - assert!( - matches!( - legacy_percentage_evidence(&config), - LegacyPercentageEvidence::Credit(_) - | LegacyPercentageEvidence::CreditZero(_) - ), - "valid percentage {valid} must survive exact range validation" - ); - } - - let config: BillingConfig = serde_json::from_str( - r#"{ "creditUsagePercent": 12.5, "productUsage": [{ "product": "GrokChat" }] }"#, + fn empty_week_shows_zero_percent_weekly_not_error() { + // THE original bug: a freshly reset week with no usage yet omits the + // percent fields. With a self-contained weekly period still present, + // that is 0% used, and the card must show a Weekly window rather than + // erroring. + let data = build_grok_data( + EMPTY_WEEK_CREDITS_BODY, + None, + &test_credentials(), + now(), + scope_none(), + None, ) .unwrap(); - assert!(matches!( - legacy_percentage_evidence(&config), - LegacyPercentageEvidence::Credit(percent) if percent == 12.5 - )); - } - - fn map_test_config(config: &str) -> Result { - let body = format!(r#"{{ "config": {config} }}"#); - let credentials = GrokCredentials { - auth_path: PathBuf::from("/tmp/unused"), - entry_key: "k".into(), - access_token: "t".into(), - refresh_token: "r".into(), - client_id: "c".into(), - expires_at: None, - email: None, - raw_json: Value::Object(Default::default()), - }; - let now = parse_timestamp("2026-07-11T12:00:00Z").unwrap(); - map_billing( - &body, - &credentials, - now, - Err(AccountScopeError::NoTrustedEvidence), - ) + assert_eq!(data.windows.len(), 1); + assert_eq!(data.windows[0].label_for_test(), "Weekly"); + assert!((data.windows[0].remaining_for_test() - 100.0).abs() < 0.01); } #[test] - fn accepts_bare_and_wrapped_numeric_amounts_for_included_and_extra_usage() { - for (label, config, card_id) in [ - ( - "bare included", - r#"{ "used": 25.0, "monthlyLimit": 100.0 }"#, - "billing.monthly.v1", - ), - ( - "wrapped included", - r#"{ "used": { "val": 25.0 }, "monthlyLimit": { "val": 100.0 } }"#, - "billing.monthly.v1", - ), + fn monthly_timeout_server_error_malformed_or_invalid_still_shows_weekly() { + // Timeout and non-success status both collapse to None at the additive + // fetch boundary; malformed or invalid success bodies are ignored here. + for (label, monthly_body) in [ + ("timeout", None), + ("server error", None), + ("malformed body", Some("not json")), ( - "bare extra", - r#"{ "onDemandUsed": 25.0, "onDemandCap": 100.0 }"#, - "extra_usage.v1", - ), - ( - "wrapped extra", - r#"{ "onDemandUsed": { "val": 25.0 }, "onDemandCap": { "val": 100.0 } }"#, - "extra_usage.v1", + "invalid meter", + Some(r#"{ "config": { "monthlyLimit": { "val": 0 }, "used": { "val": 10 } } }"#), ), ] { - let data = map_test_config(config).unwrap(); + let data = build_grok_data( + WEEKLY_CREDITS_BODY, + monthly_body, + &test_credentials(), + now(), + scope_none(), + None, + ) + .unwrap(); assert_eq!(data.windows.len(), 1, "{label}"); - assert_eq!(data.windows[0].card_id_for_test(), card_id, "{label}"); - assert!( - (data.windows[0].remaining_for_test() - 75.0).abs() < 0.01, - "{label}" - ); + assert_eq!(data.windows[0].label_for_test(), "Weekly", "{label}"); } } #[test] - fn rejects_bare_numeric_strings_for_included_and_extra_usage() { - for (label, config) in [ - ( - "included used string", - r#"{ - "used": "25", - "monthlyLimit": 100.0, - "onDemandUsed": 0.0, - "onDemandCap": 0.0 - }"#, - ), - ( - "included limit string", - r#"{ - "used": 25.0, - "monthlyLimit": "100", - "onDemandUsed": 0.0, - "onDemandCap": 0.0 - }"#, - ), - ( - "extra string", - r#"{ "onDemandUsed": "25", "onDemandCap": 100.0 }"#, - ), - ] { - assert_eq!( - map_test_config(config) - .err() - .expect("numeric strings must not become amount evidence"), - "Grok billing response has no creditUsagePercent or GrokBuild usage.", - "{label}" - ); - } + fn no_usable_windows_is_an_error() { + // Credits view with neither a percent nor a period, and no monthly data, + // has nothing to show — surface an honest error, not an empty card. + let data = build_grok_data( + r#"{ "config": {} }"#, + None, + &test_credentials(), + now(), + scope_none(), + None, + ); + assert!(data.is_err()); } #[test] - fn default_zero_credit_uses_valid_unified_included_ratio() { - let data = map_test_config( - r#"{ - "creditUsagePercent": 0.0, - "used": 25.0, - "monthlyLimit": 100.0 - }"#, - ) - .unwrap(); - - assert_eq!(data.windows.len(), 1); - assert_eq!(data.windows[0].card_id_for_test(), "billing.monthly.v1"); - assert!((data.windows[0].remaining_for_test() - 75.0).abs() < 0.01); - } + fn monthly_only_without_weekly_is_an_error() { + // Monthly is additive after weekly succeeds. A credits body that cannot + // build a weekly window must not produce a monthly-only card. + let err = match build_grok_data( + r#"{ "config": {} }"#, + Some(MONTHLY_BODY), + &test_credentials(), + now(), + scope_none(), + None, + ) { + Err(e) => e, + Ok(_) => panic!("monthly-only must not succeed"), + }; + let ProviderFetchFailure::Terminal { display } = err else { + panic!("missing weekly meter must be terminal"); + }; + assert!( + display.contains("weekly"), + "error should name the missing weekly meter: {display}" + ); - #[test] - fn fully_disabled_unified_billing_has_no_windows() { - let data = map_test_config( - r#"{ - "creditUsagePercent": 0.0, - "used": 0.0, - "monthlyLimit": 0.0, - "onDemandUsed": 0.0, - "onDemandCap": 0.0 - }"#, + // Non-weekly period on the credits view is also insufficient, even with + // a clean monthly body. + let credits_monthly_period = r#"{ + "config": { + "currentPeriod": { + "type": "USAGE_PERIOD_TYPE_MONTHLY", + "start": "2026-07-01T00:00:00+00:00", + "end": "2026-08-01T00:00:00+00:00" + } + } + }"#; + assert!(build_grok_data( + credits_monthly_period, + Some(MONTHLY_BODY), + &test_credentials(), + now(), + scope_none(), + None, ) - .unwrap(); - - assert!(data.windows.is_empty()); - } - - #[test] - fn default_zero_credit_keeps_legacy_when_included_is_absent_or_invalid() { - for (label, config) in [ - ("absent", r#"{ "creditUsagePercent": 0.0 }"#), - ( - "invalid", - r#"{ - "creditUsagePercent": 0.0, - "used": "bad", - "monthlyLimit": 100.0 - }"#, - ), - ] { - let data = map_test_config(config).unwrap(); - assert_eq!(data.windows.len(), 1, "{label}"); - assert_eq!( - data.windows[0].card_id_for_test(), - "row.billing.unknown.v1", - "{label}" - ); - assert!((data.windows[0].remaining_for_test() - 100.0).abs() < 0.01); - } + .is_err()); } #[test] - fn syntactically_nonzero_credit_still_beats_unified_when_float_underflows() { - let data = map_test_config( - r#"{ - "creditUsagePercent": 1e-400, - "used": 25.0, - "monthlyLimit": 100.0 - }"#, - ) - .unwrap(); - - assert_eq!(data.windows.len(), 1); - assert_eq!(data.windows[0].card_id_for_test(), "row.billing.unknown.v1"); - assert!((data.windows[0].remaining_for_test() - 100.0).abs() < 0.01); - } - - #[test] - fn grok_build_zero_remains_higher_confidence_than_unified_included() { - let data = map_test_config( + fn prefers_grok_build_product_percent() { + let config: BillingConfig = serde_json::from_str( r#"{ "creditUsagePercent": 50.0, - "productUsage": [{ "product": "GrokBuild", "usagePercent": 0.0 }], - "used": 25.0, - "monthlyLimit": 100.0 + "productUsage": [ + { "product": "GrokChat", "usagePercent": 10.0 }, + { "product": "GrokBuild", "usagePercent": 4.0 } + ] }"#, ) .unwrap(); - - assert_eq!(data.windows.len(), 1); - assert_eq!(data.windows[0].card_id_for_test(), "row.billing.unknown.v1"); - assert!((data.windows[0].remaining_for_test() - 100.0).abs() < 0.01); + match weekly_used_percent(&config) { + WeeklyPercent::Value(pct) => assert!((pct - 4.0).abs() < 0.01), + other => panic!("expected Value(4.0), got {other:?}"), + } } #[test] - fn maps_unified_included_usage_with_period_semantics_before_disabled_extra() { - let data = map_test_config( - r#"{ - "currentPeriod": { "type": "DAILY" }, - "billingPeriodStart": "2026-07-01T00:00:00Z", - "billingPeriodEnd": "2026-08-01T00:00:00Z", - "used": { "val": 25.0 }, - "monthlyLimit": { "val": 100.0 }, - "onDemandUsed": { "val": 0.0 }, - "onDemandCap": { "val": 0.0 } - }"#, - ) - .unwrap(); - - assert_eq!(data.windows.len(), 1); - let monthly = &data.windows[0]; - assert_eq!(monthly.label_for_test(), "Monthly"); - assert_eq!(monthly.card_id_for_test(), "billing.monthly.v1"); - assert_eq!(monthly.pace_window_key_for_test(), Some("billing.monthly.v1")); - assert_eq!(monthly.pace_reason_for_test(), None); - let wire = serde_json::to_value(monthly).unwrap(); - assert_eq!(wire["usedPercent"], 25.0); - assert_eq!(wire["paceStatus"]["durationSeconds"], 2_678_400); - - let weekly = map_test_config( + fn falls_back_to_overall_credit_percent() { + let config: BillingConfig = serde_json::from_str( r#"{ - "currentPeriod": { - "type": "WEEKLY", - "start": "2026-07-07T00:00:00Z", - "end": "2026-07-14T00:00:00Z" - }, - "used": { "val": 25.0 }, - "monthlyLimit": { "val": 100.0 } + "creditUsagePercent": 12.5, + "productUsage": [ + { "product": "GrokChat" }, + { "product": "GrokBuild" } + ] }"#, ) .unwrap(); - assert_eq!(weekly.windows[0].card_id_for_test(), "billing.weekly.v1"); + match weekly_used_percent(&config) { + WeeklyPercent::Value(pct) => assert!((pct - 12.5).abs() < 0.01), + other => panic!("expected Value(12.5), got {other:?}"), + } } #[test] - fn unrecognized_nested_period_is_not_attached_to_synthesized_monthly_card() { - let without_billing_period = map_test_config( + fn rejects_invalid_usage_percentages_before_wire() { + let invalid_product: BillingConfig = serde_json::from_str( r#"{ - "currentPeriod": { - "type": "DAILY", - "start": "2026-07-11T00:00:00Z", - "end": "2026-07-12T00:00:00Z" - }, - "used": 25.0, - "monthlyLimit": 100.0 - }"#, - ) - .unwrap(); - let wire = serde_json::to_value(&without_billing_period.windows[0]).unwrap(); - assert_eq!(wire["cardId"], "billing.monthly.v1"); - assert!(wire.get("resetsAt").is_none()); - assert!(wire["paceStatus"].get("durationSeconds").is_none()); - - let with_billing_period = map_test_config( - r#"{ - "currentPeriod": { - "type": "DAILY", - "start": "2026-07-11T00:00:00Z", - "end": "2026-07-12T00:00:00Z" - }, - "billingPeriodStart": "2026-07-01T00:00:00Z", - "billingPeriodEnd": "2026-08-01T00:00:00Z", - "used": 25.0, - "monthlyLimit": 100.0 + "creditUsagePercent": 12.5, + "productUsage": [ + { "product": "GrokBuild", "usagePercent": 150.0 } + ] }"#, ) .unwrap(); - let wire = serde_json::to_value(&with_billing_period.windows[0]).unwrap(); - assert_eq!(wire["resetsAt"], "2026-08-01T00:00:00.000Z"); - assert_eq!(wire["paceStatus"]["durationSeconds"], 2_678_400); - } + // Present-but-out-of-range GrokBuild fails that field; do not fall back. + assert_eq!( + weekly_used_percent(&invalid_product), + WeeklyPercent::Invalid + ); - #[test] - fn maps_unified_included_ratio_boundaries_and_clamps_over_limit() { - for (used, monthly_limit, expected) in [ - (0.0, 80.0, 0.0), - (80.0, 80.0, 100.0), - (160.0, 80.0, 100.0), - ] { - let config = format!( + for invalid in ["1e400", r#""NaN""#] { + let malformed_product: BillingConfig = serde_json::from_str(&format!( r#"{{ - "used": {{ "val": {used} }}, - "monthlyLimit": {{ "val": {monthly_limit} }} - }}"# - ); - let data = map_test_config(&config).unwrap(); - assert_eq!(data.windows.len(), 1); - assert_eq!(data.windows[0].card_id_for_test(), "billing.monthly.v1"); - let wire = serde_json::to_value(&data.windows[0]).unwrap(); - assert_eq!( - wire["usedPercent"], expected, - "used={used}, monthlyLimit={monthly_limit}" - ); - } - } - - #[test] - fn legacy_percentage_evidence_precedes_unified_included_usage() { - for (label, config, expected_remaining) in [ - ( - "GrokBuild", - r#"{ - "creditUsagePercent": 50.0, - "productUsage": [{ "product": "GrokBuild", "usagePercent": 4.0 }], - "used": { "val": 25.0 }, - "monthlyLimit": { "val": 100.0 } - }"#, - 96.0, - ), - ( - "creditUsagePercent", - r#"{ "creditUsagePercent": 12.5, - "used": { "val": 25.0 }, - "monthlyLimit": { "val": 100.0 } - }"#, - 87.5, - ), - ] { - let data = map_test_config(config).unwrap(); - assert_eq!(data.windows.len(), 1, "{label}"); - assert_eq!(data.windows[0].card_id_for_test(), "row.billing.unknown.v1"); - assert!( - (data.windows[0].remaining_for_test() - expected_remaining).abs() < 0.01, - "{label}" - ); - } - } - - #[test] - fn invalid_explicit_legacy_percentage_does_not_fall_through_to_unified() { - for config in [ - r#"{ - "creditUsagePercent": 12.5, - "productUsage": [{ "product": "GrokBuild", "usagePercent": "bad" }], - "used": { "val": 25.0 }, - "monthlyLimit": { "val": 100.0 } - }"#, - r#"{ - "creditUsagePercent": "bad", - "used": { "val": 25.0 }, - "monthlyLimit": { "val": 100.0 } - }"#, - ] { - assert_eq!( - map_test_config(config) - .err() - .expect("invalid legacy percentage must fail without extra usage"), - "Grok billing response has no creditUsagePercent or GrokBuild usage." - ); - } - } - - #[test] - fn malformed_legacy_is_not_masked_by_disabled_on_demand() { - for config in [ - r#"{ - "creditUsagePercent": 12.5, - "productUsage": [{ "product": "GrokBuild", "usagePercent": "bad" }], - "onDemandUsed": 0.0, - "onDemandCap": 0.0 - }"#, - r#"{ - "creditUsagePercent": "bad", - "onDemandUsed": 0.0, - "onDemandCap": 0.0 - }"#, - ] { + "productUsage": [ + {{ "product": "GrokBuild", "usagePercent": {invalid} }} + ] + }}"# + )) + .unwrap(); assert_eq!( - map_test_config(config) - .err() - .expect("disabled on-demand must not mask malformed legacy evidence"), - "Grok billing response has no creditUsagePercent or GrokBuild usage." + weekly_used_percent(&malformed_product), + WeeklyPercent::Invalid ); } - } - #[test] - fn malformed_duplicate_grok_build_evidence_fails_in_any_order() { - for config in [ - r#"{ - "productUsage": [ - { "product": "GrokBuild", "usagePercent": 4.0 }, - { "product": "GrokBuild", "usagePercent": "bad" } - ], - "onDemandUsed": 0.0, - "onDemandCap": 0.0 - }"#, + let invalid: BillingConfig = serde_json::from_str( r#"{ + "creditUsagePercent": -1.0, "productUsage": [ - { "product": "GrokBuild", "usagePercent": "bad" }, - { "product": "GrokBuild", "usagePercent": 4.0 } - ], - "onDemandUsed": 0.0, - "onDemandCap": 0.0 - }"#, - ] { - assert_eq!( - map_test_config(config) - .err() - .expect("any malformed matching GrokBuild evidence must fail closed"), - "Grok billing response has no creditUsagePercent or GrokBuild usage." - ); - } - } - - #[test] - fn malformed_unified_included_evidence_is_ignored_with_valid_legacy() { - let data = map_test_config( - r#"{ - "creditUsagePercent": 12.5, - "used": { "val": "bad", "futureField": { "large": 1e400 } }, - "monthlyLimit": "bad", - "onDemandUsed": { "val": 0.0 }, - "onDemandCap": { "val": 0.0 } + { "product": "GrokBuild", "usagePercent": 150.0 } + ] }"#, ) .unwrap(); + assert_eq!(weekly_used_percent(&invalid), WeeklyPercent::Invalid); - assert_eq!(data.windows.len(), 1); - assert!((data.windows[0].remaining_for_test() - 87.5).abs() < 0.01); - } - - #[test] - fn invalid_unified_included_evidence_is_not_masked_by_disabled_extra() { - for (label, included) in [ - ("used only", r#""used": { "val": 1.0 }"#), - ( - "monthly limit only", - r#""monthlyLimit": { "val": 10.0 }"#, - ), - ( - "wrong type", - r#""used": { "val": "1" }, "monthlyLimit": { "val": 10.0 }"#, - ), - ( - "negative used", - r#""used": { "val": -1.0 }, "monthlyLimit": { "val": 10.0 }"#, - ), - ( - "zero limit", - r#""used": { "val": 1.0 }, "monthlyLimit": { "val": 0.0 }"#, - ), - ( - "overflow", - r#""used": { "val": 1e400 }, "monthlyLimit": { "val": 10.0 }"#, - ), - ( - "used underflow", - r#""used": { "val": 1e-400 }, "monthlyLimit": { "val": 10.0 }"#, - ), - ( - "limit underflow", - r#""used": { "val": 1.0 }, "monthlyLimit": { "val": 1e-400 }"#, - ), - ] { - let config = format!( - r#"{{ - {included}, - "onDemandUsed": {{ "val": 0.0 }}, - "onDemandCap": {{ "val": 0.0 }} - }}"# - ); - assert_eq!( - map_test_config(&config) - .err() - .expect("invalid included evidence must beat disabled extra"), - "Grok billing response has no creditUsagePercent or GrokBuild usage.", - "{label}" - ); - } - } - - #[test] - fn invalid_unified_included_evidence_allows_valid_positive_extra() { - let data = map_test_config( + // Overall credit present but invalid, with no GrokBuild percent key. + let bad_credit: BillingConfig = serde_json::from_str( r#"{ - "used": { "val": 1.0 }, - "onDemandUsed": { "val": 25.0 }, - "onDemandCap": { "val": 100.0 } + "creditUsagePercent": -1.0, + "productUsage": [ { "product": "GrokBuild" } ] }"#, ) .unwrap(); + assert_eq!(weekly_used_percent(&bad_credit), WeeklyPercent::Invalid); - assert_eq!(data.windows.len(), 1); - let extra = &data.windows[0]; - assert_eq!(extra.card_id_for_test(), "extra_usage.v1"); - assert_eq!(extra.pace_window_key_for_test(), Some("extra_usage.v1")); - assert_eq!(extra.pace_reason_for_test(), Some("nonRecurring")); - assert!((extra.remaining_for_test() - 75.0).abs() < 0.01); + // Truly absent percent fields. + let absent: BillingConfig = + serde_json::from_str(r#"{ "productUsage": [ { "product": "GrokChat" } ] }"#).unwrap(); + assert_eq!(weekly_used_percent(&absent), WeeklyPercent::Absent); } #[test] - fn maps_old_and_on_demand_windows_without_changing_precedence() { - let data = map_test_config( + fn malformed_weekly_percent_with_valid_period_is_not_zero() { + // Invalid percent must not take the empty-week 0% path just because a + // self-contained weekly period is present. + let config: BillingConfig = serde_json::from_str( r#"{ "currentPeriod": { - "type": "WEEKLY", - "start": "2026-07-07T00:00:00Z", - "end": "2026-07-14T00:00:00Z" + "type": "USAGE_PERIOD_TYPE_WEEKLY", + "start": "2026-07-15T00:00:00+00:00", + "end": "2026-07-22T00:00:00+00:00" }, - "billingPeriodEnd": "2026-08-01T00:00:00Z", - "creditUsagePercent": 50.0, + "creditUsagePercent": 150.0, "productUsage": [ - { "product": "GrokChat", "usagePercent": 10.0 }, - { "product": "GrokBuild", "usagePercent": 4.0 } - ], - "onDemandUsed": { "val": 25.0 }, - "onDemandCap": { "val": 100.0 } + { "product": "GrokBuild", "usagePercent": "NaN" } + ] }"#, ) .unwrap(); + assert_eq!(weekly_used_percent(&config), WeeklyPercent::Invalid); + assert!( + weekly_window(&config, now()).is_none(), + "invalid percent must not synthesize Weekly 0%" + ); - assert_eq!(data.windows.len(), 2); - assert_eq!(data.windows[0].card_id_for_test(), "billing.weekly.v1"); - assert!((data.windows[0].remaining_for_test() - 96.0).abs() < 0.01); - - let extra = &data.windows[1]; - assert_eq!(extra.label_for_test(), "Extra usage"); - assert_eq!(extra.card_id_for_test(), "extra_usage.v1"); - assert_eq!(extra.pace_window_key_for_test(), Some("extra_usage.v1")); - assert_eq!(extra.pace_reason_for_test(), Some("nonRecurring")); - let wire = serde_json::to_value(extra).unwrap(); - assert_eq!(wire["usedPercent"], 25.0); - assert_eq!(wire["resetsAt"], "2026-08-01T00:00:00.000Z"); - } - - #[test] - fn maps_on_demand_ratio_boundaries_and_clamps_over_cap() { - for (used, cap, expected) in [ - (0.0, 80.0, 0.0), - (80.0, 80.0, 100.0), - (160.0, 80.0, 100.0), - ] { - let config = format!( - r#"{{ - "onDemandUsed": {{ "val": {used} }}, - "onDemandCap": {{ "val": {cap} }} - }}"# - ); - let data = map_test_config(&config).unwrap(); - assert_eq!(data.windows.len(), 1); - let wire = serde_json::to_value(&data.windows[0]).unwrap(); - assert_eq!(wire["usedPercent"], expected, "used={used}, cap={cap}"); - assert_eq!(data.windows[0].pace_reason_for_test(), Some("nonRecurring")); - } - } - - #[test] - fn zero_on_demand_cap_is_recognized_but_disabled() { - for cap in ["0", "0.0", "-0.0", "0e10", "-0.0e-10"] { - let config = format!( - r#"{{ - "onDemandUsed": {{ "val": 5.0 }}, - "onDemandCap": {{ "val": {cap} }} - }}"# - ); - let data = map_test_config(&config).unwrap(); - assert!(data.windows.is_empty(), "cap={cap}"); - } - } - - #[test] - fn zero_on_demand_cap_does_not_require_used() { - for (label, config) in [ - ("missing used", r#"{ "onDemandCap": 0.0 }"#), - ( - "null used", - r#"{ "onDemandUsed": null, "onDemandCap": { "val": 0.0 } }"#, - ), - ( - "fully disabled", - r#"{ - "creditUsagePercent": 0.0, - "used": 0.0, - "monthlyLimit": 0.0, - "onDemandCap": 0.0 - }"#, - ), - ] { - let data = map_test_config(config).unwrap(); - assert!(data.windows.is_empty(), "{label}"); - } - } - - #[test] - fn invalid_on_demand_pairs_do_not_create_quota() { - for (label, config) in [ - ( - "negative used", - r#"{ "onDemandUsed": { "val": -1.0 }, "onDemandCap": { "val": 10.0 } }"#, - ), - ( - "negative cap", - r#"{ "onDemandUsed": { "val": 1.0 }, "onDemandCap": { "val": -10.0 } }"#, - ), - ( - "negative used underflow", - r#"{ "onDemandUsed": { "val": -1e-400 }, "onDemandCap": { "val": 10.0 } }"#, - ), - ( - "positive used underflow", - r#"{ "onDemandUsed": { "val": 1e-400 }, "onDemandCap": { "val": 10.0 } }"#, - ), - ( - "positive cap underflow", - r#"{ "onDemandUsed": { "val": 1.0 }, "onDemandCap": { "val": 1e-400 } }"#, - ), - ("missing used", r#"{ "onDemandCap": { "val": 10.0 } }"#), - ("missing cap", r#"{ "onDemandUsed": { "val": 1.0 } }"#), - ( - "wrong outer type", - r#"{ "onDemandUsed": [1.0], "onDemandCap": { "val": 10.0 } }"#, - ), - ( - "wrong val type", - r#"{ "onDemandUsed": { "val": "1" }, "onDemandCap": { "val": 10.0 } }"#, - ), - ( - "overflow used", - r#"{ "onDemandUsed": { "val": 1e400 }, "onDemandCap": { "val": 10.0 } }"#, - ), - ( - "overflow cap", - r#"{ "onDemandUsed": { "val": 1.0 }, "onDemandCap": { "val": 1e400 } }"#, - ), - ] { - assert_eq!( - map_test_config(config) - .err() - .expect("invalid on-demand pair must fail"), - "Grok billing response has no creditUsagePercent or GrokBuild usage.", - "{label}" - ); - } - } - - #[test] - fn malformed_on_demand_does_not_break_old_subscription() { - let data = map_test_config( - r#"{ - "creditUsagePercent": 12.5, - "onDemandUsed": { - "val": "not-a-number", - "futureField": { "large": 1e400 } + // Card-level: invalid weekly + valid monthly still errors (weekly required). + let credits = r#"{ + "config": { + "currentPeriod": { + "type": "USAGE_PERIOD_TYPE_WEEKLY", + "start": "2026-07-15T00:00:00+00:00", + "end": "2026-07-22T00:00:00+00:00" }, - "onDemandCap": { "val": 100.0, "futureField": true } - }"#, + "creditUsagePercent": -5.0 + } + }"#; + assert!(build_grok_data( + credits, + Some(MONTHLY_BODY), + &test_credentials(), + now(), + scope_none(), + None, ) - .unwrap(); - - assert_eq!(data.windows.len(), 1); - assert!((data.windows[0].remaining_for_test() - 87.5).abs() < 0.01); - } - - #[test] - fn prepaid_balance_only_does_not_create_quota() { - assert_eq!( - map_test_config( - r#"{ - "prepaidBalance": { "val": 500.0 }, - "topUpMethod": "manual", - "isUnifiedBillingUser": true - }"#, - ) - .err() - .expect("prepaid balance alone must fail"), - "Grok billing response has no creditUsagePercent or GrokBuild usage." - ); + .is_err()); } #[test] @@ -1748,30 +1363,19 @@ mod tests { }, "subscriptionTiers": "X Premium+" }"#; - let credentials = GrokCredentials { - auth_path: PathBuf::from("/tmp/unused"), - entry_key: "k".into(), - access_token: "t".into(), - refresh_token: "r".into(), - client_id: "c".into(), - expires_at: None, - email: Some("user@example.com".into()), - raw_json: Value::Object(Default::default()), - }; + let credentials = test_credentials(); + assert_eq!(credentials.scope_marker(), Some("r")); let now = DateTime::parse_from_rfc3339("2026-07-11T12:00:00Z") .unwrap() .with_timezone(&Utc); - let data = map_billing( - body, - &credentials, - now, - Err(AccountScopeError::NoTrustedEvidence), - ) - .unwrap(); + let data = build_grok_data(body, None, &credentials, now, scope_none(), None).unwrap(); assert_eq!(data.windows.len(), 1); assert_eq!(data.windows[0].label_for_test(), "Weekly"); - assert_eq!(data.windows[0].card_id_for_test(), "billing.weekly.v1"); - assert_eq!(data.windows[0].pace_window_key_for_test(), Some("billing.weekly.v1")); + assert_eq!(data.windows[0].window_minutes_for_test(), Some(10_080)); + assert_eq!( + data.windows[0].pace_window_key_for_test(), + Some("billing.weekly.v1") + ); assert!((data.windows[0].remaining_for_test() - 96.0).abs() < 0.01); assert_eq!( data.identity.as_ref().and_then(|i| i.email.as_deref()), @@ -1784,219 +1388,242 @@ mod tests { } #[test] - fn stage4_period_evidence_routes_and_fails_closed() { - let credentials = GrokCredentials { - auth_path: PathBuf::from("/tmp/unused"), - entry_key: "k".into(), - access_token: "t".into(), - refresh_token: "r".into(), - client_id: "c".into(), - expires_at: None, - email: None, - raw_json: Value::Object(Default::default()), - }; - let map = |config: Value, now: DateTime| { - let body = serde_json::json!({ "config": config }).to_string(); - map_billing( - &body, - &credentials, - now, - Err(AccountScopeError::NoTrustedEvidence), - ) - .unwrap() - .windows - .into_iter() - .next() - .unwrap() - }; - - let exact_weekly = map( - serde_json::json!({ - "currentPeriod": { - "type": "usage_period_type_weekly", - "start": "2026-07-03T00:00:00.900Z", - "end": "2026-07-10T00:00:00.900Z" - }, - "creditUsagePercent": 12.0 - }), - parse_timestamp("2026-07-10T00:00:00.100Z").unwrap(), - ); - let weekly_wire = serde_json::to_value(&exact_weekly).unwrap(); - assert_eq!(weekly_wire["cardId"], "billing.weekly.v1"); - assert_eq!(weekly_wire["paceStatus"]["windowKey"], "billing.weekly.v1"); - assert_eq!(weekly_wire["paceStatus"]["state"], "learningHistory"); - assert_eq!(weekly_wire["paceStatus"]["durationSeconds"], 604_800); - assert_eq!(weekly_wire["paceStatus"]["durationSource"], "provider"); - assert_eq!(weekly_wire["resetsAt"], "2026-07-10T00:00:00.900Z"); - - for (start, end, days) in [ - ("2023-02-01T00:00:00Z", "2023-03-01T00:00:00Z", 28), - ("2024-02-01T00:00:00Z", "2024-03-01T00:00:00Z", 29), - ("2024-04-01T00:00:00Z", "2024-05-01T00:00:00Z", 30), - ("2024-05-01T00:00:00Z", "2024-06-01T00:00:00Z", 31), - ] { - let window = map( - serde_json::json!({ - "currentPeriod": { - "type": "USAGE_PERIOD_TYPE_MONTHLY", - "start": start, - "end": end - }, - "creditUsagePercent": 12.0 - }), - parse_timestamp(start).unwrap() + chrono::Duration::days(1), - ); - let wire = serde_json::to_value(&window).unwrap(); - assert_eq!(wire["cardId"], "billing.monthly.v1"); - assert_eq!(wire["paceStatus"]["windowKey"], "billing.monthly.v1"); - assert_eq!(wire["paceStatus"]["durationSeconds"], days * 86_400); - assert_eq!(wire["paceStatus"]["durationSource"], "provider"); - assert_eq!(wire["paceStatus"]["state"], "learningHistory"); - } + fn weekly_window_rejects_malformed_or_nonweekly_period_without_percent() { + // An empty period object is not a "meter exists" signal — no dates, no + // type => unknown, not 0%. + let config: BillingConfig = serde_json::from_str(r#"{ "currentPeriod": {} }"#).unwrap(); + assert!(weekly_window(&config, now()).is_none()); + + // A weekly type but no parseable start/end window is still unknown. + let config: BillingConfig = + serde_json::from_str(r#"{ "currentPeriod": { "type": "USAGE_PERIOD_TYPE_WEEKLY" } }"#) + .unwrap(); + assert!(weekly_window(&config, now()).is_none()); - let billing_fallback = map( - serde_json::json!({ - "currentPeriod": { "type": "WEEKLY" }, - "billingPeriodStart": "2026-07-03T00:00:00Z", - "billingPeriodEnd": "2026-07-10T00:00:00Z", - "creditUsagePercent": 12.0 - }), - parse_timestamp("2026-07-05T00:00:00Z").unwrap(), - ); - let fallback_wire = serde_json::to_value(&billing_fallback).unwrap(); - assert_eq!(fallback_wire["paceStatus"]["durationSeconds"], 604_800); - assert_eq!(fallback_wire["paceStatus"]["durationSource"], "provider"); + // An explicit MONTHLY period with no percent must NOT be fabricated into + // a Weekly 0%. + let config: BillingConfig = serde_json::from_str( + r#"{ "currentPeriod": { + "type": "USAGE_PERIOD_TYPE_MONTHLY", + "start": "2026-07-01T00:00:00+00:00", + "end": "2026-08-01T00:00:00+00:00" + } }"#, + ) + .unwrap(); + assert!(weekly_window(&config, now()).is_none()); - let end_only = map( - serde_json::json!({ - "currentPeriod": { - "type": "weekly", - "end": "2026-07-24T00:00:00Z" - }, - "creditUsagePercent": 12.0 - }), - parse_timestamp("2026-07-17T00:00:00Z").unwrap(), - ); - let end_only_wire = serde_json::to_value(&end_only).unwrap(); - assert_eq!(end_only_wire["paceStatus"]["state"], "learningDuration"); - assert_eq!(end_only_wire["paceStatus"]["durationSource"], "observed"); - assert!(end_only_wire["paceStatus"].get("durationSeconds").is_none()); + // MIXED SOURCE: a weekly-typed period with NO dates of its own must not + // borrow the flat billingPeriod* window to masquerade as a meter. + let config: BillingConfig = serde_json::from_str( + r#"{ "currentPeriod": { "type": "USAGE_PERIOD_TYPE_WEEKLY" }, + "billingPeriodStart": "2026-07-15T00:00:00+00:00", + "billingPeriodEnd": "2026-07-22T00:00:00+00:00" }"#, + ) + .unwrap(); + assert!(weekly_window(&config, now()).is_none()); - let unknown = map( - serde_json::json!({ - "currentPeriod": { - "type": "DAILY", - "start": "2026-07-17T00:00:00Z", - "end": "2026-07-18T00:00:00Z" - }, - "creditUsagePercent": 12.0 - }), - parse_timestamp("2026-07-17T12:00:00Z").unwrap(), - ); - let unknown_wire = serde_json::to_value(&unknown).unwrap(); - assert_eq!(unknown_wire["cardId"], "row.billing.unknown.v1"); - assert_eq!(unknown_wire["paceStatus"]["state"], "unavailable"); - assert_eq!(unknown_wire["paceStatus"]["reason"], "windowIdentity"); - assert!(unknown_wire["paceStatus"].get("windowKey").is_none()); - assert!(unknown_wire["paceStatus"].get("durationSeconds").is_none()); - for unknown_type in ["NOT_WEEKLY", "BIWEEKLY", "MONTHLYISH"] { - let window = map( - serde_json::json!({ - "currentPeriod": { - "type": unknown_type, - "start": "2026-07-17T00:00:00Z", - "end": "2026-07-18T00:00:00Z" - }, - "creditUsagePercent": 12.0 - }), - parse_timestamp("2026-07-17T12:00:00Z").unwrap(), + // SUBSTRING: types that merely CONTAIN "WEEKLY" must be rejected by the + // exact-match gate, even with a valid self-contained window. + for ty in ["USAGE_PERIOD_TYPE_BIWEEKLY", "USAGE_PERIOD_TYPE_NOT_WEEKLY"] { + let config: BillingConfig = serde_json::from_str(&format!( + r#"{{ "currentPeriod": {{ + "type": "{ty}", + "start": "2026-07-15T00:00:00+00:00", + "end": "2026-07-22T00:00:00+00:00" + }} }}"# + )) + .unwrap(); + assert!( + weekly_window(&config, now()).is_none(), + "type {ty} must not be accepted as weekly" ); - let wire = serde_json::to_value(&window).unwrap(); - assert_eq!(wire["paceStatus"]["reason"], "windowIdentity"); } + } - for (start, end) in [ - ("2026-07-18T00:00:00Z", "2026-07-17T00:00:00Z"), - ("not-a-date", "2026-07-24T00:00:00Z"), - ("2026-07-17T00:00:00Z", "not-a-date"), - ] { - let window = map( - serde_json::json!({ - "currentPeriod": { - "type": "weekly", - "start": start, - "end": end - }, - "billingPeriodStart": "2026-07-17T00:00:00Z", - "billingPeriodEnd": "2026-07-24T00:00:00Z", - "creditUsagePercent": 12.0 - }), - parse_timestamp("2026-07-17T12:00:00Z").unwrap(), - ); - let wire = serde_json::to_value(&window).unwrap(); - assert_eq!(wire["paceStatus"]["windowKey"], "billing.weekly.v1"); - assert_eq!(wire["paceStatus"]["state"], "unavailable"); - assert_eq!(wire["paceStatus"]["reason"], "invalidEvidence"); - assert!(wire["paceStatus"].get("durationSeconds").is_none()); - } + #[test] + fn monthly_window_prefers_flat_used_over_usage_total() { + let config: BillingConfig = serde_json::from_str( + r#"{ "monthlyLimit": { "val": 200 }, "used": { "val": 50 }, + "usage": { "totalUsed": { "val": 999 } } }"#, + ) + .unwrap(); + // 50/200 = 25% used -> 75% remaining. + assert!((monthly_window(&config, now()).unwrap().remaining_for_test() - 75.0).abs() < 0.01); + } - let null_primary = map( - serde_json::json!({ - "currentPeriod": { - "type": "weekly", - "start": null, - "end": "2026-07-24T00:00:00Z" - }, - "billingPeriodStart": "2026-07-17T00:00:00Z", - "creditUsagePercent": 12.0 - }), - parse_timestamp("2026-07-17T12:00:00Z").unwrap(), - ); - let null_wire = serde_json::to_value(&null_primary).unwrap(); - assert_eq!(null_wire["paceStatus"]["state"], "unavailable"); - assert_eq!(null_wire["paceStatus"]["reason"], "invalidEvidence"); + #[test] + fn monthly_window_falls_back_to_usage_total_used() { + let config: BillingConfig = serde_json::from_str( + r#"{ "monthlyLimit": { "val": 400 }, "usage": { "totalUsed": { "val": 100 } } }"#, + ) + .unwrap(); + assert!((monthly_window(&config, now()).unwrap().remaining_for_test() - 75.0).abs() < 0.01); + } - let malformed_start_only = map( - serde_json::json!({ - "currentPeriod": { - "type": "weekly", - "start": "not-a-date" - }, - "creditUsagePercent": 12.0 - }), - parse_timestamp("2026-07-17T12:00:00Z").unwrap(), - ); - let malformed_start_wire = serde_json::to_value(&malformed_start_only).unwrap(); - assert_eq!(malformed_start_wire["paceStatus"]["state"], "unavailable"); - assert_eq!( - malformed_start_wire["paceStatus"]["reason"], - "invalidEvidence" + #[test] + fn monthly_window_explicit_zero_used_is_zero_percent() { + // xAI reports a genuine zero as `{ "val": 0 }` -> 0% used, 100% remaining. + let config: BillingConfig = + serde_json::from_str(r#"{ "monthlyLimit": { "val": 15000 }, "used": { "val": 0 } }"#) + .unwrap(); + assert!( + (monthly_window(&config, now()).unwrap().remaining_for_test() - 100.0).abs() < 0.01 ); + } - let valid_start_only = map( - serde_json::json!({ - "currentPeriod": { - "type": "weekly", - "start": "2026-07-17T00:00:00Z" - }, - "creditUsagePercent": 12.0 - }), - parse_timestamp("2026-07-17T12:00:00Z").unwrap(), - ); - let start_only_wire = serde_json::to_value(&valid_start_only).unwrap(); - assert_eq!(start_only_wire["paceStatus"]["state"], "unavailable"); - assert_eq!(start_only_wire["paceStatus"]["reason"], "missingReset"); + #[test] + fn monthly_window_without_explicit_used_is_none() { + // A limit with NO explicit `used`/`totalUsed` is "unknown", not zero: + // it must not fabricate a misleading 100%-remaining across the FFI. + let config: BillingConfig = + serde_json::from_str(r#"{ "monthlyLimit": { "val": 15000 } }"#).unwrap(); + assert!(monthly_window(&config, now()).is_none()); + // An empty `{ }` wrapper (val absent) is treated the same as absent. + let config: BillingConfig = serde_json::from_str( + r#"{ "monthlyLimit": { "val": 15000 }, "used": {}, "usage": {} }"#, + ) + .unwrap(); + assert!(monthly_window(&config, now()).is_none()); } #[test] - fn unknown_period_has_no_history_identity() { + fn monthly_window_rejects_negative_used() { + // Negative consumption is invalid meter data — do not clamp to 0% healthy. + let config: BillingConfig = + serde_json::from_str(r#"{ "monthlyLimit": { "val": 15000 }, "used": { "val": -1 } }"#) + .unwrap(); + assert!(monthly_window(&config, now()).is_none()); + + let config: BillingConfig = serde_json::from_str( + r#"{ "monthlyLimit": { "val": 15000 }, "usage": { "totalUsed": { "val": -10 } } }"#, + ) + .unwrap(); + assert!(monthly_window(&config, now()).is_none()); + + // Flat used wins over usage.totalUsed; a negative flat used still rejects + // even when nested totalUsed is non-negative. let config: BillingConfig = serde_json::from_str( - r#"{ "currentPeriod": { "type": "USAGE_PERIOD_TYPE_OTHER" } }"#, + r#"{ "monthlyLimit": { "val": 200 }, "used": { "val": -5 }, + "usage": { "totalUsed": { "val": 50 } } }"#, + ) + .unwrap(); + assert!(monthly_window(&config, now()).is_none()); + } + + #[test] + fn stage4_grok_period_routes_are_exact_and_fail_closed() { + // Monthly duration routes: fixed Monthly identity from the monthly + // meter, with provider duration from the period bounds. + for (label, start, end, days) in [ + ("28-day", "2023-02-01T00:00:00Z", "2023-03-01T00:00:00Z", 28), + ("29-day", "2024-02-01T00:00:00Z", "2024-03-01T00:00:00Z", 29), + ("30-day", "2024-04-01T00:00:00Z", "2024-05-01T00:00:00Z", 30), + ("31-day", "2024-05-01T00:00:00Z", "2024-06-01T00:00:00Z", 31), + ] { + let now = parse_timestamp(start).unwrap() + chrono::Duration::days(1); + let config: BillingConfig = serde_json::from_value(serde_json::json!({ + "monthlyLimit": { "val": 100 }, + "used": { "val": 12 }, + "currentPeriod": { + "type": "USAGE_PERIOD_TYPE_MONTHLY", + "start": start, + "end": end + } + })) + .unwrap(); + let window = monthly_window(&config, now).unwrap(); + let wire = serde_json::to_value(&window).unwrap(); + assert_eq!(wire["cardId"], "billing.monthly.v1", "{label}"); + assert_eq!( + wire["paceStatus"]["windowKey"], "billing.monthly.v1", + "{label}" + ); + assert_eq!( + wire["paceStatus"]["durationSeconds"], + days * 86_400, + "{label}" + ); + assert_eq!(wire["paceStatus"]["durationSource"], "provider", "{label}"); + assert_eq!(wire["paceStatus"]["state"], "learningHistory", "{label}"); + } + + // Weekly end-only (no start): still Weekly identity, learningDuration. + let end_only: BillingConfig = serde_json::from_value(serde_json::json!({ + "creditUsagePercent": 12.0, + "currentPeriod": { + "type": "USAGE_PERIOD_TYPE_WEEKLY", + "end": "2026-07-24T00:00:00Z" + } + })) + .unwrap(); + let wire = serde_json::to_value( + weekly_window(&end_only, parse_timestamp("2026-07-17T00:00:00Z").unwrap()).unwrap(), + ) + .unwrap(); + assert_eq!(wire["cardId"], "billing.weekly.v1"); + assert_eq!(wire["paceStatus"]["windowKey"], "billing.weekly.v1"); + assert_eq!(wire["paceStatus"]["state"], "learningDuration"); + assert!(wire["resetsAt"].as_str().is_some()); + assert!(wire["paceStatus"].get("durationSeconds").is_none()); + assert!(wire["paceStatus"].get("durationSource").is_none()); + + // Credits meter is always Weekly identity even when currentPeriod type + // is not weekly — the endpoint, not the period substring, owns the key. + // Duration evidence still flows from the period bounds when valid. + let daily_typed: BillingConfig = serde_json::from_value(serde_json::json!({ + "creditUsagePercent": 12.0, + "currentPeriod": { + "type": "USAGE_PERIOD_TYPE_DAILY", + "start": "2026-07-17T00:00:00Z", + "end": "2026-07-18T00:00:00Z" + } + })) + .unwrap(); + let wire = serde_json::to_value( + weekly_window( + &daily_typed, + parse_timestamp("2026-07-17T12:00:00Z").unwrap(), + ) + .unwrap(), ) .unwrap(); - let period = period_details(&config); - assert!(period.kind.is_none()); - assert!(period.end.is_none()); + assert_eq!(wire["cardId"], "billing.weekly.v1"); + assert_eq!(wire["paceStatus"]["windowKey"], "billing.weekly.v1"); + assert_eq!(wire["paceStatus"]["durationSeconds"], 86_400); + assert_eq!(wire["paceStatus"]["durationSource"], "provider"); + + for (label, start, end) in [ + ( + "contradictory", + "2026-07-18T00:00:00Z", + "2026-07-17T00:00:00Z", + ), + ("malformed-start", "not-a-date", "2026-07-24T00:00:00Z"), + ("malformed-end", "2026-07-17T00:00:00Z", "not-a-date"), + ] { + let config: BillingConfig = serde_json::from_value(serde_json::json!({ + "creditUsagePercent": 12.0, + "currentPeriod": { + "type": "USAGE_PERIOD_TYPE_WEEKLY", + "start": start, + "end": end + } + })) + .unwrap(); + let window = + weekly_window(&config, parse_timestamp("2026-07-17T12:00:00Z").unwrap()).unwrap(); + let wire = serde_json::to_value(&window).unwrap(); + assert_eq!( + wire["paceStatus"]["windowKey"], "billing.weekly.v1", + "{label}" + ); + assert_eq!(wire["paceStatus"]["state"], "unavailable", "{label}"); + assert_eq!(wire["paceStatus"]["reason"], "invalidEvidence", "{label}"); + assert!( + wire["paceStatus"].get("durationSeconds").is_none(), + "{label}" + ); + } } #[test] @@ -2066,7 +1693,9 @@ mod tests { } }"#, ); - let creds = load_credentials_from(&path).unwrap().expect("auth.x.ai entry loads"); + let creds = load_credentials_from(&path) + .unwrap() + .expect("auth.x.ai entry loads"); assert!(creds.entry_key.contains("auth.x.ai")); assert_eq!(creds.access_token, "FAKE-XAI-ACCESS"); assert_eq!(creds.refresh_token, "FAKE-XAI-REFRESH"); @@ -2092,11 +1721,8 @@ mod tests { assert!(!is_grok_auth_entry_key( "https://auth.x.ai.evil.example::deadbeef" )); - // A foreign issuer, blank/nested client id, and shapeless key are rejected. + // A foreign issuer and a shapeless key are rejected. assert!(!is_grok_auth_entry_key("https://auth.openai.com::deadbeef")); - assert!(!is_grok_auth_entry_key("https://auth.x.ai::")); - assert!(!is_grok_auth_entry_key("https://auth.x.ai:: ")); - assert!(!is_grok_auth_entry_key("https://auth.x.ai::client::extra")); assert!(!is_grok_auth_entry_key("https://auth.x.ai")); } @@ -2144,7 +1770,10 @@ mod tests { let creds = load_credentials_from(&path) .unwrap() .expect("genuine auth.x.ai entry loads"); - assert_eq!(creds.entry_key, "https://auth.x.ai::b1a00492-073a-47ea-816f-4c329264a828"); + assert_eq!( + creds.entry_key, + "https://auth.x.ai::b1a00492-073a-47ea-816f-4c329264a828" + ); assert_eq!(creds.access_token, "FAKE-XAI-ACCESS"); assert_eq!(creds.refresh_token, "FAKE-XAI-REFRESH"); assert_ne!(creds.access_token, "FAKE-LOOKALIKE-ACCESS"); @@ -2175,6 +1804,50 @@ mod tests { ); } + #[test] + fn save_credentials_preserves_concurrent_login_and_sibling_changes() { + let original = serde_json::json!({ + (TEST_ENTRY): { + "key": "account-a-access", + "refresh_token": "account-a-refresh", + "oidc_client_id": "fixture-client" + }, + "sibling": { "value": "before" } + }); + let (dir, path) = temp_auth_json("save-race", &original.to_string()); + let mut credentials = load_credentials_from(&path).unwrap().unwrap(); + credentials.access_token = "account-a-refreshed-access".to_string(); + credentials.refresh_token = "account-a-refreshed-refresh".to_string(); + + let switched = serde_json::json!({ + (TEST_ENTRY): { + "key": "account-b-access", + "refresh_token": "account-b-refresh", + "oidc_client_id": "fixture-client" + }, + "sibling": { "value": "after-switch" } + }); + fs::write(&path, serde_json::to_vec_pretty(&switched).unwrap()).unwrap(); + assert!(save_credentials(&credentials).is_err()); + let after_switch: Value = serde_json::from_slice(&fs::read(&path).unwrap()).unwrap(); + assert_eq!(after_switch, switched, "account B must not be overwritten"); + + let sibling_changed = serde_json::json!({ + (TEST_ENTRY): original.get(TEST_ENTRY).unwrap().clone(), + "sibling": { "value": "after-sibling-update" } + }); + fs::write(&path, serde_json::to_vec_pretty(&sibling_changed).unwrap()).unwrap(); + save_credentials(&credentials).unwrap(); + let saved: Value = serde_json::from_slice(&fs::read(&path).unwrap()).unwrap(); + assert_eq!(saved["sibling"], sibling_changed["sibling"]); + assert_eq!(saved[TEST_ENTRY]["key"], "account-a-refreshed-access"); + assert_eq!( + saved[TEST_ENTRY]["refresh_token"], + "account-a-refreshed-refresh" + ); + let _ = fs::remove_dir_all(&dir); + } + #[cfg(unix)] #[test] fn atomic_write_preserves_private_permissions() { @@ -2205,120 +1878,185 @@ mod tests { const TEST_ENTRY: &str = "https://auth.x.ai::fixture-client"; - struct RecordingRefreshScope<'a> { - inner: &'a TestRefreshScope, - calls: std::cell::RefCell, Option>)>>, + #[test] + fn matching_entry_logout_remains_absent() { + let (dir, path) = temp_auth_json( + "logout", + &serde_json::json!({ + (TEST_ENTRY): { + "key": " ", + "refresh_token": "\t\n" + } + }) + .to_string(), + ); + assert!(load_credentials_from(&path).unwrap().is_none()); + let _ = fs::remove_dir_all(&dir); } - impl<'a> RecordingRefreshScope<'a> { - fn new(inner: &'a TestRefreshScope) -> Self { - Self { - inner, - calls: std::cell::RefCell::new(Vec::new()), - } - } + #[test] + fn matching_entry_tokens_are_trimmed_at_load() { + let (dir, path) = temp_auth_json( + "trimmed", + &serde_json::json!({ + (TEST_ENTRY): { + "key": " access-token\n", + "refresh_token": "\trefresh-token " + } + }) + .to_string(), + ); + let credentials = load_credentials_from(&path).unwrap().unwrap(); + assert_eq!(credentials.access_token, "access-token"); + assert_eq!(credentials.refresh_token, "refresh-token"); + assert_eq!(credentials.scope_marker(), Some("refresh-token")); + let _ = fs::remove_dir_all(&dir); } - impl RefreshScopeTransaction for RecordingRefreshScope<'_> { - fn resolve_current( - &self, - semantic_source: &str, - canonical_location: &str, - marker: &[u8], - ) -> Result { - self.calls.borrow_mut().push(( - semantic_source.to_string(), - canonical_location.to_string(), - marker.to_vec(), - None, - )); - self.inner - .resolve_current(semantic_source, canonical_location, marker) - } - - fn transfer( - &self, - semantic_source: &str, - canonical_location: &str, - old_marker: &[u8], - new_marker: &[u8], - ) -> Result { - self.calls.borrow_mut().push(( - semantic_source.to_string(), - canonical_location.to_string(), - old_marker.to_vec(), - Some(new_marker.to_vec()), - )); - self.inner - .transfer(semantic_source, canonical_location, old_marker, new_marker) + #[test] + fn matching_entry_non_string_tokens_are_terminal() { + for (tag, field, value) in [ + ("object-access", "key", serde_json::json!({ "token": "x" })), + ("number-access", "key", serde_json::json!(42)), + ( + "object-refresh", + "refresh_token", + serde_json::json!({ "token": "x" }), + ), + ("number-refresh", "refresh_token", serde_json::json!(42)), + ] { + let mut auth = serde_json::json!({ + (TEST_ENTRY): { + "key": "access-token", + "refresh_token": "refresh-token" + } + }); + auth.get_mut(TEST_ENTRY) + .unwrap() + .as_object_mut() + .unwrap() + .insert(field.to_string(), value); + let (dir, path) = temp_auth_json(tag, &auth.to_string()); + assert!( + load_credentials_from(&path).is_err(), + "{tag} must be terminal, not absent" + ); + let _ = fs::remove_dir_all(&dir); } } - fn write_test_auth(path: &Path, access_token: &str, refresh_token: &str, expires_at: &str) { - fs::create_dir_all(path.parent().unwrap()).unwrap(); - fs::write( - path, - serde_json::to_vec_pretty(&serde_json::json!({ - (TEST_ENTRY): { - "key": access_token, - "refresh_token": refresh_token, - "oidc_client_id": "fixture-client", - "expires_at": expires_at, - "email": "grok-sensitive@example.com" - }, - "https://auth.x.ai::sibling-client": { - "key": "sibling-access", - "refresh_token": "sibling-refresh", - "oidc_client_id": "sibling-client", - "expires_at": "1970-01-01T00:00:00Z" - } - })) - .unwrap(), - ) - .unwrap(); + #[test] + fn whitespace_refresh_token_is_not_a_scope_marker() { + let mut credentials = test_credentials(); + credentials.refresh_token = " \t\r\n ".to_string(); + assert_eq!(credentials.scope_marker(), None); + assert!(matches!( + credentials.resolve_account_scope(), + Err(AccountScopeError::NoTrustedEvidence) + )); } fn checkpoint_at( target: Option, - ) -> impl FnMut(RefreshCheckpoint) -> Result<(), String> { + ) -> impl FnMut(RefreshCheckpoint) -> Result<(), ProviderFetchFailure> { move |checkpoint| { if Some(checkpoint) == target { - Err("injected crash".to_string()) + Err(ProviderFetchFailure::terminal("injected crash")) } else { Ok(()) } } } + #[tokio::test] + async fn whitespace_refresh_evidence_never_binds_across_access_tokens() { + let scope = TestRefreshScope::new("grok", "grok-whitespace-binding"); + scope + .resolve_current("fixture", "baseline", b"baseline-marker") + .unwrap(); + let before = scope.metadata_bytes(); + let path = scope.root().join("grok/auth.json"); + fs::create_dir_all(path.parent().unwrap()).unwrap(); + + for access_token in ["account-a-access", "account-b-access"] { + fs::write( + &path, + serde_json::to_vec_pretty(&serde_json::json!({ + (TEST_ENTRY): { + "key": access_token, + "refresh_token": " \t ", + "oidc_client_id": "fixture-client", + "expires_at": "1970-01-01T00:00:00Z" + } + })) + .unwrap(), + ) + .unwrap(); + let loaded = load_credentials_entry_from(&path, Some(TEST_ENTRY)) + .unwrap() + .unwrap(); + assert_eq!(loaded.access_token, access_token); + assert_eq!(loaded.scope_marker(), None); + + let send_count = Arc::new(AtomicUsize::new(0)); + let request_send_count = Arc::clone(&send_count); + let failure = refresh_credentials_with( + &path, + TEST_ENTRY, + true, + &scope, + move |_, _, _| async move { + request_send_count.fetch_add(1, Ordering::SeqCst); + Ok(TokenResponse { + access_token: "unexpected-access".to_string(), + refresh_token: Some("unexpected-refresh".to_string()), + expires_in: Some(3_600), + }) + }, + save_credentials, + checkpoint_at(None), + ) + .await + .unwrap_err(); + + assert!(matches!(failure, ProviderFetchFailure::Terminal { .. })); + assert_eq!(send_count.load(Ordering::SeqCst), 0); + assert_eq!(scope.metadata_bytes(), before); + } + scope.cleanup(); + } + async fn grok_test_response( refresh_token: String, client_id: String, - ) -> Result { + _attempt_binding: ProviderCacheBinding, + ) -> Result { assert_eq!(refresh_token, "grok-old-refresh"); assert_eq!(client_id, "fixture-client"); Ok(TokenResponse { access_token: "grok-new-access".to_string(), - refresh_token: Some(" grok-new-refresh\n".to_string()), + refresh_token: Some("grok-new-refresh".to_string()), expires_in: Some(3_600), }) } - async fn unexpected_refresh_request( - _refresh_token: String, - _client_id: String, - ) -> Result { - panic!("refresh network request must be skipped") - } - fn setup_refresh(tag: &str) -> (TestRefreshScope, PathBuf, AccountScope, Vec, String) { let scope = TestRefreshScope::new("grok", tag); let path = scope.root().join("grok/auth.json"); - write_test_auth( + fs::create_dir_all(path.parent().unwrap()).unwrap(); + fs::write( &path, - "grok-old-access", - "grok-old-refresh", - "1970-01-01T00:00:00Z", - ); + serde_json::to_vec_pretty(&serde_json::json!({ + (TEST_ENTRY): { + "key": "grok-old-access", + "refresh_token": "grok-old-refresh", + "oidc_client_id": "fixture-client", + "expires_at": "1970-01-01T00:00:00Z" + } + })) + .unwrap(), + ) + .unwrap(); let credentials = load_credentials_entry_from(&path, Some(TEST_ENTRY)) .unwrap() .unwrap(); @@ -2327,7 +2065,7 @@ mod tests { .resolve_current( "grok-auth-json", &location, - credentials.scope_marker().unwrap(), + credentials.scope_marker().unwrap().as_bytes(), ) .unwrap(); let metadata = scope.metadata_bytes(); @@ -2338,7 +2076,8 @@ mod tests { scope: &TestRefreshScope, path: &Path, crash: Option, - ) -> Result<(GrokCredentials, Result), String> { + ) -> Result<(GrokCredentials, AccountScope, Option), ProviderFetchFailure> + { refresh_credentials_with( path, TEST_ENTRY, @@ -2351,211 +2090,15 @@ mod tests { .await } - fn stored_credentials(path: &Path) -> GrokCredentials { + fn stored_refresh_token(path: &Path) -> String { load_credentials_entry_from(path, Some(TEST_ENTRY)) .unwrap() .unwrap() - } - - #[test] - fn account_scope_uses_refresh_marker_and_canonical_entry_domain_without_leaks() { - let scope = TestRefreshScope::new("grok", "grok-scope-domain"); - let path = scope.root().join("grok/auth.json"); - write_test_auth( - &path, - "grok-sensitive-access-token", - "grok-sensitive-refresh-token", - "1970-01-01T00:00:00Z", - ); - let credentials = load_credentials_entry_from(&path, Some(TEST_ENTRY)) - .unwrap() - .unwrap(); - assert_eq!( - credentials.scope_marker(), - Some(b"grok-sensitive-refresh-token".as_slice()) - ); - let mut spaced_credentials = credentials.clone(); - spaced_credentials.refresh_token = " grok-sensitive-refresh-token\n".to_string(); - assert_eq!( - spaced_credentials.scope_marker(), - Some(b"grok-sensitive-refresh-token".as_slice()) - ); - let location = credentials.scope_location().unwrap(); - assert_eq!( - location, - agent_account_scope::canonical_file_location(&path, Some(TEST_ENTRY)).unwrap() - ); - - let account_scope = scope - .resolve_current( - "grok-auth-json", - &location, - credentials.scope_marker().unwrap(), - ) - .unwrap(); - let spaced_scope = scope - .resolve_current( - "grok-auth-json", - &location, - spaced_credentials.scope_marker().unwrap(), - ) - .unwrap(); - assert_eq!(spaced_scope, account_scope); - let metadata = String::from_utf8_lossy(&scope.metadata_bytes()).into_owned(); - for raw in [ - "grok-sensitive@example.com", - "grok-sensitive-access-token", - "grok-sensitive-refresh-token", - TEST_ENTRY, - location.as_str(), - ] { - assert!(!metadata.contains(raw), "metadata leaked {raw}"); - assert!(!account_scope.as_str().contains(raw), "scope leaked {raw}"); - } - let now = DateTime::parse_from_rfc3339("2026-07-18T00:00:00Z") - .unwrap() - .with_timezone(&Utc); - let error = match map_billing( - "not-json", - &credentials, - now, - Err(AccountScopeError::NoTrustedEvidence), - ) { - Err(error) => error, - Ok(_) => panic!("invalid billing payload must fail"), - }; - for raw in [ - "grok-sensitive@example.com", - "grok-sensitive-access-token", - "grok-sensitive-refresh-token", - ] { - assert!(!error.contains(raw), "error leaked {raw}"); - } - - let mut blank_marker = credentials; - blank_marker.refresh_token = " \t\n ".to_string(); - assert!(blank_marker.scope_marker().is_none()); - scope.cleanup(); - } - - #[tokio::test] - async fn refresh_reloads_exact_entry_skips_redundant_network_and_honors_force() { - let scope = TestRefreshScope::new("grok", "grok-refresh-reload"); - let path = scope.root().join("grok/auth.json"); - write_test_auth( - &path, - "stale-access", - "stale-refresh", - "1970-01-01T00:00:00Z", - ); - let stale = load_credentials_entry_from(&path, Some(TEST_ENTRY)) - .unwrap() - .unwrap(); - write_test_auth( - &path, - "grok-old-access", - "grok-old-refresh", - "2099-01-01T00:00:00Z", - ); - let current = load_credentials_entry_from(&path, Some(TEST_ENTRY)) - .unwrap() - .unwrap(); - assert_ne!(stale.refresh_token, current.refresh_token); - let location = current.scope_location().unwrap(); - let expected_scope = scope - .resolve_current("grok-auth-json", &location, current.scope_marker().unwrap()) - .unwrap(); - let before = scope.metadata_bytes(); - let recording = RecordingRefreshScope::new(&scope); - let mut checkpoints = Vec::new(); - let (reloaded, scope_outcome) = refresh_credentials_with( - &path, - TEST_ENTRY, - false, - &recording, - unexpected_refresh_request, - |_| panic!("credential save must be skipped"), - |checkpoint| { - checkpoints.push(checkpoint); - Ok(()) - }, - ) - .await - .unwrap(); - assert_eq!(reloaded.access_token, "grok-old-access"); - assert_eq!(reloaded.refresh_token, "grok-old-refresh"); - assert_eq!(scope_outcome, Ok(expected_scope.clone())); - assert_eq!(checkpoints, vec![RefreshCheckpoint::Reloaded]); - assert_eq!(scope.metadata_bytes(), before); - assert_eq!( - recording.calls.borrow().as_slice(), - &[( - "grok-auth-json".to_string(), - location.clone(), - b"grok-old-refresh".to_vec(), - None, - )] - ); - recording.calls.borrow_mut().clear(); - - let mut forced_checkpoints = Vec::new(); - let (forced, forced_scope) = refresh_credentials_with( - &path, - TEST_ENTRY, - true, - &recording, - grok_test_response, - save_credentials, - |checkpoint| { - forced_checkpoints.push(checkpoint); - Ok(()) - }, - ) - .await - .unwrap(); - assert_eq!(forced.access_token, "grok-new-access"); - assert_eq!(forced.refresh_token, "grok-new-refresh"); - assert_eq!(forced_scope, Ok(expected_scope)); - assert_eq!( - forced_checkpoints, - vec![ - RefreshCheckpoint::Reloaded, - RefreshCheckpoint::NetworkReturned, - RefreshCheckpoint::MetadataHandled, - RefreshCheckpoint::CredentialsPersisted, - ] - ); - assert_eq!( - recording.calls.borrow().as_slice(), - &[( - "grok-auth-json".to_string(), - location.clone(), - b"grok-old-refresh".to_vec(), - Some(b"grok-new-refresh".to_vec()), - )] - ); - assert_eq!(stored_credentials(&path).refresh_token, "grok-new-refresh"); - recording.calls.borrow_mut().clear(); - - write_test_auth(&path, "usable-access", " \t\n ", "2099-01-01T00:00:00Z"); - let (_, blank_scope) = refresh_credentials_with( - &path, - TEST_ENTRY, - false, - &recording, - unexpected_refresh_request, - |_| panic!("credential save must be skipped"), - |_| Ok(()), - ) - .await - .unwrap(); - assert_eq!(blank_scope, Err(AccountScopeError::NoTrustedEvidence)); - assert!(recording.calls.borrow().is_empty()); - scope.cleanup(); + .refresh_token } #[tokio::test] - async fn refresh_transfer_gate_save_failure_and_crash_boundaries_are_fail_closed() { + async fn refresh_crash_boundaries_and_scope_gate_use_production_sequence() { for boundary in [ RefreshCheckpoint::Reloaded, RefreshCheckpoint::NetworkReturned, @@ -2563,15 +2106,15 @@ mod tests { RefreshCheckpoint::CredentialsPersisted, ] { let (scope, path, old_scope, before, location) = setup_refresh("grok-crash"); + let failure = run_refresh(&scope, &path, Some(boundary)) + .await + .unwrap_err(); + assert!(matches!( + failure, + ProviderFetchFailure::Terminal { ref display } if display == "injected crash" + )); assert_eq!( - run_refresh(&scope, &path, Some(boundary)) - .await - .unwrap_err(), - "injected crash" - ); - let stored = stored_credentials(&path); - assert_eq!( - stored.refresh_token, + stored_refresh_token(&path), if boundary == RefreshCheckpoint::CredentialsPersisted { "grok-new-refresh" } else { @@ -2603,57 +2146,54 @@ mod tests { let (scope, path, old_scope, before, location) = setup_refresh("grok-metadata-fail"); scope.fail_metadata_save(); - let (refreshed, scope_outcome) = run_refresh(&scope, &path, None).await.unwrap(); - assert_eq!(refreshed.access_token, "grok-new-access"); - assert_eq!(refreshed.refresh_token, "grok-new-refresh"); - assert_eq!(scope_outcome, Err(AccountScopeError::MetadataWrite)); + let failure = run_refresh(&scope, &path, None).await.unwrap_err(); + assert!(matches!(failure, ProviderFetchFailure::Terminal { .. })); assert_eq!(scope.metadata_bytes(), before); - let stored = stored_credentials(&path); - assert_eq!(stored.access_token, "grok-old-access"); - assert_eq!(stored.refresh_token, "grok-old-refresh"); + let persisted_marker = stored_refresh_token(&path); + assert_eq!(persisted_marker, "grok-old-refresh"); + assert_eq!( + scope + .resolve_current("grok-auth-json", &location, persisted_marker.as_bytes()) + .unwrap(), + old_scope + ); + scope.cleanup(); + + let (scope, path, old_scope, _, location) = setup_refresh("grok-success"); + let (_, scope_outcome, cache_binding) = run_refresh(&scope, &path, None).await.unwrap(); + assert_eq!(scope_outcome, old_scope); + assert_eq!( + cache_binding, + Some(ProviderCacheBinding::primary(old_scope.clone())) + ); assert_eq!( scope - .resolve_current("grok-auth-json", &location, stored.refresh_token.as_bytes()) + .resolve_current("grok-auth-json", &location, b"grok-new-refresh") .unwrap(), old_scope ); scope.cleanup(); + } + #[tokio::test] + async fn refresh_persistence_failure_keeps_usage_fresh_but_uncacheable() { let (scope, path, old_scope, _, location) = setup_refresh("grok-save-fail"); - let mut checkpoints = Vec::new(); - let (refreshed, scope_outcome) = refresh_credentials_with( + let (refreshed, scope_outcome, cache_binding) = refresh_credentials_with( &path, TEST_ENTRY, true, &scope, grok_test_response, |_| Err("injected save failure".to_string()), - |checkpoint| { - checkpoints.push(checkpoint); - Ok(()) - }, + checkpoint_at(None), ) .await .unwrap(); + assert_eq!(refreshed.access_token, "grok-new-access"); - assert_eq!(refreshed.refresh_token, "grok-new-refresh"); - assert_eq!(scope_outcome, Ok(old_scope.clone())); - assert_eq!(stored_credentials(&path).refresh_token, "grok-old-refresh"); - assert_eq!( - checkpoints, - vec![ - RefreshCheckpoint::Reloaded, - RefreshCheckpoint::NetworkReturned, - RefreshCheckpoint::MetadataHandled, - RefreshCheckpoint::CredentialsPersisted, - ] - ); - assert_eq!( - scope - .resolve_current("grok-auth-json", &location, b"grok-old-refresh") - .unwrap(), - old_scope - ); + assert_eq!(scope_outcome, old_scope); + assert_eq!(cache_binding, None); + assert_eq!(stored_refresh_token(&path), "grok-old-refresh"); assert_eq!( scope .resolve_current("grok-auth-json", &location, b"grok-new-refresh") @@ -2663,55 +2203,47 @@ mod tests { scope.cleanup(); } - #[test] - fn refreshed_scope_merge_is_sticky_and_conflicts_fail_closed() { - let (scope, path, scope_a, _, location) = setup_refresh("grok-scope-merge"); - let scope_b = scope - .resolve_current("grok-auth-json", &location, b"different-refresh") - .unwrap(); - assert_ne!(scope_a, scope_b); - let credentials = load_credentials_entry_from(&path, Some(TEST_ENTRY)) - .unwrap() + #[tokio::test] + async fn refresh_transient_uses_lock_reloaded_binding_not_outer_binding() { + let (scope, path, inner_scope, _, location) = setup_refresh("grok-lock-binding"); + let outer_scope = scope + .resolve_current("grok-auth-json", &location, b"outer-refresh-a") .unwrap(); - let now = DateTime::parse_from_rfc3339("2026-07-11T12:00:00Z") - .unwrap() - .with_timezone(&Utc); - let body = r#"{ - "config": { - "creditUsagePercent": 4.0 - } - }"#; + assert_ne!(outer_scope, inner_scope); + let expected = ProviderCacheBinding::primary(inner_scope); + let request_expected = expected.clone(); - let cases = vec![ - ( - "error then success keeps first failure", - vec![Err(AccountScopeError::MetadataWrite), Ok(scope_a.clone())], - Err(AccountScopeError::MetadataWrite), - ), - ( - "success then error stays failed", - vec![Ok(scope_a.clone()), Err(AccountScopeError::MetadataRead)], - Err(AccountScopeError::MetadataRead), - ), - ( - "matching successes keep scope", - vec![Ok(scope_a.clone()), Ok(scope_a.clone())], - Ok(scope_a.clone()), - ), - ( - "different successes fail closed", - vec![Ok(scope_a.clone()), Ok(scope_b)], - Err(AccountScopeError::MetadataConflict), - ), - ]; + let failure = refresh_credentials_with( + &path, + TEST_ENTRY, + true, + &scope, + move |refresh_token, client_id, attempt_binding| async move { + assert_eq!(refresh_token, "grok-old-refresh"); + assert_eq!(client_id, "fixture-client"); + assert_eq!(attempt_binding, request_expected); + Err(ProviderFetchFailure::transient( + "Grok token refresh failed. Retrying automatically.", + Some(attempt_binding), + SafeTransportDiagnostic::from_facts(TransportErrorFacts::synthetic( + true, + false, + TransportPhase::Request, + None, + )), + )) + }, + save_credentials, + checkpoint_at(None), + ) + .await + .unwrap_err(); - for (label, outcomes, expected) in cases { - let merged = outcomes - .into_iter() - .fold(None, merge_refreshed_scope) - .unwrap(); - let mapped = map_billing(body, &credentials, now, merged).unwrap(); - assert_eq!(mapped.account_scope, expected, "{label}"); + match failure { + ProviderFetchFailure::Transient { + attempt_binding, .. + } => assert_eq!(attempt_binding, Some(expected)), + ProviderFetchFailure::Terminal { .. } => panic!("timeout must remain transient"), } scope.cleanup(); } diff --git a/crates/tb_core_ffi/src/agent_storage_windows.rs b/crates/tb_core_ffi/src/agent_storage_windows.rs index 40de8cb..06630cd 100644 --- a/crates/tb_core_ffi/src/agent_storage_windows.rs +++ b/crates/tb_core_ffi/src/agent_storage_windows.rs @@ -1931,9 +1931,7 @@ mod tests { FILE_ATTRIBUTE_NORMAL, ) .expect("reopen legacy v1 sentinel"); - assert!( - regular_file_identity(&v1_after).expect("reread v1 identity") == v1_identity - ); + assert!(regular_file_identity(&v1_after).expect("reread v1 identity") == v1_identity); assert_eq!(fs::read(&v1_path).expect("reread v1 bytes"), v1_bytes); assert_eq!( fs::metadata(&v1_path) diff --git a/crates/tb_core_ffi/src/agent_usage.rs b/crates/tb_core_ffi/src/agent_usage.rs index 8473890..96f6224 100644 --- a/crates/tb_core_ffi/src/agent_usage.rs +++ b/crates/tb_core_ffi/src/agent_usage.rs @@ -5,22 +5,23 @@ use crate::agent_account_scope::{ use crate::agent_antigravity; use crate::agent_copilot; use crate::agent_grok; -use crate::agent_quota_duration::{ - resolve_duration, valid_duration, DurationEvidence, DurationResolution, DurationSource, - DurationUnavailableReason, -}; +use crate::agent_quota_duration::{DurationEvidence, DurationSource, DurationUnavailableReason}; use crate::agent_quota_history::{ BatchObservationResult, HistoricalPace, HistoryError, HistoryOutcome, QuotaObservation, SeriesKey, }; use chrono::{DateTime, SecondsFormat, TimeZone, Utc}; +use hyper_util::client::legacy::connect::dns::{ + GaiResolver as HyperGaiResolver, Name as HyperDnsName, +}; use serde::{Deserialize, Serialize}; use serde_json::Value; use sha2::{Digest, Sha256}; -use std::collections::HashSet; +use std::collections::{HashMap, HashSet}; use std::fs; use std::path::{Path, PathBuf}; use std::sync::{LazyLock, Mutex}; +use tower_service::Service; const CODEX_USAGE_URL: &str = "https://chatgpt.com/backend-api/wham/usage"; const CODEX_REFRESH_URL: &str = "https://auth.openai.com/oauth/token"; @@ -45,6 +46,9 @@ const CLAUDE_RAW_TOKEN_KEYCHAIN_SERVICE: &str = "tokenbar-claude-oauth-token"; #[serde(rename_all = "camelCase")] pub struct AgentUsagePayload { generated_at: String, + /// Monotonic order assigned by the Rust publication gate, not by wall time + /// or provider completion order. + publication_generation: u64, agents: Vec, /// Subscription-type providers opencode is authenticated against (its /// `auth.json` `type: "oauth"` entries), e.g. ["Codex", "Copilot"]. Surfaced @@ -65,8 +69,372 @@ pub struct AgentUsageSnapshot { windows: Vec, credits: Option, error: Option, + #[serde(skip_serializing_if = "Option::is_none")] + transport_diagnostic: Option, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize)] +#[serde(rename_all = "camelCase")] +pub(crate) enum TransportCategory { + Timeout, + Dns, + Tls, + ConnectionRefused, + ConnectionReset, + Connect, + Request, + ResponseBody, + RateLimited, + ServerError, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize)] +#[serde(rename_all = "camelCase")] +pub(crate) struct SafeTransportDiagnostic { + category: TransportCategory, + #[serde(skip_serializing_if = "Option::is_none")] + status: Option, + #[serde(skip_serializing_if = "Option::is_none")] + os_code: Option, +} + +impl SafeTransportDiagnostic { + pub(crate) fn from_facts(facts: TransportErrorFacts) -> Self { + let category = if facts.is_timeout { + TransportCategory::Timeout + } else { + match facts.raw_os_code { + Some(61 | 111 | 10061) => TransportCategory::ConnectionRefused, + Some(54 | 104 | 10054) => TransportCategory::ConnectionReset, + _ if facts.is_dns => TransportCategory::Dns, + _ if facts.is_tls => TransportCategory::Tls, + _ if facts.is_connect => TransportCategory::Connect, + _ if facts.phase == TransportPhase::ResponseBody => TransportCategory::ResponseBody, + _ => TransportCategory::Request, + } + }; + Self { + category, + status: None, + os_code: facts.raw_os_code, + } + } + + fn rate_limited(status: u16) -> Self { + Self { + category: TransportCategory::RateLimited, + status: (100..=599).contains(&status).then_some(status), + os_code: None, + } + } + + fn server_error(status: u16) -> Self { + Self { + category: TransportCategory::ServerError, + status: (100..=599).contains(&status).then_some(status), + os_code: None, + } + } +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(crate) enum TransportPhase { + Request, + ResponseBody, +} + +#[derive(Debug)] +struct DnsResolutionError { + source: Box, +} + +impl DnsResolutionError { + fn new(source: impl std::error::Error + Send + Sync + 'static) -> Self { + Self { + source: Box::new(source), + } + } +} + +impl std::fmt::Display for DnsResolutionError { + fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + formatter.write_str("DNS resolution failed") + } +} + +impl std::error::Error for DnsResolutionError { + fn source(&self) -> Option<&(dyn std::error::Error + 'static)> { + Some(self.source.as_ref()) + } +} + +#[derive(Debug, Clone, Copy)] +struct TypedGaiResolver; + +impl reqwest::dns::Resolve for TypedGaiResolver { + fn resolve(&self, name: reqwest::dns::Name) -> reqwest::dns::Resolving { + let parsed_name = name.as_str().parse::(); + Box::pin(async move { + let parsed_name = parsed_name.map_err(|source| { + Box::new(DnsResolutionError::new(source)) + as Box + })?; + let addresses = HyperGaiResolver::new() + .call(parsed_name) + .await + .map_err(|source| { + Box::new(DnsResolutionError::new(source)) + as Box + })?; + Ok(Box::new(addresses) as reqwest::dns::Addrs) + }) + } +} + +pub(crate) fn provider_http_client_builder() -> reqwest::ClientBuilder { + reqwest::Client::builder().dns_resolver(TypedGaiResolver) +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(crate) struct TransportErrorFacts { + is_timeout: bool, + is_connect: bool, + is_dns: bool, + is_tls: bool, + phase: TransportPhase, + raw_os_code: Option, +} + +#[derive(Debug, Default, Clone, Copy, PartialEq, Eq)] +struct TransportSourceFacts { + is_dns: bool, + is_tls: bool, + raw_os_code: Option, +} + +fn transport_source_facts(error: &(dyn std::error::Error + 'static)) -> TransportSourceFacts { + let mut sources = vec![error]; + let mut facts = TransportSourceFacts::default(); + while let Some(current) = sources.pop() { + if let Some(io_error) = current.downcast_ref::() { + if facts.raw_os_code.is_none() { + facts.raw_os_code = io_error.raw_os_error(); + } + if let Some(inner) = io_error.get_ref() { + sources.push(inner); + } + } + facts.is_dns |= current.downcast_ref::().is_some(); + facts.is_tls |= current.downcast_ref::().is_some(); + if let Some(source) = current.source() { + sources.push(source); + } + } + facts +} + +impl TransportErrorFacts { + pub(crate) fn from_reqwest(error: &reqwest::Error, phase: TransportPhase) -> Self { + let source_facts = transport_source_facts(error); + Self { + is_timeout: error.is_timeout(), + is_connect: error.is_connect(), + is_dns: source_facts.is_dns, + is_tls: source_facts.is_tls, + phase, + raw_os_code: source_facts.raw_os_code, + } + } + + #[cfg(test)] + pub(crate) fn synthetic( + is_timeout: bool, + is_connect: bool, + phase: TransportPhase, + raw_os_code: Option, + ) -> Self { + Self { + is_timeout, + is_connect, + is_dns: false, + is_tls: false, + phase, + raw_os_code, + } + } +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub(crate) struct ProviderCacheBinding { + primary: AccountScope, + corroborating: Option, +} + +impl ProviderCacheBinding { + pub(crate) fn new(primary: AccountScope, corroborating: Option) -> Self { + Self { + primary, + corroborating, + } + } + + pub(crate) fn primary(primary: AccountScope) -> Self { + Self::new(primary, None) + } +} + +pub(crate) async fn request_after_verified_binding( + binding: Result, + request: F, +) -> Result +where + F: FnOnce(B) -> Future, + Future: std::future::Future>, +{ + request(binding?).await +} + +#[derive(Debug, Clone)] +pub(crate) enum ProviderFetchFailure { + Transient { + display: String, + attempt_binding: Option, + transport_diagnostic: SafeTransportDiagnostic, + }, + Terminal { + display: String, + }, +} + +impl ProviderFetchFailure { + pub(crate) fn transient( + display: impl Into, + attempt_binding: Option, + transport_diagnostic: SafeTransportDiagnostic, + ) -> Self { + Self::Transient { + display: display.into(), + attempt_binding, + transport_diagnostic, + } + } + + pub(crate) fn terminal(display: impl Into) -> Self { + Self::Terminal { + display: display.into(), + } + } + + pub(crate) fn from_send_error( + display: impl Into, + attempt_binding: Option, + error: &reqwest::Error, + ) -> Self { + if error.is_builder() { + return Self::terminal(display); + } + Self::transient( + display, + attempt_binding, + SafeTransportDiagnostic::from_facts(TransportErrorFacts::from_reqwest( + error, + TransportPhase::Request, + )), + ) + } +} + +#[derive(Debug, Clone)] +pub(crate) enum ProviderFetchOutcome { + Absent, + Success { + snapshot: AgentUsageSnapshot, + cache_binding: Option, + }, + Failure(ProviderFetchFailure), +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(crate) enum ResponseReadFailure { + Transient(SafeTransportDiagnostic), + Terminal(u16), +} + +pub(crate) async fn read_response_body( + status: u16, + allow_forbidden_body: bool, + read: F, +) -> Result +where + F: FnOnce() -> Future, + Future: std::future::Future>, +{ + if status == 429 { + return Err(ResponseReadFailure::Transient( + SafeTransportDiagnostic::rate_limited(status), + )); + } + if (500..=599).contains(&status) { + return Err(ResponseReadFailure::Transient( + SafeTransportDiagnostic::server_error(status), + )); + } + let may_read = (200..=299).contains(&status) || (allow_forbidden_body && status == 403); + if !may_read { + return Err(ResponseReadFailure::Terminal(status)); + } + match read().await { + Ok(body) => Ok(body), + Err(_) if status == 403 => Err(ResponseReadFailure::Terminal(status)), + Err(facts) => Err(ResponseReadFailure::Transient( + SafeTransportDiagnostic::from_facts(facts), + )), + } +} + +#[derive(Debug, Clone)] +struct LastGoodEntry { + binding: ProviderCacheBinding, + snapshot: AgentUsageSnapshot, +} + +#[derive(Debug, Default)] +struct ProviderLastGoodCache { + entries: HashMap, +} + +impl ProviderLastGoodCache { + fn clean_for( + &self, + client_id: &str, + binding: &ProviderCacheBinding, + ) -> Option { + self.entries + .get(client_id) + .filter(|entry| &entry.binding == binding) + .map(|entry| entry.snapshot.clone()) + } + + fn replace( + &mut self, + client_id: &str, + binding: ProviderCacheBinding, + mut snapshot: AgentUsageSnapshot, + ) { + snapshot.error = None; + snapshot.transport_diagnostic = None; + self.entries + .insert(client_id.to_string(), LastGoodEntry { binding, snapshot }); + } + + fn clear(&mut self, client_id: &str) { + self.entries.remove(client_id); + } } +static PROVIDER_LAST_GOOD: LazyLock> = + LazyLock::new(|| Mutex::new(ProviderLastGoodCache::default())); + #[derive(Debug, Clone, Serialize)] #[serde(rename_all = "camelCase")] pub struct AgentIdentity { @@ -116,17 +484,26 @@ pub struct UsageWindow { used_percent: f64, remaining_percent: f64, resets_at: Option, - /// Exact provider reset retained independently from millisecond wire formatting. - reset_at_evidence: Option>, reset_text: Option, - /// Resolved provider/contract evidence retained only to validate the nested wire state. - duration_evidence: Option<(DurationEvidence, DurationSource)>, - /// True only when a subsecond future reset was normalized into the next integer second. - provider_reset_normalized: bool, + /// Legacy compatibility only. Wire serialization derives this from + /// `duration_seconds`; provider adapters must never use this as identity. + window_minutes: Option, + window_key: Option, + duration_seconds: Option, + duration_source: Option, + provider_duration: Option, + contract_duration: Option, pace_status: PaceStatusPayload, historical_pace: Option, } +#[derive(Debug, Clone, Serialize)] +#[serde(rename_all = "camelCase")] +pub struct CreditsSnapshot { + remaining: Option, + unlimited: bool, +} + #[derive(Serialize)] #[serde(rename_all = "camelCase")] struct UsageWindowWire<'a> { @@ -138,7 +515,6 @@ struct UsageWindowWire<'a> { resets_at: Option<&'a str>, #[serde(skip_serializing_if = "Option::is_none")] reset_text: Option<&'a str>, - /// Compatibility mirror; it is never an independent source of duration. #[serde(skip_serializing_if = "Option::is_none")] window_minutes: Option, pace_status: &'a PaceStatusPayload, @@ -159,10 +535,7 @@ impl Serialize for UsageWindow { remaining_percent: self.remaining_percent, resets_at: self.resets_at.as_deref(), reset_text: self.reset_text.as_deref(), - window_minutes: self - .pace_status - .duration_seconds - .map(|seconds| seconds / 60), + window_minutes: self.duration_seconds.map(|seconds| seconds / 60), pace_status: &self.pace_status, historical_pace: self.historical_pace.as_ref(), } @@ -170,62 +543,53 @@ impl Serialize for UsageWindow { } } -#[derive(Debug, Clone, Serialize)] -#[serde(rename_all = "camelCase")] -pub struct CreditsSnapshot { - remaining: Option, - unlimited: bool, -} - -fn normalized_reset_at(window: &mut UsageWindow, now: DateTime) -> Option { - let Some(wire_reset) = window.resets_at.as_deref() else { - window.reset_at_evidence = None; - return None; - }; - if parse_datetime(wire_reset).is_none() { - window.resets_at = None; - window.reset_at_evidence = None; - window.reset_text = None; - return None; - } - let exact_reset = window.reset_at_evidence?; - (exact_reset > now).then(|| { - exact_reset - .timestamp() - .max(now.timestamp().saturating_add(1)) - }) -} - impl UsageWindow { - /// Builds a display window before its provider assigns stable presentation - /// identity. Stage 3A deliberately does not infer duration from a label. + /// Build a window from a "remaining fraction" (0..1) — the shape Antigravity + /// reports per model. Used-percent is derived; identity and duration are + /// attached by the provider adapter before the snapshot is emitted. pub(crate) fn from_fraction( label: String, remaining_fraction: f64, resets_at: Option>, now: DateTime, ) -> Self { - Self::from_used_percent(label, (1.0 - remaining_fraction) * 100.0, resets_at, now) + Self::from_used_percent( + label, + (1.0 - remaining_fraction) * 100.0, + resets_at, + now, + None, + ) } + /// Build a window from an absolute used-percent (0..100), with an optional + /// legacy duration hint. The hint is retained only for existing tests and + /// converted to exact seconds before any wire serialization. pub(crate) fn from_used_percent( label: String, used_percent: f64, resets_at: Option>, now: DateTime, + window_minutes: Option, ) -> Self { let used = used_percent.clamp(0.0, 100.0); let remaining = (100.0 - used).clamp(0.0, 100.0); - Self { + let duration_seconds = window_minutes + .filter(|minutes| *minutes > 0) + .and_then(|minutes| minutes.checked_mul(60)); + let mut window = UsageWindow { card_id: "row.unassigned.v1".to_string(), label, used_percent: used, remaining_percent: remaining, resets_at: resets_at.map(|d| d.to_rfc3339_opts(SecondsFormat::Millis, true)), - reset_at_evidence: resets_at, reset_text: resets_at.map(|d| reset_text(d, now)), - duration_evidence: None, - provider_reset_normalized: false, + window_minutes, + window_key: None, + duration_seconds, + duration_source: duration_seconds.map(|_| DurationSource::Contract), + provider_duration: None, + contract_duration: duration_seconds.map(DurationEvidence::contract), pace_status: PaceStatusPayload { state: PaceState::Unavailable, window_key: None, @@ -235,283 +599,220 @@ impl UsageWindow { reason: Some("windowIdentity".to_string()), }, historical_pace: None, - } + }; + window.refresh_initial_pace_status(); + window + } + + /// Preserve the raw provider reading until identity is attached so the + /// generic adapter can classify invalid evidence. `with_identity` then + /// restores finite display percentages before any wire serialization. + pub(crate) fn from_provider_used_percent( + label: String, + used_percent: f64, + resets_at: Option>, + now: DateTime, + ) -> Self { + let mut window = Self::from_used_percent(label, used_percent, resets_at, now, None); + window.used_percent = used_percent; + window.remaining_percent = 100.0 - used_percent; + window + } + + pub(crate) fn from_provider_fraction( + label: String, + remaining_fraction: f64, + resets_at: Option>, + now: DateTime, + ) -> Self { + Self::from_provider_used_percent(label, (1.0 - remaining_fraction) * 100.0, resets_at, now) + } + + pub(crate) fn try_from_provider_used_percent( + label: String, + used_percent: f64, + resets_at: Option>, + now: DateTime, + ) -> Option { + (used_percent.is_finite() && (0.0..=100.0).contains(&used_percent)) + .then(|| Self::from_provider_used_percent(label, used_percent, resets_at, now)) } + pub(crate) fn try_from_provider_fraction( + label: String, + remaining_fraction: f64, + resets_at: Option>, + now: DateTime, + ) -> Option { + (remaining_fraction.is_finite() && (0.0..=1.0).contains(&remaining_fraction)) + .then(|| Self::from_provider_fraction(label, remaining_fraction, resets_at, now)) + } + + /// Attach provider-semantic presentation and history identity plus the + /// frozen provider/contract duration evidence. pub(crate) fn with_identity( mut self, card_id: impl Into, window_key: Option, + provider_duration: Option, + contract_duration: Option, ) -> Self { + let invalid_reading = + !self.used_percent.is_finite() || !(0.0..=100.0).contains(&self.used_percent); + if invalid_reading { + self.used_percent = if self.used_percent.is_finite() { + self.used_percent.clamp(0.0, 100.0) + } else { + 0.0 + }; + self.remaining_percent = 100.0 - self.used_percent; + } self.card_id = card_id.into(); - self.duration_evidence = None; - self.provider_reset_normalized = false; - self.pace_status = match (window_key, self.resets_at.is_some()) { - (None, _) => PaceStatusPayload { + self.window_key = window_key; + self.provider_duration = provider_duration; + self.contract_duration = contract_duration; + self.duration_seconds = self + .provider_duration + .or(self.contract_duration) + .map(|evidence| evidence.duration_seconds) + .filter(|duration| *duration > 0); + self.duration_source = if self.provider_duration.is_some() { + Some(DurationSource::Provider) + } else if self.contract_duration.is_some() { + Some(DurationSource::Contract) + } else { + None + }; + self.window_minutes = self.duration_seconds.map(|seconds| seconds / 60); + self.refresh_initial_pace_status(); + if invalid_reading && self.window_key.is_some() { + self.unavailable("invalidEvidence"); + } + self + } + + fn refresh_initial_pace_status(&mut self) { + if self.window_key.is_none() { + self.duration_seconds = None; + self.duration_source = None; + self.window_minutes = None; + self.pace_status = PaceStatusPayload { state: PaceState::Unavailable, window_key: None, duration_seconds: None, duration_source: None, complete_cycles: 0, reason: Some("windowIdentity".to_string()), - }, - (Some(window_key), false) => PaceStatusPayload { + }; + self.historical_pace = None; + return; + } + if self.resets_at.is_none() { + self.duration_seconds = None; + self.duration_source = None; + self.window_minutes = None; + self.pace_status = PaceStatusPayload { state: PaceState::Unavailable, - window_key: Some(window_key), + window_key: self.window_key.clone(), duration_seconds: None, duration_source: None, complete_cycles: 0, reason: Some("missingReset".to_string()), - }, - (Some(window_key), true) => PaceStatusPayload { - state: PaceState::LearningDuration, - window_key: Some(window_key), - duration_seconds: None, - duration_source: None, - complete_cycles: 0, - reason: None, - }, - }; - self.historical_pace = None; - self - } - - /// Resolve exact provider/contract duration only for adapters with explicit - /// duration semantics. Other providers keep using `with_identity`. - fn with_duration_evidence( - mut self, - now: DateTime, - reset_was_supplied: bool, - provider: Option, - contract: Option, - ) -> Self { - if self.pace_status.window_key.is_none() { - return self; - } - self.provider_reset_normalized = false; - let had_wire_reset = self.resets_at.is_some(); - let exact_reset = self.reset_at_evidence; - let reset_at = normalized_reset_at(&mut self, now); - if reset_was_supplied != reset_at.is_some() || had_wire_reset != reset_at.is_some() { - return self.with_unavailable_reason("invalidEvidence"); - } - let provider_reset_normalized = provider.is_some_and(|evidence| { - let Some(exact_reset) = exact_reset else { - return false; }; - exact_reset > now - && exact_reset.timestamp() == now.timestamp() - && evidence.reset_at == reset_at - && evidence.reset_at == Some(exact_reset.timestamp().saturating_add(1)) - }); - match resolve_duration(now.timestamp(), reset_at, provider, contract, None) { - DurationResolution::Ready { - duration_seconds, - source, - } => { - let evidence = match source { - DurationSource::Provider => provider, - DurationSource::Contract => contract, - DurationSource::Observed => None, - }; - let Some(evidence) = evidence else { - return self.with_unavailable_reason("invalidEvidence"); - }; - self.duration_evidence = Some((evidence, source)); - self.provider_reset_normalized = - source == DurationSource::Provider && provider_reset_normalized; - self.pace_status.state = PaceState::LearningHistory; - self.pace_status.duration_seconds = Some(duration_seconds); - self.pace_status.duration_source = Some(source); - self.pace_status.complete_cycles = 0; - self.pace_status.reason = None; - self.historical_pace = None; - self - } - DurationResolution::LearningDuration => self, - DurationResolution::Unavailable(reason) => self.with_unavailable_reason(match reason { - DurationUnavailableReason::MissingReset => "missingReset", - DurationUnavailableReason::InvalidEvidence => "invalidEvidence", - }), - } - } - - pub(crate) fn with_provider_duration_evidence( - self, - now: DateTime, - reset_was_supplied: bool, - provider: Option, - ) -> Self { - self.with_duration_evidence(now, reset_was_supplied, provider, None) - } - - pub(crate) fn with_contract_duration_evidence( - self, - now: DateTime, - reset_was_supplied: bool, - contract: DurationEvidence, - ) -> Self { - self.with_duration_evidence(now, reset_was_supplied, None, Some(contract)) - } - - pub(crate) fn with_observed_duration_evidence( - self, - now: DateTime, - reset_was_supplied: bool, - ) -> Self { - let mut window = self.with_duration_evidence(now, reset_was_supplied, None, None); - if window.pace_status.state == PaceState::LearningDuration { - window.pace_status.duration_source = Some(DurationSource::Observed); + self.historical_pace = None; + return; } - window - } - - pub(crate) fn with_non_recurring(self) -> Self { - self.with_unavailable_reason("nonRecurring") - } - - fn with_unavailable_reason(mut self, reason: &str) -> Self { - self.unavailable(reason); - self + let state = if self.duration_seconds.is_some() { + PaceState::LearningHistory + } else { + PaceState::LearningDuration + }; + self.pace_status = PaceStatusPayload { + state, + window_key: self.window_key.clone(), + duration_seconds: self.duration_seconds, + duration_source: self.duration_source, + complete_cycles: 0, + reason: None, + }; + self.historical_pace = None; } - fn unavailable(&mut self, reason: &str) { - self.duration_evidence = None; - self.provider_reset_normalized = false; + pub(crate) fn unavailable(&mut self, reason: impl Into) { + let reason = reason.into(); + self.duration_seconds = None; + self.duration_source = None; + self.window_minutes = None; + self.historical_pace = None; self.pace_status = PaceStatusPayload { state: PaceState::Unavailable, - window_key: self.pace_status.window_key.clone(), + window_key: self.window_key.clone(), duration_seconds: None, duration_source: None, complete_cycles: 0, - reason: Some(reason.to_string()), + reason: Some(reason), }; - self.historical_pace = None; } fn validate_wire(&self) -> Result<(), String> { if self.card_id.trim().is_empty() { return Err("pace cardId must be non-empty".to_string()); } - if !self.used_percent.is_finite() - || !self.remaining_percent.is_finite() - || !(0.0..=100.0).contains(&self.used_percent) - || !(0.0..=100.0).contains(&self.remaining_percent) - || ((self.used_percent + self.remaining_percent) - 100.0).abs() > 1e-6 + if self.window_key != self.pace_status.window_key { + return Err("pace windowKey internal and nested values differ".to_string()); + } + if self.duration_seconds != self.pace_status.duration_seconds { + return Err("pace durationSeconds internal and nested values differ".to_string()); + } + if self.duration_source != self.pace_status.duration_source { + return Err("pace durationSource internal and nested values differ".to_string()); + } + if self.window_minutes != self.duration_seconds.map(|seconds| seconds / 60) { + return Err("pace windowMinutes must derive from durationSeconds".to_string()); + } + if self.duration_seconds.is_none() + && self.duration_source.is_some() + && !(self.pace_status.state == PaceState::LearningDuration + && self.duration_source == Some(DurationSource::Observed)) { - return Err("usage percentages must be finite, bounded, and complementary".to_string()); + return Err("pace durationSource requires a duration".to_string()); } - let has_window_key = self - .pace_status - .window_key - .as_deref() - .is_some_and(|key| !key.trim().is_empty()); - if self.pace_status.window_key.is_some() && !has_window_key { - return Err("pace windowKey must be non-empty".to_string()); + if let Some(window_key) = self.pace_status.window_key.as_deref() { + if window_key.trim().is_empty() { + return Err("pace windowKey must be non-empty".to_string()); + } } let identity_unavailable = self.pace_status.state == PaceState::Unavailable && self.pace_status.reason.as_deref() == Some("windowIdentity"); - if has_window_key == identity_unavailable { - return Err("pace window identity invariant failed".to_string()); - } - let reset_at = self - .resets_at - .as_deref() - .and_then(parse_datetime) - .map(|reset| reset.timestamp()); - let valid_reset = reset_at.is_some(); - if self.resets_at.is_some() && !valid_reset { - return Err("pace resetsAt must be a valid timestamp".to_string()); - } - let has_duration = self.pace_status.duration_seconds.is_some(); - let observed_learning_source = self.pace_status.state == PaceState::LearningDuration - && self.pace_status.duration_source == Some(DurationSource::Observed); - if self - .pace_status - .duration_seconds - .is_some_and(|duration| !valid_duration(duration)) - || (has_duration != self.pace_status.duration_source.is_some() - && !observed_learning_source) - { - return Err("pace duration invariant failed".to_string()); + if self.pace_status.window_key.is_none() != identity_unavailable { + return Err("pace windowKey identity invariant failed".to_string()); } - if self.provider_reset_normalized - && !matches!( - self.duration_evidence, - Some((_, DurationSource::Provider)) - ) - { - return Err("pace normalized provider reset lacks provider evidence".to_string()); - } - match self.duration_evidence { - Some((evidence, source)) => { - let reset_is_coherent = match source { - DurationSource::Provider => { - let Some(evidence_reset) = evidence.reset_at else { - return Err("pace provider evidence lacks reset".to_string()); - }; - let Some(wire_reset) = reset_at else { - return Err("pace provider evidence lacks wire reset".to_string()); - }; - evidence_reset == wire_reset - || (self.provider_reset_normalized - && wire_reset.checked_add(1) == Some(evidence_reset)) - } - DurationSource::Contract => evidence.reset_at.is_none(), - DurationSource::Observed => false, - }; - if !reset_is_coherent - || self.pace_status.duration_seconds != Some(evidence.duration_seconds) - || self.pace_status.duration_source != Some(source) - { - return Err("pace duration evidence and state differ".to_string()); - } + if let Some(duration) = self.pace_status.duration_seconds { + if duration <= 0 { + return Err("pace durationSeconds must be positive".to_string()); } - None if matches!( - self.pace_status.duration_source, - Some(DurationSource::Provider | DurationSource::Contract) - ) => - { - return Err("pace duration source lacks retained evidence".to_string()); + if self.pace_status.duration_source.is_none() { + return Err("pace durationSource is required with durationSeconds".to_string()); } - None => {} } match self.pace_status.state { - PaceState::LearningDuration => { - if !valid_reset - || has_duration - || self.historical_pace.is_some() - || self.pace_status.reason.is_some() - { - return Err("learningDuration pace invariant failed".to_string()); + PaceState::Available => { + if self.pace_status.duration_seconds.is_none() || self.historical_pace.is_none() { + return Err("available pace requires duration and historicalPace".to_string()); } } PaceState::LearningHistory => { - if !valid_reset - || !has_duration - || self.historical_pace.is_some() - || self.pace_status.reason.is_some() - { + if self.pace_status.duration_seconds.is_none() || self.historical_pace.is_some() { return Err("learningHistory pace invariant failed".to_string()); } } - PaceState::Available => { - if !valid_reset - || !has_duration - || self.historical_pace.is_none() - || self.pace_status.reason.is_some() - { - return Err("available pace invariant failed".to_string()); + PaceState::LearningDuration => { + if self.pace_status.duration_seconds.is_some() || self.historical_pace.is_some() { + return Err("learningDuration pace invariant failed".to_string()); } } PaceState::Unavailable => { - if has_duration - || self.pace_status.duration_source.is_some() - || self.historical_pace.is_some() - || self.pace_status.reason.is_none() - || (self.pace_status.reason.as_deref() == Some("missingReset") - && self.resets_at.is_some()) - { + if self.historical_pace.is_some() || self.pace_status.reason.as_deref().is_none() { return Err("unavailable pace invariant failed".to_string()); } } @@ -522,9 +823,9 @@ impl UsageWindow { || historical .eta_seconds .is_some_and(|eta| !eta.is_finite() || eta < 0.0) - || historical - .run_out_probability - .is_some_and(|risk| !risk.is_finite() || !(0.0..=1.0).contains(&risk)) + || historical.run_out_probability.is_some_and(|probability| { + !probability.is_finite() || !(0.0..=1.0).contains(&probability) + }) || (historical.eta_seconds.is_none() != historical.will_last_to_reset) { return Err("historicalPace contains contradictory values".to_string()); @@ -544,8 +845,13 @@ impl UsageWindow { } #[cfg(test)] - pub(crate) fn card_id_for_test(&self) -> &str { - &self.card_id + pub(crate) fn resets_at_for_test(&self) -> Option<&str> { + self.resets_at.as_deref() + } + + #[cfg(test)] + pub(crate) fn window_minutes_for_test(&self) -> Option { + self.duration_seconds.map(|seconds| seconds / 60) } #[cfg(test)] @@ -583,6 +889,13 @@ struct CodexCredentials { scope_slot: CredentialSlot, } +#[derive(Debug)] +struct CodexCredentialWriteReceipt { + path: PathBuf, + previous_root: Value, + persisted_root: Value, +} + impl CodexCredentials { fn scope_marker(&self) -> &[u8] { self.refresh_token @@ -605,9 +918,13 @@ struct ClaudeCredentials { /// Where the credentials were read from, so a rotated token can be written /// back to the same place (the Claude CLI shares this store). source: ClaudeCredentialSource, - /// Full credentials JSON as loaded, so a write-back preserves fields we - /// don't model (merge-update rather than overwrite). + /// Full credentials JSON captured at reload. The target object is the + /// optimistic write guard; top-level siblings are merged from the current + /// store at save time. raw_root: Option, + /// Exact Keychain account whose item was read at refresh reload. A later + /// write-back may validate this identity, but must never retarget it. + keychain_account: Option, scope_slot: CredentialSlot, } @@ -644,6 +961,14 @@ enum ClaudeCredentialSource { Environment, } +#[derive(Debug)] +enum ClaudeLoginResolution { + Absent, + ExplicitLogout, + Ready(ClaudeCredentials), + Terminal, +} + #[derive(Debug, Deserialize)] struct ClaudeCredentialsRoot { #[serde(default, rename = "claudeAiOauth")] @@ -675,9 +1000,9 @@ struct CodexUsageResponse { #[derive(Debug, Deserialize)] struct CodexRateLimit { - #[serde(default)] + #[serde(default, deserialize_with = "deserialize_optional_raw")] primary_window: Option, - #[serde(default)] + #[serde(default, deserialize_with = "deserialize_optional_raw")] secondary_window: Option, } @@ -706,47 +1031,53 @@ struct CodexCredits { balance: Option, } +fn finite_codex_balance(credits: Option<&CodexCredits>) -> Option { + credits + .and_then(|credits| credits.balance) + .filter(|balance| balance.is_finite()) +} + #[derive(Debug, Deserialize, Default)] struct ClaudeUsageResponse { - #[serde(default)] + #[serde(default, deserialize_with = "deserialize_optional_raw")] five_hour: Option, - #[serde(default)] + #[serde(default, deserialize_with = "deserialize_optional_raw")] seven_day: Option, - #[serde(default)] + #[serde(default, deserialize_with = "deserialize_optional_raw")] seven_day_oauth_apps: Option, - #[serde(default)] + #[serde(default, deserialize_with = "deserialize_optional_raw")] seven_day_opus: Option, - #[serde(default)] + #[serde(default, deserialize_with = "deserialize_optional_raw")] seven_day_sonnet: Option, - #[serde(default)] + #[serde(default, deserialize_with = "deserialize_optional_raw")] seven_day_design: Option, - #[serde(default)] + #[serde(default, deserialize_with = "deserialize_optional_raw")] seven_day_claude_design: Option, - #[serde(default)] + #[serde(default, deserialize_with = "deserialize_optional_raw")] claude_design: Option, - #[serde(default)] + #[serde(default, deserialize_with = "deserialize_optional_raw")] design: Option, - #[serde(default)] + #[serde(default, deserialize_with = "deserialize_optional_raw")] seven_day_omelette: Option, - #[serde(default)] + #[serde(default, deserialize_with = "deserialize_optional_raw")] omelette: Option, - #[serde(default)] + #[serde(default, deserialize_with = "deserialize_optional_raw")] omelette_promotional: Option, - #[serde(default)] + #[serde(default, deserialize_with = "deserialize_optional_raw")] seven_day_routines: Option, - #[serde(default)] + #[serde(default, deserialize_with = "deserialize_optional_raw")] seven_day_claude_routines: Option, - #[serde(default)] + #[serde(default, deserialize_with = "deserialize_optional_raw")] claude_routines: Option, - #[serde(default)] + #[serde(default, deserialize_with = "deserialize_optional_raw")] routines: Option, - #[serde(default)] + #[serde(default, deserialize_with = "deserialize_optional_raw")] routine: Option, - #[serde(default)] + #[serde(default, deserialize_with = "deserialize_optional_raw")] seven_day_cowork: Option, - #[serde(default)] + #[serde(default, deserialize_with = "deserialize_optional_raw")] cowork: Option, - #[serde(default)] + #[serde(default, deserialize_with = "deserialize_optional_raw")] extra_usage: Option, } @@ -787,7 +1118,156 @@ struct ClaudeRefreshResponse { expires_in: i64, } -pub async fn run() -> AgentUsagePayload { +fn empty_error_snapshot( + client_id: &str, + source: &str, + now: DateTime, + display: String, + transport_diagnostic: Option, +) -> AgentUsageSnapshot { + AgentUsageSnapshot { + client_id: client_id.to_string(), + source: source.to_string(), + updated_at: now.to_rfc3339_opts(SecondsFormat::Millis, true), + identity: None, + account_scope: Err(AccountScopeError::NoTrustedEvidence), + windows: Vec::new(), + credits: None, + error: Some(display), + transport_diagnostic, + } +} + +fn usable_success(snapshot: &AgentUsageSnapshot) -> bool { + match snapshot.client_id.as_str() { + "codex" => { + !snapshot.windows.is_empty() + || snapshot + .credits + .as_ref() + .and_then(|credits| credits.remaining) + .is_some_and(f64::is_finite) + } + "grok" => snapshot + .windows + .iter() + .any(|window| window.card_id == "billing.weekly.v1"), + "claude" | "copilot" | "antigravity" => !snapshot.windows.is_empty(), + _ => false, + } +} + +fn lock_last_good( + cache: &Mutex, +) -> std::sync::MutexGuard<'_, ProviderLastGoodCache> { + cache + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()) +} + +fn apply_provider_outcome_with( + cache: &Mutex, + client_id: &str, + failure_source: &str, + now: DateTime, + outcome: ProviderFetchOutcome, + mut enrich: F, +) -> Option +where + F: FnMut(&mut AgentUsageSnapshot), +{ + match outcome { + ProviderFetchOutcome::Absent => { + lock_last_good(cache).clear(client_id); + None + } + ProviderFetchOutcome::Success { + mut snapshot, + cache_binding, + } => { + match snapshot.account_scope.as_ref() { + Ok(_) | Err(AccountScopeError::NoTrustedEvidence) => {} + Err(_) => { + lock_last_good(cache).clear(client_id); + let source = snapshot.source.clone(); + return Some(empty_error_snapshot( + client_id, + &source, + now, + format!( + "{} account identity could not be verified.", + clean_plan(client_id) + ), + None, + )); + } + } + + enrich(&mut snapshot); + let cacheable = snapshot.account_scope.is_ok() + && snapshot.error.is_none() + && snapshot.transport_diagnostic.is_none() + && usable_success(&snapshot); + let mut cache = lock_last_good(cache); + match (cacheable, cache_binding) { + (true, Some(binding)) => cache.replace(client_id, binding, snapshot.clone()), + _ => cache.clear(client_id), + } + Some(snapshot) + } + ProviderFetchOutcome::Failure(ProviderFetchFailure::Terminal { display }) => { + lock_last_good(cache).clear(client_id); + Some(empty_error_snapshot( + client_id, + failure_source, + now, + display, + None, + )) + } + ProviderFetchOutcome::Failure(ProviderFetchFailure::Transient { + display, + attempt_binding, + transport_diagnostic, + }) => { + let fallback = attempt_binding + .as_ref() + .and_then(|binding| lock_last_good(cache).clean_for(client_id, binding)); + let Some(mut snapshot) = fallback else { + lock_last_good(cache).clear(client_id); + return Some(empty_error_snapshot( + client_id, + failure_source, + now, + display, + Some(transport_diagnostic), + )); + }; + snapshot.account_scope = Err(AccountScopeError::NoTrustedEvidence); + snapshot.error = Some(display); + snapshot.transport_diagnostic = Some(transport_diagnostic); + Some(snapshot) + } + } +} + +fn apply_provider_outcome( + client_id: &str, + failure_source: &str, + now: DateTime, + outcome: ProviderFetchOutcome, +) -> Option { + apply_provider_outcome_with( + &PROVIDER_LAST_GOOD, + client_id, + failure_source, + now, + outcome, + |snapshot| enrich_snapshot(snapshot, now.timestamp()), + ) +} + +pub async fn run(publication_generation: u64) -> AgentUsagePayload { let generated_at = Utc::now().to_rfc3339_opts(SecondsFormat::Millis, true); let (codex, claude, antigravity, copilot, grok) = tokio::join!( fetch_codex(), @@ -805,808 +1285,602 @@ pub async fn run() -> AgentUsagePayload { if let Some(grok) = grok { agents.push(grok); } - for snapshot in &mut agents { - retain_unique_windows(&mut snapshot.windows); - } AgentUsagePayload { generated_at, + publication_generation, agents, opencode_subscriptions: crate::opencode_integrations::detect_subscriptions(), } } -fn retain_unique_windows(windows: &mut Vec) { - let mut card_ids = HashSet::new(); - let mut window_keys = HashSet::new(); - windows.retain(|window| { - let key = window.pace_status.window_key.as_ref(); - if card_ids.contains(&window.card_id) || key.is_some_and(|key| window_keys.contains(key)) { - return false; +async fn fetch_grok() -> Option { + let now = Utc::now(); + let outcome = match agent_grok::fetch(now).await { + Ok(Some(data)) => ProviderFetchOutcome::Success { + cache_binding: data.cache_binding, + snapshot: AgentUsageSnapshot { + client_id: "grok".to_string(), + source: "oauth".to_string(), + updated_at: now.to_rfc3339_opts(SecondsFormat::Millis, true), + identity: data.identity, + account_scope: data.account_scope, + windows: data.windows, + credits: None, + error: None, + transport_diagnostic: None, + }, + }, + Ok(None) => ProviderFetchOutcome::Absent, + Err(failure) => ProviderFetchOutcome::Failure(failure), + }; + apply_provider_outcome("grok", "oauth", now, outcome) +} + +async fn fetch_copilot() -> Option { + let now = Utc::now(); + let outcome = match crate::opencode_integrations::github_copilot_credential() { + crate::opencode_integrations::GitHubCopilotCredentialLoad::Absent => { + ProviderFetchOutcome::Absent } - card_ids.insert(window.card_id.clone()); - if let Some(key) = key { - window_keys.insert(key.clone()); + crate::opencode_integrations::GitHubCopilotCredentialLoad::Terminal(display) => { + ProviderFetchOutcome::Failure(ProviderFetchFailure::terminal(display)) } - true - }); + crate::opencode_integrations::GitHubCopilotCredentialLoad::Present(credential) => { + match agent_copilot::fetch(now, credential).await { + Ok(data) => ProviderFetchOutcome::Success { + cache_binding: Some(data.cache_binding), + snapshot: AgentUsageSnapshot { + client_id: "copilot".to_string(), + source: "oauth".to_string(), + updated_at: now.to_rfc3339_opts(SecondsFormat::Millis, true), + identity: data.identity, + account_scope: data.account_scope, + windows: data.windows, + credits: None, + error: None, + transport_diagnostic: None, + }, + }, + Err(failure) => ProviderFetchOutcome::Failure(failure), + } + } + }; + apply_provider_outcome("copilot", "oauth", now, outcome) } -fn enrich_snapshot(snapshot: &mut AgentUsageSnapshot, now: DateTime) { - enrich_snapshot_with(snapshot, now, |active_keys, observations, now| { - crate::agent_quota_history::record_observations_and_evaluate(active_keys, observations, now) - }); +async fn fetch_antigravity() -> AgentUsageSnapshot { + let now = Utc::now(); + let outcome = match agent_antigravity::fetch(now).await { + Ok(fetched) => ProviderFetchOutcome::Success { + cache_binding: fetched.cache_binding, + snapshot: AgentUsageSnapshot { + client_id: "antigravity".to_string(), + source: fetched.source, + updated_at: now.to_rfc3339_opts(SecondsFormat::Millis, true), + identity: fetched.identity, + account_scope: fetched.account_scope, + windows: fetched.windows, + credits: None, + error: None, + transport_diagnostic: None, + }, + }, + Err(failure) => ProviderFetchOutcome::Failure(failure), + }; + apply_provider_outcome("antigravity", "oauth", now, outcome) + .expect("Antigravity is a required provider card") } -fn enrich_snapshot_with(snapshot: &mut AgentUsageSnapshot, now: DateTime, mut record: F) -where - F: FnMut( - &[SeriesKey], - &[QuotaObservation], - i64, - ) -> Result, HistoryError>, -{ - retain_unique_windows(&mut snapshot.windows); +async fn fetch_codex() -> AgentUsageSnapshot { + let now = Utc::now(); + apply_provider_outcome("codex", "oauth", now, fetch_codex_inner().await) + .expect("Codex is a required provider card") +} - let Ok(account_scope) = snapshot.account_scope.as_ref() else { - for window in &mut snapshot.windows { - if window.pace_status.window_key.is_some() - && window.pace_status.reason.as_deref() != Some("nonRecurring") - { - window.unavailable("accountScope"); - } - } - return; - }; - let provider_id = snapshot.client_id.clone(); - let account_scope = account_scope.as_str().to_string(); - let mut active_keys = Vec::new(); - let mut observations = Vec::new(); - let mut mapped_indices = Vec::new(); +/// Claude's `/api/oauth/usage` rate-limits aggressively. The gate stores only +/// the cooldown deadline and the exact opaque binding that triggered it; display +/// snapshots live exclusively in the provider last-good cache. +#[derive(Debug, Default)] +struct ClaudeUsageGate { + blocked_until: Option>, + binding: Option, +} - for (index, window) in snapshot.windows.iter_mut().enumerate() { - if window.pace_status.reason.as_deref() == Some("nonRecurring") { - continue; +impl ClaudeUsageGate { + fn blocked_until_for( + &mut self, + binding: &ProviderCacheBinding, + now: DateTime, + ) -> Option> { + if self.binding.as_ref() != Some(binding) { + self.blocked_until = None; + self.binding = None; + return None; } - let Some(window_key) = window.pace_status.window_key.as_deref() else { - continue; - }; - let key = SeriesKey::new(provider_id.clone(), account_scope.clone(), window_key); - active_keys.push(key.clone()); - if window.pace_status.state == PaceState::Unavailable { - continue; + match self.blocked_until { + Some(until) if until > now => Some(until), + Some(_) => { + self.blocked_until = None; + self.binding = None; + None + } + None => None, } + } - let reset_at = normalized_reset_at(window, now); - let valid_percent = - window.used_percent.is_finite() && (0.0..=100.0).contains(&window.used_percent); - if !valid_percent { - window.used_percent = if window.used_percent.is_finite() { - window.used_percent.clamp(0.0, 100.0) - } else { - 0.0 - }; - } - window.remaining_percent = 100.0 - window.used_percent; - if !valid_percent { - window.unavailable("invalidEvidence"); - continue; - } - let Some(reset_at) = reset_at else { - window.unavailable("invalidEvidence"); - continue; - }; - let (provider, contract) = match window.duration_evidence { - Some((evidence, DurationSource::Provider)) => (Some(evidence), None), - Some((evidence, DurationSource::Contract)) => (None, Some(evidence)), - Some((_, DurationSource::Observed)) => { - window.unavailable("invalidEvidence"); - continue; - } - None => (None, None), - }; - observations.push(QuotaObservation { - key, - reset_at: Some(reset_at), - used_percent: window.used_percent, - provider, - contract, - }); - mapped_indices.push(index); + fn record_rate_limit( + &mut self, + binding: ProviderCacheBinding, + retry_after: Option>, + now: DateTime, + ) { + self.blocked_until = Some( + retry_after + .filter(|until| *until > now) + .unwrap_or_else(|| now + chrono::Duration::minutes(5)), + ); + self.binding = Some(binding); } - if active_keys.is_empty() { - return; + fn clear(&mut self) { + self.blocked_until = None; + self.binding = None; } +} - let results = match record(&active_keys, &observations, now.timestamp()) { - Ok(results) if results.len() == mapped_indices.len() => results, - Ok(_) => { - for index in mapped_indices { - snapshot.windows[index].unavailable("history"); - } - return; - } - Err(error) => { - let reason = history_error_reason(error); - for index in mapped_indices { - snapshot.windows[index].unavailable(reason); - } - return; +static CLAUDE_USAGE_GATE: Mutex = Mutex::new(ClaudeUsageGate { + blocked_until: None, + binding: None, +}); + +fn lock_gate() -> std::sync::MutexGuard<'static, ClaudeUsageGate> { + CLAUDE_USAGE_GATE + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()) +} + +fn claude_gate_failure( + binding: ProviderCacheBinding, + blocked_until: DateTime, + now: DateTime, +) -> ProviderFetchFailure { + let wait_secs = (blocked_until - now).num_seconds().max(0); + ProviderFetchFailure::transient( + format!( + "Claude OAuth usage endpoint is rate limited. Retrying automatically in ~{wait_secs}s." + ), + Some(binding), + SafeTransportDiagnostic::rate_limited(429), + ) +} + +fn parse_retry_after(value: Option<&reqwest::header::HeaderValue>) -> Option> { + let raw = value?.to_str().ok()?.trim(); + if raw.is_empty() { + return None; + } + if let Ok(seconds) = raw.parse::() { + return (seconds >= 0).then(|| Utc::now() + chrono::Duration::seconds(seconds)); + } + DateTime::parse_from_rfc2822(raw) + .ok() + .map(|t| t.with_timezone(&Utc)) +} + +async fn fetch_claude() -> AgentUsageSnapshot { + let now = Utc::now(); + let (failure_source, outcome) = fetch_claude_inner().await; + apply_provider_outcome("claude", failure_source, now, outcome) + .expect("Claude is a required provider card") +} + +async fn fetch_codex_inner() -> ProviderFetchOutcome { + let loaded = match load_codex_credentials() { + Ok(credentials) => credentials, + Err(display) => { + return ProviderFetchOutcome::Failure(ProviderFetchFailure::terminal(display)); } }; - - for (index, result) in mapped_indices.into_iter().zip(results) { - let window = &mut snapshot.windows[index]; - match result { - Err(error) => window.unavailable(history_error_reason(error)), - Ok(( - HistoryOutcome::Ready { - duration_seconds, - source, - .. - }, - historical, - complete_cycles, - )) => { - let Some(reset_at) = normalized_reset_at(window, now) else { - window.unavailable("history"); - continue; - }; - if !history_duration_is_coherent( - window, - reset_at, - now, - duration_seconds, - source, - ) { - window.unavailable("history"); - continue; - } - match historical { - Some(pace) if historical_pace_is_coherent(&pace) => { - window.pace_status = PaceStatusPayload { - state: PaceState::Available, - window_key: window.pace_status.window_key.clone(), - duration_seconds: Some(duration_seconds), - duration_source: Some(source), - complete_cycles, - reason: None, - }; - window.historical_pace = Some(historical_pace_payload(pace)); - } - Some(_) => window.unavailable("history"), - None => { - window.pace_status = PaceStatusPayload { - state: PaceState::LearningHistory, - window_key: window.pace_status.window_key.clone(), - duration_seconds: Some(duration_seconds), - duration_source: Some(source), - complete_cycles, - reason: None, - }; - window.historical_pace = None; - } - } - } - Ok((HistoryOutcome::LearningDuration, None, 0)) - if window.duration_evidence.is_none() => - { - window.pace_status = PaceStatusPayload { - state: PaceState::LearningDuration, - window_key: window.pace_status.window_key.clone(), - duration_seconds: None, - duration_source: Some(DurationSource::Observed), - complete_cycles: 0, - reason: None, - }; - window.historical_pace = None; - } - Ok((HistoryOutcome::Unavailable(reason), None, 0)) => { - if reason == DurationUnavailableReason::MissingReset && window.resets_at.is_some() { - window.unavailable("history"); - } else { - window.unavailable(duration_unavailable_reason(reason)); - } + let verified = if credentials_needs_refresh(loaded.last_refresh) { + refresh_codex_credentials(&loaded.auth_path).await + } else { + resolve_codex_cache_binding(&loaded) + .map(|binding| (loaded, binding)) + .map_err(|_| { + ProviderFetchFailure::terminal("Codex account identity could not be verified.") + }) + }; + let (credentials, cache_binding, response) = + match request_after_verified_binding(verified, |(credentials, cache_binding)| async move { + let client = provider_http_client_builder() + .timeout(std::time::Duration::from_secs(30)) + .build() + .map_err(|_| { + ProviderFetchFailure::terminal("Codex usage client could not be created.") + })?; + let request_account_id = credentials + .account_id + .as_deref() + .map(str::trim) + .filter(|value| !value.is_empty()); + let mut request = client + .get(CODEX_USAGE_URL) + .bearer_auth(&credentials.access_token) + .header(reqwest::header::ACCEPT, "application/json") + .header(reqwest::header::USER_AGENT, "TokenBar"); + if let Some(account_id) = request_account_id { + request = request.header("ChatGPT-Account-Id", account_id); } - Ok(_) => window.unavailable("history"), + let response = request.send().await.map_err(|error| { + ProviderFetchFailure::from_send_error( + "Codex usage request failed. Retrying automatically.", + Some(cache_binding.clone()), + &error, + ) + })?; + Ok((credentials, cache_binding, response)) + }) + .await + { + Ok(verified) => verified, + Err(failure) => return ProviderFetchOutcome::Failure(failure), + }; + let request_account_id = credentials + .account_id + .as_deref() + .map(str::trim) + .filter(|value| !value.is_empty()); + let status = response.status().as_u16(); + let body = match read_response_body(status, false, || async { + response.text().await.map_err(|error| { + TransportErrorFacts::from_reqwest(&error, TransportPhase::ResponseBody) + }) + }) + .await + { + Ok(body) => body, + Err(ResponseReadFailure::Transient(diagnostic)) => { + return ProviderFetchOutcome::Failure(ProviderFetchFailure::transient( + "Codex usage request failed. Retrying automatically.", + Some(cache_binding), + diagnostic, + )); + } + Err(ResponseReadFailure::Terminal(401 | 403)) => { + return ProviderFetchOutcome::Failure(ProviderFetchFailure::terminal( + "Codex OAuth token expired or invalid. Run `codex` to log in again.", + )); + } + Err(ResponseReadFailure::Terminal(status)) => { + return ProviderFetchOutcome::Failure(ProviderFetchFailure::terminal(format!( + "Codex usage API rejected the request (status {status})." + ))); } - } -} - -fn history_duration_is_coherent( - window: &UsageWindow, - reset_at: i64, - now: DateTime, - duration_seconds: i64, - source: DurationSource, -) -> bool { - let (provider, contract) = match window.duration_evidence { - Some((evidence, DurationSource::Provider)) => (Some(evidence), None), - Some((evidence, DurationSource::Contract)) => (None, Some(evidence)), - Some((_, DurationSource::Observed)) => return false, - None => (None, None), }; - if source == DurationSource::Observed && window.duration_evidence.is_some() { - return false; - } - let observed = (source == DurationSource::Observed) - .then(|| DurationEvidence::observed(reset_at, duration_seconds)); - matches!( - resolve_duration(now.timestamp(), Some(reset_at), provider, contract, observed), - DurationResolution::Ready { - duration_seconds: resolved, - source: resolved_source, - } if resolved == duration_seconds && resolved_source == source - ) -} -fn history_error_reason(error: HistoryError) -> &'static str { - if error == HistoryError::StoreCapacity { - "storeCapacity" - } else { - "history" + let mut usage: CodexUsageResponse = match serde_json::from_str(&body) { + Ok(usage) => usage, + Err(_) => { + return ProviderFetchOutcome::Failure(ProviderFetchFailure::terminal( + "Codex usage response could not be decoded.", + )); + } + }; + let now = Utc::now(); + let account_scope = cache_binding + .corroborating + .clone() + .unwrap_or_else(|| cache_binding.primary.clone()); + let identity = Some(AgentIdentity { + email: credentials.id_token.as_deref().and_then(jwt_email), + plan: usage.plan_type.as_deref().map(clean_plan).or_else(|| { + credentials + .id_token + .as_deref() + .and_then(jwt_plan) + .map(clean_plan) + }), + }); + let windows = codex_windows( + usage.rate_limit.as_ref(), + usage.additional_rate_limits.as_deref(), + now, + ); + let finite_balance = finite_codex_balance(usage.credits.as_ref()); + if let Some(credits) = usage.credits.as_mut() { + credits.balance = finite_balance; } -} - -fn duration_unavailable_reason(reason: DurationUnavailableReason) -> &'static str { - match reason { - DurationUnavailableReason::MissingReset => "missingReset", - DurationUnavailableReason::InvalidEvidence => "invalidEvidence", + if windows.is_empty() && finite_balance.is_none() { + return ProviderFetchOutcome::Failure(ProviderFetchFailure::terminal( + "Codex usage API returned no usable quota data.", + )); } -} - -fn historical_pace_is_coherent(pace: &HistoricalPace) -> bool { - pace.expected_percent.is_finite() - && (0.0..=100.0).contains(&pace.expected_percent) - && pace - .eta_seconds - .is_none_or(|eta| eta.is_finite() && eta >= 0.0) - && pace - .run_out_probability - .is_none_or(|probability| probability.is_finite() && (0.0..=1.0).contains(&probability)) - && (pace.eta_seconds.is_none() == pace.will_last_to_reset) -} -fn historical_pace_payload(pace: HistoricalPace) -> HistoricalPacePayload { - HistoricalPacePayload { - expected_used_percent: pace.expected_percent, - eta_seconds: pace.eta_seconds, - will_last_to_reset: pace.will_last_to_reset, - run_out_probability: pace.run_out_probability, + if let Some(request_account_id) = request_account_id { + let _ = crate::agent_quota_history::migrate_codex_v2( + request_account_id, + account_scope.as_str(), + now.timestamp(), + ); } -} -fn finalize_grok_snapshot_with( - fetched: Result, - now: DateTime, - mut enrich: F, -) -> AgentUsageSnapshot -where - F: FnMut(&mut AgentUsageSnapshot, DateTime), -{ - let mut snapshot = match fetched { - Ok(data) => AgentUsageSnapshot { - client_id: "grok".to_string(), + ProviderFetchOutcome::Success { + snapshot: AgentUsageSnapshot { + client_id: "codex".to_string(), source: "oauth".to_string(), updated_at: now.to_rfc3339_opts(SecondsFormat::Millis, true), - identity: data.identity, - account_scope: data.account_scope, - windows: data.windows, - credits: None, + identity, + account_scope: Ok(account_scope), + windows, + credits: usage.credits.map(|credits| CreditsSnapshot { + remaining: credits.balance, + unlimited: credits.unlimited, + }), error: None, + transport_diagnostic: None, }, - Err(error) => AgentUsageSnapshot { - client_id: "grok".to_string(), - source: "oauth".to_string(), - updated_at: now.to_rfc3339_opts(SecondsFormat::Millis, true), - identity: None, - account_scope: Err(AccountScopeError::NoTrustedEvidence), - windows: Vec::new(), - credits: None, - error: Some(error), - }, - }; - enrich(&mut snapshot, now); - snapshot + cache_binding: Some(cache_binding), + } } -async fn fetch_grok() -> Option { - let now = Utc::now(); - Some(finalize_grok_snapshot_with( - agent_grok::fetch(now).await?, - now, - enrich_snapshot, - )) +fn resolve_codex_cache_binding( + credentials: &CodexCredentials, +) -> Result { + let primary = agent_account_scope::resolve_credential( + "codex", + credentials.scope_slot.semantic_source, + &credentials.scope_slot.canonical_location, + credentials.scope_marker(), + )?; + let corroborating = credentials + .account_id + .as_deref() + .map(str::trim) + .filter(|value| !value.is_empty()) + .map(|account_id| { + agent_account_scope::resolve_authoritative( + "codex", + AuthoritativeIdKind::OpaqueId, + account_id, + ) + }) + .transpose()?; + Ok(ProviderCacheBinding::new(primary, corroborating)) } -async fn fetch_copilot() -> Option { - fetch_copilot_with( - crate::opencode_integrations::github_copilot_credential(), - Utc::now(), - agent_copilot::fetch, - enrich_snapshot, +fn claude_cache_binding( + credentials: &ClaudeCredentials, +) -> Result { + credentials + .resolve_account_scope() + .map(ProviderCacheBinding::primary) +} + +fn clear_claude_gate_for_login_resolution( + login: &ClaudeLoginResolution, + gate: &mut ClaudeUsageGate, +) { + if matches!( + login, + ClaudeLoginResolution::Absent | ClaudeLoginResolution::ExplicitLogout + ) { + gate.clear(); + } +} + +async fn fetch_claude_inner() -> (&'static str, ProviderFetchOutcome) { + if let Some(token) = resolve_claude_code_oauth_token().await { + return fetch_claude_setup_token(token).await; + } + + let login = load_claude_login_credentials(); + clear_claude_gate_for_login_resolution(&login, &mut lock_gate()); + fetch_claude_login_or_setup_with( + login, + |credentials| async move { + let verified = claude_cache_binding(&credentials).map_err(|_| { + ProviderFetchFailure::terminal("Claude account identity could not be verified.") + }); + request_after_verified_binding(verified, |binding| async move { + Ok(fetch_claude_oauth_usage(credentials, binding).await) + }) + .await + .unwrap_or_else(|failure| ("oauth", ProviderFetchOutcome::Failure(failure))) + }, + resolve_claude_keychain_token, + fetch_claude_setup_token, ) .await } -async fn fetch_copilot_with( - credential: Option, - now: DateTime, - fetch: Fetch, - enrich: Enrich, -) -> Option -where - Fetch: FnOnce( - DateTime, - crate::opencode_integrations::GitHubCopilotCredential, - ) -> FetchFuture, - FetchFuture: std::future::Future>, - Enrich: FnMut(&mut AgentUsageSnapshot, DateTime), -{ - // No exact opencode github-copilot OAuth credential means no card and no network. - let credential = credential?; - Some(finalize_copilot_snapshot_with( - fetch(now, credential).await, - now, - enrich, - )) +async fn fetch_claude_setup_token( + token: ResolvedClaudeToken, +) -> (&'static str, ProviderFetchOutcome) { + let credentials = claude_credentials_from_access_token(token); + let verified = claude_cache_binding(&credentials).map_err(|_| { + ProviderFetchFailure::terminal("Claude setup-token account identity could not be verified.") + }); + let outcome = request_after_verified_binding(verified, |binding| async move { + Ok(claude_header_snapshot( + &credentials, + Utc::now(), + Ok(binding.primary.clone()), + Some(binding), + ) + .await) + }) + .await + .unwrap_or_else(ProviderFetchOutcome::Failure); + ("setup-token", outcome) } -fn finalize_copilot_snapshot_with( - fetched: Result, - now: DateTime, - mut enrich: F, -) -> AgentUsageSnapshot +async fn fetch_claude_login_or_setup_with( + login: ClaudeLoginResolution, + request_login: Login, + load_setup: LoadSetup, + request_setup: Setup, +) -> (&'static str, ProviderFetchOutcome) where - F: FnMut(&mut AgentUsageSnapshot, DateTime), + Login: FnOnce(ClaudeCredentials) -> LoginFuture, + LoginFuture: std::future::Future, + LoadSetup: FnOnce() -> Result, String>, + Setup: FnOnce(ResolvedClaudeToken) -> SetupFuture, + SetupFuture: std::future::Future, { - let mut snapshot = match fetched { - Ok(data) => AgentUsageSnapshot { - client_id: "copilot".to_string(), - source: "oauth".to_string(), - updated_at: now.to_rfc3339_opts(SecondsFormat::Millis, true), - identity: data.identity, - account_scope: data.account_scope, - windows: data.windows, - credits: None, - error: None, + match login { + ClaudeLoginResolution::Ready(credentials) => request_login(credentials).await, + ClaudeLoginResolution::Terminal => ( + "oauth", + ProviderFetchOutcome::Failure(ProviderFetchFailure::terminal( + CLAUDE_CREDENTIALS_LOAD_ERROR, + )), + ), + ClaudeLoginResolution::Absent | ClaudeLoginResolution::ExplicitLogout => { + match load_setup() { + Ok(Some(token)) => request_setup(token).await, + Ok(None) => ( + "unconfigured", + ProviderFetchOutcome::Failure(ProviderFetchFailure::terminal( + CLAUDE_UNCONFIGURED_ERROR, + )), + ), + Err(_) => ( + "setup-token", + ProviderFetchOutcome::Failure(ProviderFetchFailure::terminal( + CLAUDE_CREDENTIALS_LOAD_ERROR, + )), + ), + } + } + } +} + +async fn fetch_claude_oauth_usage( + credentials: ClaudeCredentials, + pre_binding: ProviderCacheBinding, +) -> (&'static str, ProviderFetchOutcome) { + fetch_claude_login_usage_with( + credentials, + pre_binding, + Utc::now(), + |binding, now| { + let mut gate = lock_gate(); + gate.blocked_until_for(binding, now) }, - Err(error) => AgentUsageSnapshot { - client_id: "copilot".to_string(), - source: "oauth".to_string(), - updated_at: now.to_rfc3339_opts(SecondsFormat::Millis, true), - identity: None, - account_scope: Err(AccountScopeError::NoTrustedEvidence), - windows: Vec::new(), - credits: None, - error: Some(error), + |credentials| async move { refresh_claude_credentials(&credentials).await }, + |credentials, account_scope, cache_binding| async move { + claude_header_snapshot(&credentials, Utc::now(), Ok(account_scope), cache_binding).await }, - }; - enrich(&mut snapshot, now); - snapshot + |credentials, account_scope, cache_binding, gate_binding| async move { + fetch_claude_oauth_usage_request( + &credentials, + account_scope, + cache_binding, + gate_binding, + ) + .await + }, + ) + .await } -fn finalize_antigravity_snapshot_with( - fetched: Result, +async fn fetch_claude_login_usage_with< + Gate, + Refresh, + RefreshFuture, + Header, + HeaderFuture, + Oauth, + OauthFuture, +>( + credentials: ClaudeCredentials, + pre_binding: ProviderCacheBinding, now: DateTime, - mut enrich: F, -) -> AgentUsageSnapshot + blocked_until_for: Gate, + refresh: Refresh, + header: Header, + oauth: Oauth, +) -> (&'static str, ProviderFetchOutcome) where - F: FnMut(&mut AgentUsageSnapshot, DateTime), + Gate: FnOnce(&ProviderCacheBinding, DateTime) -> Option>, + Refresh: FnOnce(ClaudeCredentials) -> RefreshFuture, + RefreshFuture: std::future::Future< + Output = Result< + ( + ClaudeCredentials, + AccountScope, + Option, + ), + ProviderFetchFailure, + >, + >, + Header: FnOnce(ClaudeCredentials, AccountScope, Option) -> HeaderFuture, + HeaderFuture: std::future::Future, + Oauth: FnOnce( + ClaudeCredentials, + AccountScope, + Option, + ProviderCacheBinding, + ) -> OauthFuture, + OauthFuture: std::future::Future, { - let mut snapshot = match fetched { - Ok(fetched) => AgentUsageSnapshot { - client_id: "antigravity".to_string(), - source: fetched.source, - updated_at: now.to_rfc3339_opts(SecondsFormat::Millis, true), - identity: fetched.identity, - account_scope: fetched.account_scope, - windows: fetched.windows, - credits: None, - error: None, - }, - Err(error) => AgentUsageSnapshot { - client_id: "antigravity".to_string(), - source: "oauth".to_string(), - updated_at: now.to_rfc3339_opts(SecondsFormat::Millis, true), - identity: None, - account_scope: Err(AccountScopeError::NoTrustedEvidence), - windows: Vec::new(), - credits: None, - error: Some(error), - }, - }; - enrich(&mut snapshot, now); - snapshot -} + let header_route = !credentials.scopes.is_empty() + && !credentials + .scopes + .iter() + .any(|scope| scope == "user:profile"); + if !header_route { + if let Some(blocked_until) = blocked_until_for(&pre_binding, now) { + return ( + "oauth", + ProviderFetchOutcome::Failure(claude_gate_failure(pre_binding, blocked_until, now)), + ); + } + } -async fn fetch_antigravity() -> AgentUsageSnapshot { - let now = Utc::now(); - finalize_antigravity_snapshot_with(agent_antigravity::fetch(now).await, now, enrich_snapshot) -} + let (credentials, account_scope, cache_binding) = if claude_credentials_expired(&credentials) { + match refresh(credentials).await { + Ok(refreshed) => refreshed, + Err(failure) => return ("oauth", ProviderFetchOutcome::Failure(failure)), + } + } else { + let account_scope = pre_binding.primary.clone(); + (credentials, account_scope, Some(pre_binding)) + }; -async fn fetch_codex() -> AgentUsageSnapshot { - match fetch_codex_inner().await { - Ok(snapshot) => snapshot, - Err(error) => AgentUsageSnapshot { - client_id: "codex".to_string(), - source: "oauth".to_string(), - updated_at: Utc::now().to_rfc3339_opts(SecondsFormat::Millis, true), - identity: None, - account_scope: Err(AccountScopeError::NoTrustedEvidence), - windows: Vec::new(), - credits: None, - error: Some(error), - }, + if header_route { + return ( + "setup-token", + header(credentials, account_scope, cache_binding).await, + ); } -} -/// Claude's `/api/oauth/usage` rate-limits aggressively (and the budget is -/// shared with any other monitor on the account, e.g. codexbar). Modeled on -/// codexbar's ClaudeOAuthUsageRateLimitGate: after a 429, stop hitting the -/// endpoint until Retry-After (default 5 min) and serve the last good -/// snapshot so the card keeps its data instead of flashing an error. -struct ClaudeUsageGate { - blocked_until: Option>, - last_good: Option, + let gate_binding = ProviderCacheBinding::primary(account_scope.clone()); + oauth(credentials, account_scope, cache_binding, gate_binding).await } -static CLAUDE_USAGE_GATE: Mutex = Mutex::new(ClaudeUsageGate { - blocked_until: None, - last_good: None, -}); +async fn fetch_claude_oauth_usage_request( + credentials: &ClaudeCredentials, + account_scope: AccountScope, + cache_binding: Option, + gate_binding: ProviderCacheBinding, +) -> (&'static str, ProviderFetchOutcome) { + let client = match provider_http_client_builder() + .timeout(std::time::Duration::from_secs(30)) + .build() + { + Ok(client) => client, + Err(_) => { + return ( + "oauth", + ProviderFetchOutcome::Failure(ProviderFetchFailure::terminal( + "Claude usage client could not be created.", + )), + ); + } + }; -/// Lock the gate, recovering from a poisoned mutex instead of panicking. Under -/// the release profile's unwind + FFI-boundary `catch_unwind` (see `guarded` in -/// lib.rs), a panic caught mid-section poisons this static; `into_inner()` keeps -/// the 429 gate working for the rest of the process instead of wedging every -/// later `tb_agent_usage` call — same stance as the live-tail lock in lib.rs. -fn lock_gate() -> std::sync::MutexGuard<'static, ClaudeUsageGate> { - CLAUDE_USAGE_GATE - .lock() - .unwrap_or_else(|poisoned| poisoned.into_inner()) -} - -fn claude_gate_blocked_until(now: DateTime) -> Option> { - let mut gate = lock_gate(); - match gate.blocked_until { - Some(until) if until > now => Some(until), - Some(_) => { - gate.blocked_until = None; - None - } - None => None, - } -} - -fn claude_gate_record_rate_limit(retry_after: Option>, now: DateTime) { - let blocked_until = retry_after - .filter(|until| *until > now) - .unwrap_or_else(|| now + chrono::Duration::minutes(5)); - lock_gate().blocked_until = Some(blocked_until); -} - -fn claude_gate_record_success(snapshot: &AgentUsageSnapshot) { - let mut gate = lock_gate(); - gate.blocked_until = None; - gate.last_good = Some(snapshot.clone()); -} - -/// While the gate is closed, prefer the cached snapshot (its `updated_at` -/// stays honest); with nothing cached yet, surface a countdown error. -fn claude_gate_fallback(blocked_until: DateTime, now: DateTime) -> AgentUsageSnapshot { - if let Some(mut snapshot) = lock_gate().last_good.clone() { - // A cached 429 response is not current account-scope evidence. Keeping - // the stale scope would attribute a later poll to an unauthenticated account. - snapshot.account_scope = Err(AccountScopeError::NoTrustedEvidence); - return snapshot; - } - let wait_secs = (blocked_until - now).num_seconds().max(0); - AgentUsageSnapshot { - client_id: "claude".to_string(), - source: "oauth".to_string(), - updated_at: now.to_rfc3339_opts(SecondsFormat::Millis, true), - identity: None, - account_scope: Err(AccountScopeError::NoTrustedEvidence), - windows: Vec::new(), - credits: None, - error: Some(format!( - "Claude OAuth usage endpoint is rate limited. Retrying automatically in ~{}s.", - wait_secs - )), - } -} - -fn parse_retry_after(value: Option<&reqwest::header::HeaderValue>) -> Option> { - let raw = value?.to_str().ok()?.trim(); - if raw.is_empty() { - return None; - } - if let Ok(seconds) = raw.parse::() { - return (seconds >= 0).then(|| Utc::now() + chrono::Duration::seconds(seconds)); - } - DateTime::parse_from_rfc2822(raw) - .ok() - .map(|t| t.with_timezone(&Utc)) -} - -async fn fetch_claude() -> AgentUsageSnapshot { - let now = Utc::now(); - if let Some(blocked_until) = claude_gate_blocked_until(now) { - return claude_gate_fallback(blocked_until, now); - } - match fetch_claude_inner().await { - Ok(mut snapshot) => { - enrich_snapshot(&mut snapshot, now); - // Cache the display-ready snapshot. A later 429 fallback returns it - // without another enrichment pass or history write. - claude_gate_record_success(&snapshot); - snapshot - } - Err(error) => { - // A 429 inside fetch_claude_inner arms the gate; fall back to the - // cached, already-enriched snapshot rather than blanking the card. - let now = Utc::now(); - if let Some(blocked_until) = claude_gate_blocked_until(now) { - return claude_gate_fallback(blocked_until, now); - } - // "unconfigured" == no credential at all, so the UI shows a setup - // prompt; every other error is a real failure of a present credential. - let source = if error.as_str() == CLAUDE_UNCONFIGURED_ERROR { - "unconfigured" - } else { - "oauth" - }; - AgentUsageSnapshot { - client_id: "claude".to_string(), - source: source.to_string(), - updated_at: now.to_rfc3339_opts(SecondsFormat::Millis, true), - identity: None, - account_scope: Err(AccountScopeError::NoTrustedEvidence), - windows: Vec::new(), - credits: None, - error: Some(error), - } - } - } -} - -async fn fetch_codex_inner() -> Result { - let mut credentials = load_codex_credentials()?; - let mut refreshed_scope = None; - if credentials_needs_refresh(credentials.last_refresh) { - if credentials - .refresh_token - .as_deref() - .unwrap_or("") - .is_empty() - { - return Err( - "Codex OAuth token needs refresh but auth.json has no refresh token.".to_string(), - ); - } - let refreshed = refresh_codex_credentials(&credentials.auth_path).await?; - credentials = refreshed.0; - refreshed_scope = Some(refreshed.1); - } - - let client = reqwest::Client::builder() - .timeout(std::time::Duration::from_secs(30)) - .build() - .map_err(|e| format!("build Codex OAuth client: {}", e))?; - - let mut request = client - .get(CODEX_USAGE_URL) - .bearer_auth(&credentials.access_token) - .header(reqwest::header::ACCEPT, "application/json") - .header(reqwest::header::USER_AGENT, "TokenBar"); - let request_account_id = credentials - .account_id - .as_deref() - .map(str::trim) - .filter(|value| !value.is_empty()); - if let Some(account_id) = request_account_id { - request = request.header("ChatGPT-Account-Id", account_id); - } - - let response = request - .send() - .await - .map_err(|e| format!("Codex OAuth request failed: {}", e))?; - let status = response.status(); - let body = response - .text() - .await - .map_err(|e| format!("read Codex OAuth response: {}", e))?; - - if status == reqwest::StatusCode::UNAUTHORIZED || status == reqwest::StatusCode::FORBIDDEN { - return Err( - "Codex OAuth token expired or invalid. Run `codex` to log in again.".to_string(), - ); - } - if !status.is_success() { - return Err(format!("Codex usage API returned {}.", status.as_u16())); - } - - let usage: CodexUsageResponse = - serde_json::from_str(&body).map_err(|e| format!("decode Codex usage response: {}", e))?; - let now = Utc::now(); - let account_scope = resolve_codex_account_scope( - refreshed_scope, - request_account_id, - |account_id| { - agent_account_scope::resolve_authoritative( - "codex", - AuthoritativeIdKind::OpaqueId, - account_id, - ) - }, - || { - agent_account_scope::resolve_credential( - "codex", - credentials.scope_slot.semantic_source, - &credentials.scope_slot.canonical_location, - credentials.scope_marker(), - ) - }, - ); - let identity = Some(AgentIdentity { - email: credentials.id_token.as_deref().and_then(jwt_email), - plan: usage.plan_type.as_deref().map(clean_plan).or_else(|| { - credentials - .id_token - .as_deref() - .and_then(jwt_plan) - .map(clean_plan) - }), - }); - let windows = codex_windows( - usage.rate_limit.as_ref(), - usage.additional_rate_limits.as_deref(), - now, - ); - if windows.is_empty() && usage.credits.as_ref().and_then(|c| c.balance).is_none() { - return Err("Codex usage API returned no rate-limit windows.".to_string()); - } - - // Only the non-empty account ID actually attached to this successful request - // may authorize legacy migration. Migration remains best-effort so the live - // v3 observation still records if import fails. - maybe_migrate_codex_v2_with( - request_account_id, - &account_scope, - now.timestamp(), - crate::agent_quota_history::migrate_codex_v2, - ); - - let mut snapshot = AgentUsageSnapshot { - client_id: "codex".to_string(), - source: "oauth".to_string(), - updated_at: now.to_rfc3339_opts(SecondsFormat::Millis, true), - identity, - account_scope, - windows, - credits: usage.credits.map(|credits| CreditsSnapshot { - remaining: credits.balance, - unlimited: credits.unlimited, - }), - error: None, - }; - enrich_snapshot(&mut snapshot, now); - Ok(snapshot) -} - -fn resolve_codex_account_scope( - refreshed_scope: Option>, - request_account_id: Option<&str>, - resolve_authoritative: ResolveAuthoritative, - resolve_credential: ResolveCredential, -) -> Result -where - ResolveAuthoritative: FnOnce(&str) -> Result, - ResolveCredential: FnOnce() -> Result, -{ - if let Some(Err(error)) = refreshed_scope.as_ref() { - return Err(*error); - } - if let Some(account_id) = request_account_id { - return resolve_authoritative(account_id); - } - refreshed_scope.unwrap_or_else(resolve_credential) -} - -fn maybe_migrate_codex_v2_with( - request_account_id: Option<&str>, - account_scope: &Result, - now: i64, - migrate: F, -) where - F: FnOnce(&str, &str, i64) -> Result, -{ - let Some(request_account_id) = request_account_id - .map(str::trim) - .filter(|account_id| !account_id.is_empty()) - else { - return; - }; - let Ok(account_scope) = account_scope else { - return; - }; - let _ = migrate(request_account_id, account_scope.as_str(), now); -} - -async fn fetch_claude_inner() -> Result { - // Mirror Claude Code's auth precedence: CLAUDE_CODE_OAUTH_TOKEN (our env, or - // harvested from the user's ~/.zshrc) outranks a stored subscription /login, - // because Claude Code itself consumes that token first. So TokenBar reports - // the account Claude Code is actually spending against, read from the - // ratelimit headers. (This is why the harvest runs even for /login users.) - if let Some(token) = resolve_claude_code_oauth_token().await { - return claude_header_snapshot( - &claude_credentials_from_access_token(token), - Utc::now(), - None, - ) - .await; - } - - // A stored full login (TokenBar env override / Keychain / file) uses the - // richer oauth/usage endpoint. Any failure -- a login that can't refresh, or - // a credentials file that exists but can't be read (permissions / I/O) -- is - // deferred: we still try the tokenbar Keychain setup-token below, and surface - // the error only if that misses too. So a stale login / read error never - // strands a working setup-token, yet a genuine failure isn't masked by the - // generic "unconfigured" setup prompt. - let deferred_error: Option = match load_claude_login_credentials() { - Ok(Some(credentials)) => match fetch_claude_oauth_usage(credentials).await { - Ok(snapshot) => return Ok(snapshot), - Err(login_error) => Some(login_error), - }, - Ok(None) => None, - Err(read_error) => Some(read_error), - }; - - // Last resort: the tokenbar-claude-oauth-token Keychain item reads limits - // straight from the ratelimit headers (no oauth/usage GET, no 429 gate). - if let Some(token) = resolve_claude_keychain_token() { - return claude_header_snapshot( - &claude_credentials_from_access_token(token), - Utc::now(), - None, - ) - .await; - } - - Err(deferred_error.unwrap_or_else(|| CLAUDE_UNCONFIGURED_ERROR.to_string())) -} - -async fn fetch_claude_oauth_usage( - mut credentials: ClaudeCredentials, -) -> Result { - let mut refreshed_scope = None; - if claude_credentials_expired(&credentials) { - let refreshed = refresh_claude_credentials(&credentials).await?; - credentials = refreshed.0; - refreshed_scope = Some(refreshed.1); - } - - if !credentials.scopes.is_empty() - && !credentials - .scopes - .iter() - .any(|scope| scope == "user:profile") - { - // Inference-only token declared explicit non-user:profile scopes — skip - // the (guaranteed-403) oauth/usage GET and read limits from headers. - return claude_header_snapshot(&credentials, Utc::now(), refreshed_scope).await; - } - - let client = reqwest::Client::builder() - .timeout(std::time::Duration::from_secs(30)) - .build() - .map_err(|e| format!("build Claude OAuth client: {}", e))?; - - let response = client + let response = match client .get(CLAUDE_USAGE_URL) .bearer_auth(&credentials.access_token) .header(reqwest::header::ACCEPT, "application/json") @@ -1615,73 +1889,133 @@ async fn fetch_claude_oauth_usage( .header("anthropic-beta", "oauth-2025-04-20") .send() .await - .map_err(|e| format!("Claude OAuth request failed: {}", e))?; - let status = response.status(); - let retry_after = if status == reqwest::StatusCode::TOO_MANY_REQUESTS { - parse_retry_after(response.headers().get(reqwest::header::RETRY_AFTER)) - } else { - None + { + Ok(response) => response, + Err(error) => { + return ( + "oauth", + ProviderFetchOutcome::Failure(ProviderFetchFailure::from_send_error( + "Claude usage request failed. Retrying automatically.", + cache_binding, + &error, + )), + ); + } + }; + let status = response.status().as_u16(); + let retry_after = (status == 429) + .then(|| parse_retry_after(response.headers().get(reqwest::header::RETRY_AFTER))) + .flatten(); + let body = match read_response_body(status, true, || async { + response.text().await.map_err(|error| { + TransportErrorFacts::from_reqwest(&error, TransportPhase::ResponseBody) + }) + }) + .await + { + Ok(body) => body, + Err(ResponseReadFailure::Transient(diagnostic)) => { + if status == 429 { + lock_gate().record_rate_limit(gate_binding, retry_after, Utc::now()); + } + return ( + "oauth", + ProviderFetchOutcome::Failure(ProviderFetchFailure::transient( + "Claude usage request failed. Retrying automatically.", + cache_binding, + diagnostic, + )), + ); + } + Err(ResponseReadFailure::Terminal(401)) => { + return ( + "oauth", + ProviderFetchOutcome::Failure(ProviderFetchFailure::terminal( + "Claude OAuth token expired or invalid. Run `claude` to re-authenticate.", + )), + ); + } + Err(ResponseReadFailure::Terminal(403)) => { + return ( + "oauth", + ProviderFetchOutcome::Failure(ProviderFetchFailure::terminal( + "Claude OAuth usage was denied. Run `claude logout && claude login` to grant user:profile.", + )), + ); + } + Err(ResponseReadFailure::Terminal(status)) => { + return ( + "oauth", + ProviderFetchOutcome::Failure(ProviderFetchFailure::terminal(format!( + "Claude usage API rejected the request (status {status})." + ))), + ); + } }; - let body = response - .text() - .await - .map_err(|e| format!("read Claude OAuth response: {}", e))?; - if status == reqwest::StatusCode::UNAUTHORIZED { - return Err( - "Claude OAuth token expired or invalid. Run `claude` to re-authenticate.".to_string(), - ); - } - if status == reqwest::StatusCode::FORBIDDEN { - // oauth/usage requires user:profile. An inference-only token (e.g. - // `claude setup-token`) is denied *specifically* for that scope — fall - // back to the unified rate-limit headers, which it *is* allowed to read. - // Any other 403 keeps the actionable re-auth error (and skips the probe, - // so we don't spend an inference call on an unrelated denial). + if status == 403 { if body.contains("user:profile") { - return claude_header_snapshot(&credentials, Utc::now(), refreshed_scope).await; + return ( + "setup-token", + claude_header_snapshot(credentials, Utc::now(), Ok(account_scope), cache_binding) + .await, + ); } - return Err( - "Claude OAuth usage was denied. Run `claude logout && claude login` to grant user:profile." - .to_string(), - ); - } - if status == reqwest::StatusCode::TOO_MANY_REQUESTS { - claude_gate_record_rate_limit(retry_after, Utc::now()); - return Err( - "Claude OAuth usage endpoint is rate limited. Backing off automatically.".to_string(), + return ( + "oauth", + ProviderFetchOutcome::Failure(ProviderFetchFailure::terminal( + "Claude OAuth usage was denied. Run `claude logout && claude login` to grant user:profile.", + )), ); } - if !status.is_success() { - return Err(format!("Claude usage API returned {}.", status.as_u16())); - } - let usage: ClaudeUsageResponse = - serde_json::from_str(&body).map_err(|e| format!("decode Claude usage response: {}", e))?; + let usage: ClaudeUsageResponse = match serde_json::from_str(&body) { + Ok(usage) => usage, + Err(_) => { + return ( + "oauth", + ProviderFetchOutcome::Failure(ProviderFetchFailure::terminal( + "Claude usage response could not be decoded.", + )), + ); + } + }; let now = Utc::now(); let windows = claude_windows(&usage, now); if windows.is_empty() { - return Err("Claude usage API returned no rate-limit windows.".to_string()); + return ( + "oauth", + ProviderFetchOutcome::Failure(ProviderFetchFailure::terminal( + "Claude usage API returned no usable quota windows.", + )), + ); } - let account_scope = refreshed_scope.unwrap_or_else(|| credentials.resolve_account_scope()); + lock_gate().clear(); - Ok(AgentUsageSnapshot { - client_id: "claude".to_string(), - source: "oauth".to_string(), - updated_at: now.to_rfc3339_opts(SecondsFormat::Millis, true), - identity: Some(AgentIdentity { - email: None, - plan: first_non_empty([ - credentials.subscription_type.as_deref(), - credentials.rate_limit_tier.as_deref(), - ]) - .map(clean_plan), - }), - account_scope, - windows, - credits: claude_credits(usage.extra_usage.as_ref()), - error: None, - }) + ( + "oauth", + ProviderFetchOutcome::Success { + snapshot: AgentUsageSnapshot { + client_id: "claude".to_string(), + source: "oauth".to_string(), + updated_at: now.to_rfc3339_opts(SecondsFormat::Millis, true), + identity: Some(AgentIdentity { + email: None, + plan: first_non_empty([ + credentials.subscription_type.as_deref(), + credentials.rate_limit_tier.as_deref(), + ]) + .map(clean_plan), + }), + account_scope: Ok(account_scope), + windows, + credits: claude_credits(usage.extra_usage.as_ref()), + error: None, + transport_diagnostic: None, + }, + cache_binding, + }, + ) } /// Fallback for inference-only tokens (`claude setup-token`): the oauth/usage @@ -1707,9 +2041,10 @@ fn refresh_cached_windows(windows: &[UsageWindow], now: DateTime) -> Option let mut refreshed = Vec::with_capacity(windows.len()); for window in windows { let mut window = window.clone(); - if window.resets_at.is_some() { - normalized_reset_at(&mut window, now)?; - let reset = window.reset_at_evidence?; + if let Some(reset) = window.resets_at.as_deref().and_then(parse_datetime) { + if now >= reset { + return None; + } window.reset_text = Some(reset_text(reset, now)); } refreshed.push(window); @@ -1717,7 +2052,10 @@ fn refresh_cached_windows(windows: &[UsageWindow], now: DateTime) -> Option Some(refreshed) } -async fn fetch_claude_via_headers(access_token: &str) -> Result, String> { +async fn fetch_claude_via_headers( + access_token: &str, + attempt_binding: Option, +) -> Result, ProviderFetchFailure> { { let now = Utc::now(); let guard = CLAUDE_HEADER_CACHE @@ -1732,10 +2070,12 @@ async fn fetch_claude_via_headers(access_token: &str) -> Result } } - let client = reqwest::Client::builder() + let client = provider_http_client_builder() .timeout(std::time::Duration::from_secs(30)) .build() - .map_err(|e| format!("build Claude header-probe client: {}", e))?; + .map_err(|_| { + ProviderFetchFailure::terminal("Claude header-probe client could not be created.") + })?; let response = client .post(CLAUDE_MESSAGES_URL) @@ -1752,63 +2092,83 @@ async fn fetch_claude_via_headers(access_token: &str) -> Result })) .send() .await - .map_err(|e| format!("Claude header probe failed: {}", e))?; + .map_err(|error| { + ProviderFetchFailure::from_send_error( + "Claude header probe failed. Retrying automatically.", + attempt_binding.clone(), + &error, + ) + })?; - let status = response.status(); - // Read headers before consuming the body — this returns an owned Vec, ending - // the borrow of `response`. + let status = response.status().as_u16(); let windows = parse_unified_ratelimit_windows(response.headers(), Utc::now()); - - if status.is_success() || status == reqwest::StatusCode::TOO_MANY_REQUESTS { + if (200..=299).contains(&status) || status == 429 { if windows.is_empty() { - return Err("Claude header probe returned no unified rate-limit headers.".to_string()); - } - { - let mut guard = CLAUDE_HEADER_CACHE - .lock() - .unwrap_or_else(|e| e.into_inner()); - *guard = Some((Utc::now(), access_token.to_string(), windows.clone())); + if status == 429 { + return Err(ProviderFetchFailure::transient( + "Claude header probe is rate limited. Retrying automatically.", + attempt_binding, + SafeTransportDiagnostic::rate_limited(status), + )); + } + return Err(ProviderFetchFailure::terminal( + "Claude header probe returned no usable rate-limit headers.", + )); } + let mut guard = CLAUDE_HEADER_CACHE + .lock() + .unwrap_or_else(|e| e.into_inner()); + *guard = Some((Utc::now(), access_token.to_string(), windows.clone())); return Ok(windows); } - - let body = response.text().await.unwrap_or_default(); - Err(format!( - "Claude header probe returned {} ({}).", - status.as_u16(), - body.chars().take(200).collect::() + if (500..=599).contains(&status) { + return Err(ProviderFetchFailure::transient( + "Claude header probe failed. Retrying automatically.", + attempt_binding, + SafeTransportDiagnostic::server_error(status), + )); + } + Err(ProviderFetchFailure::terminal( + if matches!(status, 401 | 403) { + "Claude setup-token expired or lacks access.".to_string() + } else { + format!("Claude header probe rejected the request (status {status}).") + }, )) } -/// Build a Claude snapshot from the unified rate-limit headers. Shared by the -/// scope-guard and HTTP-403 branches of `fetch_claude_inner`. `source` is -/// `"setup-token"` — it doubles as the limits-card badge, so it names the auth -/// method the user recognizes rather than the fetch mechanism, and still lets -/// telemetry tell it apart from the richer oauth/usage path. async fn claude_header_snapshot( credentials: &ClaudeCredentials, now: DateTime, - account_scope: Option>, -) -> Result { - let windows = fetch_claude_via_headers(&credentials.access_token).await?; - let account_scope = account_scope.unwrap_or_else(|| credentials.resolve_account_scope()); - Ok(AgentUsageSnapshot { - client_id: "claude".to_string(), - source: "setup-token".to_string(), - updated_at: now.to_rfc3339_opts(SecondsFormat::Millis, true), - identity: Some(AgentIdentity { - email: None, - plan: first_non_empty([ - credentials.subscription_type.as_deref(), - credentials.rate_limit_tier.as_deref(), - ]) - .map(clean_plan), - }), - account_scope, - windows, - credits: None, - error: None, - }) + account_scope: Result, + cache_binding: Option, +) -> ProviderFetchOutcome { + let windows = + match fetch_claude_via_headers(&credentials.access_token, cache_binding.clone()).await { + Ok(windows) => windows, + Err(failure) => return ProviderFetchOutcome::Failure(failure), + }; + ProviderFetchOutcome::Success { + snapshot: AgentUsageSnapshot { + client_id: "claude".to_string(), + source: "setup-token".to_string(), + updated_at: now.to_rfc3339_opts(SecondsFormat::Millis, true), + identity: Some(AgentIdentity { + email: None, + plan: first_non_empty([ + credentials.subscription_type.as_deref(), + credentials.rate_limit_tier.as_deref(), + ]) + .map(clean_plan), + }), + account_scope, + windows, + credits: None, + error: None, + transport_diagnostic: None, + }, + cache_binding, + } } fn load_codex_credentials() -> Result { @@ -1869,45 +2229,72 @@ fn load_codex_credentials_from(auth_path: &Path) -> Result Result, String> { - if let Some(credentials) = load_claude_credentials_from_environment()? { - return Ok(Some(credentials)); - } - if let Some(raw) = load_claude_credentials_from_keychain()? { - if let Ok(credentials) = - parse_claude_credentials_data(&raw, ClaudeCredentialSource::Keychain) - { - return Ok(Some(credentials)); - } - } - match fs::read_to_string(claude_credentials_path()) { - Ok(raw) => { - if let Ok(credentials) = - parse_claude_credentials_data(&raw, ClaudeCredentialSource::File) - { - return Ok(Some(credentials)); - } - // Parsed but unusable (logged-out / no accessToken): fall through. - Ok(None) +/// TokenBar env override. Only a genuinely missing higher-priority store falls +/// through; the explicit #26 logout shape stops full-login precedence. +fn load_claude_login_credentials() -> ClaudeLoginResolution { + match load_claude_credentials_from_environment() { + Ok(Some(credentials)) => return ClaudeLoginResolution::Ready(credentials), + Ok(None) => {} + Err(_) => return ClaudeLoginResolution::Terminal, + } + load_stored_claude_login_with( + load_claude_credentials_from_keychain, + || match fs::read_to_string(claude_credentials_path()) { + Ok(raw) => Ok(Some(raw)), + Err(error) if error.kind() == std::io::ErrorKind::NotFound => Ok(None), + Err(_) => Err("Claude credentials file could not be read.".to_string()), + }, + ) +} + +fn load_stored_claude_login_with( + load_keychain: LoadKeychain, + load_file: LoadFile, +) -> ClaudeLoginResolution +where + LoadKeychain: FnOnce() -> Result, String>, + LoadFile: FnOnce() -> Result, String>, +{ + match load_keychain() { + Ok(Some(raw)) => { + return resolve_stored_claude_login(&raw, ClaudeCredentialSource::Keychain); } - // Absent is normal (no file login). A genuine read failure (permissions / - // I/O) is a real problem — return it so the caller can surface the - // actionable error after setup-token fallbacks miss, rather than the - // generic "unconfigured" setup prompt. - Err(error) if error.kind() == std::io::ErrorKind::NotFound => Ok(None), - Err(error) => Err(format!( - "read Claude credentials file {}: {}", - claude_credentials_path().display(), - error - )), + Ok(None) => {} + Err(_) => return ClaudeLoginResolution::Terminal, + } + match load_file() { + Ok(Some(raw)) => resolve_stored_claude_login(&raw, ClaudeCredentialSource::File), + Ok(None) => ClaudeLoginResolution::Absent, + Err(_) => ClaudeLoginResolution::Terminal, + } +} + +fn resolve_stored_claude_login(raw: &str, source: ClaudeCredentialSource) -> ClaudeLoginResolution { + let raw_root: Value = match serde_json::from_str(raw) { + Ok(root) => root, + Err(_) => return ClaudeLoginResolution::Terminal, + }; + let explicitly_logged_out = raw_root + .get("claudeAiOauth") + .and_then(Value::as_object) + .is_some_and(|oauth| { + oauth.contains_key("refreshToken") + && match oauth.get("accessToken") { + None | Some(Value::Null) => true, + Some(Value::String(token)) => token.trim().is_empty(), + _ => false, + } + }); + if explicitly_logged_out { + return ClaudeLoginResolution::ExplicitLogout; + } + match parse_claude_credentials_data(raw, source) { + Ok(credentials) => ClaudeLoginResolution::Ready(credentials), + Err(_) => ClaudeLoginResolution::Terminal, } } @@ -1939,17 +2326,16 @@ async fn resolve_claude_code_oauth_token() -> Option { /// The `tokenbar-claude-oauth-token` Keychain item (a TokenBar-specific setup /// token). A last-resort fallback, below the stored `/login`. -fn resolve_claude_keychain_token() -> Option { - load_claude_raw_token_from_keychain() - .ok() - .flatten() - .map(|access_token| ResolvedClaudeToken { +fn resolve_claude_keychain_token() -> Result, String> { + load_claude_raw_token_from_keychain().map(|token| { + token.map(|access_token| ResolvedClaudeToken { access_token, scope_slot: CredentialSlot { semantic_source: "claude-setup-keychain", canonical_location: CLAUDE_RAW_TOKEN_KEYCHAIN_SERVICE.to_string(), }, }) + }) } fn load_claude_credentials_from_environment() -> Result, String> { @@ -1987,6 +2373,7 @@ fn load_claude_credentials_from_environment() -> Result Result bool { + status.code() == Some(44) +} + fn load_claude_credentials_from_keychain() -> Result, String> { - let output = std::process::Command::new("/usr/bin/security") - .args(["find-generic-password", "-s", CLAUDE_KEYCHAIN_SERVICE, "-w"]) + load_claude_credentials_from_keychain_item(None) +} + +#[cfg(target_os = "macos")] +fn load_claude_credentials_from_keychain_item( + account: Option<&str>, +) -> Result, String> { + let mut command = std::process::Command::new("/usr/bin/security"); + command.args(["find-generic-password", "-s", CLAUDE_KEYCHAIN_SERVICE]); + if let Some(account) = account { + command.args(["-a", account]); + } + let output = command + .arg("-w") .output() .map_err(|e| format!("read Claude Keychain credentials: {}", e))?; if !output.status.success() { - return Ok(None); + return if keychain_item_not_found(&output.status) { + Ok(None) + } else { + Err("Claude Keychain credentials could not be read.".to_string()) + }; } let raw = String::from_utf8(output.stdout) .map_err(|_| "Claude Keychain credentials are not UTF-8 JSON.".to_string())?; let raw = raw.trim_matches(['\r', '\n']).to_string(); if raw.trim().is_empty() { - return Ok(None); + return Err("Claude Keychain credentials are empty.".to_string()); } Ok(Some(raw)) } #[cfg(not(target_os = "macos"))] -fn load_claude_credentials_from_keychain() -> Result, String> { +fn load_claude_credentials_from_keychain_item( + _account: Option<&str>, +) -> Result, String> { Ok(None) } /// Build credentials from a bare access token (no refresh/expiry/scope metadata). -/// Used by the setup-token delivery paths (env var, shell harvest, raw keychain); -/// empty scopes make `fetch_claude_inner` skip the scope guard and reach the -/// header fallback on the resulting oauth/usage 403. +/// Setup-token delivery paths use these credentials only for the header probe. fn claude_credentials_from_access_token(token: ResolvedClaudeToken) -> ClaudeCredentials { ClaudeCredentials { access_token: token.access_token, @@ -2088,6 +2496,7 @@ fn claude_credentials_from_access_token(token: ResolvedClaudeToken) -> ClaudeCre // to, so treat it as read-only — save_claude_credentials skips it. source: ClaudeCredentialSource::Environment, raw_root: None, + keychain_account: None, scope_slot: token.scope_slot, } } @@ -2247,13 +2656,17 @@ fn load_claude_raw_token_from_keychain() -> Result, String> { .output() .map_err(|e| format!("read TokenBar Claude token from Keychain: {}", e))?; if !output.status.success() { - return Ok(None); + return if keychain_item_not_found(&output.status) { + Ok(None) + } else { + Err("TokenBar Claude Keychain token could not be read.".to_string()) + }; } let raw = String::from_utf8(output.stdout) .map_err(|_| "TokenBar Claude Keychain token is not UTF-8.".to_string())?; let raw = raw.trim().to_string(); if raw.is_empty() { - return Ok(None); + return Err("TokenBar Claude Keychain token is empty.".to_string()); } Ok(Some(raw)) } @@ -2265,9 +2678,10 @@ fn load_claude_raw_token_from_keychain() -> Result, String> { async fn refresh_codex_credentials( auth_path: &Path, -) -> Result<(CodexCredentials, Result), String> { - let refresh = agent_account_scope::begin_refresh("codex") - .map_err(|_| "Codex credential refresh lock is unavailable.".to_string())?; +) -> Result<(CodexCredentials, ProviderCacheBinding), ProviderFetchFailure> { + let refresh = agent_account_scope::begin_refresh("codex").map_err(|_| { + ProviderFetchFailure::terminal("Codex credential refresh lock is unavailable.") + })?; refresh_codex_credentials_with( auth_path, &refresh, @@ -2278,11 +2692,16 @@ async fn refresh_codex_credentials( .await } -async fn request_codex_refresh(refresh_token: String) -> Result { - let client = reqwest::Client::builder() +async fn request_codex_refresh( + refresh_token: String, + attempt_binding: ProviderCacheBinding, +) -> Result { + let client = provider_http_client_builder() .timeout(std::time::Duration::from_secs(30)) .build() - .map_err(|e| format!("build Codex refresh client: {}", e))?; + .map_err(|_| { + ProviderFetchFailure::terminal("Codex refresh client could not be created.") + })?; let body = serde_json::json!({ "client_id": CODEX_CLIENT_ID, "grant_type": "refresh_token", @@ -2295,16 +2714,58 @@ async fn request_codex_refresh(refresh_token: String) -> Result { .json(&body) .send() .await - .map_err(|e| format!("Codex token refresh failed: {}", e))?; - let status = response.status(); - let body = response - .text() - .await - .map_err(|e| format!("read Codex refresh response: {}", e))?; - if !status.is_success() { - return Err("Codex OAuth refresh failed. Run `codex` to log in again.".to_string()); - } - serde_json::from_str(&body).map_err(|e| format!("decode Codex refresh response: {}", e)) + .map_err(|error| { + ProviderFetchFailure::from_send_error( + "Codex token refresh failed. Retrying automatically.", + Some(attempt_binding.clone()), + &error, + ) + })?; + let status = response.status().as_u16(); + let body = read_response_body(status, false, || async { + response.text().await.map_err(|error| { + TransportErrorFacts::from_reqwest(&error, TransportPhase::ResponseBody) + }) + }) + .await + .map_err(|failure| match failure { + ResponseReadFailure::Transient(diagnostic) => ProviderFetchFailure::transient( + "Codex token refresh failed. Retrying automatically.", + Some(attempt_binding), + diagnostic, + ), + ResponseReadFailure::Terminal(_) => ProviderFetchFailure::terminal( + "Codex OAuth refresh failed. Run `codex` to log in again.", + ), + })?; + serde_json::from_str(&body).map_err(|_| { + ProviderFetchFailure::terminal("Codex OAuth refresh response could not be decoded.") + }) +} + +fn resolve_codex_cache_binding_with( + credentials: &CodexCredentials, + refresh: &R, +) -> Result { + let primary = refresh.resolve_current( + credentials.scope_slot.semantic_source, + &credentials.scope_slot.canonical_location, + credentials.scope_marker(), + )?; + let corroborating = credentials + .account_id + .as_deref() + .map(str::trim) + .filter(|value| !value.is_empty()) + .map(|account_id| { + agent_account_scope::resolve_authoritative( + "codex", + AuthoritativeIdKind::OpaqueId, + account_id, + ) + }) + .transpose()?; + Ok(ProviderCacheBinding::new(primary, corroborating)) } async fn refresh_codex_credentials_with( @@ -2313,25 +2774,22 @@ async fn refresh_codex_credentials_with Result<(CodexCredentials, Result), String> +) -> Result<(CodexCredentials, ProviderCacheBinding), ProviderFetchFailure> where R: RefreshScopeTransaction + ?Sized, - Request: FnOnce(String) -> RequestFuture, - RequestFuture: std::future::Future>, - Save: FnOnce(&CodexCredentials) -> Result<(), String>, - Checkpoint: FnMut(RefreshCheckpoint) -> Result<(), String>, + Request: FnOnce(String, ProviderCacheBinding) -> RequestFuture, + RequestFuture: std::future::Future>, + Save: FnOnce(&CodexCredentials) -> Result, + Checkpoint: FnMut(RefreshCheckpoint) -> Result<(), ProviderFetchFailure>, { - // Another TokenBar process may have refreshed while this caller waited. - // Reload the request-bearing record only after the refresh lock is held. - let credentials = load_codex_credentials_from(auth_path)?; + let credentials = + load_codex_credentials_from(auth_path).map_err(ProviderFetchFailure::terminal)?; checkpoint(RefreshCheckpoint::Reloaded)?; + let pre_binding = resolve_codex_cache_binding_with(&credentials, refresh).map_err(|_| { + ProviderFetchFailure::terminal("Codex account identity could not be verified.") + })?; if !credentials_needs_refresh(credentials.last_refresh) { - let scope = refresh.resolve_current( - credentials.scope_slot.semantic_source, - &credentials.scope_slot.canonical_location, - credentials.scope_marker(), - ); - return Ok((credentials, scope)); + return Ok((credentials, pre_binding)); } let refresh_token = credentials @@ -2339,52 +2797,70 @@ where .as_deref() .map(str::trim) .filter(|token| !token.is_empty()) - .ok_or_else(|| "Codex auth.json has no refresh token.".to_string())? + .ok_or_else(|| ProviderFetchFailure::terminal("Codex auth.json has no refresh token."))? .to_string(); let old_marker = credentials.scope_marker().to_vec(); - let json = request(refresh_token).await?; + let json = request(refresh_token, pre_binding.clone()).await?; checkpoint(RefreshCheckpoint::NetworkReturned)?; - let response = json.as_object(); + let response = json.as_object().ok_or_else(|| { + ProviderFetchFailure::terminal("Codex OAuth refresh response was not a JSON object.") + })?; + let access_token = string_key(response, "access_token", "accessToken").ok_or_else(|| { + ProviderFetchFailure::terminal( + "Codex OAuth refresh response contained no usable access token.", + ) + })?; let refreshed = CodexCredentials { - access_token: response - .and_then(|tokens| string_key(tokens, "access_token", "accessToken")) - .unwrap_or(credentials.access_token), - refresh_token: response - .and_then(|tokens| string_key(tokens, "refresh_token", "refreshToken")) + access_token, + refresh_token: string_key(response, "refresh_token", "refreshToken") .or(credentials.refresh_token), - id_token: response - .and_then(|tokens| string_key(tokens, "id_token", "idToken")) - .or(credentials.id_token), + id_token: string_key(response, "id_token", "idToken").or(credentials.id_token), account_id: credentials.account_id, last_refresh: Some(Utc::now()), auth_path: credentials.auth_path, raw_json: credentials.raw_json, scope_slot: credentials.scope_slot, }; - let marker_rotated = refreshed.scope_marker() != old_marker.as_slice(); - let scope = refresh.transfer( + let write_receipt = save(&refreshed).map_err(|_| { + ProviderFetchFailure::terminal("Codex refreshed credentials could not be saved.") + })?; + checkpoint(RefreshCheckpoint::CredentialsPersisted)?; + let post_primary = match refresh.transfer( refreshed.scope_slot.semantic_source, &refreshed.scope_slot.canonical_location, &old_marker, refreshed.scope_marker(), - ); + ) { + Ok(account_scope) => account_scope, + Err(_) => { + let _ = rollback_codex_credentials_if_unchanged(&write_receipt); + return Err(ProviderFetchFailure::terminal( + "Codex credential lineage could not be preserved.", + )); + } + }; checkpoint(RefreshCheckpoint::MetadataHandled)?; - // A rotated marker may reach disk only after its lineage transfer is durable. - // The refreshed access token remains usable in memory for this poll. - if marker_rotated && scope.is_err() { - return Ok((refreshed, scope)); - } - save(&refreshed)?; - checkpoint(RefreshCheckpoint::CredentialsPersisted)?; - Ok((refreshed, scope)) + // The account ID is unchanged by refresh and was already verified in the + // pre-request binding. Reuse it instead of performing a second fallible + // metadata resolution after credentials and lineage are durable. + let post_binding = ProviderCacheBinding::new(post_primary, pre_binding.corroborating); + Ok((refreshed, post_binding)) } async fn refresh_claude_credentials( original: &ClaudeCredentials, -) -> Result<(ClaudeCredentials, Result), String> { - let refresh = agent_account_scope::begin_refresh("claude") - .map_err(|_| "Claude credential refresh lock is unavailable.".to_string())?; +) -> Result< + ( + ClaudeCredentials, + AccountScope, + Option, + ), + ProviderFetchFailure, +> { + let refresh = agent_account_scope::begin_refresh("claude").map_err(|_| { + ProviderFetchFailure::terminal("Claude credential refresh lock is unavailable.") + })?; refresh_claude_credentials_with( original, &refresh, @@ -2396,11 +2872,16 @@ async fn refresh_claude_credentials( .await } -async fn request_claude_refresh(refresh_token: String) -> Result { - let client = reqwest::Client::builder() +async fn request_claude_refresh( + refresh_token: String, + attempt_binding: ProviderCacheBinding, +) -> Result { + let client = provider_http_client_builder() .timeout(std::time::Duration::from_secs(30)) .build() - .map_err(|e| format!("build Claude refresh client: {}", e))?; + .map_err(|_| { + ProviderFetchFailure::terminal("Claude refresh client could not be created.") + })?; let response = client .post(CLAUDE_REFRESH_URL) .header(reqwest::header::ACCEPT, "application/json") @@ -2415,16 +2896,33 @@ async fn request_claude_refresh(refresh_token: String) -> Result ProviderFetchFailure::transient( + "Claude OAuth refresh failed. Retrying automatically.", + Some(attempt_binding), + diagnostic, + ), + ResponseReadFailure::Terminal(_) => ProviderFetchFailure::terminal( + "Claude OAuth refresh failed. Run `claude` to re-authenticate.", + ), + })?; + serde_json::from_str(&body).map_err(|_| { + ProviderFetchFailure::terminal("Claude OAuth refresh response could not be decoded.") + }) } async fn refresh_claude_credentials_with( @@ -2434,27 +2932,40 @@ async fn refresh_claude_credentials_with Result<(ClaudeCredentials, Result), String> +) -> Result< + ( + ClaudeCredentials, + AccountScope, + Option, + ), + ProviderFetchFailure, +> where R: RefreshScopeTransaction + ?Sized, Reload: FnOnce(&ClaudeCredentials) -> Result, - Request: FnOnce(String) -> RequestFuture, - RequestFuture: std::future::Future>, + Request: FnOnce(String, ProviderCacheBinding) -> RequestFuture, + RequestFuture: + std::future::Future>, Save: FnOnce(&ClaudeCredentials) -> Result<(), String>, - Checkpoint: FnMut(RefreshCheckpoint) -> Result<(), String>, + Checkpoint: FnMut(RefreshCheckpoint) -> Result<(), ProviderFetchFailure>, { - let credentials = reload(original)?; + let credentials = reload(original).map_err(ProviderFetchFailure::terminal)?; checkpoint(RefreshCheckpoint::Reloaded)?; + let marker = credentials.scope_marker().ok_or_else(|| { + ProviderFetchFailure::terminal("Claude credential has no trusted account marker.") + })?; + let pre_scope = refresh + .resolve_current( + credentials.scope_slot.semantic_source, + &credentials.scope_slot.canonical_location, + marker, + ) + .map_err(|_| { + ProviderFetchFailure::terminal("Claude account identity could not be verified.") + })?; + let pre_binding = ProviderCacheBinding::primary(pre_scope.clone()); if !claude_credentials_expired(&credentials) { - let scope = match credentials.scope_marker() { - Some(marker) => refresh.resolve_current( - credentials.scope_slot.semantic_source, - &credentials.scope_slot.canonical_location, - marker, - ), - None => Err(AccountScopeError::NoTrustedEvidence), - }; - return Ok((credentials, scope)); + return Ok((credentials, pre_scope, Some(pre_binding))); } let refresh_token = credentials @@ -2462,14 +2973,22 @@ where .as_deref() .filter(|token| !token.is_empty()) .ok_or_else(|| { - "Claude OAuth token is expired and has no refresh token. Run `claude`.".to_string() + ProviderFetchFailure::terminal( + "Claude OAuth token is expired and has no refresh token. Run `claude`.", + ) })? .to_string(); let old_marker = refresh_token.as_bytes().to_vec(); - let token_response = request(refresh_token).await?; + let token_response = request(refresh_token, pre_binding).await?; checkpoint(RefreshCheckpoint::NetworkReturned)?; + let access_token = token_response.access_token.trim(); + if access_token.is_empty() { + return Err(ProviderFetchFailure::terminal( + "Claude OAuth refresh response has no access token.", + )); + } let refreshed = ClaudeCredentials { - access_token: token_response.access_token, + access_token: access_token.to_string(), refresh_token: token_response .refresh_token .as_deref() @@ -2483,40 +3002,59 @@ where subscription_type: credentials.subscription_type.clone(), source: credentials.source, raw_root: credentials.raw_root.clone(), + keychain_account: credentials.keychain_account.clone(), scope_slot: credentials.scope_slot.clone(), }; - let new_marker = refreshed.scope_marker(); - let marker_rotated = new_marker.is_some_and(|marker| marker != old_marker.as_slice()); - let scope = match new_marker { - Some(new_marker) => refresh.transfer( + let new_marker = refreshed.scope_marker().ok_or_else(|| { + ProviderFetchFailure::terminal("Claude refreshed credential has no trusted marker.") + })?; + let scope = refresh + .transfer( refreshed.scope_slot.semantic_source, &refreshed.scope_slot.canonical_location, &old_marker, new_marker, - ), - None => Err(AccountScopeError::NoTrustedEvidence), - }; + ) + .map_err(|_| { + ProviderFetchFailure::terminal("Claude credential lineage could not be preserved.") + })?; checkpoint(RefreshCheckpoint::MetadataHandled)?; - // A rotated marker may reach the shared provider store only after its - // lineage transfer is durable. The new access token remains usable in - // memory for this poll. - if marker_rotated && scope.is_err() { - return Ok((refreshed, scope)); - } - if let Err(error) = save(&refreshed) { - eprintln!("tb_core_ffi: failed to persist refreshed Claude credentials: {error}"); - } + let persisted = save(&refreshed).is_ok(); checkpoint(RefreshCheckpoint::CredentialsPersisted)?; - Ok((refreshed, scope)) + let cache_binding = if persisted { + Some(ProviderCacheBinding::primary( + refresh + .resolve_current( + refreshed.scope_slot.semantic_source, + &refreshed.scope_slot.canonical_location, + new_marker, + ) + .map_err(|_| { + ProviderFetchFailure::terminal( + "Claude account identity could not be verified after refresh.", + ) + })?, + )) + } else { + None + }; + Ok((refreshed, scope, cache_binding)) } fn reload_claude_credentials(original: &ClaudeCredentials) -> Result { match original.source { ClaudeCredentialSource::Keychain => { - let raw = load_claude_credentials_from_keychain()?.ok_or_else(|| { - "Claude Keychain credentials disappeared during refresh.".to_string() + let account = claude_keychain_account().ok_or_else(|| { + "Claude Keychain account could not be captured during refresh.".to_string() })?; - parse_claude_credentials_data(&raw, ClaudeCredentialSource::Keychain) + let raw = + load_claude_credentials_from_keychain_item(Some(&account))?.ok_or_else(|| { + "Claude Keychain credentials disappeared during refresh.".to_string() + })?; + let mut credentials = + parse_claude_credentials_data(&raw, ClaudeCredentialSource::Keychain)?; + credentials.keychain_account = Some(account); + Ok(credentials) } ClaudeCredentialSource::File => { let raw = fs::read_to_string(claude_credentials_path()) @@ -2533,9 +3071,7 @@ fn reload_claude_credentials(original: &ClaudeCredentials) -> Result Result<(), String> { match credentials.source { - ClaudeCredentialSource::Keychain => { - save_claude_credentials_to_keychain(&merge_claude_credentials_json(credentials)?) - } + ClaudeCredentialSource::Keychain => save_claude_credentials_to_keychain(credentials), ClaudeCredentialSource::File => { save_claude_credentials_to_file(credentials, &claude_credentials_path()) } @@ -2547,7 +3083,10 @@ fn save_claude_credentials_to_file( credentials: &ClaudeCredentials, path: &Path, ) -> Result<(), String> { - atomic_write(path, &merge_claude_credentials_json(credentials)?) + let current_raw = fs::read_to_string(path) + .map_err(|e| format!("read current Claude credentials file: {e}"))?; + let data = merge_claude_credentials_json(credentials, ¤t_raw)?; + atomic_write(path, &data) } /// Replace `path` atomically: write a sibling temp file, then rename over the @@ -2575,7 +3114,7 @@ fn atomic_write(path: &Path, data: &str) -> Result<(), String> { let seq = TMP_SEQ.fetch_add(1, std::sync::atomic::Ordering::Relaxed); let tmp = parent.join(format!(".{}.tmp.{}.{}", file_name, std::process::id(), seq)); - // Stage into the temp, fsync it, then rename over the target. Create with + // Stage into the temp, fsync it, then atomically replace the target. Create with // O_EXCL + 0600 up front: the mode-at-creation closes the umask-default // window a write-then-chmod leaves the secret readable in, and O_EXCL // refuses to follow a symlink pre-seeded at the temp path. @@ -2605,7 +3144,7 @@ fn atomic_write(path: &Path, data: &str) -> Result<(), String> { let _ = fs::remove_file(&tmp); return Err(error); } - if let Err(error) = fs::rename(&tmp, path) { + if let Err(error) = tokscale_core::fs_atomic::replace_file(&tmp, path) { let _ = fs::remove_file(&tmp); return Err(format!("replace {}: {}", path.display(), error)); } @@ -2617,17 +3156,36 @@ fn atomic_write(path: &Path, data: &str) -> Result<(), String> { Ok(()) } -/// Merge the rotated tokens into the loaded credentials JSON, preserving any -/// other fields, and return it serialized. Pure so it's unit-testable. -fn merge_claude_credentials_json(credentials: &ClaudeCredentials) -> Result { - let mut root = credentials +/// Merge rotated tokens into the current credentials JSON only when the +/// `claudeAiOauth` object still matches the one captured at refresh reload. +/// Top-level siblings come from `current_raw`, so unrelated concurrent writes +/// survive. Pure so both file and Keychain decisions are fixture-testable. +fn merge_claude_credentials_json( + credentials: &ClaudeCredentials, + current_raw: &str, +) -> Result { + let expected_oauth = credentials .raw_root - .clone() - .unwrap_or_else(|| serde_json::json!({ "claudeAiOauth": {} })); - let oauth = root + .as_ref() + .and_then(|root| root.get("claudeAiOauth")) + .and_then(Value::as_object) + .ok_or_else(|| "Reloaded Claude credentials have no claudeAiOauth object.".to_string())?; + let mut current_root: Value = serde_json::from_str(current_raw) + .map_err(|e| format!("decode current Claude credentials: {e}"))?; + let current_oauth = current_root + .get("claudeAiOauth") + .and_then(Value::as_object) + .ok_or_else(|| "Current Claude credentials have no claudeAiOauth object.".to_string())?; + if current_oauth != expected_oauth { + return Err( + "Claude credentials changed during refresh; refusing stale write-back.".to_string(), + ); + } + + let oauth = current_root .get_mut("claudeAiOauth") .and_then(Value::as_object_mut) - .ok_or_else(|| "Claude credentials JSON has no claudeAiOauth object.".to_string())?; + .ok_or_else(|| "Current Claude credentials have no claudeAiOauth object.".to_string())?; oauth.insert( "accessToken".to_string(), Value::String(credentials.access_token.clone()), @@ -2641,27 +3199,44 @@ fn merge_claude_credentials_json(credentials: &ClaudeCredentials) -> Result( + credentials: &'a ClaudeCredentials, + current_account: Option<&str>, + current_raw: &str, +) -> Result<(&'a str, String), String> { + let captured_account = credentials.keychain_account.as_deref().ok_or_else(|| { + "Claude Keychain refresh has no captured account; refusing write-back.".to_string() + })?; + if current_account != Some(captured_account) { + return Err( + "Claude Keychain account changed during refresh; refusing write-back.".to_string(), + ); + } + let data = merge_claude_credentials_json(credentials, current_raw)?; + Ok((captured_account, data)) } #[cfg(target_os = "macos")] -fn save_claude_credentials_to_keychain(data: &str) -> Result<(), String> { - // Fail closed: only update the item once we can confirm the exact account - // the Claude CLI stored it under. `add-generic-password -U` matches on - // (service, account), so updating with the wrong or an empty account would - // create a SECOND "Claude Code-credentials" item and confuse the store the - // CLI shares — worse than not persisting. If the account can't be read, - // skip the write-back (the caller logs it); the next refresh retries. - let account = claude_keychain_account().ok_or_else(|| { - "could not resolve the Claude Keychain account; skipping write-back to avoid a duplicate item" - .to_string() +fn save_claude_credentials_to_keychain(credentials: &ClaudeCredentials) -> Result<(), String> { + let captured_account = credentials.keychain_account.as_deref().ok_or_else(|| { + "Claude Keychain refresh has no captured account; refusing write-back.".to_string() })?; - // NOTE: `-w ` puts the credential JSON on the argv, briefly visible via - // `ps` to same-user processes. security(1) has no stdin form for - // add-generic-password (only an interactive `-w` prompt, unusable from a - // background app) and the item is already same-user-readable once the - // keychain is unlocked, so on a single-user Mac this narrow window is an - // accepted trade-off; move to the SecItem API if that assumption changes. + let current_raw = load_claude_credentials_from_keychain_item(Some(captured_account))? + .ok_or_else(|| "Claude Keychain credentials disappeared during refresh.".to_string())?; + let current_account = claude_keychain_account(); + let (account, data) = + prepare_claude_keychain_write(credentials, current_account.as_deref(), ¤t_raw)?; + + // NOTE: security(1) has no compare-and-swap operation. The exact-item read, + // account guard, and target comparison close the network-wait race and the + // write always stays pinned to the captured account. A same-item mutation + // between this check and `-U` remains the existing CLI platform limitation. + // `-w ` also puts the JSON on argv briefly; the item is already + // same-user-readable while the Keychain is unlocked. Move to SecItem only if + // either platform assumption changes. let status = std::process::Command::new("/usr/bin/security") .args([ "add-generic-password", @@ -2669,9 +3244,9 @@ fn save_claude_credentials_to_keychain(data: &str) -> Result<(), String> { "-s", CLAUDE_KEYCHAIN_SERVICE, "-a", - &account, + account, "-w", - data, + &data, ]) .status() .map_err(|e| format!("write Claude Keychain credentials: {}", e))?; @@ -2682,7 +3257,7 @@ fn save_claude_credentials_to_keychain(data: &str) -> Result<(), String> { } #[cfg(not(target_os = "macos"))] -fn save_claude_credentials_to_keychain(_data: &str) -> Result<(), String> { +fn save_claude_credentials_to_keychain(_credentials: &ClaudeCredentials) -> Result<(), String> { Err("Keychain writes are only supported on macOS.".to_string()) } @@ -2722,136 +3297,432 @@ fn claude_keychain_account() -> Option { None } -fn save_codex_credentials(credentials: &CodexCredentials) -> Result<(), String> { - let mut raw = credentials.raw_json.clone(); - raw["tokens"]["access_token"] = Value::String(credentials.access_token.clone()); +#[cfg(not(target_os = "macos"))] +fn claude_keychain_account() -> Option { + None +} + +fn save_codex_credentials( + credentials: &CodexCredentials, +) -> Result { + let expected_tokens = credentials + .raw_json + .get("tokens") + .ok_or_else(|| "Codex tokens missing from the loaded credentials.".to_string())?; + let mut raw = load_codex_credentials_from(&credentials.auth_path) + .map_err(|e| format!("reload Codex auth.json before saving: {}", e))? + .raw_json; + let current_tokens = raw + .get("tokens") + .ok_or_else(|| "Codex tokens disappeared before saving.".to_string())?; + if current_tokens != expected_tokens { + return Err("Codex tokens changed during refresh.".to_string()); + } + let previous_root = raw.clone(); + let tokens = raw + .get_mut("tokens") + .and_then(Value::as_object_mut) + .ok_or_else(|| "Codex tokens are not an object while saving.".to_string())?; + + tokens.insert( + "access_token".to_string(), + Value::String(credentials.access_token.clone()), + ); if let Some(refresh_token) = &credentials.refresh_token { - raw["tokens"]["refresh_token"] = Value::String(refresh_token.clone()); + tokens.insert( + "refresh_token".to_string(), + Value::String(refresh_token.clone()), + ); } if let Some(id_token) = &credentials.id_token { - raw["tokens"]["id_token"] = Value::String(id_token.clone()); + tokens.insert("id_token".to_string(), Value::String(id_token.clone())); } if let Some(account_id) = &credentials.account_id { - raw["tokens"]["account_id"] = Value::String(account_id.clone()); + tokens.insert("account_id".to_string(), Value::String(account_id.clone())); } raw["last_refresh"] = Value::String(Utc::now().to_rfc3339_opts(SecondsFormat::Millis, true)); let data = - serde_json::to_vec_pretty(&raw).map_err(|e| format!("encode Codex auth.json: {}", e))?; - fs::write(&credentials.auth_path, data).map_err(|e| format!("save Codex auth.json: {}", e)) + serde_json::to_string_pretty(&raw).map_err(|e| format!("encode Codex auth.json: {}", e))?; + atomic_write(&credentials.auth_path, &data) + .map_err(|e| format!("save Codex auth.json: {}", e))?; + Ok(CodexCredentialWriteReceipt { + path: credentials.auth_path.clone(), + previous_root, + persisted_root: raw, + }) } -fn codex_windows( - rate_limit: Option<&CodexRateLimit>, - additional_rate_limits: Option<&[CodexAdditionalRateLimit]>, - now: DateTime, -) -> Vec { - let mut windows = Vec::new(); - if let Some(rate_limit) = rate_limit { - let mut main = [ - ("primary", rate_limit.primary_window.clone()), - ("secondary", rate_limit.secondary_window.clone()), - ]; - main.sort_by_key(|(_, window)| { - match window.as_ref().map(|window| window.limit_window_seconds) { - Some(18_000) => 0, - Some(604_800) => 1, - _ => 2, - } - }); - for (slot, window) in main - .into_iter() - .filter_map(|(slot, window)| window.map(|window| (slot, window))) - { - let (label, card_id, window_key) = match window.limit_window_seconds { - 18_000 => ( - "Session", - "main.session.v1".to_string(), - Some("main.session.v1".to_string()), - ), - 604_800 => ( - "Weekly", - "main.weekly.v1".to_string(), - Some("main.weekly.v1".to_string()), - ), - _ => ("Unknown", format!("row.main.{slot}.v1"), None), - }; - if let Some(window) = map_window_with_identity(label, window, now, card_id, window_key) - { - windows.push(window); +/// Restore the pre-refresh Codex root only while this refresh still owns the +/// exact root it persisted. External Codex writers do not share TokenBar's +/// refresh lock, so the compare-to-rename interval remains a known residual +/// window rather than a filesystem compare-and-swap. +fn rollback_codex_credentials_if_unchanged( + receipt: &CodexCredentialWriteReceipt, +) -> Result { + let current_raw = fs::read_to_string(&receipt.path) + .map_err(|e| format!("read Codex auth.json before rollback: {}", e))?; + let current_root: Value = serde_json::from_str(¤t_raw) + .map_err(|e| format!("decode Codex auth.json before rollback: {}", e))?; + if current_root != receipt.persisted_root { + return Ok(false); + } + let previous_data = serde_json::to_string_pretty(&receipt.previous_root) + .map_err(|e| format!("encode Codex auth.json rollback: {}", e))?; + atomic_write(&receipt.path, &previous_data) + .map_err(|e| format!("rollback Codex auth.json: {}", e))?; + Ok(true) +} + +fn enrich_snapshot(snapshot: &mut AgentUsageSnapshot, now: i64) { + enrich_snapshot_with(snapshot, now, |active_keys, observations, now| { + crate::agent_quota_history::record_observations_and_evaluate(active_keys, observations, now) + }); +} + +fn enrich_snapshot_with(snapshot: &mut AgentUsageSnapshot, now: i64, mut record: F) +where + F: FnMut( + &[SeriesKey], + &[QuotaObservation], + i64, + ) -> Result, HistoryError>, +{ + let mut card_ids = HashSet::new(); + let mut window_keys = HashSet::new(); + snapshot.windows.retain(|window| { + let card_is_unique = !card_ids.contains(&window.card_id); + let key_is_unique = window + .window_key + .as_ref() + .is_none_or(|window_key| !window_keys.contains(window_key)); + if !card_is_unique || !key_is_unique { + return false; + } + card_ids.insert(window.card_id.clone()); + if let Some(window_key) = window.window_key.as_ref() { + window_keys.insert(window_key.clone()); + } + true + }); + + let Ok(account_scope) = snapshot.account_scope.as_ref() else { + for window in &mut snapshot.windows { + if window.window_key.is_some() { + window.unavailable("accountScope"); } } - } + return; + }; + let account_scope = account_scope.as_str(); + let mut active_keys = Vec::new(); + let mut observations = Vec::new(); + let mut mapped_indices = Vec::new(); - for extra in additional_rate_limits.unwrap_or(&[]) { - let Some(rate_limit) = extra.rate_limit.as_ref() else { + for (index, window) in snapshot.windows.iter_mut().enumerate() { + let Some(window_key) = window.window_key.as_deref() else { + // The provider already classified this card as windowIdentity. continue; }; - let (slot, window) = match ( - rate_limit.primary_window.clone(), - rate_limit.secondary_window.clone(), - ) { - (Some(window), _) => ("primary", window), - (None, Some(window)) => ("secondary", window), - (None, None) => continue, - }; - let source = additional_limit_source(extra); - let (label, card_id, window_key) = match source { - Some(source) => { - let key = format!("additional.{}.{slot}.v1", sha256_hex(&source)); - (additional_limit_label(extra), key.clone(), Some(key)) - } - None => ( - "Unknown".to_string(), - format!("row.additional.unknown.{slot}.v1"), - None, - ), + let key = SeriesKey::new(snapshot.client_id.clone(), account_scope, window_key); + active_keys.push(key.clone()); + if matches!(window.pace_status.state, PaceState::Unavailable) { + // Emission protects existing history from capacity eviction, but + // missing reset and other typed early rejects never record a sample. + continue; + } + let Some(reset_at) = window + .resets_at + .as_deref() + .and_then(parse_datetime) + .map(|reset| reset.timestamp()) + else { + window.unavailable("invalidEvidence"); + continue; }; - if let Some(window) = map_window_with_identity(&label, window, now, card_id, window_key) { - windows.push(window); + if reset_at <= now + || !window.used_percent.is_finite() + || !(0.0..=100.0).contains(&window.used_percent) + { + window.unavailable("invalidEvidence"); + continue; } + observations.push(QuotaObservation { + key, + reset_at: Some(reset_at), + used_percent: window.used_percent, + provider: window.provider_duration, + contract: window.contract_duration, + }); + mapped_indices.push(index); } - windows -} -fn claude_windows(usage: &ClaudeUsageResponse, now: DateTime) -> Vec { - let mut windows = Vec::new(); - push_claude_window( - &mut windows, - "Session", - "session.v1", - DurationEvidence::contract(18_000), - usage.five_hour.as_ref(), - now, - ); - push_claude_window( - &mut windows, - "Weekly", - "weekly.v1", - DurationEvidence::contract(604_800), - usage.seven_day.as_ref(), - now, - ); - push_claude_window( - &mut windows, - "OAuth Apps", - "oauth_apps.weekly.v1", - DurationEvidence::contract(604_800), - usage.seven_day_oauth_apps.as_ref(), - now, - ); - push_claude_window( - &mut windows, - "Sonnet", - "sonnet.weekly.v1", - DurationEvidence::contract(604_800), - usage.seven_day_sonnet.as_ref(), - now, + if active_keys.is_empty() { + return; + } + + let results = match record(&active_keys, &observations, now) { + Ok(results) if results.len() == mapped_indices.len() => results, + Ok(_) => { + for index in mapped_indices { + snapshot.windows[index].unavailable("history"); + } + return; + } + Err(error) => { + let reason = if error == HistoryError::StoreCapacity { + "storeCapacity" + } else { + "history" + }; + for index in mapped_indices { + snapshot.windows[index].unavailable(reason); + } + return; + } + }; + + for (index, result) in mapped_indices.into_iter().zip(results) { + let window = &mut snapshot.windows[index]; + match result { + Ok(( + HistoryOutcome::Ready { + duration_seconds, + source, + .. + }, + historical, + complete_cycles, + )) => { + window.duration_seconds = Some(duration_seconds); + window.duration_source = Some(source); + window.window_minutes = Some(duration_seconds / 60); + match historical { + Some(pace) if historical_pace_is_coherent(&pace) => { + window.pace_status = PaceStatusPayload { + state: PaceState::Available, + window_key: window.window_key.clone(), + duration_seconds: Some(duration_seconds), + duration_source: Some(source), + complete_cycles, + reason: None, + }; + window.historical_pace = Some(historical_pace_payload(pace)); + } + Some(_) => { + window.unavailable("history"); + } + None => { + window.pace_status = PaceStatusPayload { + state: PaceState::LearningHistory, + window_key: window.window_key.clone(), + duration_seconds: Some(duration_seconds), + duration_source: Some(source), + complete_cycles, + reason: None, + }; + window.historical_pace = None; + } + } + } + Ok((HistoryOutcome::LearningDuration, None, _)) => { + window.duration_seconds = None; + window.duration_source = Some(DurationSource::Observed); + window.window_minutes = None; + window.pace_status = PaceStatusPayload { + state: PaceState::LearningDuration, + window_key: window.window_key.clone(), + duration_seconds: None, + duration_source: Some(DurationSource::Observed), + complete_cycles: 0, + reason: None, + }; + window.historical_pace = None; + } + Ok((HistoryOutcome::Unavailable(reason), _, _)) => { + window.unavailable(duration_unavailable_reason(reason)); + } + Err(error) => { + window.unavailable(if error == HistoryError::StoreCapacity { + "storeCapacity" + } else { + "history" + }); + } + Ok((HistoryOutcome::LearningDuration, Some(_), _)) => { + window.unavailable("history"); + } + } + } +} + +fn duration_unavailable_reason(reason: DurationUnavailableReason) -> &'static str { + match reason { + DurationUnavailableReason::MissingReset => "missingReset", + DurationUnavailableReason::InvalidEvidence => "invalidEvidence", + } +} + +fn historical_pace_is_coherent(pace: &HistoricalPace) -> bool { + pace.expected_percent.is_finite() + && (0.0..=100.0).contains(&pace.expected_percent) + && pace + .eta_seconds + .is_none_or(|eta| eta.is_finite() && eta >= 0.0) + && pace + .run_out_probability + .is_none_or(|probability| probability.is_finite() && (0.0..=1.0).contains(&probability)) + && (pace.eta_seconds.is_none() == pace.will_last_to_reset) +} + +fn historical_pace_payload(pace: HistoricalPace) -> HistoricalPacePayload { + HistoricalPacePayload { + expected_used_percent: pace.expected_percent, + eta_seconds: pace.eta_seconds, + will_last_to_reset: pace.will_last_to_reset, + run_out_probability: pace.run_out_probability, + } +} + +fn codex_windows( + rate_limit: Option<&CodexRateLimit>, + additional_rate_limits: Option<&[CodexAdditionalRateLimit]>, + now: DateTime, +) -> Vec { + let mut windows = Vec::new(); + let mut emitted_card_ids = HashSet::new(); + if let Some(rate_limit) = rate_limit { + let mut main = [ + ("primary", rate_limit.primary_window.clone()), + ("secondary", rate_limit.secondary_window.clone()), + ]; + main.sort_by_key(|(_, window)| { + window + .as_ref() + .map_or(2, |window| match window.limit_window_seconds { + 18_000 => 0, + 604_800 => 1, + _ => 2, + }) + }); + for (slot, window) in main + .into_iter() + .filter_map(|(slot, window)| window.map(|window| (slot, window))) + { + let semantic = match window.limit_window_seconds { + 18_000 => Some(("Session", "main.session.v1")), + 604_800 => Some(("Weekly", "main.weekly.v1")), + _ => None, + }; + let (label, window_key) = semantic.unwrap_or(("Unknown", "")); + let card_id = if window_key.is_empty() { + format!("row.main.{slot}.v1") + } else { + window_key.to_string() + }; + let Some(mapped) = map_window_with_identity( + label, + window, + now, + card_id.clone(), + (!window_key.is_empty()).then(|| window_key.to_string()), + ) else { + continue; + }; + if !emitted_card_ids.insert(card_id) { + continue; + } + windows.push(mapped); + } + } + + let mut anonymous_slots = HashSet::new(); + for extra in additional_rate_limits.unwrap_or(&[]) { + let source = additional_limit_source(extra); + let digest = source.map(sha256_hex); + let Some(rate_limit) = extra.rate_limit.as_ref() else { + continue; + }; + for (slot, window) in [ + ("primary", rate_limit.primary_window.clone()), + ("secondary", rate_limit.secondary_window.clone()), + ] + .into_iter() + .filter_map(|(slot, window)| window.map(|window| (slot, window))) + { + let Some(digest) = digest.as_deref() else { + let Some(mapped) = map_window_with_identity( + "Unknown", + window, + now, + format!("row.additional.unknown.{slot}.v1"), + None, + ) else { + continue; + }; + if anonymous_slots.insert(slot) { + windows.push(mapped); + } + continue; + }; + let label = additional_limit_label(extra); + let window_key = format!("additional.{digest}.{slot}.v1"); + let Some(mapped) = map_window_with_identity( + &label, + window, + now, + window_key.clone(), + Some(window_key.clone()), + ) else { + continue; + }; + if !emitted_card_ids.insert(window_key) { + continue; + } + windows.push(mapped); + } + } + windows +} + +fn claude_windows(usage: &ClaudeUsageResponse, now: DateTime) -> Vec { + let mut windows = Vec::new(); + push_claude_window( + &mut windows, + "Session", + "session.v1", + DurationEvidence::contract(300 * 60), + usage.five_hour.as_ref(), + now, + ); + push_claude_window( + &mut windows, + "Weekly", + "weekly.v1", + DurationEvidence::contract(7 * 24 * 60 * 60), + usage.seven_day.as_ref(), + now, + ); + push_claude_window( + &mut windows, + "OAuth Apps", + "oauth_apps.weekly.v1", + DurationEvidence::contract(7 * 24 * 60 * 60), + usage.seven_day_oauth_apps.as_ref(), + now, + ); + push_claude_window( + &mut windows, + "Sonnet", + "sonnet.weekly.v1", + DurationEvidence::contract(7 * 24 * 60 * 60), + usage.seven_day_sonnet.as_ref(), + now, ); push_claude_window( &mut windows, "Opus", "opus.weekly.v1", - DurationEvidence::contract(604_800), + DurationEvidence::contract(7 * 24 * 60 * 60), usage.seven_day_opus.as_ref(), now, ); @@ -2859,7 +3730,7 @@ fn claude_windows(usage: &ClaudeUsageResponse, now: DateTime) -> Vec) -> Vec, label: &str, window_key: &str, - contract: DurationEvidence, + contract_duration: DurationEvidence, window: Option<&ClaudeWindow>, now: DateTime, ) { - if let Some(mapped) = - window.and_then(|window| map_claude_window(label, window_key, contract, window, now)) + if let Some(mapped) = window + .and_then(|window| map_claude_window(label, window_key, contract_duration, window, now)) { windows.push(mapped); } @@ -2927,20 +3798,21 @@ fn push_claude_window( fn map_claude_window( label: &str, window_key: &str, - contract: DurationEvidence, + contract_duration: DurationEvidence, window: &ClaudeWindow, now: DateTime, ) -> Option { - if !window.has_valid_utilization() { - return None; - } let used = window.utilization?; - let reset_was_supplied = window.resets_at.is_some(); let resets_at = window.resets_at.as_deref().and_then(parse_datetime); - Some( - UsageWindow::from_used_percent(label.to_string(), used, resets_at, now) - .with_identity(window_key, Some(window_key.to_string())) - .with_duration_evidence(now, reset_was_supplied, None, Some(contract)), + UsageWindow::try_from_provider_used_percent(label.to_string(), used, resets_at, now).map( + |window| { + window.with_identity( + window_key, + Some(window_key.to_string()), + None, + Some(contract_duration), + ) + }, ) } @@ -2964,10 +3836,9 @@ fn parse_unified_ratelimit_windows( if let Some(window) = unified_ratelimit_window_with_identity( "Session", "session.v1", - DurationEvidence::contract(18_000), + DurationEvidence::contract(300 * 60), read_f64("anthropic-ratelimit-unified-5h-utilization"), read_i64("anthropic-ratelimit-unified-5h-reset"), - headers.contains_key("anthropic-ratelimit-unified-5h-reset"), now, ) { windows.push(window); @@ -2975,10 +3846,9 @@ fn parse_unified_ratelimit_windows( if let Some(window) = unified_ratelimit_window_with_identity( "Weekly", "weekly.v1", - DurationEvidence::contract(604_800), + DurationEvidence::contract(7 * 24 * 60 * 60), read_f64("anthropic-ratelimit-unified-7d-utilization"), read_i64("anthropic-ratelimit-unified-7d-reset"), - headers.contains_key("anthropic-ratelimit-unified-7d-reset"), now, ) { windows.push(window); @@ -2993,24 +3863,24 @@ fn parse_unified_ratelimit_windows( fn unified_ratelimit_window_with_identity( label: &str, window_key: &str, - contract: DurationEvidence, + contract_duration: DurationEvidence, utilization_fraction: Option, reset_epoch_seconds: Option, - reset_was_supplied: bool, now: DateTime, ) -> Option { - let fraction = utilization_fraction?; - if !fraction.is_finite() || !(0.0..=1.0).contains(&fraction) { - return None; - } - let used = fraction * 100.0; + let used = utilization_fraction? * 100.0; let resets_at = reset_epoch_seconds .filter(|seconds| *seconds > 0) .and_then(|seconds| Utc.timestamp_opt(seconds, 0).single()); - Some( - UsageWindow::from_used_percent(label.to_string(), used, resets_at, now) - .with_identity(window_key, Some(window_key.to_string())) - .with_duration_evidence(now, reset_was_supplied, None, Some(contract)), + UsageWindow::try_from_provider_used_percent(label.to_string(), used, resets_at, now).map( + |window| { + window.with_identity( + window_key, + Some(window_key.to_string()), + None, + Some(contract_duration), + ) + }, ) } @@ -3021,26 +3891,22 @@ fn unified_ratelimit_window( reset_epoch_seconds: Option, now: DateTime, ) -> Option { - let (key, contract) = if label.eq_ignore_ascii_case("Session") { - ("session.v1", DurationEvidence::contract(18_000)) + let (window_key, duration) = if label.eq_ignore_ascii_case("Session") { + ("session.v1", DurationEvidence::contract(300 * 60)) } else { - ("weekly.v1", DurationEvidence::contract(604_800)) + ("weekly.v1", DurationEvidence::contract(7 * 24 * 60 * 60)) }; unified_ratelimit_window_with_identity( label, - key, - contract, + window_key, + duration, utilization_fraction, reset_epoch_seconds, - reset_epoch_seconds.is_some(), now, ) } -fn claude_extra_usage_window( - extra: Option<&ClaudeExtraUsage>, - now: DateTime, -) -> Option { +fn claude_extra_usage_window(extra: Option<&ClaudeExtraUsage>) -> Option { let extra = extra?; if !extra.is_enabled { return None; @@ -3054,9 +3920,6 @@ fn claude_extra_usage_window( None } })?; - if !used.is_finite() || !(0.0..=100.0).contains(&used) { - return None; - } let reset_text = match (extra.used_credits, extra.monthly_limit) { (Some(used), Some(limit)) => Some(format!( "Monthly cap: {} / {}", @@ -3065,9 +3928,18 @@ fn claude_extra_usage_window( )), _ => None, }; - let mut window = UsageWindow::from_used_percent("Extra usage".to_string(), used, None, now) - .with_identity("extra_usage.v1", Some("extra_usage.v1".to_string())) - .with_unavailable_reason("nonRecurring"); + let mut window = UsageWindow::try_from_provider_used_percent( + "Extra usage".to_string(), + used, + None, + Utc::now(), + )? + .with_identity( + "extra_usage.v1", + Some("extra_usage.v1".to_string()), + None, + None, + ); window.reset_text = reset_text; Some(window) } @@ -3145,15 +4017,18 @@ fn map_window_with_identity( card_id: impl Into, window_key: Option, ) -> Option { - let resets_at = (window.reset_at > 0) + let resets_at = (window.reset_at != 0) .then(|| Utc.timestamp_opt(window.reset_at, 0).single()) .flatten(); - let provider = DurationEvidence::provider(window.reset_at, window.limit_window_seconds); - (window.used_percent.is_finite() && (0.0..=100.0).contains(&window.used_percent)).then(|| { - UsageWindow::from_used_percent(label.to_string(), window.used_percent, resets_at, now) - .with_identity(card_id, window_key) - .with_duration_evidence(now, true, Some(provider), None) - }) + let provider_duration = (window.limit_window_seconds != 0) + .then(|| DurationEvidence::provider(window.reset_at, window.limit_window_seconds)); + UsageWindow::try_from_provider_used_percent( + label.to_string(), + window.used_percent, + resets_at, + now, + ) + .map(|window| window.with_identity(card_id, window_key, provider_duration, None)) } fn additional_limit_source(limit: &CodexAdditionalRateLimit) -> Option { @@ -3164,11 +4039,9 @@ fn additional_limit_source(limit: &CodexAdditionalRateLimit) -> Option { .map(str::to_string) } -fn sha256_hex(value: &str) -> String { - Sha256::digest(value.trim().as_bytes()) - .iter() - .map(|byte| format!("{byte:02x}")) - .collect() +fn sha256_hex(value: String) -> String { + let digest = Sha256::digest(value.trim().as_bytes()); + digest.iter().map(|byte| format!("{byte:02x}")).collect() } pub(crate) fn reset_text(reset: DateTime, now: DateTime) -> String { @@ -3362,6 +4235,15 @@ pub(crate) fn clean_plan(value: impl AsRef) -> String { .join(" ") } +pub(crate) fn deserialize_optional_raw<'de, D, T>(deserializer: D) -> Result, D::Error> +where + D: serde::Deserializer<'de>, + T: serde::de::DeserializeOwned, +{ + let raw = Option::>::deserialize(deserializer)?; + Ok(raw.and_then(|raw| serde_json::from_str(raw.get()).ok())) +} + fn deserialize_optional_non_empty_string<'de, D>( deserializer: D, ) -> Result, D::Error> @@ -3394,27 +4276,6 @@ mod tests { use super::*; use crate::agent_account_scope::test_support::TestRefreshScope; - static COPILOT_TEMP_COUNTER: std::sync::atomic::AtomicU64 = - std::sync::atomic::AtomicU64::new(0); - - fn copilot_auth_path(tag: &str) -> PathBuf { - let root = std::env::temp_dir().join(format!( - "tb-usage-copilot-{tag}-{}-{}", - std::process::id(), - COPILOT_TEMP_COUNTER.fetch_add(1, std::sync::atomic::Ordering::Relaxed) - )); - std::fs::create_dir_all(&root).unwrap(); - root.join("auth.json") - } - - fn copilot_credential( - path: &Path, - json: &serde_json::Value, - ) -> Option { - std::fs::write(path, serde_json::to_vec(json).unwrap()).unwrap(); - crate::opencode_integrations::github_copilot_credential_from(path, json) - } - #[test] fn claude_user_agent_uses_first_version_token() { assert_eq!( @@ -3425,2786 +4286,2461 @@ mod tests { assert_eq!(claude_user_agent_from_stdout(&[0xff]), None); } - fn enrichment_scope(tag: &str) -> (TestRefreshScope, AccountScope) { - let scope = TestRefreshScope::new("fixture", tag); - let account_scope = scope - .resolve_current("fixture", tag, tag.as_bytes()) - .unwrap(); - (scope, account_scope) - } + #[test] + fn parses_retry_after_seconds_and_http_date() { + let header = reqwest::header::HeaderValue::from_static("120"); + let parsed = parse_retry_after(Some(&header)).unwrap(); + let delta = (parsed - Utc::now()).num_seconds(); + assert!((118..=120).contains(&delta), "delta was {}", delta); - fn enrichment_snapshot( - account_scope: Result, - windows: Vec, - ) -> AgentUsageSnapshot { - AgentUsageSnapshot { - client_id: "fixture".to_string(), - source: "fixture".to_string(), - updated_at: String::new(), - identity: None, - account_scope, - windows, - credits: None, - error: None, - } + let header = reqwest::header::HeaderValue::from_static("Fri, 21 Nov 2025 09:00:00 GMT"); + let parsed = parse_retry_after(Some(&header)).unwrap(); + assert_eq!(parsed.timestamp(), 1_763_715_600); + + let header = reqwest::header::HeaderValue::from_static("bogus"); + assert!(parse_retry_after(Some(&header)).is_none()); + assert!(parse_retry_after(None).is_none()); } - fn enrichment_window( - now: DateTime, - card_id: &str, - window_key: &str, - used_percent: f64, - duration_source: Option, - ) -> UsageWindow { - let reset = now + chrono::Duration::days(1); - let window = - UsageWindow::from_used_percent(card_id.to_string(), used_percent, Some(reset), now) - .with_identity(card_id, Some(window_key.to_string())); - match duration_source { - Some(DurationSource::Provider) => window.with_duration_evidence( - now, - true, - Some(DurationEvidence::provider(reset.timestamp(), 86_400)), - None, + #[test] + fn string_key_uses_first_valid_snake_or_camel_alias() { + let cases = [ + ( + "snake priority", + serde_json::json!({ + "snake_key": " snake-value ", + "camelKey": "camel-value" + }), + Some("snake-value"), ), - Some(DurationSource::Contract) => window.with_duration_evidence( - now, - true, + ( + "snake missing", + serde_json::json!({ "camelKey": " camel-value " }), + Some("camel-value"), + ), + ( + "snake null", + serde_json::json!({ "snake_key": null, "camelKey": "camel-value" }), + Some("camel-value"), + ), + ( + "snake empty", + serde_json::json!({ "snake_key": "", "camelKey": "camel-value" }), + Some("camel-value"), + ), + ( + "snake whitespace", + serde_json::json!({ "snake_key": " \t\n ", "camelKey": "camel-value" }), + Some("camel-value"), + ), + ( + "snake non-string", + serde_json::json!({ + "snake_key": { "unexpected": true }, + "camelKey": "camel-value" + }), + Some("camel-value"), + ), + ( + "both invalid", + serde_json::json!({ "snake_key": false, "camelKey": " " }), None, - Some(DurationEvidence::contract(86_400)), ), - Some(DurationSource::Observed) => panic!("observed duration is never retained"), - None => window, + ]; + + for (label, value, expected) in cases { + let map = value.as_object().unwrap(); + assert_eq!( + string_key(map, "snake_key", "camelKey").as_deref(), + expected, + "{label}" + ); } } - fn enrichment_failure_windows(now: DateTime) -> Vec { - vec![ - enrichment_window(now, "first.v1", "first.v1", 10.0, None), - enrichment_window(now, "second.v1", "second.v1", 20.0, None), - UsageWindow::from_used_percent("Missing".to_string(), 30.0, None, now) - .with_identity("missing.v1", Some("missing.v1".to_string())), - UsageWindow::from_used_percent("Non-recurring".to_string(), 40.0, None, now) - .with_identity("non-recurring.v1", Some("non-recurring.v1".to_string())) - .with_unavailable_reason("nonRecurring"), - ] + #[test] + fn claude_refresh_response_ignores_invalid_optional_refresh_token() { + let cases = [ + ( + "valid", + serde_json::json!({ + "access_token": "new-access", + "refresh_token": " new-refresh ", + "expires_in": 3600 + }), + Some("new-refresh"), + ), + ( + "missing", + serde_json::json!({ "access_token": "new-access", "expires_in": 3600 }), + None, + ), + ( + "null", + serde_json::json!({ + "access_token": "new-access", + "refresh_token": null, + "expires_in": 3600 + }), + None, + ), + ( + "empty", + serde_json::json!({ + "access_token": "new-access", + "refresh_token": "", + "expires_in": 3600 + }), + None, + ), + ( + "whitespace", + serde_json::json!({ + "access_token": "new-access", + "refresh_token": " \t\n ", + "expires_in": 3600 + }), + None, + ), + ( + "non-string", + serde_json::json!({ + "access_token": "new-access", + "refresh_token": { "unexpected": true }, + "expires_in": 3600 + }), + None, + ), + ]; + + for (label, value, expected) in cases { + let response: ClaudeRefreshResponse = serde_json::from_value(value).unwrap(); + assert_eq!(response.access_token, "new-access", "{label}"); + assert_eq!(response.expires_in, 3_600, "{label}"); + assert_eq!(response.refresh_token.as_deref(), expected, "{label}"); + } + assert!(serde_json::from_value::( + serde_json::json!({ "expires_in": 3600 }) + ) + .is_err()); + assert!(serde_json::from_value::( + serde_json::json!({ "access_token": "new-access" }) + ) + .is_err()); } - #[tokio::test] - async fn copilot_blank_or_missing_credential_skips_fetch() { - let path = copilot_auth_path("no-credential"); + #[test] + fn claude_gate_is_binding_scoped_and_expires() { let now = Utc.timestamp_opt(1_700_000_000, 0).single().unwrap(); - let fetch_calls = std::cell::Cell::new(0); - let enrich_calls = std::cell::Cell::new(0); - let record_calls = std::cell::Cell::new(0); - for json in [ - serde_json::json!({}), - serde_json::json!({"github-copilot": { - "type": "oauth", "refresh": " ", "access": "\t\n" - }}), - serde_json::json!({ - "github-copilot": {"type": "oauth"}, - "copilot": {"type": "oauth", "refresh": "foreign-token"} - }), - ] { - let snapshot = fetch_copilot_with( - copilot_credential(&path, &json), - now, - |_, _| { - fetch_calls.set(fetch_calls.get() + 1); - std::future::ready(Err::( - "network must not run".to_string(), - )) - }, - |snapshot, callback_now| { - enrich_calls.set(enrich_calls.get() + 1); - enrich_snapshot_with(snapshot, callback_now, |_, _, _| { - record_calls.set(record_calls.get() + 1); - Ok(Vec::new()) - }); - }, - ) - .await; - assert!(snapshot.is_none()); - } - assert_eq!(fetch_calls.get(), 0); - assert_eq!(enrich_calls.get(), 0); - assert_eq!(record_calls.get(), 0); - let _ = std::fs::remove_dir_all(path.parent().unwrap()); + let scope = TestRefreshScope::new("claude", "local-gate"); + let binding_a = ProviderCacheBinding::primary( + scope + .resolve_current("fixture", "account-a", b"marker-a") + .unwrap(), + ); + let binding_b = ProviderCacheBinding::primary( + scope + .resolve_current("fixture", "account-b", b"marker-b") + .unwrap(), + ); + let mut gate = ClaudeUsageGate::default(); + + gate.record_rate_limit(binding_a.clone(), None, now); + let until = gate.blocked_until_for(&binding_a, now).unwrap(); + assert_eq!((until - now).num_seconds(), 300); + + assert!(gate.blocked_until_for(&binding_b, now).is_none()); + assert!(gate.blocked_until_for(&binding_a, now).is_none()); + + gate.record_rate_limit( + binding_a.clone(), + Some(now + chrono::Duration::seconds(60)), + now, + ); + assert!(gate + .blocked_until_for(&binding_a, now + chrono::Duration::seconds(61)) + .is_none()); + gate.clear(); + assert!(gate.blocked_until_for(&binding_a, now).is_none()); + scope.cleanup(); } #[tokio::test] - async fn copilot_success_enriches_distinct_series_once_without_wire_metadata() { - let path = copilot_auth_path("snapshot"); - let credential = copilot_credential( - &path, - &serde_json::json!({"github-copilot": { - "type": "oauth", - "refresh": " fake-copilot-marker ", - "access": "unused-access-token" - }}), + async fn claude_login_usage_gates_current_binding_before_refresh_and_request() { + let now = Utc.timestamp_opt(1_700_000_000, 0).single().unwrap(); + let scope = TestRefreshScope::new("claude", "routed-gate"); + let binding_a = ProviderCacheBinding::primary( + scope + .resolve_current("fixture", "account-a", b"marker-a") + .unwrap(), ); - let location = crate::agent_account_scope::canonical_file_location( - &path, - Some("github-copilot"), - ) - .unwrap(); - let scope_store = TestRefreshScope::new("copilot", "copilot-central-snapshot"); - let now = "2026-06-15T12:00:00Z".parse::>().unwrap(); - let premium_reset = "2026-07-01T00:00:00Z".parse::>().unwrap(); - let chat_reset = premium_reset; - let premium_duration = 30 * 86_400; - let chat_duration = premium_duration; - let premium = UsageWindow::from_fraction( - "Premium".to_string(), - 0.30, - Some(premium_reset), + let binding_b = ProviderCacheBinding::primary( + scope + .resolve_current("fixture", "account-b", b"marker-b") + .unwrap(), + ); + let mut credentials = ClaudeCredentials { + access_token: "claude-access".to_string(), + refresh_token: Some("claude-refresh".to_string()), + expires_at: Some(Utc::now() - chrono::Duration::minutes(1)), + scopes: vec!["user:profile".to_string()], + rate_limit_tier: None, + subscription_type: None, + source: ClaudeCredentialSource::File, + raw_root: None, + keychain_account: None, + scope_slot: CredentialSlot { + semantic_source: "fixture", + canonical_location: "fixture".to_string(), + }, + }; + let mut gate = ClaudeUsageGate::default(); + gate.record_rate_limit(binding_a.clone(), None, now); + let refresh_calls = std::cell::Cell::new(0); + let header_calls = std::cell::Cell::new(0); + let usage_calls = std::cell::Cell::new(0); + + let (source, outcome) = fetch_claude_login_usage_with( + credentials.clone(), + binding_a.clone(), now, + |binding, at| gate.blocked_until_for(binding, at), + |credentials| { + refresh_calls.set(refresh_calls.get() + 1); + let binding = binding_a.clone(); + async move { Ok((credentials, binding.primary.clone(), Some(binding))) } + }, + |_, _, _| async { + header_calls.set(header_calls.get() + 1); + claude_test_success_outcome() + }, + |_, _, _, _| async { + usage_calls.set(usage_calls.get() + 1); + ("oauth", claude_test_success_outcome()) + }, ) - .with_identity( - "premium_interactions.v1", - Some("premium_interactions.v1".to_string()), - ) - .with_contract_duration_evidence( + .await; + assert_eq!(source, "oauth"); + assert!(matches!( + outcome, + ProviderFetchOutcome::Failure(ProviderFetchFailure::Transient { + attempt_binding: Some(ref binding), + transport_diagnostic: SafeTransportDiagnostic { + category: TransportCategory::RateLimited, + status: Some(429), + .. + }, + .. + }) if binding == &binding_a + )); + assert_eq!(refresh_calls.get(), 0); + assert_eq!(header_calls.get(), 0); + assert_eq!(usage_calls.get(), 0); + + credentials.expires_at = None; + credentials.scopes = vec!["org:create_api_key".to_string()]; + gate.record_rate_limit(binding_a.clone(), None, now); + let (source, outcome) = fetch_claude_login_usage_with( + credentials.clone(), + binding_a.clone(), now, - true, - DurationEvidence::contract(premium_duration), - ); - let chat = UsageWindow::from_fraction("Chat".to_string(), 0.75, Some(chat_reset), now) - .with_identity("chat.v1", Some("chat.v1".to_string())) - .with_contract_duration_evidence( - now, - true, - DurationEvidence::contract(chat_duration), - ); - let enrich_calls = std::cell::Cell::new(0); - let record_calls = std::cell::Cell::new(0); - let snapshot = fetch_copilot_with( - credential, + |binding, at| gate.blocked_until_for(binding, at), + |credentials| { + refresh_calls.set(refresh_calls.get() + 1); + let binding = binding_a.clone(); + async move { Ok((credentials, binding.primary.clone(), Some(binding))) } + }, + |_, _, _| async { + header_calls.set(header_calls.get() + 1); + claude_test_success_outcome() + }, + |_, _, _, _| async { + usage_calls.set(usage_calls.get() + 1); + ("oauth", claude_test_success_outcome()) + }, + ) + .await; + assert_eq!(source, "setup-token"); + assert!(matches!(outcome, ProviderFetchOutcome::Success { .. })); + assert_eq!(refresh_calls.get(), 0); + assert_eq!(header_calls.get(), 1); + assert_eq!(usage_calls.get(), 0); + assert!(gate.blocked_until_for(&binding_a, now).is_some()); + + credentials.scopes = vec!["user:profile".to_string()]; + gate.record_rate_limit(binding_a.clone(), None, now); + let (source, outcome) = fetch_claude_login_usage_with( + credentials, + binding_b.clone(), now, - |request_now, credential| { - assert_eq!(request_now, now); - assert_eq!(credential.request_token, "fake-copilot-marker"); - assert_eq!(credential.marker, b"fake-copilot-marker"); - assert_eq!(credential.semantic_source, "opencode-auth-json"); - assert_eq!(credential.canonical_location, location); - let account_scope = scope_store.resolve_current( - credential.semantic_source, - &credential.canonical_location, - &credential.marker, - ); - std::future::ready(Ok(agent_copilot::CopilotData { - identity: None, - account_scope, - windows: vec![premium, chat], - })) + |binding, at| gate.blocked_until_for(binding, at), + |credentials| { + refresh_calls.set(refresh_calls.get() + 1); + let binding = binding_b.clone(); + async move { Ok((credentials, binding.primary.clone(), Some(binding))) } }, - |snapshot, callback_now| { - enrich_calls.set(enrich_calls.get() + 1); - let account_scope = snapshot.account_scope.as_ref().unwrap().clone(); - enrich_snapshot_with(snapshot, callback_now, |active, observations, history_now| { - record_calls.set(record_calls.get() + 1); - assert_eq!(history_now, now.timestamp()); - assert_eq!( - active, - &[ - SeriesKey::new( - "copilot", - account_scope.as_str(), - "premium_interactions.v1", - ), - SeriesKey::new("copilot", account_scope.as_str(), "chat.v1"), - ] - ); - assert_eq!(observations.len(), 2); - assert_eq!(observations[0].key, active[0]); - assert_eq!(observations[0].reset_at, Some(premium_reset.timestamp())); - assert_eq!(observations[0].used_percent, 70.0); - assert_eq!( - observations[0].contract, - Some(DurationEvidence::contract(premium_duration)) - ); - assert!(observations[0].provider.is_none()); - assert_eq!(observations[1].key, active[1]); - assert_eq!(observations[1].reset_at, Some(chat_reset.timestamp())); - assert_eq!(observations[1].used_percent, 25.0); - assert!(observations[1].provider.is_none()); - assert_eq!( - observations[1].contract, - Some(DurationEvidence::contract(chat_duration)) - ); - Ok(vec![ - Ok(( - HistoryOutcome::Ready { - duration_seconds: premium_duration, - source: DurationSource::Contract, - sampled: true, - }, - Some(HistoricalPace { - expected_percent: 35.0, - eta_seconds: Some(123.0), - will_last_to_reset: false, - run_out_probability: None, - }), - 3, - )), - Ok(( - HistoryOutcome::Ready { - duration_seconds: chat_duration, - source: DurationSource::Contract, - sampled: true, - }, - Some(HistoricalPace { - expected_percent: 26.0, - eta_seconds: Some(456.0), - will_last_to_reset: false, - run_out_probability: None, - }), - 3, - )), - ]) - }); + |_, _, _| async { + header_calls.set(header_calls.get() + 1); + claude_test_success_outcome() + }, + |_, _, _, _| async { + usage_calls.set(usage_calls.get() + 1); + ("oauth", claude_test_success_outcome()) }, ) - .await - .unwrap(); - assert_eq!(enrich_calls.get(), 1); - assert_eq!(record_calls.get(), 1); - let account_scope = snapshot.account_scope.as_ref().unwrap(); - let opaque_scope = account_scope.as_str().to_string(); - assert_eq!(snapshot.windows.len(), 2); - assert_eq!( - snapshot.windows[0].pace_window_key_for_test(), - Some("premium_interactions.v1") - ); - assert_eq!(snapshot.windows[1].pace_window_key_for_test(), Some("chat.v1")); - assert_eq!(snapshot.windows[0].pace_status.state, PaceState::Available); - assert_eq!(snapshot.windows[0].pace_status.complete_cycles, 3); - assert_eq!( - snapshot.windows[0].pace_status.duration_source, - Some(DurationSource::Contract) - ); - assert!(snapshot.windows[0].historical_pace.is_some()); - assert_eq!(snapshot.windows[1].pace_status.state, PaceState::Available); - assert_eq!(snapshot.windows[1].pace_status.complete_cycles, 3); - assert_eq!( - snapshot.windows[1].pace_status.duration_seconds, - Some(chat_duration) - ); - assert_eq!( - snapshot.windows[1].pace_status.duration_source, - Some(DurationSource::Contract) - ); - assert_eq!( - snapshot.windows[1] - .historical_pace - .as_ref() - .map(|pace| pace.expected_used_percent), - Some(26.0) - ); - - let wire = serde_json::to_string(&snapshot).unwrap(); - assert!(!wire.contains("accountScope")); - assert!(!wire.contains("fake-copilot-marker")); - assert!(!wire.contains("unused-access-token")); - assert!(!wire.contains(path.to_string_lossy().as_ref())); - assert!(!wire.contains(&location)); - assert!(!wire.contains(&opaque_scope)); - scope_store.cleanup(); - let _ = std::fs::remove_dir_all(path.parent().unwrap()); + .await; + assert_eq!(source, "oauth"); + assert!(matches!(outcome, ProviderFetchOutcome::Success { .. })); + assert_eq!(refresh_calls.get(), 0); + assert_eq!(header_calls.get(), 1); + assert_eq!(usage_calls.get(), 1); + scope.cleanup(); } - #[test] - fn copilot_one_cycle_stays_in_learning_history_without_historical_pace() { - let (scope, account_scope) = enrichment_scope("copilot-one-cycle"); - let now = "2026-06-15T12:00:00Z".parse::>().unwrap(); - let reset = "2026-07-01T00:00:00Z".parse::>().unwrap(); - let duration = 30 * 86_400; - let window = UsageWindow::from_fraction("Premium".to_string(), 0.30, Some(reset), now) - .with_identity( - "premium_interactions.v1", - Some("premium_interactions.v1".to_string()), - ) - .with_contract_duration_evidence( + fn cache_test_snapshot( + client_id: &str, + account_scope: Result, + now: DateTime, + ) -> AgentUsageSnapshot { + AgentUsageSnapshot { + client_id: client_id.to_string(), + source: "oauth".to_string(), + updated_at: now.to_rfc3339_opts(SecondsFormat::Millis, true), + identity: Some(AgentIdentity { + email: Some("fixture@example.invalid".to_string()), + plan: Some("Fixture".to_string()), + }), + account_scope, + windows: vec![UsageWindow::from_provider_used_percent( + "Session".to_string(), + 20.0, + Some(now + chrono::Duration::hours(5)), now, - true, - DurationEvidence::contract(duration), - ); - let enrich_calls = std::cell::Cell::new(0); - let record_calls = std::cell::Cell::new(0); - let snapshot = finalize_copilot_snapshot_with( - Ok(agent_copilot::CopilotData { - identity: None, - account_scope: Ok(account_scope), - windows: vec![window], + ) + .with_identity( + "main.session.v1", + Some("main.session.v1".to_string()), + None, + Some(DurationEvidence::contract(300 * 60)), + )], + credits: Some(CreditsSnapshot { + remaining: Some(-2.5), + unlimited: false, }), - now, - |snapshot, callback_now| { - enrich_calls.set(enrich_calls.get() + 1); - enrich_snapshot_with(snapshot, callback_now, |_, observations, _| { - record_calls.set(record_calls.get() + 1); - assert_eq!(observations.len(), 1); - Ok(vec![Ok(( - HistoryOutcome::Ready { - duration_seconds: duration, - source: DurationSource::Contract, - sampled: true, - }, - None, - 1, - ))]) - }); - }, - ); - assert_eq!(enrich_calls.get(), 1); - assert_eq!(record_calls.get(), 1); - assert_eq!(snapshot.windows[0].pace_status.state, PaceState::LearningHistory); - assert_eq!(snapshot.windows[0].pace_status.complete_cycles, 1); - assert_eq!(snapshot.windows[0].pace_status.duration_seconds, Some(duration)); - assert_eq!( - snapshot.windows[0].pace_status.duration_source, - Some(DurationSource::Contract) - ); - assert!(snapshot.windows[0].historical_pace.is_none()); - scope.cleanup(); + error: None, + transport_diagnostic: None, + } } - #[test] - fn copilot_non_month_reset_enters_observed_history_learning() { - let (scope, account_scope) = enrichment_scope("copilot-observed"); - let now = "2026-07-01T12:00:00Z".parse::>().unwrap(); - let reset = "2026-07-15T00:00:00Z".parse::>().unwrap(); - let window = UsageWindow::from_fraction("Chat".to_string(), 0.40, Some(reset), now) - .with_identity("chat.v1", Some("chat.v1".to_string())) - .with_observed_duration_evidence(now, true); - let record_calls = std::cell::Cell::new(0); - let snapshot = finalize_copilot_snapshot_with( - Ok(agent_copilot::CopilotData { - identity: None, - account_scope: Ok(account_scope), - windows: vec![window], - }), - now, - |snapshot, callback_now| { - enrich_snapshot_with(snapshot, callback_now, |_, observations, _| { - record_calls.set(record_calls.get() + 1); - assert_eq!(observations.len(), 1); - assert_eq!(observations[0].reset_at, Some(reset.timestamp())); - assert_eq!(observations[0].used_percent, 60.0); - assert!(observations[0].provider.is_none()); - assert!(observations[0].contract.is_none()); - Ok(vec![Ok((HistoryOutcome::LearningDuration, None, 0))]) - }); - }, - ); - assert_eq!(record_calls.get(), 1); - assert_eq!(snapshot.windows[0].used_percent, 60.0); - assert_eq!(snapshot.windows[0].reset_at_evidence, Some(reset)); - assert_eq!(snapshot.windows[0].pace_status.state, PaceState::LearningDuration); - assert_eq!( - snapshot.windows[0].pace_status.duration_source, - Some(DurationSource::Observed) - ); - assert!(snapshot.windows[0].pace_status.duration_seconds.is_none()); - assert!(snapshot.windows[0].historical_pace.is_none()); - scope.cleanup(); + fn claude_test_login_credentials() -> ClaudeCredentials { + ClaudeCredentials { + access_token: "claude-access".to_string(), + refresh_token: Some("claude-refresh".to_string()), + expires_at: None, + scopes: vec!["user:profile".to_string()], + rate_limit_tier: None, + subscription_type: None, + source: ClaudeCredentialSource::File, + raw_root: None, + keychain_account: None, + scope_slot: CredentialSlot { + semantic_source: "fixture", + canonical_location: "fixture".to_string(), + }, + } } - #[test] - fn copilot_scope_error_enriches_once_without_recording() { - let (scope_store, _) = enrichment_scope("copilot-scope-error"); - let now = "2026-07-01T00:00:00Z".parse::>().unwrap(); - let reset = now + chrono::Duration::days(1); - let windows = vec![ - UsageWindow::from_used_percent("Premium".to_string(), 30.0, Some(reset), now) - .with_identity( - "premium_interactions.v1", - Some("premium_interactions.v1".to_string()), - ), - UsageWindow::from_used_percent("Chat".to_string(), 40.0, Some(reset), now) - .with_identity("chat.v1", Some("chat.v1".to_string())), - ]; - let enrich_calls = std::cell::Cell::new(0); - let record_calls = std::cell::Cell::new(0); - let snapshot = finalize_copilot_snapshot_with( - Ok(agent_copilot::CopilotData { - identity: None, - account_scope: Err(AccountScopeError::MetadataWrite), - windows, - }), - now, - |snapshot, callback_now| { - enrich_calls.set(enrich_calls.get() + 1); - enrich_snapshot_with(snapshot, callback_now, |_, _, _| { - record_calls.set(record_calls.get() + 1); - Ok(Vec::new()) - }); + fn claude_test_setup_token() -> ResolvedClaudeToken { + ResolvedClaudeToken { + access_token: "setup-access".to_string(), + scope_slot: CredentialSlot { + semantic_source: "fixture-setup", + canonical_location: "fixture-setup".to_string(), }, - ); - assert_eq!(enrich_calls.get(), 1); - assert_eq!(record_calls.get(), 0); - assert_eq!(snapshot.windows.len(), 2); - assert_eq!(snapshot.windows[0].used_percent, 30.0); - assert_eq!(snapshot.windows[0].remaining_percent, 70.0); - assert_eq!(snapshot.windows[1].used_percent, 40.0); - assert_eq!(snapshot.windows[1].remaining_percent, 60.0); - assert!(snapshot - .windows - .iter() - .all(|window| window.pace_reason_for_test() == Some("accountScope"))); - scope_store.cleanup(); + } } - #[test] - fn copilot_api_error_and_empty_success_enrich_once_without_recording() { + fn claude_test_success_outcome() -> ProviderFetchOutcome { let now = Utc.timestamp_opt(1_700_000_000, 0).single().unwrap(); - let api_enrich_calls = std::cell::Cell::new(0); - let api_record_calls = std::cell::Cell::new(0); - let api_error = finalize_copilot_snapshot_with( - Err("Copilot usage API returned 503.".to_string()), - now, - |snapshot, callback_now| { - api_enrich_calls.set(api_enrich_calls.get() + 1); - enrich_snapshot_with(snapshot, callback_now, |_, _, _| { - api_record_calls.set(api_record_calls.get() + 1); - Ok(Vec::new()) - }); + ProviderFetchOutcome::Success { + snapshot: cache_test_snapshot("claude", Err(AccountScopeError::NoTrustedEvidence), now), + cache_binding: None, + } + } + + #[tokio::test] + async fn claude_login_precedence_falls_through_only_for_absent_credentials() { + { + let primary_calls = std::cell::Cell::new(0); + let setup_loads = std::cell::Cell::new(0); + let setup_calls = std::cell::Cell::new(0); + let (source, outcome) = fetch_claude_login_or_setup_with( + resolve_stored_claude_login("{", ClaudeCredentialSource::File), + |_| async { + primary_calls.set(primary_calls.get() + 1); + ("oauth", claude_test_success_outcome()) + }, + || { + setup_loads.set(setup_loads.get() + 1); + Ok(Some(claude_test_setup_token())) + }, + |_| async { + setup_calls.set(setup_calls.get() + 1); + ("setup-token", claude_test_success_outcome()) + }, + ) + .await; + assert_eq!(source, "oauth"); + assert!(matches!( + outcome, + ProviderFetchOutcome::Failure(ProviderFetchFailure::Terminal { .. }) + )); + assert_eq!(primary_calls.get(), 0); + assert_eq!(setup_loads.get(), 0); + assert_eq!(setup_calls.get(), 0); + } + + for (label, primary_outcome) in [ + ( + "401", + ProviderFetchOutcome::Failure(ProviderFetchFailure::terminal( + "Claude OAuth token expired or invalid. Run `claude` to re-authenticate.", + )), + ), + ( + "transient", + ProviderFetchOutcome::Failure(ProviderFetchFailure::transient( + "Claude usage request failed. Retrying automatically.", + None, + SafeTransportDiagnostic::server_error(503), + )), + ), + ] { + let primary_calls = std::cell::Cell::new(0); + let setup_loads = std::cell::Cell::new(0); + let setup_calls = std::cell::Cell::new(0); + let (source, outcome) = fetch_claude_login_or_setup_with( + ClaudeLoginResolution::Ready(claude_test_login_credentials()), + |_| { + primary_calls.set(primary_calls.get() + 1); + async move { ("oauth", primary_outcome) } + }, + || { + setup_loads.set(setup_loads.get() + 1); + Ok(Some(claude_test_setup_token())) + }, + |_| async { + setup_calls.set(setup_calls.get() + 1); + ("setup-token", claude_test_success_outcome()) + }, + ) + .await; + assert_eq!(source, "oauth", "{label}"); + match label { + "401" => assert!(matches!( + outcome, + ProviderFetchOutcome::Failure(ProviderFetchFailure::Terminal { .. }) + )), + _ => assert!(matches!( + outcome, + ProviderFetchOutcome::Failure(ProviderFetchFailure::Transient { .. }) + )), + } + assert_eq!(primary_calls.get(), 1, "{label}"); + assert_eq!(setup_loads.get(), 0, "{label}"); + assert_eq!(setup_calls.get(), 0, "{label}"); + } + + { + let primary_calls = std::cell::Cell::new(0); + let setup_loads = std::cell::Cell::new(0); + let setup_calls = std::cell::Cell::new(0); + let logged_out = resolve_stored_claude_login( + r#"{"claudeAiOauth":{"refreshToken":"stale"}}"#, + ClaudeCredentialSource::File, + ); + assert!(matches!(logged_out, ClaudeLoginResolution::ExplicitLogout)); + let (source, outcome) = fetch_claude_login_or_setup_with( + logged_out, + |_| async { + primary_calls.set(primary_calls.get() + 1); + ("oauth", claude_test_success_outcome()) + }, + || { + setup_loads.set(setup_loads.get() + 1); + Ok(Some(claude_test_setup_token())) + }, + |_| async { + setup_calls.set(setup_calls.get() + 1); + ("setup-token", claude_test_success_outcome()) + }, + ) + .await; + assert_eq!(source, "setup-token"); + assert!(matches!(outcome, ProviderFetchOutcome::Success { .. })); + assert_eq!(primary_calls.get(), 0); + assert_eq!(setup_loads.get(), 1); + assert_eq!(setup_calls.get(), 1); + } + } + + #[tokio::test] + async fn claude_keychain_logout_blocks_stale_file_login_and_allows_setup_token() { + const VALID_FILE_LOGIN: &str = + r#"{"claudeAiOauth":{"accessToken":"file-access","refreshToken":"file-refresh"}}"#; + + for malformed in [ + r#"{"claudeAiOauth":{}}"#, + r#"{"claudeAiOauth":{"accessToken":null}}"#, + ] { + assert!(matches!( + resolve_stored_claude_login(malformed, ClaudeCredentialSource::Keychain), + ClaudeLoginResolution::Terminal + )); + } + + let missing_file_loads = std::cell::Cell::new(0); + let missing = load_stored_claude_login_with( + || Ok(None), + || { + missing_file_loads.set(missing_file_loads.get() + 1); + Ok(Some(VALID_FILE_LOGIN.to_string())) }, ); - assert_eq!(api_enrich_calls.get(), 1); - assert_eq!(api_record_calls.get(), 0); - assert_eq!(api_error.account_scope, Err(AccountScopeError::NoTrustedEvidence)); - assert!(api_error.windows.is_empty()); - assert_eq!(api_error.error.as_deref(), Some("Copilot usage API returned 503.")); - - let (scope_store, account_scope) = enrichment_scope("copilot-empty-success"); - let empty_enrich_calls = std::cell::Cell::new(0); - let empty_record_calls = std::cell::Cell::new(0); - let empty = finalize_copilot_snapshot_with( - Ok(agent_copilot::CopilotData { - identity: None, - account_scope: Ok(account_scope), - windows: Vec::new(), - }), - now, - |snapshot, callback_now| { - empty_enrich_calls.set(empty_enrich_calls.get() + 1); - enrich_snapshot_with(snapshot, callback_now, |_, _, _| { - empty_record_calls.set(empty_record_calls.get() + 1); - Ok(Vec::new()) - }); + assert!(matches!( + missing, + ClaudeLoginResolution::Ready(ClaudeCredentials { + source: ClaudeCredentialSource::File, + .. + }) + )); + assert_eq!(missing_file_loads.get(), 1); + + let file_loads = std::cell::Cell::new(0); + let login = load_stored_claude_login_with( + || { + Ok(Some( + r#"{"claudeAiOauth":{"refreshToken":"stale-file-shape"}}"#.to_string(), + )) + }, + || { + file_loads.set(file_loads.get() + 1); + Ok(Some(VALID_FILE_LOGIN.to_string())) }, ); - assert_eq!(empty_enrich_calls.get(), 1); - assert_eq!(empty_record_calls.get(), 0); - assert!(empty.windows.is_empty()); - scope_store.cleanup(); + assert!(matches!(login, ClaudeLoginResolution::ExplicitLogout)); + assert_eq!(file_loads.get(), 0); + + let now = Utc.timestamp_opt(1_700_000_000, 0).single().unwrap(); + let scope = TestRefreshScope::new("claude", "keychain-explicit-logout"); + let binding = ProviderCacheBinding::primary( + scope + .resolve_current("fixture", "logged-in-account", b"marker") + .unwrap(), + ); + let mut gate = ClaudeUsageGate::default(); + gate.record_rate_limit(binding.clone(), None, now); + assert!(gate.blocked_until_for(&binding, now).is_some()); + clear_claude_gate_for_login_resolution(&login, &mut gate); + assert!(gate.blocked_until_for(&binding, now).is_none()); + + let oauth_calls = std::cell::Cell::new(0); + let setup_loads = std::cell::Cell::new(0); + let setup_calls = std::cell::Cell::new(0); + let (source, outcome) = fetch_claude_login_or_setup_with( + login, + |_| async { + oauth_calls.set(oauth_calls.get() + 1); + ("oauth", claude_test_success_outcome()) + }, + || { + setup_loads.set(setup_loads.get() + 1); + Ok(Some(claude_test_setup_token())) + }, + |_| async { + setup_calls.set(setup_calls.get() + 1); + ("setup-token", claude_test_success_outcome()) + }, + ) + .await; + assert_eq!(source, "setup-token"); + assert!(matches!(outcome, ProviderFetchOutcome::Success { .. })); + assert_eq!(oauth_calls.get(), 0); + assert_eq!(setup_loads.get(), 1); + assert_eq!(setup_calls.get(), 1); + scope.cleanup(); + } + + fn timeout_diagnostic() -> SafeTransportDiagnostic { + SafeTransportDiagnostic::from_facts(TransportErrorFacts::synthetic( + true, + false, + TransportPhase::Request, + None, + )) } #[test] - fn copilot_unknown_and_non_recurring_successes_do_not_record_history() { - let (scope, account_scope) = enrichment_scope("copilot-no-stable-window"); - let now = "2026-06-15T12:00:00Z".parse::>().unwrap(); - let reset = "2026-07-01T00:00:00Z".parse::>().unwrap(); - let windows = [ - vec![UsageWindow::from_used_percent("Unknown".to_string(), 10.0, Some(reset), now) - .with_identity("row.copilot.unknown.v1", None)], - vec![UsageWindow::from_used_percent( - "Non-recurring".to_string(), - 20.0, - Some(reset), - now, - ) - .with_identity("non-recurring.v1", Some("non-recurring.v1".to_string())) - .with_unavailable_reason("nonRecurring")], - ]; - let enrich_calls = std::cell::Cell::new(0); - let record_calls = std::cell::Cell::new(0); - for windows in windows { - let snapshot = finalize_copilot_snapshot_with( - Ok(agent_copilot::CopilotData { - identity: None, - account_scope: Ok(account_scope.clone()), + fn copilot_malformed_optional_reset_remains_success_and_keeps_last_good() { + let scope = TestRefreshScope::new("copilot", "lossy-optional-reset"); + let account_scope = scope + .resolve_current("fixture", "account-a", b"marker-a") + .unwrap(); + let binding = ProviderCacheBinding::primary(account_scope.clone()); + let cache = Mutex::new(ProviderLastGoodCache::default()); + let fresh_at = Utc.timestamp_opt(1_700_000_000, 0).single().unwrap(); + + apply_provider_outcome_with( + &cache, + "copilot", + "oauth", + fresh_at, + ProviderFetchOutcome::Success { + snapshot: cache_test_snapshot("copilot", Ok(account_scope.clone()), fresh_at), + cache_binding: Some(binding.clone()), + }, + |_| {}, + ) + .unwrap(); + + let response_at = fresh_at + chrono::Duration::minutes(1); + let decoded = agent_copilot::decode_usage_response( + r#"{ + "quota_reset_date": {"credential":"token-secret"}, + "quota_snapshots": { + "premium_interactions": { + "entitlement": 100, + "remaining": 60, + "percent_remaining": 60 + } + } + }"#, + response_at, + ); + let outcome = match decoded { + Ok((plan, windows)) => ProviderFetchOutcome::Success { + snapshot: AgentUsageSnapshot { + client_id: "copilot".to_string(), + source: "oauth".to_string(), + updated_at: response_at.to_rfc3339_opts(SecondsFormat::Millis, true), + identity: Some(AgentIdentity { email: None, plan }), + account_scope: Ok(account_scope), windows, - }), - now, - |snapshot, callback_now| { - enrich_calls.set(enrich_calls.get() + 1); - enrich_snapshot_with(snapshot, callback_now, |active, observations, _| { - record_calls.set(record_calls.get() + 1); - assert!(active.is_empty()); - assert!(observations.is_empty()); - Ok(Vec::new()) - }); + credits: None, + error: None, + transport_diagnostic: None, }, - ); - assert_eq!(snapshot.windows.len(), 1); - } - assert_eq!(enrich_calls.get(), 2); - assert_eq!(record_calls.get(), 0); + cache_binding: Some(binding), + }, + Err(failure) => ProviderFetchOutcome::Failure(failure), + }; + let snapshot = + apply_provider_outcome_with(&cache, "copilot", "oauth", response_at, outcome, |_| {}) + .unwrap(); + + assert!(snapshot.error.is_none()); + assert_eq!(snapshot.windows.len(), 1); + assert!((snapshot.windows[0].remaining_percent - 60.0).abs() < 0.01); + assert!(snapshot.windows[0].resets_at.is_none()); + let cached = lock_last_good(&cache).entries["copilot"].snapshot.clone(); + assert_eq!(cached.updated_at, snapshot.updated_at); + assert_eq!(cached.windows.len(), 1); + assert!(cached.error.is_none()); + assert!(cached.transport_diagnostic.is_none()); scope.cleanup(); } #[test] - fn finalizes_grok_scope_and_enriches_success_but_errors_fail_closed() { - let scope_store = TestRefreshScope::new("grok", "grok-finalize"); - let marker = b"grok-sensitive-refresh-marker"; - let account_scope = scope_store - .resolve_current("grok-auth-json", "fixture-location", marker) + fn last_good_same_binding_fallback_preserves_clean_snapshot_without_enrichment() { + let scope = TestRefreshScope::new("codex", "last-good-same-binding"); + let account_scope = scope + .resolve_current("fixture", "account-a", b"marker-a") .unwrap(); - let opaque_scope = account_scope.as_str().to_string(); - let now = DateTime::parse_from_rfc3339("2026-07-18T00:00:00Z") - .unwrap() - .with_timezone(&Utc); - let reset = now + chrono::Duration::days(7); - let weekly = UsageWindow::from_used_percent( - "Weekly".to_string(), - 25.0, - Some(reset), - now, + let binding = ProviderCacheBinding::primary(account_scope.clone()); + let cache = Mutex::new(ProviderLastGoodCache::default()); + let fresh_at = Utc.timestamp_opt(1_700_000_000, 0).single().unwrap(); + let failure_at = fresh_at + chrono::Duration::minutes(1); + let enrich_calls = std::cell::Cell::new(0); + + let fresh = apply_provider_outcome_with( + &cache, + "codex", + "oauth", + fresh_at, + ProviderFetchOutcome::Success { + snapshot: cache_test_snapshot("codex", Ok(account_scope), fresh_at), + cache_binding: Some(binding.clone()), + }, + |snapshot| { + enrich_calls.set(enrich_calls.get() + 1); + snapshot.windows[0].pace_status = PaceStatusPayload { + state: PaceState::Available, + window_key: Some("main.session.v1".to_string()), + duration_seconds: Some(300 * 60), + duration_source: Some(DurationSource::Contract), + complete_cycles: 6, + reason: None, + }; + snapshot.windows[0].historical_pace = Some(HistoricalPacePayload { + expected_used_percent: 35.0, + eta_seconds: Some(1_800.0), + will_last_to_reset: false, + run_out_probability: Some(0.42), + }); + }, ) - .with_identity( - "billing.weekly.v1", - Some("billing.weekly.v1".to_string()), + .unwrap(); + assert_eq!(enrich_calls.get(), 1); + + let fallback = apply_provider_outcome_with( + &cache, + "codex", + "oauth", + failure_at, + ProviderFetchOutcome::Failure(ProviderFetchFailure::transient( + "Codex usage request failed. Retrying automatically.", + Some(binding.clone()), + timeout_diagnostic(), + )), + |_| enrich_calls.set(enrich_calls.get() + 1), ) - .with_provider_duration_evidence( - now, - true, - Some(DurationEvidence::provider(reset.timestamp(), 604_800)), - ); - let callback_calls = std::cell::Cell::new(0); - let record_calls = std::cell::Cell::new(0); - let success = finalize_grok_snapshot_with( - Ok(agent_grok::GrokData { - identity: None, - account_scope: Ok(account_scope.clone()), - windows: vec![weekly], - }), - now, - |snapshot, callback_now| { - callback_calls.set(callback_calls.get() + 1); - assert_eq!(callback_now, now); - assert_eq!(snapshot.client_id, "grok"); - assert_eq!(snapshot.account_scope.as_ref().unwrap(), &account_scope); - enrich_snapshot_with( - snapshot, - callback_now, - |active, observations, history_now| { - record_calls.set(record_calls.get() + 1); - assert_eq!(history_now, now.timestamp()); - assert_eq!( - active, - &[SeriesKey::new( - "grok", - account_scope.as_str(), - "billing.weekly.v1", - )] - ); - assert_eq!(observations.len(), 1); - assert_eq!( - observations[0].key, - SeriesKey::new( - "grok", - account_scope.as_str(), - "billing.weekly.v1", - ) - ); - assert_eq!(observations[0].reset_at, Some(reset.timestamp())); - assert_eq!(observations[0].used_percent, 25.0); - assert_eq!( - observations[0].provider, - Some(DurationEvidence::provider(reset.timestamp(), 604_800)) - ); - assert!(observations[0].contract.is_none()); - Ok(vec![Ok(( - HistoryOutcome::Ready { - duration_seconds: 604_800, - source: DurationSource::Provider, - sampled: true, - }, - Some(HistoricalPace { - expected_percent: 27.5, - eta_seconds: Some(123.0), - will_last_to_reset: false, - run_out_probability: None, - }), - 3, - ))]) - }, - ); - }, - ); - assert_eq!(callback_calls.get(), 1); - assert_eq!(record_calls.get(), 1); - assert_eq!(success.account_scope, Ok(account_scope)); - assert_eq!(success.windows[0].pace_status.state, PaceState::Available); - assert_eq!(success.windows[0].pace_status.duration_seconds, Some(604_800)); + .unwrap(); + assert_eq!(enrich_calls.get(), 1); + assert_eq!(fallback.updated_at, fresh.updated_at); + assert_eq!(fallback.source, fresh.source); assert_eq!( - success.windows[0].pace_status.duration_source, - Some(DurationSource::Provider) + fallback.identity.as_ref().unwrap().plan.as_deref(), + Some("Fixture") ); - assert_eq!(success.windows[0].pace_status.complete_cycles, 3); + assert_eq!(fallback.windows.len(), 1); + assert_eq!(fallback.windows[0].pace_status.complete_cycles, 6); assert_eq!( - success.windows[0] + fallback.windows[0] .historical_pace .as_ref() .map(|pace| pace.expected_used_percent), - Some(27.5) - ); - let wire = serde_json::to_string(&success).unwrap(); - assert!(!wire.contains("accountScope")); - assert!(!wire.contains(String::from_utf8_lossy(marker).as_ref())); - assert!(!wire.contains(&opaque_scope)); - - let learning = UsageWindow::from_used_percent( - "Weekly".to_string(), - 10.0, - Some(now + chrono::Duration::days(7)), - now, - ) - .with_identity("billing.weekly.v1", Some("billing.weekly.v1".to_string())); - let unscoped = finalize_grok_snapshot_with( - Ok(agent_grok::GrokData { - identity: None, - account_scope: Err(AccountScopeError::NoTrustedEvidence), - windows: vec![learning], - }), - now, - enrich_snapshot, - ); - assert_eq!( - unscoped.windows[0].pace_status.state, - PaceState::Unavailable + Some(35.0) ); assert_eq!( - unscoped.windows[0].pace_status.reason.as_deref(), - Some("accountScope") - ); - - let callback_calls = std::cell::Cell::new(0); - let record_calls = std::cell::Cell::new(0); - let failed = finalize_grok_snapshot_with( - Err("Grok billing failed".to_string()), - now, - |snapshot, callback_now| { - callback_calls.set(callback_calls.get() + 1); - assert_eq!(callback_now, now); - assert!(snapshot.windows.is_empty()); - assert!(matches!( - snapshot.account_scope, - Err(AccountScopeError::NoTrustedEvidence) - )); - enrich_snapshot_with(snapshot, callback_now, |_, _, _| { - record_calls.set(record_calls.get() + 1); - Ok(Vec::new()) - }); - }, + fallback + .credits + .as_ref() + .and_then(|credits| credits.remaining), + Some(-2.5) ); - assert_eq!(callback_calls.get(), 1); - assert_eq!(record_calls.get(), 0); - assert_eq!( - failed.account_scope, + assert!(matches!( + fallback.account_scope, Err(AccountScopeError::NoTrustedEvidence) + )); + assert!(fallback.error.is_some()); + assert_eq!( + fallback + .transport_diagnostic + .map(|diagnostic| diagnostic.category), + Some(TransportCategory::Timeout) ); - assert!(failed.windows.is_empty()); - assert_eq!(failed.error.as_deref(), Some("Grok billing failed")); - scope_store.cleanup(); - } - #[test] - fn finalizes_grok_same_second_provider_reset_as_coherent_history_observation() { - let (scope, account_scope) = enrichment_scope("grok-finalize-normalized-reset"); - let now = DateTime::parse_from_rfc3339("2026-07-18T00:00:00.000500Z") + let cached = lock_last_good(&cache) + .entries + .get("codex") .unwrap() - .with_timezone(&Utc); - let reset = DateTime::parse_from_rfc3339("2026-07-18T00:00:00.000900Z") - .unwrap() - .with_timezone(&Utc); - let window = UsageWindow::from_used_percent( - "Weekly".to_string(), - 33.0, - Some(reset), - now, + .snapshot + .clone(); + assert!(cached.error.is_none()); + assert!(cached.transport_diagnostic.is_none()); + drop(cached); + + let fallback_again = apply_provider_outcome_with( + &cache, + "codex", + "oauth", + failure_at + chrono::Duration::minutes(1), + ProviderFetchOutcome::Failure(ProviderFetchFailure::transient( + "Codex usage request failed. Retrying automatically.", + Some(binding.clone()), + SafeTransportDiagnostic::server_error(503), + )), + |_| enrich_calls.set(enrich_calls.get() + 1), ) - .with_identity( - "billing.weekly.v1", - Some("billing.weekly.v1".to_string()), + .unwrap(); + assert_eq!(enrich_calls.get(), 1); + assert_eq!(fallback_again.updated_at, fresh.updated_at); + assert_eq!(fallback_again.windows[0].pace_status.complete_cycles, 6); + + let dns_fallback = apply_provider_outcome_with( + &cache, + "codex", + "oauth", + failure_at + chrono::Duration::minutes(2), + ProviderFetchOutcome::Failure(ProviderFetchFailure::transient( + "Codex usage request failed. Retrying automatically.", + Some(binding), + SafeTransportDiagnostic::from_facts(TransportErrorFacts { + is_timeout: false, + is_connect: true, + is_dns: true, + is_tls: false, + phase: TransportPhase::Request, + raw_os_code: None, + }), + )), + |_| enrich_calls.set(enrich_calls.get() + 1), ) - .with_provider_duration_evidence( - now, - true, - Some(DurationEvidence::provider(now.timestamp() + 1, 604_800)), - ); - assert!(window.provider_reset_normalized); - - let record_calls = std::cell::Cell::new(0); - let snapshot = finalize_grok_snapshot_with( - Ok(agent_grok::GrokData { - identity: None, - account_scope: Ok(account_scope.clone()), - windows: vec![window], - }), - now, - |snapshot, callback_now| { - enrich_snapshot_with( - snapshot, - callback_now, - |_, observations, _| { - record_calls.set(record_calls.get() + 1); - assert_eq!(observations.len(), 1); - assert_eq!(observations[0].reset_at, Some(now.timestamp() + 1)); - assert_eq!(observations[0].used_percent, 33.0); - assert_eq!( - observations[0].provider, - Some(DurationEvidence::provider(now.timestamp() + 1, 604_800)) - ); - Ok(vec![Ok(( - HistoryOutcome::Ready { - duration_seconds: 604_800, - source: DurationSource::Provider, - sampled: true, - }, - None, - 1, - ))]) - }, - ); - }, - ); - assert_eq!(record_calls.get(), 1); + .unwrap(); + assert_eq!(enrich_calls.get(), 1); + assert_eq!(dns_fallback.updated_at, fresh.updated_at); + assert_eq!(dns_fallback.windows[0].pace_status.complete_cycles, 6); assert_eq!( - snapshot.windows[0].pace_status.state, - PaceState::LearningHistory + dns_fallback + .transport_diagnostic + .map(|diagnostic| diagnostic.category), + Some(TransportCategory::Dns) ); - assert_eq!(snapshot.windows[0].pace_status.reason, None); - assert!(snapshot.windows[0].provider_reset_normalized); - assert_eq!(snapshot.windows[0].pace_status.duration_seconds, Some(604_800)); - assert_eq!(snapshot.windows[0].pace_status.complete_cycles, 1); - assert!(snapshot.windows[0].historical_pace.is_none()); - let wire = serde_json::to_value(&snapshot.windows[0]).unwrap(); - assert_eq!(wire["resetsAt"], "2026-07-18T00:00:00.000Z"); scope.cleanup(); } #[test] - fn finalizes_antigravity_success_and_enriches_once() { - let (scope, account_scope) = enrichment_scope("antigravity-finalize-success"); - let now = DateTime::parse_from_rfc3339("2026-07-10T00:00:00.000500Z") - .unwrap() - .with_timezone(&Utc); - let reset = DateTime::parse_from_rfc3339("2026-07-10T00:00:00.000900Z") - .unwrap() - .with_timezone(&Utc); - let window = UsageWindow::from_used_percent( - "Gemini".to_string(), - 25.0, - Some(reset), - now, - ) - .with_identity( - "model.gemini.v1", - Some("model.gemini.v1".to_string()), - ) - .with_observed_duration_evidence(now, true); - assert_eq!(window.pace_status.state, PaceState::LearningDuration); - assert_eq!(window.pace_status.duration_source, Some(DurationSource::Observed)); + fn last_good_mismatch_unbound_terminal_and_absent_clear_cache() { + let scope = TestRefreshScope::new("codex", "last-good-clear"); + let scope_a = scope + .resolve_current("fixture", "account-a", b"marker-a") + .unwrap(); + let scope_b = scope + .resolve_current("fixture", "account-b", b"marker-b") + .unwrap(); + let binding_a = ProviderCacheBinding::primary(scope_a.clone()); + let binding_b = ProviderCacheBinding::primary(scope_b); + let now = Utc.timestamp_opt(1_700_000_000, 0).single().unwrap(); - let fetched = agent_antigravity::Fetched { - source: "cli".to_string(), - identity: Some(AgentIdentity { - email: Some("fixture@example.com".to_string()), - plan: Some("Pro".to_string()), - }), - account_scope: Ok(account_scope.clone()), - windows: vec![window], - }; - let callback_calls = std::cell::Cell::new(0); - let snapshot = finalize_antigravity_snapshot_with( - Ok(fetched), + for failure in [ + ProviderFetchFailure::transient("mismatch", Some(binding_b), timeout_diagnostic()), + ProviderFetchFailure::transient("unbound", None, timeout_diagnostic()), + ProviderFetchFailure::terminal("terminal"), + ] { + let cache = Mutex::new(ProviderLastGoodCache::default()); + apply_provider_outcome_with( + &cache, + "codex", + "oauth", + now, + ProviderFetchOutcome::Success { + snapshot: cache_test_snapshot("codex", Ok(scope_a.clone()), now), + cache_binding: Some(binding_a.clone()), + }, + |_| {}, + ); + let result = apply_provider_outcome_with( + &cache, + "codex", + "oauth", + now + chrono::Duration::seconds(1), + ProviderFetchOutcome::Failure(failure), + |_| panic!("failure must not enrich"), + ) + .unwrap(); + assert!(result.windows.is_empty()); + assert!(!lock_last_good(&cache).entries.contains_key("codex")); + } + + let cache = Mutex::new(ProviderLastGoodCache::default()); + apply_provider_outcome_with( + &cache, + "codex", + "oauth", now, - |snapshot, callback_now| { - callback_calls.set(callback_calls.get() + 1); - assert_eq!(callback_now, now); - assert_eq!(snapshot.client_id, "antigravity"); - assert_eq!(snapshot.source, "cli"); - assert_eq!(snapshot.account_scope.as_ref().unwrap(), &account_scope); - enrich_snapshot_with( - snapshot, - callback_now, - |active, observations, history_now| { - assert_eq!( - active, - &[SeriesKey::new( - "antigravity", - account_scope.as_str(), - "model.gemini.v1", - )] - ); - assert_eq!(history_now, now.timestamp()); - assert_eq!(observations.len(), 1); - assert_eq!( - observations[0].key, - SeriesKey::new( - "antigravity", - account_scope.as_str(), - "model.gemini.v1", - ) - ); - assert_eq!(observations[0].reset_at, Some(now.timestamp() + 1)); - assert_eq!(observations[0].used_percent, 25.0); - Ok(vec![Ok((HistoryOutcome::LearningDuration, None, 0))]) - }, - ); + ProviderFetchOutcome::Success { + snapshot: cache_test_snapshot("codex", Ok(scope_a), now), + cache_binding: Some(binding_a), }, + |_| {}, ); - - assert_eq!(callback_calls.get(), 1); - assert_eq!(snapshot.windows.len(), 1); - assert_eq!( - snapshot.windows[0].pace_status.state, - PaceState::LearningDuration - ); - assert_eq!( - snapshot.windows[0].pace_status.duration_source, - Some(DurationSource::Observed) - ); - assert!(snapshot.error.is_none()); + assert!(apply_provider_outcome_with( + &cache, + "codex", + "oauth", + now, + ProviderFetchOutcome::Absent, + |_| panic!("absent must not enrich"), + ) + .is_none()); + assert!(!lock_last_good(&cache).entries.contains_key("codex")); scope.cleanup(); } #[test] - fn finalizes_antigravity_error_and_skips_history_recording() { + fn uncacheable_or_invalid_success_clears_prior_last_good() { + let scope = TestRefreshScope::new("antigravity", "last-good-uncacheable"); + let account_scope = scope + .resolve_current("fixture", "account-a", b"marker-a") + .unwrap(); + let binding = ProviderCacheBinding::primary(account_scope.clone()); let now = Utc.timestamp_opt(1_700_000_000, 0).single().unwrap(); - let callback_calls = std::cell::Cell::new(0); - let record_calls = std::cell::Cell::new(0); - let snapshot = finalize_antigravity_snapshot_with( - Err("synthetic Antigravity failure".to_string()), + + let cache = Mutex::new(ProviderLastGoodCache::default()); + apply_provider_outcome_with( + &cache, + "antigravity", + "oauth", now, - |snapshot, callback_now| { - callback_calls.set(callback_calls.get() + 1); - assert_eq!(callback_now, now); - assert!(snapshot.windows.is_empty()); - assert!(matches!( - snapshot.account_scope, - Err(AccountScopeError::NoTrustedEvidence) - )); - enrich_snapshot_with(snapshot, callback_now, |_, _, _| { - record_calls.set(record_calls.get() + 1); - Ok(Vec::new()) - }); + ProviderFetchOutcome::Success { + snapshot: cache_test_snapshot("antigravity", Ok(account_scope.clone()), now), + cache_binding: Some(binding.clone()), }, + |_| {}, ); + let anonymous = apply_provider_outcome_with( + &cache, + "antigravity", + "local", + now, + ProviderFetchOutcome::Success { + snapshot: cache_test_snapshot( + "antigravity", + Err(AccountScopeError::NoTrustedEvidence), + now, + ), + cache_binding: None, + }, + |_| {}, + ) + .unwrap(); + assert_eq!(anonymous.windows.len(), 1); + assert!(!lock_last_good(&cache).entries.contains_key("antigravity")); - assert_eq!(callback_calls.get(), 1); - assert_eq!(record_calls.get(), 0); - assert_eq!(snapshot.error.as_deref(), Some("synthetic Antigravity failure")); - assert!(snapshot.windows.is_empty()); - } - - #[test] - fn exact_reset_normalization_fails_closed_and_keeps_wire_milliseconds() { - let (scope, account_scope) = enrichment_scope("antigravity-exact-reset"); - let now = DateTime::parse_from_rfc3339("2026-07-10T00:00:00.000500Z") - .unwrap() - .with_timezone(&Utc); - let future = DateTime::parse_from_rfc3339("2026-07-10T00:00:00.000900Z") - .unwrap() - .with_timezone(&Utc); - let future_window = UsageWindow::from_used_percent( - "Future".to_string(), - 20.0, - Some(future), + apply_provider_outcome_with( + &cache, + "antigravity", + "oauth", + now, + ProviderFetchOutcome::Success { + snapshot: cache_test_snapshot("antigravity", Ok(account_scope.clone()), now), + cache_binding: Some(binding.clone()), + }, + |_| {}, + ); + let mut empty = cache_test_snapshot("antigravity", Ok(account_scope.clone()), now); + empty.windows.clear(); + let live_empty = apply_provider_outcome_with( + &cache, + "antigravity", + "oauth", now, + ProviderFetchOutcome::Success { + snapshot: empty, + cache_binding: Some(binding.clone()), + }, + |_| {}, ) - .with_identity("future.v1", Some("future.v1".to_string())); - let mut future_snapshot = enrichment_snapshot(Ok(account_scope.clone()), vec![future_window]); - enrich_snapshot_with( - &mut future_snapshot, + .unwrap(); + assert!(live_empty.windows.is_empty()); + assert!(!lock_last_good(&cache).entries.contains_key("antigravity")); + + apply_provider_outcome_with( + &cache, + "antigravity", + "oauth", now, - |_, observations, history_now| { - assert_eq!(history_now, now.timestamp()); - assert_eq!(observations.len(), 1); - assert_eq!(observations[0].reset_at, Some(now.timestamp() + 1)); - Ok(vec![Ok((HistoryOutcome::LearningDuration, None, 0))]) + ProviderFetchOutcome::Success { + snapshot: cache_test_snapshot("antigravity", Ok(account_scope), now), + cache_binding: Some(binding), }, + |_| {}, ); - assert_eq!( - future_snapshot.windows[0].pace_status.state, - PaceState::LearningDuration - ); - assert_eq!( - future_snapshot.windows[0].pace_status.duration_source, - Some(DurationSource::Observed) - ); - let wire = serde_json::to_value(&future_snapshot.windows[0]).unwrap(); - assert_eq!( - wire["resetsAt"], - "2026-07-10T00:00:00.000Z" - ); - let serialized = serde_json::to_string(&future_snapshot.windows[0]).unwrap(); - assert!(!serialized.contains("reset_at_evidence")); - assert!(!serialized.contains("resetAtEvidence")); - assert!(!serialized.contains("000900Z")); - - let past = DateTime::parse_from_rfc3339("2026-07-10T00:00:00.000400Z") - .unwrap() - .with_timezone(&Utc); - let past_window = UsageWindow::from_used_percent( - "Past".to_string(), - 20.0, - Some(past), + let enrich_calls = std::cell::Cell::new(0); + let invalid = apply_provider_outcome_with( + &cache, + "antigravity", + "oauth", now, + ProviderFetchOutcome::Success { + snapshot: cache_test_snapshot( + "antigravity", + Err(AccountScopeError::MetadataRead), + now, + ), + cache_binding: None, + }, + |_| enrich_calls.set(enrich_calls.get() + 1), ) - .with_identity("past.v1", Some("past.v1".to_string())) - .with_observed_duration_evidence(now, true); - assert_eq!(past_window.pace_reason_for_test(), Some("invalidEvidence")); - assert!(past_window.reset_at_evidence.is_some()); + .unwrap(); + assert!(invalid.windows.is_empty()); + assert_eq!(enrich_calls.get(), 0); + assert!(!lock_last_good(&cache).entries.contains_key("antigravity")); scope.cleanup(); } - #[test] - fn provider_duration_wrapper_accepts_subsecond_future_reset() { - let now = DateTime::parse_from_rfc3339("2026-07-10T00:00:00.100Z") - .unwrap() - .with_timezone(&Utc); - let reset = DateTime::parse_from_rfc3339("2026-07-10T00:00:00.900Z") - .unwrap() - .with_timezone(&Utc); - let window = UsageWindow::from_used_percent( - "Weekly".to_string(), - 20.0, - Some(reset), - now, - ) - .with_identity("weekly.v1", Some("weekly.v1".to_string())) - .with_provider_duration_evidence( - now, - true, - Some(DurationEvidence::provider(now.timestamp() + 1, 604_800)), - ); - let wire = serde_json::to_value(&window).unwrap(); - assert_eq!(wire["resetsAt"], "2026-07-10T00:00:00.900Z"); - assert_eq!(wire["paceStatus"]["state"], "learningHistory"); - assert_eq!(wire["paceStatus"]["durationSeconds"], 604_800); - assert_eq!(wire["paceStatus"]["durationSource"], "provider"); + #[tokio::test] + async fn status_before_body_enforces_terminal_transient_and_claude_exception() { + use std::cell::Cell; + + for status in [401, 403, 418, 429, 500, 503] { + let reads = Cell::new(0); + let result = read_response_body(status, false, || async { + reads.set(reads.get() + 1); + Ok("sensitive body".to_string()) + }) + .await; + assert_eq!(reads.get(), 0, "status {status} must not read body"); + match status { + 429 | 500 | 503 => { + assert!(matches!(result, Err(ResponseReadFailure::Transient(_)))) + } + _ => assert_eq!(result, Err(ResponseReadFailure::Terminal(status))), + } + } - let earlier_now = DateTime::parse_from_rfc3339("2026-07-09T00:00:00.100Z") - .unwrap() - .with_timezone(&Utc); - let mismatched = UsageWindow::from_used_percent( - "Mismatched".to_string(), - 20.0, - Some(reset), - earlier_now, - ) - .with_identity("mismatched.v1", Some("mismatched.v1".to_string())) - .with_provider_duration_evidence( - earlier_now, - true, - Some(DurationEvidence::provider(reset.timestamp() + 1, 604_800)), - ); - let mismatched_wire = serde_json::to_value(&mismatched).unwrap(); - assert_eq!(mismatched_wire["paceStatus"]["state"], "unavailable"); - assert_eq!( - mismatched_wire["paceStatus"]["reason"], - "invalidEvidence" - ); + let reads = Cell::new(0); + let body = read_response_body(200, false, || async { + reads.set(reads.get() + 1); + Ok("success".to_string()) + }) + .await + .unwrap(); + assert_eq!(body, "success"); + assert_eq!(reads.get(), 1); + + let reads = Cell::new(0); + let failure = read_response_body(200, false, || async { + reads.set(reads.get() + 1); + Err(TransportErrorFacts::synthetic( + false, + false, + TransportPhase::ResponseBody, + Some(54), + )) + }) + .await + .unwrap_err(); + assert_eq!(reads.get(), 1); + assert!(matches!( + failure, + ResponseReadFailure::Transient(SafeTransportDiagnostic { + category: TransportCategory::ConnectionReset, + os_code: Some(54), + .. + }) + )); + + let reads = Cell::new(0); + let body = read_response_body(403, true, || async { + reads.set(reads.get() + 1); + Ok("missing user:profile".to_string()) + }) + .await + .unwrap(); + assert_eq!(body, "missing user:profile"); + assert_eq!(reads.get(), 1); + + let reads = Cell::new(0); + let failure = read_response_body(403, true, || async { + reads.set(reads.get() + 1); + Err(TransportErrorFacts::synthetic( + true, + false, + TransportPhase::ResponseBody, + None, + )) + }) + .await + .unwrap_err(); + assert_eq!(reads.get(), 1); + assert_eq!(failure, ResponseReadFailure::Terminal(403)); } - #[test] - fn contract_duration_wrapper_uses_shared_resolver_and_fails_closed() { - let now = "2026-07-10T00:00:00Z".parse::>().unwrap(); - let reset = now + chrono::Duration::days(1); - let valid = UsageWindow::from_used_percent( - "Contract".to_string(), - 20.0, - Some(reset), - now, - ) - .with_identity("contract.v1", Some("contract.v1".to_string())) - .with_contract_duration_evidence( - now, - true, - DurationEvidence::contract(86_400), - ); - let wire = serde_json::to_value(&valid).unwrap(); - assert_eq!(wire["paceStatus"]["state"], "learningHistory"); - assert_eq!(wire["paceStatus"]["durationSource"], "contract"); - assert_eq!(wire["paceStatus"]["durationSeconds"], 86_400); - assert_eq!(wire["windowMinutes"], 1_440); - - let invalid = UsageWindow::from_used_percent( - "Invalid contract".to_string(), - 20.0, - Some(reset), - now, - ) - .with_identity("invalid-contract.v1", Some("invalid-contract.v1".to_string())) - .with_contract_duration_evidence(now, true, DurationEvidence::contract(0)); - let invalid_wire = serde_json::to_value(&invalid).unwrap(); - assert_eq!(invalid_wire["paceStatus"]["state"], "unavailable"); - assert_eq!(invalid_wire["paceStatus"]["reason"], "invalidEvidence"); - assert!(invalid_wire["paceStatus"].get("durationSeconds").is_none()); - assert!(invalid_wire.get("windowMinutes").is_none()); - assert!(invalid_wire.get("historicalPace").is_none()); + #[tokio::test] + async fn verified_binding_failure_prevents_every_provider_request() { + for provider in ["codex", "claude", "grok", "copilot", "antigravity"] { + let sends = std::cell::Cell::new(0); + let result: Result<(), &str> = + request_after_verified_binding(Err::<(), _>("scope unavailable"), |()| async { + sends.set(sends.get() + 1); + Ok(()) + }) + .await; + assert_eq!(result, Err("scope unavailable"), "{provider}"); + assert_eq!(sends.get(), 0, "{provider}"); + } } - #[test] - fn serializes_stage3a_pace_states_without_legacy_scalars() { - let now = Utc.timestamp_opt(1_700_000_000, 0).single().unwrap(); - let ready = UsageWindow::from_used_percent( - "Known".to_string(), - 20.0, - Some(now + chrono::Duration::hours(1)), - now, - ) - .with_identity("known.v1", Some("known.v1".to_string())); - let missing_reset = UsageWindow::from_used_percent("No reset".to_string(), 20.0, None, now) - .with_identity("no_reset.v1", Some("no_reset.v1".to_string())); - let unknown = UsageWindow::from_used_percent( - "Unknown".to_string(), - 20.0, - Some(now + chrono::Duration::hours(1)), - now, - ) - .with_identity("row.unknown.v1", None); + #[derive(Debug)] + struct SensitiveTestError(&'static str); - let values = [ready, missing_reset, unknown] - .into_iter() - .map(|window| serde_json::to_value(window).unwrap()) - .collect::>(); - assert_eq!(values[0]["cardId"], "known.v1"); - assert_eq!(values[0]["paceStatus"]["state"], "learningDuration"); - assert_eq!(values[1]["paceStatus"]["state"], "unavailable"); - assert_eq!(values[1]["paceStatus"]["reason"], "missingReset"); - assert_eq!(values[2]["paceStatus"]["state"], "unavailable"); - assert_eq!(values[2]["paceStatus"]["reason"], "windowIdentity"); - for value in values { - assert!(value.get("paceStatus").is_some()); - assert!(value.get("historicalPace").is_none()); - assert!(value.get("windowMinutes").is_none()); - assert!(value.get("historicalExpectedPercent").is_none()); - assert!(value.get("runOutProbability").is_none()); + impl std::fmt::Display for SensitiveTestError { + fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + formatter.write_str(self.0) } } - #[test] - fn observed_duration_evidence_preserves_reset_presence() { - let now = Utc.timestamp_opt(1_700_000_000, 0).single().unwrap(); - let future_reset = now + chrono::Duration::hours(1); - let past_reset = now - chrono::Duration::hours(1); - - let absent = UsageWindow::from_used_percent("Absent".to_string(), 20.0, None, now) - .with_identity("absent.v1", Some("absent.v1".to_string())) - .with_observed_duration_evidence(now, false); - assert_eq!(absent.pace_reason_for_test(), Some("missingReset")); - - let mut malformed = UsageWindow::from_used_percent( - "Malformed".to_string(), - 20.0, - Some(future_reset), - now, - ) - .with_identity("malformed.v1", Some("malformed.v1".to_string())); - malformed.resets_at = Some("bogus".to_string()); - malformed = malformed.with_observed_duration_evidence(now, true); - assert_eq!(malformed.pace_reason_for_test(), Some("invalidEvidence")); - assert!(malformed.resets_at.is_none()); - assert!(malformed.reset_at_evidence.is_none()); - - let past = UsageWindow::from_used_percent("Past".to_string(), 20.0, Some(past_reset), now) - .with_identity("past.v1", Some("past.v1".to_string())) - .with_observed_duration_evidence(now, true); - assert_eq!(past.pace_reason_for_test(), Some("invalidEvidence")); - - let future = UsageWindow::from_used_percent( - "Future".to_string(), - 20.0, - Some(future_reset), - now, - ) - .with_identity("future.v1", Some("future.v1".to_string())) - .with_observed_duration_evidence(now, true); - assert_eq!(future.pace_reason_for_test(), None); - assert_eq!(future.pace_status.state, PaceState::LearningDuration); - assert_eq!( - future.pace_status.duration_source, - Some(DurationSource::Observed) - ); + impl std::error::Error for SensitiveTestError {} - let subsecond_now = DateTime::parse_from_rfc3339("2026-07-10T00:00:00.000500Z") - .unwrap() - .with_timezone(&Utc); - let subsecond_reset = DateTime::parse_from_rfc3339("2026-07-10T00:00:00.000900Z") + #[derive(Debug)] + struct NestedTestError { + source: Box, + } + + impl NestedTestError { + fn new(source: impl std::error::Error + Send + Sync + 'static) -> Self { + Self { + source: Box::new(source), + } + } + } + + impl std::fmt::Display for NestedTestError { + fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + formatter.write_str("nested transport failure") + } + } + + impl std::error::Error for NestedTestError { + fn source(&self) -> Option<&(dyn std::error::Error + 'static)> { + Some(self.source.as_ref()) + } + } + + #[derive(Debug, Clone, Copy)] + struct InjectedDnsFailureResolver; + + impl reqwest::dns::Resolve for InjectedDnsFailureResolver { + fn resolve(&self, _name: reqwest::dns::Name) -> reqwest::dns::Resolving { + Box::pin(async { + Err(Box::new(DnsResolutionError::new(SensitiveTestError( + "token-secret user@example.invalid /private/credential/path", + ))) + as Box) + }) + } + } + + #[tokio::test] + async fn typed_gai_adapter_preserves_loopback_addresses_without_network_io() { + let name: reqwest::dns::Name = "127.0.0.1".parse().unwrap(); + let addresses = reqwest::dns::Resolve::resolve(&TypedGaiResolver, name) + .await .unwrap() - .with_timezone(&Utc); - let subsecond = UsageWindow::from_used_percent( - "Subsecond".to_string(), - 20.0, - Some(subsecond_reset), - subsecond_now, - ) - .with_identity("subsecond.v1", Some("subsecond.v1".to_string())) - .with_observed_duration_evidence(subsecond_now, true); - assert_eq!(subsecond.pace_status.state, PaceState::LearningDuration); - assert_eq!( - subsecond.pace_status.duration_source, - Some(DurationSource::Observed) - ); + .collect::>(); + assert!(!addresses.is_empty()); + assert!(addresses.iter().all(|address| address.ip().is_loopback())); + assert!(provider_http_client_builder().build().is_ok()); + } - let mut missing_identity = UsageWindow::from_used_percent( - "Missing identity".to_string(), - 20.0, - Some(future_reset), - now, - ) - .with_identity("row.missing.v1", None); - missing_identity.resets_at = Some("bogus".to_string()); - missing_identity = missing_identity.with_observed_duration_evidence(now, true); - assert_eq!(missing_identity.pace_reason_for_test(), Some("windowIdentity")); + #[tokio::test] + async fn injected_typed_dns_failure_is_classified_without_source_disclosure() { + let client = reqwest::Client::builder() + .no_proxy() + .dns_resolver(InjectedDnsFailureResolver) + .timeout(std::time::Duration::from_secs(2)) + .build() + .unwrap(); + let error = client + .get("http://account-123.example.invalid/private/path?token=token-secret") + .send() + .await + .unwrap_err(); + let diagnostic = SafeTransportDiagnostic::from_facts(TransportErrorFacts::from_reqwest( + &error, + TransportPhase::Request, + )); + assert_eq!(diagnostic.category, TransportCategory::Dns); + let wire = serde_json::to_string(&diagnostic).unwrap(); + assert_eq!(wire, r#"{"category":"dns"}"#); + for secret in [ + "token-secret", + "user@example.invalid", + "account-123", + "example.invalid", + "/private/path", + "/private/credential/path", + ] { + assert!(!wire.contains(secret)); + } } #[test] - fn rejects_malformed_usage_window_wire() { - let now = Utc.timestamp_opt(1_700_000_000, 0).single().unwrap(); - let mut window = UsageWindow::from_used_percent( - "Known".to_string(), - 20.0, - Some(now + chrono::Duration::hours(1)), - now, - ) - .with_identity("known.v1", Some("known.v1".to_string())); - window.card_id.clear(); - assert!(serde_json::to_value(&window).is_err()); - window.card_id = "known.v1".to_string(); - window.used_percent = f64::NAN; - assert!(serde_json::to_value(&window).is_err()); + fn nested_typed_sources_are_found_without_text_classification() { + let dns_error = NestedTestError::new(std::io::Error::other(DnsResolutionError::new( + SensitiveTestError("token-secret"), + ))); + let dns_facts = transport_source_facts(&dns_error); + assert!(dns_facts.is_dns); + assert!(!dns_facts.is_tls); + assert_eq!(dns_facts.raw_os_code, None); + + let tls_error = NestedTestError::new(std::io::Error::other(rustls::Error::General( + "token-secret".to_string(), + ))); + let tls_facts = transport_source_facts(&tls_error); + assert!(!tls_facts.is_dns); + assert!(tls_facts.is_tls); + assert_eq!(tls_facts.raw_os_code, None); + + let os_error = + NestedTestError::new(std::io::Error::other(std::io::Error::from_raw_os_error(61))); + assert_eq!(transport_source_facts(&os_error).raw_os_code, Some(61)); } - #[test] - fn serialized_duration_mirror_is_nested_and_state_coherent() { - let now = Utc.timestamp_opt(1_700_000_000, 0).single().unwrap(); - let window = map_window_with_identity( - "Session", - CodexWindow { - used_percent: 20.0, - reset_at: now.timestamp() + 3_600, - limit_window_seconds: 18_000, - }, - now, - "main.session.v1", - Some("main.session.v1".to_string()), + #[tokio::test] + async fn loopback_plaintext_on_tls_endpoint_is_classified_as_tls() { + use tokio::io::{AsyncReadExt, AsyncWriteExt}; + + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let address = listener.local_addr().unwrap(); + let server = tokio::spawn(async move { + let (mut stream, _) = listener.accept().await.unwrap(); + let mut client_hello = [0_u8; 1024]; + let _ = stream.read(&mut client_hello).await; + let _ = stream + .write_all(b"HTTP/1.1 200 OK\r\ncontent-length: 0\r\n\r\n") + .await; + tokio::time::sleep(std::time::Duration::from_millis(100)).await; + }); + let client = provider_http_client_builder() + .no_proxy() + .timeout(std::time::Duration::from_secs(3)) + .build() + .unwrap(); + let error = tokio::time::timeout( + std::time::Duration::from_secs(5), + client.get(format!("https://{address}/")).send(), ) - .unwrap(); - let wire = serde_json::to_value(&window).unwrap(); - assert_eq!(wire["paceStatus"]["durationSeconds"], 18_000); - assert_eq!(wire["paceStatus"]["durationSource"], "provider"); - assert_eq!(wire["windowMinutes"], 300); + .await + .expect("loopback TLS request timed out") + .unwrap_err(); + let facts = TransportErrorFacts::from_reqwest(&error, TransportPhase::Request); + assert!(facts.is_tls); assert_eq!( - wire["windowMinutes"].as_i64(), - wire["paceStatus"]["durationSeconds"] - .as_i64() - .map(|seconds| seconds / 60) + SafeTransportDiagnostic::from_facts(facts).category, + TransportCategory::Tls ); - - let mut contradictory = window; - contradictory.pace_status.state = PaceState::Unavailable; - contradictory.pace_status.reason = Some("invalidEvidence".to_string()); - assert!(serde_json::to_value(&contradictory).is_err()); - contradictory.pace_status.state = PaceState::LearningHistory; - contradictory.pace_status.reason = None; - contradictory.pace_status.duration_source = None; - assert!(serde_json::to_value(&contradictory).is_err()); - contradictory.pace_status.duration_source = Some(DurationSource::Provider); - contradictory.pace_status.duration_seconds = Some(604_800); - assert!(serde_json::to_value(&contradictory).is_err()); + server.await.unwrap(); } #[test] - fn retain_unique_windows_drops_later_card_or_window_identity() { - let now = Utc.timestamp_opt(1_700_000_000, 0).single().unwrap(); - let first = UsageWindow::from_used_percent( - "First".to_string(), - 20.0, - Some(now + chrono::Duration::hours(1)), - now, - ) - .with_identity("card.a.v1", Some("window.a.v1".to_string())); - let duplicate_card = UsageWindow::from_used_percent( - "Later card".to_string(), - 30.0, - Some(now + chrono::Duration::hours(1)), - now, - ) - .with_identity("card.a.v1", Some("window.b.v1".to_string())); - let duplicate_key = UsageWindow::from_used_percent( - "Later key".to_string(), - 40.0, - Some(now + chrono::Duration::hours(1)), - now, - ) - .with_identity("card.c.v1", Some("window.a.v1".to_string())); - let mut windows = vec![first, duplicate_card, duplicate_key]; - retain_unique_windows(&mut windows); - assert_eq!(windows.len(), 1); - assert_eq!(windows[0].label_for_test(), "First"); - } + fn transport_diagnostic_precedence_and_generic_categories_are_stable() { + let facts = + |is_timeout, is_connect, is_dns, is_tls, phase, raw_os_code| TransportErrorFacts { + is_timeout, + is_connect, + is_dns, + is_tls, + phase, + raw_os_code, + }; + let category = |facts| SafeTransportDiagnostic::from_facts(facts).category; - #[test] - fn enrichment_scope_failure_preserves_identity_and_non_recurring_rows() { - let (scope, _) = enrichment_scope("enrichment-scope-failure"); - let now = Utc.timestamp_opt(1_700_000_000, 0).single().unwrap(); - let stable = enrichment_window( - now, - "stable.v1", - "stable.v1", - 20.0, - Some(DurationSource::Contract), + assert_eq!( + category(facts( + true, + true, + true, + true, + TransportPhase::Request, + Some(61), + )), + TransportCategory::Timeout ); - let identity = UsageWindow::from_used_percent( - "Identity".to_string(), - 30.0, - Some(now + chrono::Duration::days(1)), - now, - ) - .with_identity("identity.v1", None); - let non_recurring = - UsageWindow::from_used_percent("Non-recurring".to_string(), 40.0, None, now) - .with_identity("non-recurring.v1", Some("non-recurring.v1".to_string())) - .with_unavailable_reason("nonRecurring"); - let mut snapshot = enrichment_snapshot( - Err(AccountScopeError::MetadataWrite), - vec![stable, identity, non_recurring], + assert_eq!( + category(facts( + false, + true, + true, + true, + TransportPhase::Request, + Some(61), + )), + TransportCategory::ConnectionRefused ); - let calls = std::cell::Cell::new(0); - - enrich_snapshot_with(&mut snapshot, now, |_, _, _| { - calls.set(calls.get() + 1); - Ok(Vec::new()) - }); - - assert_eq!(calls.get(), 0); assert_eq!( - snapshot.windows[0].pace_reason_for_test(), - Some("accountScope") + category(facts( + false, + true, + true, + true, + TransportPhase::Request, + Some(54), + )), + TransportCategory::ConnectionReset ); assert_eq!( - snapshot.windows[1].pace_reason_for_test(), - Some("windowIdentity") + category(facts( + false, + true, + true, + true, + TransportPhase::Request, + None, + )), + TransportCategory::Dns + ); + assert_eq!( + category(facts( + false, + true, + false, + true, + TransportPhase::Request, + None, + )), + TransportCategory::Tls ); assert_eq!( - snapshot.windows[2].pace_reason_for_test(), - Some("nonRecurring") + category(facts( + false, + true, + false, + false, + TransportPhase::Request, + None, + )), + TransportCategory::Connect + ); + assert_eq!( + category(facts( + false, + false, + false, + false, + TransportPhase::Request, + None, + )), + TransportCategory::Request + ); + assert_eq!( + category(facts( + false, + false, + false, + false, + TransportPhase::ResponseBody, + None, + )), + TransportCategory::ResponseBody ); - assert!(serde_json::to_value(&snapshot).is_ok()); - scope.cleanup(); } #[test] - fn enrichment_filters_duplicate_card_and_window_keys_before_batching() { - let (scope, account_scope) = enrichment_scope("enrichment-duplicates"); - let now = Utc.timestamp_opt(1_700_000_000, 0).single().unwrap(); - let mut snapshot = enrichment_snapshot( - Ok(account_scope), - vec![ - enrichment_window(now, "shared-card.v1", "first.v1", 10.0, None), - enrichment_window(now, "second-card.v1", "first.v1", 20.0, None), - enrichment_window(now, "shared-card.v1", "third.v1", 30.0, None), - ], + fn structured_transport_diagnostic_serializes_only_allowlisted_fields() { + let diagnostic = SafeTransportDiagnostic::from_facts(TransportErrorFacts::synthetic( + false, + true, + TransportPhase::Request, + Some(61), + )); + let wire = serde_json::to_string(&diagnostic).unwrap(); + assert_eq!(wire, r#"{"category":"connectionRefused","osCode":61}"#); + for secret in [ + "token-secret", + "Authorization", + "https://example.invalid/path?query=secret#fragment", + "user@example.invalid", + "account-123", + "/private/credential/path", + ] { + assert!(!wire.contains(secret)); + } + assert_eq!( + serde_json::to_value(SafeTransportDiagnostic::rate_limited(429)).unwrap(), + serde_json::json!({ "category": "rateLimited", "status": 429 }) ); - - enrich_snapshot_with( - &mut snapshot, - now, - |active, observations, batch_now| { - assert_eq!(batch_now, now.timestamp()); - assert_eq!(active.len(), 1); - assert_eq!(active[0].window_key, "first.v1"); - assert_eq!(observations.len(), 1); - assert_eq!(observations[0].used_percent, 10.0); - Ok(vec![Ok((HistoryOutcome::LearningDuration, None, 0))]) - }, + assert_eq!( + serde_json::to_value(SafeTransportDiagnostic::server_error(503)).unwrap(), + serde_json::json!({ "category": "serverError", "status": 503 }) ); + } + + #[test] + fn codex_credit_is_usable_only_when_balance_is_finite() { + let credits = |balance, unlimited| CodexCredits { balance, unlimited }; - assert_eq!(snapshot.windows.len(), 1); - assert_eq!(snapshot.windows[0].card_id_for_test(), "shared-card.v1"); assert_eq!( - snapshot.windows[0].pace_status.duration_source, - Some(DurationSource::Observed) + finite_codex_balance(Some(&credits(Some(12.5), false))), + Some(12.5) ); - assert!(serde_json::to_value(&snapshot).is_ok()); - scope.cleanup(); + assert_eq!( + finite_codex_balance(Some(&credits(Some(-2.5), false))), + Some(-2.5), + "finite negative balances preserve the existing present-credit semantics" + ); + assert_eq!( + finite_codex_balance(Some(&credits(Some(0.0), false))), + Some(0.0) + ); + assert_eq!( + finite_codex_balance(Some(&credits(Some(f64::NAN), false))), + None + ); + assert_eq!( + finite_codex_balance(Some(&credits(Some(f64::INFINITY), false))), + None + ); + assert_eq!( + finite_codex_balance(Some(&credits(None, true))), + None, + "unlimited without a balance is not usable credit" + ); + assert_eq!(finite_codex_balance(None), None); } #[test] - fn enrichment_builds_active_keys_observations_and_coherent_results() { - let (scope, account_scope) = enrichment_scope("enrichment-batch-map"); - let expected_scope = account_scope.as_str().to_string(); + fn maps_codex_primary_and_secondary_windows() { let now = Utc.timestamp_opt(1_700_000_000, 0).single().unwrap(); - let reset = now + chrono::Duration::days(1); - let mut invalid_reset = - enrichment_window(now, "invalid-reset.v1", "invalid-reset.v1", 50.0, None); - invalid_reset.resets_at = Some("not-rfc3339".to_string()); - invalid_reset.used_percent = f64::NAN; - invalid_reset.remaining_percent = f64::NAN; - let expired = UsageWindow::from_used_percent( - "Expired".to_string(), - 60.0, - Some(now - chrono::Duration::seconds(1)), - now, - ) - .with_identity("expired.v1", Some("expired.v1".to_string())); - let mut invalid_percent = - enrichment_window(now, "invalid-percent.v1", "invalid-percent.v1", 70.0, None); - invalid_percent.used_percent = f64::NAN; - invalid_percent.remaining_percent = f64::NAN; - let missing = UsageWindow::from_used_percent("Missing".to_string(), 30.0, None, now) - .with_identity("missing.v1", Some("missing.v1".to_string())); - let non_recurring = - UsageWindow::from_used_percent("Non-recurring".to_string(), 40.0, None, now) - .with_identity("non-recurring.v1", Some("non-recurring.v1".to_string())) - .with_unavailable_reason("nonRecurring"); - let mut snapshot = enrichment_snapshot( - Ok(account_scope), - vec![ - enrichment_window( - now, - "provider.v1", - "provider.v1", - 10.0, - Some(DurationSource::Provider), - ), - enrichment_window( - now, - "contract.v1", - "contract.v1", - 20.0, - Some(DurationSource::Contract), - ), - enrichment_window(now, "observed.v1", "observed.v1", 25.0, None), - missing, - non_recurring, - invalid_reset, - expired, - invalid_percent, - ], - ); - - enrich_snapshot_with(&mut snapshot, now, |active, observations, _| { - assert_eq!( - active, - &[ - SeriesKey::new("fixture", &expected_scope, "provider.v1"), - SeriesKey::new("fixture", &expected_scope, "contract.v1"), - SeriesKey::new("fixture", &expected_scope, "observed.v1"), - SeriesKey::new("fixture", &expected_scope, "missing.v1"), - SeriesKey::new("fixture", &expected_scope, "invalid-reset.v1"), - SeriesKey::new("fixture", &expected_scope, "expired.v1"), - SeriesKey::new("fixture", &expected_scope, "invalid-percent.v1"), - ] - ); - assert_eq!(observations.len(), 3); - assert_eq!(observations[0].reset_at, Some(reset.timestamp())); - assert_eq!( - observations[0].provider, - Some(DurationEvidence::provider(reset.timestamp(), 86_400)) - ); - assert_eq!(observations[0].contract, None); - assert_eq!(observations[1].provider, None); - assert_eq!( - observations[1].contract, - Some(DurationEvidence::contract(86_400)) - ); - assert_eq!(observations[2].provider, None); - assert_eq!(observations[2].contract, None); - Ok(vec![ - Ok(( - HistoryOutcome::Ready { - duration_seconds: 86_400, - source: DurationSource::Provider, - sampled: true, - }, - None, - 2, - )), - Ok(( - HistoryOutcome::Ready { - duration_seconds: 86_400, - source: DurationSource::Contract, - sampled: true, - }, - Some(HistoricalPace { - expected_percent: 42.0, - eta_seconds: Some(900.0), - will_last_to_reset: false, - run_out_probability: None, - }), - 4, - )), - Ok((HistoryOutcome::LearningDuration, None, 0)), - ]) - }); + let rate_limit = CodexRateLimit { + primary_window: Some(CodexWindow { + used_percent: 8.0, + reset_at: 1_700_005_400, + limit_window_seconds: 18_000, + }), + secondary_window: Some(CodexWindow { + used_percent: 35.0, + reset_at: 1_700_172_800, + limit_window_seconds: 604_800, + }), + }; + let windows = codex_windows(Some(&rate_limit), None, now); + assert_eq!(windows.len(), 2); + assert_eq!(windows[0].label, "Session"); + assert_eq!(windows[0].remaining_percent, 92.0); + assert_eq!(windows[0].window_minutes, Some(300)); + assert_eq!(windows[1].label, "Weekly"); + assert_eq!(windows[1].remaining_percent, 65.0); + assert_eq!(windows[1].window_minutes, Some(10_080)); + } + #[test] + fn stage0_freezes_codex_duration_roles_and_unknown_window_baseline() { + let now = Utc.timestamp_opt(1_700_000_000, 0).single().unwrap(); + let reversed = CodexRateLimit { + primary_window: Some(CodexWindow { + used_percent: 35.0, + reset_at: 1_700_172_800, + limit_window_seconds: 604_800, + }), + secondary_window: Some(CodexWindow { + used_percent: 8.0, + reset_at: 1_700_005_400, + limit_window_seconds: 18_000, + }), + }; + let windows = codex_windows(Some(&reversed), None, now); + assert_eq!(windows.len(), 2); + assert_eq!(windows[0].label, "Session", "codex.main.18000.session"); + assert_eq!(windows[0].card_id, "main.session.v1"); + assert_eq!(windows[0].window_key.as_deref(), Some("main.session.v1")); + assert_eq!(windows[0].window_minutes, Some(300)); + assert_eq!(windows[1].label, "Weekly", "codex.main.604800.weekly"); + assert_eq!(windows[1].card_id, "main.weekly.v1"); + assert_eq!(windows[1].window_key.as_deref(), Some("main.weekly.v1")); + assert_eq!(windows[1].window_minutes, Some(10_080)); + + let unknown_rate_limit = CodexRateLimit { + primary_window: Some(CodexWindow { + used_percent: 10.0, + reset_at: now.timestamp() + 3_600, + limit_window_seconds: 3_600, + }), + secondary_window: None, + }; + let unknown = codex_windows(Some(&unknown_rate_limit), None, now); + assert_eq!(unknown.len(), 1); + let unknown = &unknown[0]; + assert_eq!(unknown.card_id, "row.main.primary.v1"); + assert_eq!(unknown.window_key, None); + assert_eq!(unknown.window_minutes, None); + assert_eq!(unknown.pace_status.state, PaceState::Unavailable); assert_eq!( - snapshot.windows[0].pace_status.state, - PaceState::LearningHistory - ); - assert_eq!(snapshot.windows[0].pace_status.complete_cycles, 2); - assert_eq!(snapshot.windows[1].pace_status.state, PaceState::Available); - assert_eq!(snapshot.windows[1].pace_status.complete_cycles, 4); - let historical = snapshot.windows[1].historical_pace.as_ref().unwrap(); - assert_eq!(historical.expected_used_percent, 42.0); - assert_eq!(historical.eta_seconds, Some(900.0)); - assert!(!historical.will_last_to_reset); - assert_eq!(historical.run_out_probability, None); - assert_eq!( - snapshot.windows[2].pace_status.state, - PaceState::LearningDuration - ); - assert_eq!( - snapshot.windows[2].pace_status.duration_source, - Some(DurationSource::Observed) - ); - assert_eq!( - snapshot.windows[3].pace_reason_for_test(), - Some("missingReset") - ); - assert_eq!( - snapshot.windows[4].pace_reason_for_test(), - Some("nonRecurring") + unknown.pace_status.reason.as_deref(), + Some("windowIdentity") ); - for window in &snapshot.windows[5..] { - assert_eq!(window.pace_reason_for_test(), Some("invalidEvidence")); - } - let wire = serde_json::to_value(&snapshot).unwrap(); - assert_eq!(wire["windows"][0]["windowMinutes"], 1_440); - assert_eq!(wire["windows"][1]["windowMinutes"], 1_440); - assert!(wire["windows"][2].get("windowMinutes").is_none()); - scope.cleanup(); + let wire = serde_json::to_value(unknown).unwrap(); + assert_eq!(wire["cardId"], "row.main.primary.v1"); + assert!(wire["paceStatus"].get("windowKey").is_none()); + assert_eq!(wire["paceStatus"]["state"], "unavailable"); + assert_eq!(wire["paceStatus"]["reason"], "windowIdentity"); } #[test] - fn enrichment_maps_unavailable_and_rejects_contradictory_results() { - let (scope, account_scope) = enrichment_scope("enrichment-result-validation"); + fn serializes_nested_historical_pace_without_legacy_scalars() { let now = Utc.timestamp_opt(1_700_000_000, 0).single().unwrap(); - let mut snapshot = enrichment_snapshot( - Ok(account_scope), - vec![ - enrichment_window(now, "missing.v1", "missing.v1", 10.0, None), - enrichment_window(now, "invalid.v1", "invalid.v1", 20.0, None), - enrichment_window( - now, - "learning-conflict.v1", - "learning-conflict.v1", - 30.0, - Some(DurationSource::Contract), - ), - enrichment_window( - now, - "historical-conflict.v1", - "historical-conflict.v1", - 40.0, - Some(DurationSource::Contract), - ), - enrichment_window( - now, - "source-conflict.v1", - "source-conflict.v1", - 50.0, - Some(DurationSource::Contract), - ), - enrichment_window( - now, - "unavailable-conflict.v1", - "unavailable-conflict.v1", - 60.0, - None, - ), - enrichment_window( - now, - "nonfinite-history.v1", - "nonfinite-history.v1", - 70.0, - Some(DurationSource::Contract), - ), - ], + let mut window = UsageWindow::from_used_percent( + "Weekly".to_string(), + 60.0, + Some(now + chrono::Duration::hours(12)), + now, + Some(10_080), + ) + .with_identity( + "weekly.v1", + Some("weekly.v1".to_string()), + None, + Some(DurationEvidence::contract(10_080 * 60)), ); - let historical = HistoricalPace { - expected_percent: 42.0, - eta_seconds: Some(900.0), + window.pace_status.state = PaceState::Available; + window.historical_pace = Some(HistoricalPacePayload { + expected_used_percent: 55.0, + eta_seconds: Some(3_600.0), will_last_to_reset: false, - run_out_probability: None, + run_out_probability: Some(0.8), + }); + + let value = serde_json::to_value(&window).unwrap(); + assert!(value.get("historicalPace").is_some()); + assert!(value.get("historicalExpectedPercent").is_none()); + assert!(value.get("runOutProbability").is_none()); + let historical = value.get("historicalPace").unwrap(); + assert_eq!(historical["expectedUsedPercent"], 55.0); + assert_eq!(historical["etaSeconds"], 3_600.0); + assert_eq!(historical["willLastToReset"], false); + assert_eq!(historical["runOutProbability"], 0.8); + } + + #[test] + fn stage1_credential_markers_follow_the_frozen_provider_routes() { + let slot = CredentialSlot { + semantic_source: "fixture", + canonical_location: "fixture".to_string(), }; + let codex = CodexCredentials { + access_token: "codex-access".to_string(), + refresh_token: Some("codex-refresh".to_string()), + id_token: None, + account_id: None, + last_refresh: None, + auth_path: PathBuf::new(), + raw_json: Value::Null, + scope_slot: slot.clone(), + }; + assert_eq!(codex.scope_marker(), b"codex-refresh"); + let mut codex_access_only = codex.clone(); + codex_access_only.refresh_token = None; + assert_eq!(codex_access_only.scope_marker(), b"codex-access"); - enrich_snapshot_with(&mut snapshot, now, |_, observations, _| { - assert_eq!(observations.len(), 7); - Ok(vec![ - Ok(( - HistoryOutcome::Unavailable(DurationUnavailableReason::MissingReset), - None, - 0, - )), - Ok(( - HistoryOutcome::Unavailable(DurationUnavailableReason::InvalidEvidence), - None, - 0, - )), - Ok((HistoryOutcome::LearningDuration, None, 0)), - Ok(( - HistoryOutcome::Ready { - duration_seconds: 86_400, - source: DurationSource::Contract, - sampled: true, - }, - Some(HistoricalPace { - expected_percent: f64::NAN, - ..historical.clone() - }), - 3, - )), - Ok(( - HistoryOutcome::Ready { - duration_seconds: 86_400, - source: DurationSource::Provider, - sampled: true, - }, - None, - 2, - )), - Ok(( - HistoryOutcome::Unavailable(DurationUnavailableReason::InvalidEvidence), - None, - 3, - )), - Ok(( - HistoryOutcome::Ready { - duration_seconds: 86_400, - source: DurationSource::Contract, - sampled: true, - }, - Some(HistoricalPace { - expected_percent: f64::NAN, - ..historical.clone() - }), - 3, - )), - ]) - }); + let claude_login = ClaudeCredentials { + access_token: "claude-access".to_string(), + refresh_token: Some("claude-refresh".to_string()), + expires_at: None, + scopes: Vec::new(), + rate_limit_tier: None, + subscription_type: None, + source: ClaudeCredentialSource::File, + raw_root: None, + keychain_account: None, + scope_slot: slot.clone(), + }; + assert_eq!( + claude_login.scope_marker(), + Some(b"claude-refresh".as_slice()) + ); + let mut login_without_refresh = claude_login.clone(); + login_without_refresh.refresh_token = None; + assert_eq!(login_without_refresh.scope_marker(), None); - assert_eq!(snapshot.windows[0].pace_reason_for_test(), Some("history")); + let claude_setup = ClaudeCredentials { + source: ClaudeCredentialSource::Environment, + scope_slot: slot, + ..login_without_refresh + }; assert_eq!( - snapshot.windows[1].pace_reason_for_test(), - Some("invalidEvidence") + claude_setup.scope_marker(), + Some(b"claude-access".as_slice()) ); - for window in &snapshot.windows[2..] { - assert_eq!(window.pace_reason_for_test(), Some("history")); - } - assert!(serde_json::to_value(&snapshot).is_ok()); - scope.cleanup(); } #[test] - fn enrichment_maps_global_row_and_count_failures_only_to_observations() { - let (scope, account_scope) = enrichment_scope("enrichment-errors"); - let now = Utc.timestamp_opt(1_700_000_000, 0).single().unwrap(); + fn provider_cache_binding_requires_structural_exact_match() { + let scope_store = TestRefreshScope::new("codex", "binding-exact-match"); + let primary_a = scope_store + .resolve_current("fixture", "primary-a", b"primary-a") + .unwrap(); + let primary_b = scope_store + .resolve_current("fixture", "primary-b", b"primary-b") + .unwrap(); + let corroborating_a = scope_store + .resolve_current("fixture", "corroborating-a", b"corroborating-a") + .unwrap(); + let corroborating_b = scope_store + .resolve_current("fixture", "corroborating-b", b"corroborating-b") + .unwrap(); - let mut global_capacity = - enrichment_snapshot(Ok(account_scope.clone()), enrichment_failure_windows(now)); - enrich_snapshot_with( - &mut global_capacity, - now, - |active, observations, _| { - assert_eq!(active.len(), 3); - assert_eq!(observations.len(), 2); - Err(HistoryError::StoreCapacity) - }, - ); - assert_eq!( - global_capacity.windows[0].pace_reason_for_test(), - Some("storeCapacity") + let full = ProviderCacheBinding::new(primary_a.clone(), Some(corroborating_a.clone())); + assert_eq!(full, full.clone()); + assert_ne!( + full, + ProviderCacheBinding::new(primary_b, Some(corroborating_a.clone())) ); - assert_eq!( - global_capacity.windows[1].pace_reason_for_test(), - Some("storeCapacity") + assert_ne!( + full, + ProviderCacheBinding::new(primary_a.clone(), Some(corroborating_b)) ); + assert_ne!(full, ProviderCacheBinding::primary(primary_a)); + scope_store.cleanup(); + } + + #[test] + fn maps_codex_additional_model_limits() { + let now = Utc.timestamp_opt(1_700_000_000, 0).single().unwrap(); + let extra = CodexAdditionalRateLimit { + limit_name: Some("gpt-5.2-codex-spark".to_string()), + metered_feature: None, + rate_limit: Some(CodexRateLimit { + primary_window: Some(CodexWindow { + used_percent: 41.0, + reset_at: 1_700_003_600, + limit_window_seconds: 18_000, + }), + secondary_window: None, + }), + }; + let windows = codex_windows(None, Some(&[extra]), now); + assert_eq!(windows.len(), 1); + assert_eq!(windows[0].label, "Codex Spark"); + assert_eq!(windows[0].remaining_percent, 59.0); + } + + #[test] + fn stage0_freezes_codex_additional_identity_baseline() { + let metered_only = CodexAdditionalRateLimit { + limit_name: None, + metered_feature: Some("gpt-5.2-codex-spark".to_string()), + rate_limit: None, + }; assert_eq!( - global_capacity.windows[2].pace_reason_for_test(), - Some("missingReset") + additional_limit_label(&metered_only), + "Codex Spark", + "codex.additional.metered-feature.primary" ); + + let named = CodexAdditionalRateLimit { + limit_name: Some("named-limit".to_string()), + metered_feature: Some("metered-feature".to_string()), + rate_limit: None, + }; assert_eq!( - global_capacity.windows[3].pace_reason_for_test(), - Some("nonRecurring") + additional_limit_label(&named), + "Named Limit", + "display label remains separate from the metered-feature identity" ); - let mut global_history = - enrichment_snapshot(Ok(account_scope.clone()), enrichment_failure_windows(now)); - enrich_snapshot_with(&mut global_history, now, |_, _, _| { - Err(HistoryError::Read) - }); + let anonymous = CodexAdditionalRateLimit { + limit_name: None, + metered_feature: None, + rate_limit: None, + }; + assert_eq!(additional_limit_source(&anonymous), None); + + let now = Utc.timestamp_opt(1_700_000_000, 0).single().unwrap(); + let both_slots = CodexAdditionalRateLimit { + limit_name: Some("named-limit".to_string()), + metered_feature: Some(" metered-feature ".to_string()), + rate_limit: Some(CodexRateLimit { + primary_window: Some(CodexWindow { + used_percent: 10.0, + reset_at: 1_700_003_600, + limit_window_seconds: 18_000, + }), + secondary_window: Some(CodexWindow { + used_percent: 20.0, + reset_at: 1_700_086_400, + limit_window_seconds: 604_800, + }), + }), + }; assert_eq!( - global_history.windows[0].pace_reason_for_test(), - Some("history") + additional_limit_source(&both_slots).as_deref(), + Some("metered-feature") ); + let windows = codex_windows(None, Some(&[both_slots]), now); assert_eq!( - global_history.windows[1].pace_reason_for_test(), - Some("history") - ); + windows.len(), + 2, + "codex.additional.primary-secondary emits both semantic slots" + ); + let digest = sha256_hex("metered-feature".to_string()); + let primary_key = format!("additional.{digest}.primary.v1"); + let secondary_key = format!("additional.{digest}.secondary.v1"); + assert_eq!(windows[0].card_id, primary_key); assert_eq!( - global_history.windows[2].pace_reason_for_test(), - Some("missingReset") + windows[0].window_key.as_deref(), + Some(windows[0].card_id.as_str()) ); + assert_eq!(windows[1].card_id, secondary_key); assert_eq!( - global_history.windows[3].pace_reason_for_test(), - Some("nonRecurring") + windows[1].window_key.as_deref(), + Some(windows[1].card_id.as_str()) ); + } - let mut count_mismatch = - enrichment_snapshot(Ok(account_scope.clone()), enrichment_failure_windows(now)); - enrich_snapshot_with(&mut count_mismatch, now, |_, _, _| { - Ok(vec![Ok((HistoryOutcome::LearningDuration, None, 0))]) - }); - assert_eq!( - count_mismatch.windows[0].pace_reason_for_test(), - Some("history") - ); - assert_eq!( - count_mismatch.windows[1].pace_reason_for_test(), - Some("history") + #[test] + fn codex_unknown_and_anonymous_windows_are_structural_and_skip_history() { + let now = Utc.timestamp_opt(1_700_000_000, 0).single().unwrap(); + let unknown_main = CodexRateLimit { + primary_window: Some(CodexWindow { + used_percent: 5.0, + reset_at: now.timestamp() + 3_600, + limit_window_seconds: 3_600, + }), + secondary_window: None, + }; + let anonymous = |primary_used: f64, secondary_used: f64| CodexAdditionalRateLimit { + limit_name: None, + metered_feature: None, + rate_limit: Some(CodexRateLimit { + primary_window: Some(CodexWindow { + used_percent: primary_used, + reset_at: now.timestamp() + 7_200, + limit_window_seconds: 7_200, + }), + secondary_window: Some(CodexWindow { + used_percent: secondary_used, + reset_at: now.timestamp() + 86_400, + limit_window_seconds: 86_400, + }), + }), + }; + let windows = codex_windows( + Some(&unknown_main), + Some(&[anonymous(10.0, 20.0), anonymous(30.0, 40.0)]), + now, ); + assert_eq!(windows.len(), 3); assert_eq!( - count_mismatch.windows[2].pace_reason_for_test(), - Some("missingReset") + windows + .iter() + .map(|window| window.card_id.as_str()) + .collect::>(), + vec![ + "row.main.primary.v1", + "row.additional.unknown.primary.v1", + "row.additional.unknown.secondary.v1" + ] ); assert_eq!( - count_mismatch.windows[3].pace_reason_for_test(), - Some("nonRecurring") + windows + .iter() + .map(|window| window.used_percent) + .collect::>(), + vec![5.0, 10.0, 20.0], + "duplicate anonymous slots keep the provider-order first row" ); + for window in &windows[1..] { + assert_eq!(window.label_for_test(), "Unknown"); + assert_ne!(window.label_for_test(), "Codex extra limit"); + } + for window in &windows { + assert_eq!(window.window_key, None); + assert_eq!(window.pace_status.state, PaceState::Unavailable); + assert_eq!(window.pace_status.reason.as_deref(), Some("windowIdentity")); + } - let mut row_errors = - enrichment_snapshot(Ok(account_scope), enrichment_failure_windows(now)); - enrich_snapshot_with(&mut row_errors, now, |_, _, _| { - Ok(vec![ - Err(HistoryError::StoreCapacity), - Err(HistoryError::Read), - ]) + let scope = TestRefreshScope::new("codex", "unknown-windows"); + let account_scope = scope + .resolve_current("fixture", "unknown-windows", b"marker") + .unwrap(); + let mut snapshot = AgentUsageSnapshot { + client_id: "codex".to_string(), + source: "fixture".to_string(), + updated_at: String::new(), + identity: None, + account_scope: Ok(account_scope), + windows, + credits: None, + error: None, + transport_diagnostic: None, + }; + let history_calls = std::cell::Cell::new(0); + enrich_snapshot_with(&mut snapshot, now.timestamp(), |_, _, _| { + history_calls.set(history_calls.get() + 1); + Ok(Vec::new()) }); - assert_eq!( - row_errors.windows[0].pace_reason_for_test(), - Some("storeCapacity") - ); - assert_eq!( - row_errors.windows[1].pace_reason_for_test(), - Some("history") - ); - assert_eq!( - row_errors.windows[2].pace_reason_for_test(), - Some("missingReset") - ); - assert_eq!( - row_errors.windows[3].pace_reason_for_test(), - Some("nonRecurring") - ); - for snapshot in [global_capacity, global_history, count_mismatch, row_errors] { - assert!(serde_json::to_value(snapshot).is_ok()); + assert_eq!(history_calls.get(), 0); + + let wire = serde_json::to_value(&snapshot).unwrap(); + let rows = wire["windows"].as_array().unwrap(); + assert_eq!(rows.len(), 3); + for row in rows { + assert!(row["paceStatus"].get("windowKey").is_none()); + assert_eq!(row["paceStatus"]["state"], "unavailable"); + assert_eq!(row["paceStatus"]["reason"], "windowIdentity"); } scope.cleanup(); } #[test] - fn parses_retry_after_seconds_and_http_date() { - let header = reqwest::header::HeaderValue::from_static("120"); - let parsed = parse_retry_after(Some(&header)).unwrap(); - let delta = (parsed - Utc::now()).num_seconds(); - assert!((118..=120).contains(&delta), "delta was {}", delta); - - let header = reqwest::header::HeaderValue::from_static("Fri, 21 Nov 2025 09:00:00 GMT"); - let parsed = parse_retry_after(Some(&header)).unwrap(); - assert_eq!(parsed.timestamp(), 1_763_715_600); - - let header = reqwest::header::HeaderValue::from_static("bogus"); - assert!(parse_retry_after(Some(&header)).is_none()); - assert!(parse_retry_after(None).is_none()); + fn parses_claude_credentials_file() { + let raw = r#"{ + "claudeAiOauth": { + "accessToken": "access", + "refreshToken": "refresh", + "expiresAt": 1700000000000, + "scopes": ["user:profile"], + "rateLimitTier": "max", + "subscriptionType": "pro" + } + }"#; + let credentials = parse_claude_credentials_data(raw, ClaudeCredentialSource::File).unwrap(); + assert_eq!(credentials.access_token, "access"); + assert_eq!(credentials.refresh_token.as_deref(), Some("refresh")); + assert_eq!(credentials.scopes, vec!["user:profile"]); + assert_eq!(credentials.subscription_type.as_deref(), Some("pro")); } #[test] - fn string_key_uses_first_valid_snake_or_camel_alias() { - let cases = [ - ( - "snake priority", - serde_json::json!({ - "snake_key": " snake-value ", - "camelKey": "camel-value" - }), - Some("snake-value"), - ), - ( - "snake missing", - serde_json::json!({ "camelKey": " camel-value " }), - Some("camel-value"), - ), - ( - "snake null", - serde_json::json!({ "snake_key": null, "camelKey": "camel-value" }), - Some("camel-value"), - ), - ( - "snake empty", - serde_json::json!({ "snake_key": "", "camelKey": "camel-value" }), - Some("camel-value"), - ), - ( - "snake whitespace", - serde_json::json!({ "snake_key": " \t\n ", "camelKey": "camel-value" }), - Some("camel-value"), - ), - ( - "snake non-string", - serde_json::json!({ - "snake_key": { "unexpected": true }, - "camelKey": "camel-value" - }), - Some("camel-value"), - ), - ( - "both invalid", - serde_json::json!({ "snake_key": false, "camelKey": " " }), - None, - ), - ]; + fn merge_claude_credentials_rotates_tokens_and_preserves_other_fields() { + let raw = r#"{ + "claudeAiOauth": { + "accessToken": "old-access", + "refreshToken": "old-refresh", + "expiresAt": 1700000000000, + "scopes": ["user:profile"], + "subscriptionType": "pro" + } + }"#; + let mut credentials = + parse_claude_credentials_data(raw, ClaudeCredentialSource::File).unwrap(); + credentials.access_token = "new-access".to_string(); + credentials.refresh_token = Some("new-refresh".to_string()); + credentials.expires_at = Utc.timestamp_millis_opt(1_700_009_999_000).single(); - for (label, value, expected) in cases { - let map = value.as_object().unwrap(); - assert_eq!( - string_key(map, "snake_key", "camelKey").as_deref(), - expected, - "{label}" - ); - } + let merged = merge_claude_credentials_json(&credentials, raw).unwrap(); + let reparsed = + parse_claude_credentials_data(&merged, ClaudeCredentialSource::File).unwrap(); + assert_eq!(reparsed.access_token, "new-access"); + assert_eq!(reparsed.refresh_token.as_deref(), Some("new-refresh")); + assert_eq!( + reparsed.expires_at, + Utc.timestamp_millis_opt(1_700_009_999_000).single() + ); + // Untouched fields the Claude CLI wrote survive the merge. + assert_eq!(reparsed.subscription_type.as_deref(), Some("pro")); + assert_eq!(reparsed.scopes, vec!["user:profile"]); } #[test] - fn claude_refresh_response_ignores_invalid_optional_refresh_token() { - let cases = [ - ( - "valid", - serde_json::json!({ - "access_token": "new-access", - "refresh_token": " new-refresh ", - "expires_in": 3600 - }), - Some("new-refresh"), - ), - ( - "missing", - serde_json::json!({ "access_token": "new-access", "expires_in": 3600 }), - None, - ), - ( - "null", - serde_json::json!({ - "access_token": "new-access", - "refresh_token": null, - "expires_in": 3600 - }), - None, - ), - ( - "empty", - serde_json::json!({ - "access_token": "new-access", - "refresh_token": "", - "expires_in": 3600 - }), - None, - ), - ( - "whitespace", - serde_json::json!({ - "access_token": "new-access", - "refresh_token": " \t\n ", - "expires_in": 3600 - }), - None, - ), - ( - "non-string", - serde_json::json!({ - "access_token": "new-access", - "refresh_token": { "unexpected": true }, - "expires_in": 3600 - }), - None, - ), - ]; + fn claude_keychain_write_decision_pins_account_and_rejects_target_mismatch() { + let raw_a = r#"{ + "claudeAiOauth": { + "accessToken": "old-access", + "refreshToken": "old-refresh", + "expiresAt": 0 + }, + "sibling": "a" + }"#; + let mut credentials = + parse_claude_credentials_data(raw_a, ClaudeCredentialSource::Keychain).unwrap(); + credentials.keychain_account = Some("account-a".to_string()); + credentials.access_token = "new-access".to_string(); + credentials.refresh_token = Some("new-refresh".to_string()); - for (label, value, expected) in cases { - let response: ClaudeRefreshResponse = serde_json::from_value(value).unwrap(); - assert_eq!(response.access_token, "new-access", "{label}"); - assert_eq!(response.expires_in, 3_600, "{label}"); - assert_eq!(response.refresh_token.as_deref(), expected, "{label}"); - } - } + let (account, merged) = + prepare_claude_keychain_write(&credentials, Some("account-a"), raw_a).unwrap(); + assert_eq!(account, "account-a"); + assert_eq!( + serde_json::from_str::(&merged).unwrap()["claudeAiOauth"]["accessToken"], + "new-access" + ); - #[test] - fn account_scope_and_credential_markers_never_reach_the_wire() { - let scope_store = TestRefreshScope::new("codex", "agent-usage-wire-privacy"); - let marker = b"sensitive-refresh-token-marker"; - let account_scope = scope_store - .resolve_current("codex-auth-json", "fixture-location", marker) - .unwrap(); - let opaque_scope = account_scope.as_str().to_string(); - let snapshot = AgentUsageSnapshot { - client_id: "codex".to_string(), - source: "oauth".to_string(), - updated_at: "2026-07-18T00:00:00.000Z".to_string(), - identity: None, - account_scope: Ok(account_scope), - windows: Vec::new(), - credits: None, - error: None, - }; + assert!(prepare_claude_keychain_write(&credentials, Some("account-b"), raw_a).is_err()); + assert!(prepare_claude_keychain_write(&credentials, None, raw_a).is_err()); - let wire = serde_json::to_string(&snapshot).unwrap(); - assert!(!wire.contains("accountScope")); - assert!(!wire.contains(String::from_utf8_lossy(marker).as_ref())); - assert!(!wire.contains(&opaque_scope)); - scope_store.cleanup(); + let raw_changed_target = r#"{ + "claudeAiOauth": { + "accessToken": "account-b-access", + "refreshToken": "account-b-refresh", + "expiresAt": 0 + }, + "sibling": "b" + }"#; + assert!( + prepare_claude_keychain_write(&credentials, Some("account-a"), raw_changed_target,) + .is_err() + ); } #[test] - fn credential_markers_and_locations_follow_the_canonical_routes() { - let scope_store = TestRefreshScope::new("codex", "agent-usage-locations"); - let auth_path = scope_store.root().join("codex/auth.json"); - fs::create_dir_all(auth_path.parent().unwrap()).unwrap(); - fs::write( - &auth_path, - serde_json::json!({ - "tokens": { - "access_token": " codex-access ", - "refresh_token": " codex-refresh " - } - }) - .to_string(), - ) - .unwrap(); - let codex = load_codex_credentials_from(&auth_path).unwrap(); - assert_eq!(codex.scope_slot.semantic_source, "codex-auth-json"); - assert_eq!( - codex.scope_slot.canonical_location, - agent_account_scope::canonical_file_location(&auth_path, Some("tokens")).unwrap() - ); - assert_eq!(codex.scope_marker(), b"codex-refresh"); - let mut codex_access_only = codex.clone(); - codex_access_only.refresh_token = None; - assert_eq!(codex_access_only.scope_marker(), b"codex-access"); + fn atomic_write_replaces_existing_file_contents() { + let dir = std::env::temp_dir().join(format!("tb_atomic_{}", std::process::id())); + fs::create_dir_all(&dir).unwrap(); + let path = dir.join(".credentials.json"); + fs::write(&path, "old").unwrap(); - let claude_file_slot = claude_login_scope_slot(ClaudeCredentialSource::File).unwrap(); - assert_eq!(claude_file_slot.semantic_source, "claude-login-file"); - assert_eq!( - claude_file_slot.canonical_location, - agent_account_scope::canonical_file_location( - &claude_credentials_path(), - Some("claudeAiOauth") - ) + atomic_write(&path, "new").unwrap(); + assert_eq!(fs::read_to_string(&path).unwrap(), "new"); + // No temp turds left in the directory. + let leftovers: Vec<_> = fs::read_dir(&dir) .unwrap() - ); - let claude_keychain_slot = - claude_login_scope_slot(ClaudeCredentialSource::Keychain).unwrap(); - assert_eq!( - claude_keychain_slot.semantic_source, - "claude-login-keychain" - ); - assert_eq!( - claude_keychain_slot.canonical_location, - CLAUDE_KEYCHAIN_SERVICE - ); - - let claude_login = ClaudeCredentials { - access_token: "claude-access".to_string(), - refresh_token: Some("claude-refresh".to_string()), - expires_at: None, - scopes: Vec::new(), - rate_limit_tier: None, - subscription_type: None, - source: ClaudeCredentialSource::File, - raw_root: None, - scope_slot: claude_file_slot, - }; - assert_eq!( - claude_login.scope_marker(), - Some(b"claude-refresh".as_slice()) - ); - let mut login_without_refresh = claude_login.clone(); - login_without_refresh.refresh_token = None; - assert_eq!(login_without_refresh.scope_marker(), None); + .filter_map(|e| e.ok()) + .filter(|e| e.file_name().to_string_lossy().contains(".tmp.")) + .collect(); + assert!(leftovers.is_empty(), "temp file not cleaned up"); - let claude_setup = claude_credentials_from_access_token(ResolvedClaudeToken { - access_token: "claude-setup-access".to_string(), - scope_slot: CredentialSlot { - semantic_source: "claude-code-environment", - canonical_location: "CLAUDE_CODE_OAUTH_TOKEN".to_string(), - }, - }); - assert_eq!( - claude_setup.scope_marker(), - Some(b"claude-setup-access".as_slice()) - ); - assert_eq!( - claude_setup.scope_slot.semantic_source, - "claude-code-environment" - ); - assert_eq!( - claude_setup.scope_slot.canonical_location, - "CLAUDE_CODE_OAUTH_TOKEN" - ); - scope_store.cleanup(); + let _ = fs::remove_dir_all(&dir); } - #[test] - fn codex_scope_precedence_keeps_refresh_failure_sticky() { - let scope_store = TestRefreshScope::new("codex", "codex-scope-precedence"); - let refresh_scope = scope_store - .resolve_current("fixture", "refresh", b"refresh-marker") - .unwrap(); - let authoritative_scope = scope_store - .resolve_current("fixture", "authoritative", b"authoritative-marker") - .unwrap(); - let credential_scope = scope_store - .resolve_current("fixture", "credential", b"credential-marker") - .unwrap(); - let authoritative_calls = std::cell::Cell::new(0); - let credential_calls = std::cell::Cell::new(0); - - let resolved = resolve_codex_account_scope( - Some(Err(AccountScopeError::MetadataWrite)), - Some("acct-id"), - |_| { - authoritative_calls.set(authoritative_calls.get() + 1); - Ok(authoritative_scope.clone()) - }, - || { - credential_calls.set(credential_calls.get() + 1); - Ok(credential_scope.clone()) - }, - ); - assert_eq!(resolved, Err(AccountScopeError::MetadataWrite)); - assert_eq!(authoritative_calls.get(), 0); - assert_eq!(credential_calls.get(), 0); + #[cfg(target_os = "windows")] + fn open_without_delete_sharing(path: &Path) -> fs::File { + use std::os::windows::fs::OpenOptionsExt as _; + use windows_sys::Win32::Storage::FileSystem::{FILE_SHARE_READ, FILE_SHARE_WRITE}; - let resolved = resolve_codex_account_scope( - Some(Err(AccountScopeError::MetadataRead)), - None, - |_| { - authoritative_calls.set(authoritative_calls.get() + 1); - Ok(authoritative_scope.clone()) - }, - || { - credential_calls.set(credential_calls.get() + 1); - Ok(credential_scope.clone()) - }, - ); - assert_eq!(resolved, Err(AccountScopeError::MetadataRead)); - assert_eq!(authoritative_calls.get(), 0); - assert_eq!(credential_calls.get(), 0); + let mut options = fs::OpenOptions::new(); + options + .read(true) + .share_mode(FILE_SHARE_READ | FILE_SHARE_WRITE); + options.open(path).unwrap() + } - let resolved = resolve_codex_account_scope( - Some(Ok(refresh_scope.clone())), - Some("acct-id"), - |_| { - authoritative_calls.set(authoritative_calls.get() + 1); - Ok(authoritative_scope.clone()) - }, - || { - credential_calls.set(credential_calls.get() + 1); - Ok(credential_scope.clone()) - }, - ); - assert_eq!(resolved.unwrap(), authoritative_scope); - assert_eq!(authoritative_calls.get(), 1); - assert_eq!(credential_calls.get(), 0); + #[cfg(target_os = "windows")] + fn atomic_temp_path(dir: &Path) -> Option { + fs::read_dir(dir) + .unwrap() + .filter_map(|entry| entry.ok()) + .find(|entry| entry.file_name().to_string_lossy().contains(".tmp.")) + .map(|entry| entry.path()) + } - let resolved = resolve_codex_account_scope( - Some(Ok(refresh_scope.clone())), - None, - |_| { - authoritative_calls.set(authoritative_calls.get() + 1); - Ok(authoritative_scope.clone()) - }, - || { - credential_calls.set(credential_calls.get() + 1); - Ok(credential_scope.clone()) - }, - ); - assert_eq!(resolved.unwrap(), refresh_scope); - assert_eq!(authoritative_calls.get(), 1); - assert_eq!(credential_calls.get(), 0); + #[cfg(target_os = "windows")] + fn lock_staged_atomic_temp(dir: &Path) -> Option { + use std::os::windows::fs::OpenOptionsExt as _; - let resolved = resolve_codex_account_scope( - None, - None, - |_| { - authoritative_calls.set(authoritative_calls.get() + 1); - Ok(authoritative_scope.clone()) - }, - || { - credential_calls.set(credential_calls.get() + 1); - Ok(credential_scope.clone()) - }, - ); - assert_eq!(resolved.unwrap(), credential_scope); - assert_eq!(authoritative_calls.get(), 1); - assert_eq!(credential_calls.get(), 1); - scope_store.cleanup(); + let mut options = fs::OpenOptions::new(); + options.read(true).share_mode(0); + options.open(atomic_temp_path(dir)?).ok() } + #[cfg(target_os = "windows")] #[test] - fn codex_v2_migration_requires_request_id_and_scope_and_is_best_effort() { - let scope_store = TestRefreshScope::new("codex", "codex-v2-migration-gate"); - let account_scope = scope_store - .resolve_current("fixture", "codex-v2", b"codex-v2-marker") - .unwrap(); - let opaque_scope = account_scope.as_str().to_string(); - let now = Utc.timestamp_opt(1_700_000_000, 0).single().unwrap(); - let calls = std::cell::RefCell::new(Vec::new()); - - maybe_migrate_codex_v2_with( - Some("request-account"), - &Ok(account_scope.clone()), - now.timestamp(), - |request_account_id, scope, call_now| { - calls.borrow_mut().push(( - request_account_id.to_string(), - scope.to_string(), - call_now, - )); - Ok(()) - }, - ); - assert_eq!( - calls.borrow().as_slice(), - &[( - "request-account".to_string(), - opaque_scope.clone(), - now.timestamp(), - )] - ); - - let skipped_calls = std::cell::Cell::new(0); - maybe_migrate_codex_v2_with( - None, - &Ok(account_scope.clone()), - now.timestamp(), - |_, _, _| { - skipped_calls.set(skipped_calls.get() + 1); - Ok(()) - }, - ); - maybe_migrate_codex_v2_with( - Some(" \t"), - &Ok(account_scope.clone()), - now.timestamp(), - |_, _, _| { - skipped_calls.set(skipped_calls.get() + 1); - Ok(()) - }, - ); - maybe_migrate_codex_v2_with( - Some("request-account"), - &Err(AccountScopeError::MetadataRead), - now.timestamp(), - |_, _, _| { - skipped_calls.set(skipped_calls.get() + 1); - Ok(()) - }, - ); - assert_eq!(skipped_calls.get(), 0); + fn atomic_write_retries_transient_windows_destination_lock() { + use std::time::{Duration, Instant}; - let migration_error_calls = std::cell::Cell::new(0); - maybe_migrate_codex_v2_with( - Some("request-account"), - &Ok(account_scope.clone()), - now.timestamp(), - |_, _, _| { - migration_error_calls.set(migration_error_calls.get() + 1); - Err::<(), _>(HistoryError::AtomicSave) - }, - ); + let dir = std::env::temp_dir().join(format!( + "tb_atomic_windows_transient_{}", + std::process::id() + )); + let _ = fs::remove_dir_all(&dir); + fs::create_dir_all(&dir).unwrap(); + let path = dir.join("auth.json"); + fs::write(&path, "old").unwrap(); + let destination_lock = open_without_delete_sharing(&path); + + let writer_path = path.clone(); + let writer = std::thread::spawn(move || atomic_write(&writer_path, "new")); + let deadline = Instant::now() + Duration::from_secs(1); + let staged_temp_lock = loop { + if let Some(file) = lock_staged_atomic_temp(&dir) { + break Some(file); + } + if writer.is_finished() || Instant::now() >= deadline { + break None; + } + std::thread::sleep(Duration::from_millis(1)); + }; + let staged_temp_locked = staged_temp_lock.is_some(); + if staged_temp_locked { + std::thread::sleep(Duration::from_millis(20)); + } + let waited_for_retry = !writer.is_finished(); + drop(staged_temp_lock); + drop(destination_lock); + let result = writer.join().expect("atomic writer thread panicked"); - let mut snapshot = enrichment_snapshot( - Ok(account_scope), - vec![enrichment_window( - now, - "weekly.v1", - "weekly.v1", - 20.0, - Some(DurationSource::Contract), - )], + assert!( + staged_temp_locked, + "atomic write never completed temp-file staging" ); - let record_calls = std::cell::Cell::new(0); - enrich_snapshot_with(&mut snapshot, now, |_, observations, _| { - record_calls.set(record_calls.get() + 1); - assert_eq!(observations.len(), 1); - Ok(vec![Ok(( - HistoryOutcome::Ready { - duration_seconds: 86_400, - source: DurationSource::Contract, - sampled: true, - }, - None, - 1, - ))]) - }); - assert_eq!(migration_error_calls.get(), 1); - assert_eq!(record_calls.get(), 1); - assert_eq!( - snapshot.windows[0].pace_status.state, - PaceState::LearningHistory + assert!( + waited_for_retry, + "atomic write did not retry the sharing denial" ); - scope_store.cleanup(); + result.unwrap(); + assert_eq!(fs::read_to_string(&path).unwrap(), "new"); + assert!(atomic_temp_path(&dir).is_none(), "temp file not cleaned up"); + let _ = fs::remove_dir_all(&dir); } - // Single test for the whole gate lifecycle — the gate is a process-wide - // static, so split tests would race under the parallel test runner. + #[cfg(target_os = "windows")] #[test] - fn claude_gate_blocks_then_clears() { - let now = Utc.timestamp_opt(1_700_000_000, 0).single().unwrap(); - assert!(claude_gate_blocked_until(now).is_none()); - - // 429 with no Retry-After → default 5-minute cooldown. - claude_gate_record_rate_limit(None, now); - let until = claude_gate_blocked_until(now).unwrap(); - assert_eq!((until - now).num_seconds(), 300); + fn atomic_write_exhausts_windows_retry_budget_without_losing_original() { + use std::time::{Duration, Instant}; - // No cached snapshot yet → countdown error. - let fallback = claude_gate_fallback(until, now); - assert!(fallback.error.unwrap().contains("~300s")); - assert!(fallback.windows.is_empty()); - - // Cooldown expiry clears the gate lazily. - let later = now + chrono::Duration::seconds(301); - assert!(claude_gate_blocked_until(later).is_none()); - - // Success caches the display-ready snapshot; a later 429 returns those - // rows unchanged, without another enrichment/history pass, while dropping - // stale account evidence from the earlier authenticated poll. - let scope_store = TestRefreshScope::new("claude", "cached-429"); - let account_scope = scope_store - .resolve_current("fixture", "cached-429", b"cached-429-marker") - .unwrap(); - let reset = now + chrono::Duration::days(1); - let mut snapshot = AgentUsageSnapshot { - client_id: "claude".to_string(), - source: "oauth".to_string(), - updated_at: now.to_rfc3339_opts(SecondsFormat::Millis, true), - identity: None, - account_scope: Ok(account_scope), - windows: vec![UsageWindow::from_used_percent( - "Session".to_string(), - 20.0, - Some(reset), - now, - ) - .with_identity("session.v1", Some("session.v1".to_string())) - .with_duration_evidence( - now, - true, - None, - Some(DurationEvidence::contract(86_400)), - )], - credits: None, - error: None, - }; - let record_calls = std::cell::Cell::new(0); - enrich_snapshot_with( - &mut snapshot, - now, - |active, observations, batch_now| { - record_calls.set(record_calls.get() + 1); - assert_eq!(batch_now, now.timestamp()); - assert_eq!(active.len(), 1); - assert_eq!(observations.len(), 1); - Ok(vec![Ok(( - HistoryOutcome::Ready { - duration_seconds: 86_400, - source: DurationSource::Contract, - sampled: true, - }, - Some(HistoricalPace { - expected_percent: 35.0, - eta_seconds: None, - will_last_to_reset: true, - run_out_probability: Some(0.42), - }), - 6, - ))]) - }, - ); - assert_eq!(record_calls.get(), 1); - assert_eq!(snapshot.windows[0].pace_status.state, PaceState::Available); - claude_gate_record_success(&snapshot); - assert!(claude_gate_blocked_until(later).is_none()); - claude_gate_record_rate_limit(Some(later + chrono::Duration::seconds(60)), later); - let until = claude_gate_blocked_until(later).unwrap(); - let fallback = claude_gate_fallback(until, later); - assert_eq!(record_calls.get(), 1); - assert!(fallback.error.is_none()); - assert_eq!(fallback.windows.len(), 1); - assert_eq!(fallback.windows[0].label, "Session"); - assert_eq!(fallback.windows[0].card_id, "session.v1"); - assert_eq!(fallback.windows[0].used_percent, 20.0); - assert_eq!(fallback.windows[0].remaining_percent, 80.0); - assert_eq!(fallback.windows[0].pace_status.state, PaceState::Available); - assert_eq!(fallback.windows[0].pace_status.complete_cycles, 6); - assert_eq!( - fallback.windows[0] - .historical_pace - .as_ref() - .map(|pace| pace.expected_used_percent), - Some(35.0) - ); - assert!(matches!( - &fallback.account_scope, - Err(AccountScopeError::NoTrustedEvidence) + let dir = std::env::temp_dir().join(format!( + "tb_atomic_windows_persistent_{}", + std::process::id() )); + let _ = fs::remove_dir_all(&dir); + fs::create_dir_all(&dir).unwrap(); + let path = dir.join("auth.json"); + fs::write(&path, "old").unwrap(); + let destination_lock = open_without_delete_sharing(&path); - // Leave the gate clean for any other test touching the static. - claude_gate_record_success(&snapshot); - scope_store.cleanup(); + let started = Instant::now(); + let result = atomic_write(&path, "new"); + let elapsed = started.elapsed(); + drop(destination_lock); + + assert!(result.is_err(), "persistent sharing denial must fail"); + assert!( + elapsed >= Duration::from_millis(80), + "atomic write returned before exhausting the retry budget: {elapsed:?}" + ); + assert_eq!(fs::read_to_string(&path).unwrap(), "old"); + assert!(atomic_temp_path(&dir).is_none(), "temp file not cleaned up"); + let _ = fs::remove_dir_all(&dir); } #[test] - fn maps_codex_primary_and_secondary_windows() { + fn maps_claude_oauth_windows() { let now = Utc.timestamp_opt(1_700_000_000, 0).single().unwrap(); - let rate_limit = CodexRateLimit { - primary_window: Some(CodexWindow { - used_percent: 8.0, - reset_at: 1_700_005_400, - limit_window_seconds: 18_000, + let usage = ClaudeUsageResponse { + five_hour: Some(ClaudeWindow { + utilization: Some(8.0), + resets_at: Some("2023-11-14T23:13:20Z".to_string()), }), - secondary_window: Some(CodexWindow { - used_percent: 35.0, - reset_at: 1_700_172_800, - limit_window_seconds: 604_800, + seven_day: Some(ClaudeWindow { + utilization: Some(23.0), + resets_at: Some("2023-11-17T22:13:20Z".to_string()), + }), + seven_day_oauth_apps: None, + seven_day_opus: None, + seven_day_sonnet: Some(ClaudeWindow { + utilization: Some(3.0), + resets_at: None, + }), + seven_day_design: Some(ClaudeWindow { + utilization: Some(0.0), + resets_at: None, }), + seven_day_routines: None, + extra_usage: None, + ..Default::default() }; - let windows = codex_windows(Some(&rate_limit), None, now); - assert_eq!(windows.len(), 2); + let windows = claude_windows(&usage, now); + assert_eq!(windows.len(), 4); assert_eq!(windows[0].label, "Session"); - assert_eq!(windows[0].card_id_for_test(), "main.session.v1"); - assert_eq!( - windows[0].pace_window_key_for_test(), - Some("main.session.v1") - ); assert_eq!(windows[0].remaining_percent, 92.0); - assert_eq!(windows[0].pace_status.state, PaceState::LearningHistory); - assert_eq!(windows[0].pace_status.duration_seconds, Some(18_000)); - assert_eq!( - windows[0].pace_status.duration_source, - Some(DurationSource::Provider) - ); assert_eq!(windows[1].label, "Weekly"); - assert_eq!(windows[1].card_id_for_test(), "main.weekly.v1"); - assert_eq!( - windows[1].pace_window_key_for_test(), - Some("main.weekly.v1") - ); - assert_eq!(windows[1].remaining_percent, 65.0); - assert_eq!(windows[1].pace_status.state, PaceState::LearningHistory); - assert_eq!(windows[1].pace_status.duration_seconds, Some(604_800)); - assert_eq!( - windows[1].pace_status.duration_source, - Some(DurationSource::Provider) - ); + assert_eq!(windows[1].remaining_percent, 77.0); + assert_eq!(windows[2].label, "Sonnet"); + assert_eq!(windows[2].remaining_percent, 97.0); + assert_eq!(windows[3].label, "Designs"); + assert_eq!(windows[3].remaining_percent, 100.0); } #[test] - fn agent_usage_payload_omits_legacy_history_fields() { + fn stage4_claude_json_and_headers_share_canonical_duration_contracts() { let now = Utc.timestamp_opt(1_700_000_000, 0).single().unwrap(); - let rate_limit = CodexRateLimit { - primary_window: Some(CodexWindow { - used_percent: 35.0, - reset_at: 1_700_172_800, - limit_window_seconds: 604_800, - }), - secondary_window: None, + let reset = Some("2026-07-24T00:00:00Z".to_string()); + let window = |utilization| ClaudeWindow { + utilization: Some(utilization), + resets_at: reset.clone(), }; - let payload = AgentUsagePayload { - generated_at: now.to_rfc3339_opts(SecondsFormat::Millis, true), - agents: vec![AgentUsageSnapshot { - client_id: "codex".to_string(), - source: "oauth".to_string(), - updated_at: now.to_rfc3339_opts(SecondsFormat::Millis, true), - identity: None, - account_scope: Err(AccountScopeError::NoTrustedEvidence), - windows: codex_windows(Some(&rate_limit), None, now), - credits: None, - error: None, - }], - opencode_subscriptions: Vec::new(), + let usage = ClaudeUsageResponse { + five_hour: Some(window(5.0)), + seven_day: Some(window(10.0)), + seven_day_oauth_apps: Some(window(15.0)), + seven_day_sonnet: Some(window(20.0)), + seven_day_opus: Some(window(25.0)), + ..Default::default() }; - let serialized = serde_json::to_value(payload).unwrap(); - let weekly = serialized["agents"][0]["windows"] - .as_array() - .unwrap() + let windows = claude_windows(&usage, now); + let contracts = windows .iter() - .find(|window| window["label"] == "Weekly") - .expect("normal Codex Weekly mapping"); - let object = weekly.as_object().unwrap(); - assert_eq!(object["cardId"], "main.weekly.v1"); - assert_eq!(object["usedPercent"], 35.0); - assert_eq!(object["remainingPercent"], 65.0); - assert_eq!(object["paceStatus"]["state"], "learningHistory"); - assert_eq!(object["paceStatus"]["windowKey"], "main.weekly.v1"); - assert_eq!(object["paceStatus"]["durationSeconds"], 604_800); - assert_eq!(object["paceStatus"]["durationSource"], "provider"); - assert_eq!(object["windowMinutes"], 10_080); - assert!(!object.contains_key("historicalExpectedPercent")); - assert!(!object.contains_key("runOutProbability")); - assert!(!object.contains_key("historicalPace")); + .map(|window| { + ( + window.card_id.as_str(), + window.pace_status.window_key.as_deref(), + window.duration_seconds, + window.duration_source, + window.pace_status.state, + ) + }) + .collect::>(); + assert_eq!( + contracts, + vec![ + ( + "session.v1", + Some("session.v1"), + Some(18_000), + Some(DurationSource::Contract), + PaceState::LearningHistory, + ), + ( + "weekly.v1", + Some("weekly.v1"), + Some(604_800), + Some(DurationSource::Contract), + PaceState::LearningHistory, + ), + ( + "oauth_apps.weekly.v1", + Some("oauth_apps.weekly.v1"), + Some(604_800), + Some(DurationSource::Contract), + PaceState::LearningHistory, + ), + ( + "sonnet.weekly.v1", + Some("sonnet.weekly.v1"), + Some(604_800), + Some(DurationSource::Contract), + PaceState::LearningHistory, + ), + ( + "opus.weekly.v1", + Some("opus.weekly.v1"), + Some(604_800), + Some(DurationSource::Contract), + PaceState::LearningHistory, + ), + ] + ); + + let headers = header_map(&[ + ("anthropic-ratelimit-unified-5h-utilization", "0.11"), + ("anthropic-ratelimit-unified-5h-reset", "1783111200"), + ("anthropic-ratelimit-unified-7d-utilization", "0.6"), + ("anthropic-ratelimit-unified-7d-reset", "1783504800"), + ]); + let header_windows = parse_unified_ratelimit_windows(&headers, now); + assert_eq!(header_windows.len(), 2); + for (window, expected_key, expected_duration) in [ + (&header_windows[0], "session.v1", 18_000), + (&header_windows[1], "weekly.v1", 604_800), + ] { + assert_eq!(window.card_id, expected_key); + assert_eq!(window.pace_status.window_key.as_deref(), Some(expected_key)); + assert_eq!(window.duration_seconds, Some(expected_duration)); + assert_eq!(window.duration_source, Some(DurationSource::Contract)); + assert_eq!(window.pace_status.state, PaceState::LearningHistory); + } } #[test] - fn maps_codex_additional_model_limits() { + fn decodes_claude_alias_windows_without_duplicate_error() { + let raw = r#"{ + "five_hour": { "utilization": 5, "resets_at": "2026-05-28T14:00:00Z" }, + "seven_day": { "utilization": 23, "resets_at": "2026-05-31T14:00:00Z" }, + "seven_day_sonnet": { "utilization": 3, "resets_at": null }, + "seven_day_omelette": { "utilization": 0, "resets_at": null }, + "omelette_promotional": { "utilization": 0, "resets_at": null }, + "seven_day_cowork": { "utilization": 0, "resets_at": null } + }"#; + let usage: ClaudeUsageResponse = serde_json::from_str(raw).unwrap(); let now = Utc.timestamp_opt(1_700_000_000, 0).single().unwrap(); - let extra = CodexAdditionalRateLimit { - limit_name: Some("gpt-5.2-codex-spark".to_string()), - metered_feature: None, - rate_limit: Some(CodexRateLimit { - primary_window: Some(CodexWindow { - used_percent: 41.0, - reset_at: 1_700_003_600, - limit_window_seconds: 18_000, - }), - secondary_window: None, - }), - }; - let windows = codex_windows(None, Some(&[extra]), now); - assert_eq!(windows.len(), 1); - assert_eq!(windows[0].label, "Codex Spark"); - assert_eq!( - windows[0].card_id_for_test(), - format!( - "additional.{}.primary.v1", - sha256_hex("gpt-5.2-codex-spark") - ) - ); - assert_eq!(windows[0].remaining_percent, 59.0); - assert_eq!(windows[0].pace_status.state, PaceState::LearningHistory); - assert_eq!(windows[0].pace_status.duration_seconds, Some(18_000)); + let windows = claude_windows(&usage, now); assert_eq!( - windows[0].pace_status.duration_source, - Some(DurationSource::Provider) + windows.iter().map(|w| w.label.as_str()).collect::>(), + vec!["Session", "Weekly", "Sonnet", "Designs", "Daily Routines"] ); } #[test] - fn invalid_codex_duration_or_reset_evidence_fails_closed() { + fn stage4_claude_weekly_alias_groups_share_canonical_contracts() { let now = Utc.timestamp_opt(1_700_000_000, 0).single().unwrap(); - for (case, reset_at, duration_seconds) in [ - ("duration", now.timestamp() + 3_600, 0), - ("reset", 0, 3_600), - ] { - let window = map_window_with_identity( - "Additional", - CodexWindow { - used_percent: 25.0, - reset_at, - limit_window_seconds: duration_seconds, - }, - now, - "additional.stable.primary.v1", - Some("additional.stable.primary.v1".to_string()), - ) - .unwrap(); - assert_eq!(window.pace_status.state, PaceState::Unavailable, "{case}"); + let design_aliases = [ + "seven_day_design", + "seven_day_claude_design", + "claude_design", + "design", + "seven_day_omelette", + "omelette", + "omelette_promotional", + ]; + for alias in design_aliases { + let raw = + format!(r#"{{"{alias}":{{"utilization":12,"resets_at":"2026-07-24T00:00:00Z"}}}}"#); + let usage: ClaudeUsageResponse = serde_json::from_str(&raw).unwrap(); + let windows = claude_windows(&usage, now); + assert_eq!(windows.len(), 1, "claude.design.aliases: {alias}"); assert_eq!( - window.pace_status.reason.as_deref(), - Some("invalidEvidence"), - "{case}" + windows[0].label, "Designs", + "claude.design.aliases: {alias}" ); - assert!(window.pace_status.duration_seconds.is_none(), "{case}"); - assert!(window.pace_status.duration_source.is_none(), "{case}"); - assert!( - serde_json::to_value(window) - .unwrap() - .get("windowMinutes") - .is_none(), - "{case}" + assert_eq!(windows[0].card_id, "design.weekly.v1", "{alias}"); + assert_eq!( + windows[0].pace_status.window_key.as_deref(), + Some("design.weekly.v1"), + "{alias}" + ); + assert_eq!(windows[0].duration_seconds, Some(604_800), "{alias}"); + assert_eq!( + windows[0].duration_source, + Some(DurationSource::Contract), + "{alias}" ); + assert_eq!(windows[0].pace_status.state, PaceState::LearningHistory); } - let unknown = map_window_with_identity( - "Unknown", - CodexWindow { - used_percent: 25.0, - reset_at: 0, - limit_window_seconds: 0, - }, - now, - "row.additional.unknown.primary.v1", - None, - ) - .unwrap(); - assert_eq!(unknown.pace_reason_for_test(), Some("windowIdentity")); - - let reset = now + chrono::Duration::hours(1); - let mismatched = - UsageWindow::from_used_percent("Additional".to_string(), 25.0, Some(reset), now) - .with_identity( - "additional.stable.primary.v1", - Some("additional.stable.primary.v1".to_string()), - ) - .with_duration_evidence( - now, - true, - Some(DurationEvidence::provider(reset.timestamp() + 1, 3_600)), - None, - ); - assert_eq!( - mismatched.pace_status.reason.as_deref(), - Some("invalidEvidence") - ); - } - - #[test] - fn parses_claude_credentials_file() { - let raw = r#"{ - "claudeAiOauth": { - "accessToken": "access", - "refreshToken": "refresh", - "expiresAt": 1700000000000, - "scopes": ["user:profile"], - "rateLimitTier": "max", - "subscriptionType": "pro" - } - }"#; - let credentials = parse_claude_credentials_data(raw, ClaudeCredentialSource::File).unwrap(); - assert_eq!(credentials.access_token, "access"); - assert_eq!(credentials.refresh_token.as_deref(), Some("refresh")); - assert_eq!(credentials.scopes, vec!["user:profile"]); - assert_eq!(credentials.subscription_type.as_deref(), Some("pro")); + let routines_aliases = [ + "seven_day_routines", + "seven_day_claude_routines", + "claude_routines", + "routines", + "routine", + "seven_day_cowork", + "cowork", + ]; + for alias in routines_aliases { + let raw = + format!(r#"{{"{alias}":{{"utilization":12,"resets_at":"2026-07-24T00:00:00Z"}}}}"#); + let usage: ClaudeUsageResponse = serde_json::from_str(&raw).unwrap(); + let windows = claude_windows(&usage, now); + assert_eq!(windows.len(), 1, "claude.routines.aliases: {alias}"); + assert_eq!( + windows[0].label, "Daily Routines", + "claude.routines.aliases: {alias}" + ); + assert_eq!(windows[0].card_id, "routines.weekly.v1", "{alias}"); + assert_eq!( + windows[0].pace_status.window_key.as_deref(), + Some("routines.weekly.v1"), + "{alias}" + ); + assert_eq!(windows[0].duration_seconds, Some(604_800), "{alias}"); + assert_eq!( + windows[0].duration_source, + Some(DurationSource::Contract), + "{alias}" + ); + assert_eq!(windows[0].pace_status.state, PaceState::LearningHistory); + } } #[test] - fn merge_claude_credentials_rotates_tokens_and_preserves_other_fields() { + fn stage0_freezes_claude_named_windows_and_invalid_baseline() { + let now = Utc.timestamp_opt(1_700_000_000, 0).single().unwrap(); let raw = r#"{ - "claudeAiOauth": { - "accessToken": "old-access", - "refreshToken": "old-refresh", - "expiresAt": 1700000000000, - "scopes": ["user:profile"], - "subscriptionType": "pro" - } + "five_hour": { "utilization": 5, "resets_at": "2026-07-18T00:00:00Z" }, + "seven_day": { "utilization": 10, "resets_at": "2026-07-19T00:00:00Z" }, + "seven_day_oauth_apps": { "utilization": 15, "resets_at": "2026-07-20T00:00:00Z" }, + "seven_day_sonnet": { "utilization": 20, "resets_at": "2026-07-21T00:00:00Z" }, + "seven_day_opus": { "utilization": 25, "resets_at": "2026-07-22T00:00:00Z" } }"#; - let mut credentials = - parse_claude_credentials_data(raw, ClaudeCredentialSource::File).unwrap(); - credentials.access_token = "new-access".to_string(); - credentials.refresh_token = Some("new-refresh".to_string()); - credentials.expires_at = Utc.timestamp_millis_opt(1_700_009_999_000).single(); - - let merged = merge_claude_credentials_json(&credentials).unwrap(); - let reparsed = - parse_claude_credentials_data(&merged, ClaudeCredentialSource::File).unwrap(); - assert_eq!(reparsed.access_token, "new-access"); - assert_eq!(reparsed.refresh_token.as_deref(), Some("new-refresh")); + let usage: ClaudeUsageResponse = serde_json::from_str(raw).unwrap(); + let windows = claude_windows(&usage, now); + let mapped: Vec<_> = windows + .iter() + .map(|window| (window.label.as_str(), window.window_minutes)) + .collect(); assert_eq!( - reparsed.expires_at, - Utc.timestamp_millis_opt(1_700_009_999_000).single() + mapped, + vec![ + ("Session", Some(300)), + ("Weekly", Some(10_080)), + ("OAuth Apps", Some(10_080)), + ("Sonnet", Some(10_080)), + ("Opus", Some(10_080)), + ], + "claude.named-window-contracts" ); - // Untouched fields the Claude CLI wrote survive the merge. - assert_eq!(reparsed.subscription_type.as_deref(), Some("pro")); - assert_eq!(reparsed.scopes, vec!["user:profile"]); - } - - #[test] - fn atomic_write_replaces_existing_file_contents() { - let dir = std::env::temp_dir().join(format!("tb_atomic_{}", std::process::id())); - fs::create_dir_all(&dir).unwrap(); - let path = dir.join(".credentials.json"); - fs::write(&path, "old").unwrap(); - - atomic_write(&path, "new").unwrap(); - assert_eq!(fs::read_to_string(&path).unwrap(), "new"); - // No temp turds left in the directory. - let leftovers: Vec<_> = fs::read_dir(&dir) - .unwrap() - .filter_map(|e| e.ok()) - .filter(|e| e.file_name().to_string_lossy().contains(".tmp.")) - .collect(); - assert!(leftovers.is_empty(), "temp file not cleaned up"); - - let _ = fs::remove_dir_all(&dir); - } - #[test] - fn maps_claude_oauth_windows() { - let now = Utc.timestamp_opt(1_700_000_000, 0).single().unwrap(); - let usage = ClaudeUsageResponse { - five_hour: Some(ClaudeWindow { - utilization: Some(8.0), - resets_at: Some("2023-11-14T23:13:20Z".to_string()), - }), - seven_day: Some(ClaudeWindow { - utilization: Some(23.0), - resets_at: Some("2023-11-17T22:13:20Z".to_string()), - }), - seven_day_oauth_apps: None, - seven_day_opus: None, - seven_day_sonnet: Some(ClaudeWindow { - utilization: Some(3.0), - resets_at: None, - }), - seven_day_design: Some(ClaudeWindow { - utilization: Some(0.0), - resets_at: None, - }), - seven_day_routines: None, - extra_usage: None, - ..Default::default() - }; - let windows = claude_windows(&usage, now); - assert_eq!(windows.len(), 4); - assert_eq!(windows[0].label, "Session"); - assert_eq!(windows[0].card_id_for_test(), "session.v1"); - assert_eq!(windows[0].pace_window_key_for_test(), Some("session.v1")); - assert_eq!(windows[0].remaining_percent, 92.0); - assert_eq!(windows[0].pace_status.state, PaceState::LearningHistory); - assert_eq!(windows[0].pace_status.duration_seconds, Some(18_000)); + let out_of_range = UsageWindow::from_used_percent( + "Out of range".to_string(), + 150.0, + Some(now - chrono::Duration::seconds(1)), + now, + Some(-1), + ); assert_eq!( - windows[0].pace_status.duration_source, - Some(DurationSource::Contract) + out_of_range.used_percent, 100.0, + "invalid.out-of-range captures the current clamping baseline" + ); + assert!( + out_of_range.resets_at.is_some(), + "invalid.expired-reset captures the current emitted baseline" ); - assert_eq!(windows[1].label, "Weekly"); - assert_eq!(windows[1].card_id_for_test(), "weekly.v1"); - assert_eq!(windows[1].pace_window_key_for_test(), Some("weekly.v1")); - assert_eq!(windows[1].remaining_percent, 77.0); - assert_eq!(windows[1].pace_status.state, PaceState::LearningHistory); - assert_eq!(windows[1].pace_status.duration_seconds, Some(604_800)); assert_eq!( - windows[1].pace_status.duration_source, - Some(DurationSource::Contract) + out_of_range.window_minutes, None, + "invalid.contradictory-duration is not emitted as legacy duration" + ); + + let non_finite = + UsageWindow::from_used_percent("Non-finite".to_string(), f64::NAN, None, now, None); + assert!( + non_finite.used_percent.is_nan(), + "invalid.non-finite captures the current emitted baseline" ); - assert_eq!(windows[2].label, "Sonnet"); - assert_eq!(windows[2].card_id_for_test(), "sonnet.weekly.v1"); - assert_eq!(windows[2].pace_reason_for_test(), Some("missingReset")); - assert_eq!(windows[2].remaining_percent, 97.0); - assert_eq!(windows[3].label, "Designs"); - assert_eq!(windows[3].card_id_for_test(), "design.weekly.v1"); - assert_eq!(windows[3].remaining_percent, 100.0); } #[test] - fn claude_body_and_aliases_reject_invalid_utilization() { - let now = Utc.timestamp_opt(1_700_000_000, 0).single().unwrap(); - let reset = (now + chrono::Duration::hours(1)).to_rfc3339_opts(SecondsFormat::Secs, true); - for invalid in [-1.0, 150.0, f64::NAN, f64::INFINITY] { - assert!(map_claude_window( - "Session", - "session.v1", - DurationEvidence::contract(18_000), - &ClaudeWindow { - utilization: Some(invalid), - resets_at: Some(reset.clone()), - }, - now, - ) - .is_none()); - } - - let usage = ClaudeUsageResponse { - five_hour: Some(ClaudeWindow { - utilization: Some(150.0), - resets_at: Some(reset.clone()), - }), - seven_day: Some(ClaudeWindow { - utilization: Some(20.0), - resets_at: Some(reset.clone()), - }), - seven_day_design: Some(ClaudeWindow { - utilization: Some(f64::NAN), - resets_at: Some(reset.clone()), - }), - design: Some(ClaudeWindow { - utilization: Some(30.0), - resets_at: Some(reset.clone()), - }), - seven_day_routines: Some(ClaudeWindow { - utilization: Some(150.0), - resets_at: Some(reset.clone()), - }), - routines: Some(ClaudeWindow { - utilization: Some(40.0), - resets_at: Some(reset), - }), - ..Default::default() - }; - let windows = claude_windows(&usage, now); - assert_eq!( - windows - .iter() - .map(|window| (window.label.as_str(), window.used_percent)) - .collect::>(), - vec![ - ("Weekly", 20.0), - ("Designs", 30.0), - ("Daily Routines", 40.0) - ] - ); - } - - #[test] - fn claude_extra_usage_is_explicitly_non_recurring() { - let now = Utc.timestamp_opt(1_700_000_000, 0).single().unwrap(); - let window = claude_extra_usage_window( - Some(&ClaudeExtraUsage { - is_enabled: true, - monthly_limit: Some(10_000.0), - used_credits: Some(2_500.0), - utilization: None, - currency: Some("USD".to_string()), - }), - now, - ) + fn stage4_claude_extra_usage_is_active_without_recording_an_observation() { + let window = claude_extra_usage_window(Some(&ClaudeExtraUsage { + is_enabled: true, + monthly_limit: Some(10_000.0), + used_credits: Some(2_500.0), + utilization: None, + currency: Some("USD".to_string()), + })) .unwrap(); + assert_eq!(window.label, "Extra usage"); assert_eq!(window.card_id, "extra_usage.v1"); assert_eq!(window.used_percent, 25.0); + assert!(window.resets_at.is_none()); assert_eq!( - window.reset_text.as_deref(), - Some("Monthly cap: $25.00 / $100.00") + window.pace_status.window_key.as_deref(), + Some("extra_usage.v1") ); assert_eq!(window.pace_status.state, PaceState::Unavailable); - assert_eq!(window.pace_reason_for_test(), Some("nonRecurring")); - let wire = serde_json::to_value(window).unwrap(); - assert_eq!(wire["paceStatus"]["reason"], "nonRecurring"); - assert!(wire.get("windowMinutes").is_none()); - } + assert_eq!(window.pace_status.reason.as_deref(), Some("missingReset")); + assert!(window.duration_seconds.is_none()); + assert!(window.historical_pace.is_none()); - #[test] - fn invalid_claude_extra_usage_does_not_poison_valid_windows() { - let now = Utc.timestamp_opt(1_700_000_000, 0).single().unwrap(); - let usage = ClaudeUsageResponse { - five_hour: Some(ClaudeWindow { - utilization: Some(20.0), - resets_at: Some( - (now + chrono::Duration::hours(1)).to_rfc3339_opts(SecondsFormat::Secs, true), - ), - }), - extra_usage: Some(ClaudeExtraUsage { - is_enabled: true, - monthly_limit: Some(10_000.0), - used_credits: Some(2_500.0), - utilization: Some(f64::NAN), - currency: Some("USD".to_string()), - }), - ..Default::default() + let scope = TestRefreshScope::new("claude", "extra-usage"); + let account_scope = scope + .resolve_current("fixture", "extra-usage", b"extra-usage-marker") + .unwrap(); + let expected_scope = account_scope.as_str().to_string(); + let mut snapshot = AgentUsageSnapshot { + client_id: "claude".to_string(), + source: "oauth".to_string(), + updated_at: String::new(), + identity: None, + account_scope: Ok(account_scope), + windows: vec![window], + credits: None, + error: None, + transport_diagnostic: None, }; - - let windows = claude_windows(&usage, now); - assert_eq!(windows.len(), 1); - assert_eq!(windows[0].card_id, "session.v1"); - assert!(serde_json::to_value(windows).is_ok()); - } - - #[test] - fn decodes_claude_alias_windows_without_duplicate_error() { - let raw = r#"{ - "five_hour": { "utilization": 5, "resets_at": "2026-05-28T14:00:00Z" }, - "seven_day": { "utilization": 23, "resets_at": "2026-05-31T14:00:00Z" }, - "seven_day_sonnet": { "utilization": 3, "resets_at": null }, - "seven_day_omelette": { "utilization": 0, "resets_at": null }, - "omelette_promotional": { "utilization": 0, "resets_at": null }, - "seven_day_cowork": { "utilization": 0, "resets_at": null } - }"#; - let usage: ClaudeUsageResponse = serde_json::from_str(raw).unwrap(); - let now = Utc.timestamp_opt(1_700_000_000, 0).single().unwrap(); - let windows = claude_windows(&usage, now); + let calls = std::cell::Cell::new(0); + enrich_snapshot_with(&mut snapshot, 1_700_000_000, |active, observations, _| { + calls.set(calls.get() + 1); + assert_eq!( + active, + &[SeriesKey::new("claude", &expected_scope, "extra_usage.v1")] + ); + assert!(observations.is_empty()); + Ok(Vec::new()) + }); + assert_eq!(calls.get(), 1); assert_eq!( - windows.iter().map(|w| w.label.as_str()).collect::>(), - vec!["Session", "Weekly", "Sonnet", "Designs", "Daily Routines"] + snapshot.windows[0].pace_status.state, + PaceState::Unavailable + ); + assert_eq!( + snapshot.windows[0].pace_status.reason.as_deref(), + Some("missingReset") ); + scope.cleanup(); } #[test] - fn claude_named_and_alias_windows_use_exact_contract_durations() { - let now = Utc.timestamp_opt(1_700_000_000, 0).single().unwrap(); - let reset = (now + chrono::Duration::hours(1)).to_rfc3339_opts(SecondsFormat::Secs, true); - let window = |utilization| ClaudeWindow { - utilization: Some(utilization), - resets_at: Some(reset.clone()), - }; - let usage = ClaudeUsageResponse { - five_hour: Some(window(5.0)), - seven_day: Some(window(10.0)), - seven_day_oauth_apps: Some(window(15.0)), - seven_day_sonnet: Some(window(20.0)), - seven_day_opus: Some(window(25.0)), - seven_day_omelette: Some(window(30.0)), - seven_day_cowork: Some(window(35.0)), - ..Default::default() - }; - let windows = claude_windows(&usage, now); - let expected = [ - ("session.v1", 18_000), - ("weekly.v1", 604_800), - ("oauth_apps.weekly.v1", 604_800), - ("sonnet.weekly.v1", 604_800), - ("opus.weekly.v1", 604_800), - ("design.weekly.v1", 604_800), - ("routines.weekly.v1", 604_800), - ]; - assert_eq!(windows.len(), expected.len()); - for (window, (key, duration)) in windows.iter().zip(expected) { - assert_eq!(window.card_id, key); - assert_eq!(window.pace_status.state, PaceState::LearningHistory); - assert_eq!(window.pace_status.duration_seconds, Some(duration)); - assert_eq!( - window.pace_status.duration_source, - Some(DurationSource::Contract) - ); + fn stage4_emitted_unavailable_series_survives_capacity_admission() { + let scope = TestRefreshScope::new("claude", "emitted-capacity"); + let account_scope = scope + .resolve_current("fixture", "capacity", b"capacity-marker") + .unwrap(); + let account_scope_value = account_scope.as_str().to_string(); + let history_path = scope + .root() + .join(crate::agent_quota_history::HISTORY_FILE_NAME); + let seed_now = 1_800_000_000_i64; + let seed_reset = seed_now + 86_400; + let weekly_key = SeriesKey::new("claude", &account_scope_value, "weekly.v1"); + let mut seeded_keys = vec![weekly_key.clone()]; + seeded_keys.extend( + (0..crate::agent_quota_history::MAX_SERIES - 1).map(|index| { + SeriesKey::new( + "claude", + &account_scope_value, + format!("zzzz.{index:04}.v1"), + ) + }), + ); + for (sample_index, sampled_at) in [ + seed_now, + seed_now + 86_400 / 5, + seed_now + 2 * 86_400 / 5, + seed_now + 3 * 86_400 / 5, + seed_now + 4 * 86_400 / 5, + seed_reset - 1, + ] + .into_iter() + .enumerate() + { + let seeded_observations = seeded_keys + .iter() + .cloned() + .map(|key| QuotaObservation { + key, + reset_at: Some(seed_reset), + used_percent: 10.0 + sample_index as f64 * 10.0, + provider: None, + contract: Some(DurationEvidence::contract(86_400)), + }) + .collect::>(); + let seeded = crate::agent_quota_history::record_observations_at_path_and_evaluate( + &seeded_keys, + &seeded_observations, + sampled_at, + &history_path, + ) + .unwrap(); + assert_eq!(seeded.len(), crate::agent_quota_history::MAX_SERIES); } - for (alias, label, key) in [ - ("seven_day_design", "Designs", "design.weekly.v1"), - ("seven_day_claude_design", "Designs", "design.weekly.v1"), - ("claude_design", "Designs", "design.weekly.v1"), - ("design", "Designs", "design.weekly.v1"), - ("seven_day_omelette", "Designs", "design.weekly.v1"), - ("omelette", "Designs", "design.weekly.v1"), - ("omelette_promotional", "Designs", "design.weekly.v1"), - ("seven_day_routines", "Daily Routines", "routines.weekly.v1"), - ( - "seven_day_claude_routines", - "Daily Routines", - "routines.weekly.v1", - ), - ("claude_routines", "Daily Routines", "routines.weekly.v1"), - ("routines", "Daily Routines", "routines.weekly.v1"), - ("routine", "Daily Routines", "routines.weekly.v1"), - ("seven_day_cowork", "Daily Routines", "routines.weekly.v1"), - ("cowork", "Daily Routines", "routines.weekly.v1"), - ] { - let raw = format!(r#"{{"{alias}":{{"utilization":12,"resets_at":"{reset}"}}}}"#); - let usage: ClaudeUsageResponse = serde_json::from_str(&raw).unwrap(); - let windows = claude_windows(&usage, now); - assert_eq!(windows.len(), 1, "{alias}"); - assert_eq!(windows[0].label, label, "{alias}"); - assert_eq!(windows[0].card_id, key, "{alias}"); - assert_eq!( - windows[0].pace_status.duration_seconds, - Some(604_800), - "{alias}" - ); - assert_eq!( - windows[0].pace_status.duration_source, - Some(DurationSource::Contract), - "{alias}" - ); - assert_eq!( - windows[0].pace_status.state, - PaceState::LearningHistory, - "{alias}" - ); - } + let now = seed_reset + 15 * 60 + 1; + let now_date = Utc.timestamp_opt(now, 0).single().unwrap(); + let mut weekly = + UsageWindow::from_provider_used_percent("Weekly".to_string(), 20.0, None, now_date) + .with_identity( + "weekly.v1", + Some("weekly.v1".to_string()), + None, + Some(DurationEvidence::contract(86_400)), + ); + weekly.unavailable("missingReset"); + let new_window = UsageWindow::from_provider_used_percent( + "New quota".to_string(), + 5.0, + Some(Utc.timestamp_opt(now + 86_400, 0).single().unwrap()), + now_date, + ) + .with_identity( + "new.v1", + Some("new.v1".to_string()), + None, + Some(DurationEvidence::contract(86_400)), + ); + let mut snapshot = AgentUsageSnapshot { + client_id: "claude".to_string(), + source: "oauth".to_string(), + updated_at: String::new(), + identity: None, + account_scope: Ok(account_scope), + windows: vec![weekly, new_window], + credits: None, + error: None, + transport_diagnostic: None, + }; + + enrich_snapshot_with( + &mut snapshot, + now, + |active, observations, transaction_now| { + assert_eq!(active.len(), 2); + assert!(active.contains(&weekly_key)); + assert_eq!(observations.len(), 1); + assert_eq!(observations[0].key.window_key, "new.v1"); + crate::agent_quota_history::record_observations_at_path_and_evaluate( + active, + observations, + transaction_now, + &history_path, + ) + }, + ); + + let store: Value = serde_json::from_slice(&fs::read(&history_path).unwrap()).unwrap(); + let series = store["series"].as_array().unwrap(); + assert_eq!(series.len(), crate::agent_quota_history::MAX_SERIES); + assert!(series.iter().any(|entry| { + entry["providerId"] == "claude" + && entry["accountScope"] == account_scope_value + && entry["windowKey"] == "weekly.v1" + })); + assert_eq!( + snapshot.windows[0].pace_status.reason.as_deref(), + Some("missingReset") + ); + scope.cleanup(); } fn header_map(pairs: &[(&'static str, &'static str)]) -> reqwest::header::HeaderMap { @@ -6223,36 +6759,22 @@ mod tests { let now = Utc.timestamp_opt(1_700_000_000, 0).single().unwrap(); let headers = header_map(&[ ("anthropic-ratelimit-unified-5h-utilization", "0.11"), - ("anthropic-ratelimit-unified-5h-reset", "1700003600"), + ("anthropic-ratelimit-unified-5h-reset", "1783111200"), ("anthropic-ratelimit-unified-7d-utilization", "0.6"), - ("anthropic-ratelimit-unified-7d-reset", "1700172800"), + ("anthropic-ratelimit-unified-7d-reset", "1783504800"), ]); let windows = parse_unified_ratelimit_windows(&headers, now); assert_eq!(windows.len(), 2); assert_eq!(windows[0].label, "Session"); - assert_eq!(windows[0].card_id_for_test(), "session.v1"); - assert_eq!(windows[0].pace_window_key_for_test(), Some("session.v1")); assert!((windows[0].used_percent - 11.0).abs() < 1e-9); assert!((windows[0].remaining_percent - 89.0).abs() < 1e-9); + assert_eq!(windows[0].window_minutes, Some(300)); assert!(windows[0].resets_at.is_some()); assert!(windows[0].reset_text.is_some()); - assert_eq!(windows[0].pace_status.state, PaceState::LearningHistory); - assert_eq!(windows[0].pace_status.duration_seconds, Some(18_000)); - assert_eq!( - windows[0].pace_status.duration_source, - Some(DurationSource::Contract) - ); assert_eq!(windows[1].label, "Weekly"); - assert_eq!(windows[1].card_id_for_test(), "weekly.v1"); - assert_eq!(windows[1].pace_window_key_for_test(), Some("weekly.v1")); assert!((windows[1].used_percent - 60.0).abs() < 1e-9); assert!((windows[1].remaining_percent - 40.0).abs() < 1e-9); - assert_eq!(windows[1].pace_status.state, PaceState::LearningHistory); - assert_eq!(windows[1].pace_status.duration_seconds, Some(604_800)); - assert_eq!( - windows[1].pace_status.duration_source, - Some(DurationSource::Contract) - ); + assert_eq!(windows[1].window_minutes, Some(10_080)); } #[test] @@ -6293,24 +6815,10 @@ mod tests { let window = unified_ratelimit_window("Weekly", Some(0.4), None, now).unwrap(); assert!(window.resets_at.is_none()); assert!(window.reset_text.is_none()); - assert_eq!(window.pace_reason_for_test(), Some("missingReset")); - - let invalid_reset = parse_unified_ratelimit_windows( - &header_map(&[ - ("anthropic-ratelimit-unified-5h-utilization", "0.2"), - ("anthropic-ratelimit-unified-5h-reset", "bogus"), - ]), - now, - ); - assert_eq!(invalid_reset.len(), 1); - assert_eq!( - invalid_reset[0].pace_reason_for_test(), - Some("invalidEvidence") - ); } #[test] - fn unified_window_accepts_boundaries_and_rejects_invalid_fraction() { + fn unified_window_rejects_invalid_fraction_before_wire() { let now = Utc.timestamp_opt(1_700_000_000, 0).single().unwrap(); let zero = unified_ratelimit_window("Session", Some(0.0), None, now).unwrap(); assert!((zero.used_percent - 0.0).abs() < 1e-9); @@ -6318,20 +6826,151 @@ mod tests { let full = unified_ratelimit_window("Session", Some(1.0), None, now).unwrap(); assert!((full.used_percent - 100.0).abs() < 1e-9); assert!((full.remaining_percent - 0.0).abs() < 1e-9); - for invalid in [-0.1, 1.5, f64::NAN, f64::INFINITY] { - assert!(unified_ratelimit_window("Session", Some(invalid), None, now).is_none()); - } + + assert!(unified_ratelimit_window("Session", Some(1.5), None, now).is_none()); + assert!(unified_ratelimit_window("Session", Some(f64::NAN), None, now).is_none()); assert!(parse_unified_ratelimit_windows( &header_map(&[ - ("anthropic-ratelimit-unified-5h-utilization", "1.5"), - ("anthropic-ratelimit-unified-7d-utilization", "NaN"), + ("anthropic-ratelimit-unified-5h-utilization", "NaN"), + ("anthropic-ratelimit-unified-5h-reset", "1700003600"), ]), now, ) .is_empty()); + + // None utilization -> no window assert!(unified_ratelimit_window("Session", None, Some(1_783_111_200), now).is_none()); } + #[test] + fn provider_adapters_reject_invalid_percentages_before_wire() { + let now = Utc.timestamp_opt(1_700_000_000, 0).single().unwrap(); + assert!(map_claude_window( + "Session", + "session.v1", + DurationEvidence::contract(300 * 60), + &ClaudeWindow { + utilization: Some(150.0), + resets_at: None, + }, + now, + ) + .is_none()); + assert!(claude_extra_usage_window(Some(&ClaudeExtraUsage { + is_enabled: true, + monthly_limit: None, + used_credits: None, + utilization: Some(f64::NAN), + currency: None, + })) + .is_none()); + assert!(map_window_with_identity( + "Weekly", + CodexWindow { + used_percent: -1.0, + reset_at: 1_700_003_600, + limit_window_seconds: 604_800, + }, + now, + "main.weekly.v1", + Some("main.weekly.v1".to_string()), + ) + .is_none()); + + let valid_duplicate = codex_windows( + Some(&CodexRateLimit { + primary_window: Some(CodexWindow { + used_percent: 150.0, + reset_at: 1_700_003_600, + limit_window_seconds: 18_000, + }), + secondary_window: Some(CodexWindow { + used_percent: 20.0, + reset_at: 1_700_003_600, + limit_window_seconds: 18_000, + }), + }), + None, + now, + ); + assert_eq!(valid_duplicate.len(), 1); + assert_eq!(valid_duplicate[0].card_id, "main.session.v1"); + assert_eq!(valid_duplicate[0].used_percent, 20.0); + } + + #[test] + fn provider_payloads_isolate_malformed_percentage_rows() { + let now = Utc.timestamp_opt(1_700_000_000, 0).single().unwrap(); + for invalid in ["1e400", r#""NaN""#] { + let codex: CodexUsageResponse = serde_json::from_str(&format!( + r#"{{ + "rate_limit": {{ + "primary_window": {{ + "used_percent": {invalid}, + "reset_at": 1700003600, + "limit_window_seconds": 18000 + }}, + "secondary_window": {{ + "used_percent": 20, + "reset_at": 1700003600, + "limit_window_seconds": 18000 + }} + }} + }}"# + )) + .unwrap(); + let codex_windows = codex_windows(codex.rate_limit.as_ref(), None, now); + assert_eq!(codex_windows.len(), 1); + assert_eq!(codex_windows[0].card_id, "main.session.v1"); + assert_eq!(codex_windows[0].used_percent, 20.0); + + let claude: ClaudeUsageResponse = serde_json::from_str(&format!( + r#"{{ + "five_hour": {{ + "utilization": {invalid}, + "resets_at": "2023-11-15T00:13:20Z" + }}, + "seven_day": {{ + "utilization": 20, + "resets_at": "2023-11-21T22:13:20Z" + }}, + "seven_day_design": {{ + "utilization": {invalid}, + "resets_at": "2023-11-21T22:13:20Z" + }}, + "design": {{ + "utilization": 30, + "resets_at": "2023-11-21T22:13:20Z" + }}, + "seven_day_routines": {{ + "utilization": {invalid}, + "resets_at": "2023-11-21T22:13:20Z" + }}, + "routines": {{ + "utilization": 40, + "resets_at": "2023-11-21T22:13:20Z" + }}, + "extra_usage": {{ + "is_enabled": true, + "utilization": {invalid} + }} + }}"# + )) + .unwrap(); + let claude_windows = claude_windows(&claude, now); + assert_eq!(claude_windows.len(), 3); + assert!(claude_windows + .iter() + .any(|window| window.card_id == "weekly.v1" && window.used_percent == 20.0)); + assert!(claude_windows + .iter() + .any(|window| window.card_id == "design.weekly.v1" && window.used_percent == 30.0)); + assert!(claude_windows.iter().any( + |window| window.card_id == "routines.weekly.v1" && window.used_percent == 40.0 + )); + } + } + #[test] fn reads_claude_code_oauth_token_via_lookup() { let token = claude_token_from_lookup(|key| match key { @@ -6343,28 +6982,130 @@ mod tests { assert!(claude_token_from_lookup(|_| Some(" ".to_string())).is_none()); } - fn checkpoint_at( - target: Option, - ) -> impl FnMut(RefreshCheckpoint) -> Result<(), String> { - move |checkpoint| { - if Some(checkpoint) == target { - Err("injected crash".to_string()) - } else { - Ok(()) + #[test] + fn refreshes_or_expires_cached_windows() { + let base = Utc.timestamp_opt(1_700_000_000, 0).single().unwrap(); + let window = + unified_ratelimit_window("Session", Some(0.2), Some(1_700_000_000 + 3600), base) + .unwrap(); + + // 30 min later, still before the reset: reset_text recomputed to the + // shorter countdown (not the frozen original). + let later = base + chrono::Duration::seconds(1800); + let refreshed = refresh_cached_windows(std::slice::from_ref(&window), later).unwrap(); + assert_eq!(refreshed.len(), 1); + assert!(refreshed[0].reset_text.as_deref().unwrap().contains("30m")); + + // Past the reset: stale -> expire (None) so the caller re-probes. + let after = base + chrono::Duration::seconds(3700); + assert!(refresh_cached_windows(std::slice::from_ref(&window), after).is_none()); + } + + struct RecordingRefreshScope<'a> { + inner: &'a TestRefreshScope, + resolves: Mutex, + transfers: Mutex, Vec)>>, + } + + impl<'a> RecordingRefreshScope<'a> { + fn new(inner: &'a TestRefreshScope) -> Self { + Self { + inner, + resolves: Mutex::new(0), + transfers: Mutex::new(Vec::new()), } } - } - async fn codex_test_response(refresh_token: String) -> Result { - assert_eq!(refresh_token, "codex-old-refresh"); - Ok(serde_json::json!({ - "access_token": "codex-new-access", - "refresh_token": "codex-new-refresh" - })) + fn resolve_count(&self) -> usize { + *self.resolves.lock().unwrap() + } + + fn transfers(&self) -> Vec<(Vec, Vec)> { + self.transfers.lock().unwrap().clone() + } } - fn setup_codex_refresh( - tag: &str, + impl RefreshScopeTransaction for RecordingRefreshScope<'_> { + fn resolve_current( + &self, + semantic_source: &str, + canonical_location: &str, + marker: &[u8], + ) -> Result { + *self.resolves.lock().unwrap() += 1; + self.inner + .resolve_current(semantic_source, canonical_location, marker) + } + + fn transfer( + &self, + semantic_source: &str, + canonical_location: &str, + old_marker: &[u8], + new_marker: &[u8], + ) -> Result { + self.transfers + .lock() + .unwrap() + .push((old_marker.to_vec(), new_marker.to_vec())); + self.inner + .transfer(semantic_source, canonical_location, old_marker, new_marker) + } + } + + struct MetadataFailingRefreshScope<'a> { + inner: &'a TestRefreshScope, + } + + impl RefreshScopeTransaction for MetadataFailingRefreshScope<'_> { + fn resolve_current( + &self, + semantic_source: &str, + canonical_location: &str, + marker: &[u8], + ) -> Result { + self.inner + .resolve_current(semantic_source, canonical_location, marker) + } + + fn transfer( + &self, + semantic_source: &str, + canonical_location: &str, + old_marker: &[u8], + new_marker: &[u8], + ) -> Result { + self.inner.fail_metadata_save(); + self.inner + .transfer(semantic_source, canonical_location, old_marker, new_marker) + } + } + + fn checkpoint_at( + target: Option, + ) -> impl FnMut(RefreshCheckpoint) -> Result<(), ProviderFetchFailure> { + move |checkpoint| { + if Some(checkpoint) == target { + Err(ProviderFetchFailure::terminal("injected crash")) + } else { + Ok(()) + } + } + } + + async fn codex_test_response( + refresh_token: String, + _attempt_binding: ProviderCacheBinding, + ) -> Result { + assert_eq!(refresh_token, "codex-old-refresh"); + Ok(serde_json::json!({ + "access_token": "codex-new-access", + "refresh_token": "codex-new-refresh" + })) + } + + fn setup_codex_refresh( + tag: &str, ) -> (TestRefreshScope, PathBuf, AccountScope, Vec, String) { let scope = TestRefreshScope::new("codex", tag); let path = scope.root().join("codex/auth.json"); @@ -6394,11 +7135,11 @@ mod tests { (scope, path, old_scope, metadata, location) } - async fn run_codex_refresh( - scope: &TestRefreshScope, + async fn run_codex_refresh( + scope: &R, path: &Path, crash: Option, - ) -> Result<(CodexCredentials, Result), String> { + ) -> Result<(CodexCredentials, ProviderCacheBinding), ProviderFetchFailure> { refresh_codex_credentials_with( path, scope, @@ -6410,35 +7151,343 @@ mod tests { } #[tokio::test] - async fn codex_refresh_transfer_and_crash_boundaries_use_production_sequence() { + async fn codex_refresh_rejects_concurrent_account_switch_without_touching_b() { + const B_BYTES: &[u8] = br#"{ + "tokens": { + "access_token": "account-b-access", + "refresh_token": "account-b-refresh", + "id_token": "account-b-id", + "account_id": "account-b" + }, + "sibling": {"writer": "b", "revision": 2} +} +"#; + let (scope, path, _, metadata_before, _) = setup_codex_refresh("codex-target-switch"); + let recording = RecordingRefreshScope::new(&scope); + let request_path = path.clone(); + + let failure = refresh_codex_credentials_with( + &path, + &recording, + move |refresh_token, _attempt_binding| async move { + assert_eq!(refresh_token, "codex-old-refresh"); + fs::write(&request_path, B_BYTES).unwrap(); + Ok(serde_json::json!({ + "access_token": "codex-new-access", + "refresh_token": "codex-new-refresh" + })) + }, + save_codex_credentials, + checkpoint_at(None), + ) + .await + .unwrap_err(); + + assert!(matches!(failure, ProviderFetchFailure::Terminal { .. })); + assert!(recording.transfers().is_empty()); + assert_eq!(scope.metadata_bytes(), metadata_before); + let stored_bytes = fs::read(&path).unwrap(); + assert_eq!(stored_bytes, B_BYTES); + assert!(!String::from_utf8_lossy(&stored_bytes).contains("codex-new")); + let stored = load_codex_credentials_from(&path).unwrap(); + assert_eq!(stored.access_token, "account-b-access"); + assert_eq!(stored.refresh_token.as_deref(), Some("account-b-refresh")); + assert_eq!(stored.account_id.as_deref(), Some("account-b")); + scope.cleanup(); + } + + #[tokio::test] + async fn codex_refresh_patches_unchanged_target_and_preserves_siblings() { + let (scope, path, old_scope, _, location) = setup_codex_refresh("codex-target-unchanged"); + let original = serde_json::json!({ + "tokens": { + "access_token": "codex-old-access", + "refresh_token": "codex-old-refresh", + "id_token": "codex-old-id", + "token_sibling": {"keep": true} + }, + "sibling": {"writer": "before", "revision": 1} + }); + fs::write(&path, serde_json::to_vec_pretty(&original).unwrap()).unwrap(); + let current = serde_json::json!({ + "tokens": original["tokens"].clone(), + "sibling": {"writer": "codex-cli", "revision": 2}, + "unrelated": [1, 2, 3] + }); + let current_bytes = serde_json::to_vec_pretty(¤t).unwrap(); + let request_path = path.clone(); + + let (refreshed, post_binding) = refresh_codex_credentials_with( + &path, + &scope, + move |refresh_token, _attempt_binding| async move { + assert_eq!(refresh_token, "codex-old-refresh"); + fs::write(&request_path, current_bytes).unwrap(); + Ok(serde_json::json!({ + "access_token": "codex-new-access", + "refresh_token": "codex-new-refresh" + })) + }, + save_codex_credentials, + checkpoint_at(None), + ) + .await + .unwrap(); + + assert_eq!(refreshed.access_token, "codex-new-access"); + assert_eq!(post_binding.primary, old_scope); + let stored: Value = serde_json::from_slice(&fs::read(&path).unwrap()).unwrap(); + assert_eq!(stored["tokens"]["access_token"], "codex-new-access"); + assert_eq!(stored["tokens"]["refresh_token"], "codex-new-refresh"); + assert_eq!(stored["tokens"]["id_token"], "codex-old-id"); + assert_eq!(stored["tokens"]["token_sibling"]["keep"], true); + assert_eq!(stored["sibling"]["writer"], "codex-cli"); + assert_eq!(stored["sibling"]["revision"], 2); + assert_eq!(stored["unrelated"], serde_json::json!([1, 2, 3])); + assert_eq!( + scope + .resolve_current("codex-auth-json", &location, b"codex-old-refresh") + .unwrap(), + old_scope + ); + assert_eq!( + scope + .resolve_current("codex-auth-json", &location, b"codex-new-refresh") + .unwrap(), + old_scope + ); + scope.cleanup(); + } + + #[tokio::test] + async fn codex_refresh_rejects_concurrent_logout_without_restoring_a() { + const LOGGED_OUT_BYTES: &[u8] = br#"{ + "sibling": {"writer": "logout", "revision": 2} +} +"#; + let (scope, path, _, metadata_before, _) = setup_codex_refresh("codex-target-logout"); + let recording = RecordingRefreshScope::new(&scope); + let request_path = path.clone(); + + let failure = refresh_codex_credentials_with( + &path, + &recording, + move |refresh_token, _attempt_binding| async move { + assert_eq!(refresh_token, "codex-old-refresh"); + fs::write(&request_path, LOGGED_OUT_BYTES).unwrap(); + Ok(serde_json::json!({ + "access_token": "codex-new-access", + "refresh_token": "codex-new-refresh" + })) + }, + save_codex_credentials, + checkpoint_at(None), + ) + .await + .unwrap_err(); + + assert!(matches!(failure, ProviderFetchFailure::Terminal { .. })); + assert!(recording.transfers().is_empty()); + assert_eq!(scope.metadata_bytes(), metadata_before); + let stored_bytes = fs::read(&path).unwrap(); + assert_eq!(stored_bytes, LOGGED_OUT_BYTES); + assert!(!String::from_utf8_lossy(&stored_bytes).contains("codex-new")); + let stored: Value = serde_json::from_slice(&stored_bytes).unwrap(); + assert!(stored.get("tokens").is_none()); + scope.cleanup(); + } + + #[tokio::test] + async fn codex_refresh_canonicalizes_tokens_and_preserves_unrotated_marker() { + for (tag, refresh_value) in [ + ("missing", None), + ("null", Some(Value::Null)), + ("empty", Some(Value::String(String::new()))), + ("whitespace", Some(Value::String(" \t\n ".to_string()))), + ( + "non-string", + Some(serde_json::json!({ "unexpected": true })), + ), + ] { + let (scope, path, old_scope, _, _) = + setup_codex_refresh(&format!("codex-canonical-{tag}")); + let recording = RecordingRefreshScope::new(&scope); + let mut response = serde_json::json!({ + "access_token": { "unexpected": true }, + "accessToken": " codex-new-access ", + "id_token": " \t\n ", + "idToken": " codex-new-id " + }); + if let Some(refresh_value) = refresh_value { + response + .as_object_mut() + .unwrap() + .insert("refresh_token".to_string(), refresh_value); + } + + let (refreshed, post_binding) = refresh_codex_credentials_with( + &path, + &recording, + move |refresh_token, _attempt_binding| async move { + assert_eq!(refresh_token, "codex-old-refresh"); + Ok(response) + }, + save_codex_credentials, + checkpoint_at(None), + ) + .await + .unwrap(); + + assert_eq!(refreshed.access_token, "codex-new-access", "{tag}"); + assert_eq!( + refreshed.refresh_token.as_deref(), + Some("codex-old-refresh"), + "{tag}" + ); + assert_eq!(refreshed.id_token.as_deref(), Some("codex-new-id"), "{tag}"); + assert_eq!(post_binding.primary, old_scope, "{tag}"); + assert_eq!(post_binding.corroborating, None, "{tag}"); + assert_eq!( + recording.transfers(), + vec![(b"codex-old-refresh".to_vec(), b"codex-old-refresh".to_vec())], + "{tag}" + ); + + let stored: Value = serde_json::from_str(&fs::read_to_string(&path).unwrap()).unwrap(); + assert_eq!( + stored["tokens"]["refresh_token"], + Value::String("codex-old-refresh".to_string()), + "{tag}" + ); + let reloaded = load_codex_credentials_from(&path).unwrap(); + assert_eq!(reloaded.access_token, refreshed.access_token, "{tag}"); + assert_eq!(reloaded.refresh_token, refreshed.refresh_token, "{tag}"); + assert_eq!(reloaded.id_token, refreshed.id_token, "{tag}"); + assert_eq!(reloaded.scope_marker(), refreshed.scope_marker(), "{tag}"); + assert_eq!( + scope + .resolve_current( + reloaded.scope_slot.semantic_source, + &reloaded.scope_slot.canonical_location, + reloaded.scope_marker(), + ) + .unwrap(), + old_scope, + "{tag}" + ); + scope.cleanup(); + } + } + + #[tokio::test] + async fn codex_refresh_rejects_invalid_success_schema_before_state_or_usage() { + for (tag, response) in [ + ("array", serde_json::json!([])), + ("empty-object", serde_json::json!({})), + ("null", Value::Null), + ("bool", Value::Bool(true)), + ("string", Value::String("codex-new-access".to_string())), + ( + "blank-access", + serde_json::json!({ "access_token": " \t\n " }), + ), + ( + "invalid-aliases", + serde_json::json!({ + "access_token": false, + "accessToken": [], + "refreshToken": " codex-new-refresh " + }), + ), + ] { + let (scope, path, old_scope, metadata_before, location) = + setup_codex_refresh(&format!("codex-invalid-schema-{tag}")); + let credentials_before = fs::read(&path).unwrap(); + let recording = RecordingRefreshScope::new(&scope); + let save_calls = std::cell::Cell::new(0); + let usage_calls = std::cell::Cell::new(0); + + let refresh_result = refresh_codex_credentials_with( + &path, + &recording, + move |refresh_token, _attempt_binding| async move { + assert_eq!(refresh_token, "codex-old-refresh"); + Ok(response) + }, + |_| -> Result { + save_calls.set(save_calls.get() + 1); + Err("unexpected save".to_string()) + }, + checkpoint_at(None), + ) + .await; + let result: Result<(), ProviderFetchFailure> = + request_after_verified_binding(refresh_result, |_| async { + usage_calls.set(usage_calls.get() + 1); + Ok(()) + }) + .await; + + assert!( + matches!(result, Err(ProviderFetchFailure::Terminal { .. })), + "{tag}" + ); + assert!(recording.transfers().is_empty(), "{tag}"); + assert_eq!(save_calls.get(), 0, "{tag}"); + assert_eq!(usage_calls.get(), 0, "{tag}"); + assert_eq!(fs::read(&path).unwrap(), credentials_before, "{tag}"); + assert_eq!(scope.metadata_bytes(), metadata_before, "{tag}"); + let reloaded = load_codex_credentials_from(&path).unwrap(); + assert_eq!(reloaded.access_token, "codex-old-access", "{tag}"); + assert_eq!( + reloaded.refresh_token.as_deref(), + Some("codex-old-refresh"), + "{tag}" + ); + assert!(reloaded.last_refresh.is_none(), "{tag}"); + assert_eq!( + scope + .resolve_current("codex-auth-json", &location, reloaded.scope_marker()) + .unwrap(), + old_scope, + "{tag}" + ); + scope.cleanup(); + } + } + + #[tokio::test] + async fn codex_refresh_crash_boundaries_and_scope_gate_use_production_sequence() { + // These checkpoints model process stops, not a cross-resource transaction: + // after credential persistence, metadata may still be the pre-refresh bytes. for boundary in [ RefreshCheckpoint::Reloaded, RefreshCheckpoint::NetworkReturned, - RefreshCheckpoint::MetadataHandled, RefreshCheckpoint::CredentialsPersisted, + RefreshCheckpoint::MetadataHandled, ] { let (scope, path, old_scope, before, location) = setup_codex_refresh("codex-crash"); - assert_eq!( - run_codex_refresh(&scope, &path, Some(boundary)) - .await - .unwrap_err(), - "injected crash" + let failure = run_codex_refresh(&scope, &path, Some(boundary)) + .await + .unwrap_err(); + assert!(matches!( + failure, + ProviderFetchFailure::Terminal { ref display } if display == "injected crash" + )); + let credentials_persisted = matches!( + boundary, + RefreshCheckpoint::CredentialsPersisted | RefreshCheckpoint::MetadataHandled ); let stored = load_codex_credentials_from(&path).unwrap(); assert_eq!( stored.refresh_token.as_deref(), - Some(if boundary == RefreshCheckpoint::CredentialsPersisted { + Some(if credentials_persisted { "codex-new-refresh" } else { "codex-old-refresh" }) ); - if matches!( - boundary, - RefreshCheckpoint::Reloaded | RefreshCheckpoint::NetworkReturned - ) { - assert_eq!(scope.metadata_bytes(), before); - } else { + if boundary == RefreshCheckpoint::MetadataHandled { assert_ne!(scope.metadata_bytes(), before); assert_eq!( scope @@ -6452,16 +7501,19 @@ mod tests { .unwrap(), old_scope ); + } else { + assert_eq!(scope.metadata_bytes(), before); } scope.cleanup(); } let (scope, path, old_scope, before, location) = setup_codex_refresh("codex-metadata-fail"); - scope.fail_metadata_save(); - let (refreshed, scope_outcome) = run_codex_refresh(&scope, &path, None).await.unwrap(); - assert_eq!(refreshed.access_token, "codex-new-access"); - assert_eq!(scope_outcome, Err(AccountScopeError::MetadataWrite)); + let auth_before = fs::read(&path).unwrap(); + let failing = MetadataFailingRefreshScope { inner: &scope }; + let failure = run_codex_refresh(&failing, &path, None).await.unwrap_err(); + assert!(matches!(failure, ProviderFetchFailure::Terminal { .. })); assert_eq!(scope.metadata_bytes(), before); + assert_eq!(fs::read(&path).unwrap(), auth_before); let persisted = load_codex_credentials_from(&path).unwrap(); assert_eq!(persisted.access_token, "codex-old-access"); assert_eq!( @@ -6476,58 +7528,129 @@ mod tests { ); scope.cleanup(); - let (scope, path, _old_scope, before, _) = - setup_codex_refresh("codex-metadata-fail-unchanged"); - scope.fail_metadata_save(); - let (refreshed, scope_outcome) = refresh_codex_credentials_with( + const CONCURRENT_LOGIN_BYTES: &[u8] = br#"{ + "tokens": { + "access_token": "concurrent-access", + "refresh_token": "concurrent-refresh", + "id_token": "concurrent-id", + "account_id": "concurrent-account" + }, + "sibling": {"writer": "codex-cli", "revision": 3} +} +"#; + let (scope, path, _, metadata_before, _) = + setup_codex_refresh("codex-metadata-fail-concurrent-login"); + let failing = MetadataFailingRefreshScope { inner: &scope }; + let save_path = path.clone(); + let failure = refresh_codex_credentials_with( &path, - &scope, - |refresh_token| async move { - assert_eq!(refresh_token, "codex-old-refresh"); - Ok(serde_json::json!({ "access_token": "codex-new-access" })) + &failing, + codex_test_response, + move |credentials| { + let receipt = save_codex_credentials(credentials)?; + fs::write(&save_path, CONCURRENT_LOGIN_BYTES) + .map_err(|error| format!("inject concurrent Codex login: {error}"))?; + Ok(receipt) }, - save_codex_credentials, checkpoint_at(None), ) .await - .unwrap(); - assert_eq!(scope_outcome, Err(AccountScopeError::MetadataWrite)); - assert_eq!(scope.metadata_bytes(), before); - assert_eq!( - refreshed.refresh_token.as_deref(), - Some("codex-old-refresh") - ); - let persisted = load_codex_credentials_from(&path).unwrap(); - assert_eq!(persisted.access_token, "codex-new-access"); - assert_eq!( - persisted.refresh_token.as_deref(), - Some("codex-old-refresh") - ); + .unwrap_err(); + assert!(matches!(failure, ProviderFetchFailure::Terminal { .. })); + assert_eq!(scope.metadata_bytes(), metadata_before); + assert_eq!(fs::read(&path).unwrap(), CONCURRENT_LOGIN_BYTES); + scope.cleanup(); + + let (scope, path, _, metadata_before, _) = setup_codex_refresh("codex-save-fail"); + let auth_before = fs::read(&path).unwrap(); + let recording = RecordingRefreshScope::new(&scope); + let failure = refresh_codex_credentials_with( + &path, + &recording, + codex_test_response, + |_| Err("injected save failure".to_string()), + checkpoint_at(None), + ) + .await + .unwrap_err(); + assert!(matches!(failure, ProviderFetchFailure::Terminal { .. })); + assert!(recording.transfers().is_empty()); + assert_eq!(scope.metadata_bytes(), metadata_before); + assert_eq!(fs::read(&path).unwrap(), auth_before); scope.cleanup(); let (scope, path, old_scope, _, location) = setup_codex_refresh("codex-success"); - let (_, scope_outcome) = run_codex_refresh(&scope, &path, None).await.unwrap(); - assert_eq!(scope_outcome.unwrap(), old_scope); + let recording = RecordingRefreshScope::new(&scope); + let (_, post_binding) = refresh_codex_credentials_with( + &path, + &recording, + codex_test_response, + save_codex_credentials, + checkpoint_at(None), + ) + .await + .unwrap(); + assert_eq!(recording.resolve_count(), 1); + assert_eq!(post_binding.primary, old_scope); + assert_eq!(post_binding.corroborating, None); assert_eq!( scope .resolve_current("codex-auth-json", &location, b"codex-new-refresh") .unwrap(), old_scope ); - assert_eq!( - load_codex_credentials_from(&path) - .unwrap() - .refresh_token - .as_deref(), - Some("codex-new-refresh") - ); scope.cleanup(); } - async fn claude_test_response(refresh_token: String) -> Result { + #[tokio::test] + async fn codex_refresh_transient_uses_lock_reloaded_binding_not_outer_binding() { + let (scope, path, inner_scope, _, location) = setup_codex_refresh("codex-lock-binding"); + let outer_scope = scope + .resolve_current("codex-auth-json", &location, b"outer-refresh-a") + .unwrap(); + assert_ne!(outer_scope, inner_scope); + let expected = ProviderCacheBinding::primary(inner_scope); + let request_expected = expected.clone(); + + let failure = refresh_codex_credentials_with( + &path, + &scope, + move |refresh_token, attempt_binding| async move { + assert_eq!(refresh_token, "codex-old-refresh"); + assert_eq!(attempt_binding, request_expected); + Err(ProviderFetchFailure::transient( + "Codex token refresh failed. Retrying automatically.", + Some(attempt_binding), + SafeTransportDiagnostic::from_facts(TransportErrorFacts::synthetic( + true, + false, + TransportPhase::Request, + None, + )), + )) + }, + save_codex_credentials, + checkpoint_at(None), + ) + .await + .unwrap_err(); + + match failure { + ProviderFetchFailure::Transient { + attempt_binding, .. + } => assert_eq!(attempt_binding, Some(expected)), + ProviderFetchFailure::Terminal { .. } => panic!("timeout must remain transient"), + } + scope.cleanup(); + } + + async fn claude_test_response( + refresh_token: String, + _attempt_binding: ProviderCacheBinding, + ) -> Result { assert_eq!(refresh_token, "claude-old-refresh"); Ok(ClaudeRefreshResponse { - access_token: "claude-new-access".to_string(), + access_token: " claude-new-access ".to_string(), refresh_token: Some("claude-new-refresh".to_string()), expires_in: 3_600, }) @@ -6582,7 +7705,14 @@ mod tests { path: &Path, original: &ClaudeCredentials, crash: Option, - ) -> Result<(ClaudeCredentials, Result), String> { + ) -> Result< + ( + ClaudeCredentials, + AccountScope, + Option, + ), + ProviderFetchFailure, + > { let reload_path = path.to_path_buf(); let save_path = path.to_path_buf(); refresh_claude_credentials_with( @@ -6603,78 +7733,254 @@ mod tests { .await } - fn stored_claude_credentials(path: &Path) -> ClaudeCredentials { + fn stored_claude_refresh_token(path: &Path) -> Option { parse_claude_credentials_data( &fs::read_to_string(path).unwrap(), ClaudeCredentialSource::File, ) .unwrap() + .refresh_token } #[tokio::test] - async fn claude_refresh_invalid_new_marker_preserves_old_lineage_and_store() { - for (tag, refresh_value) in [ - ("empty", serde_json::json!("")), - ("non-string", serde_json::json!({ "unexpected": true })), - ] { - let (scope, path, original, old_scope, _, location) = - setup_claude_refresh(&format!("claude-invalid-refresh-{tag}")); - let response: ClaudeRefreshResponse = serde_json::from_value(serde_json::json!({ - "access_token": "claude-new-access", - "refresh_token": refresh_value, - "expires_in": 3600 - })) - .unwrap(); - let reload_path = path.clone(); - let save_path = path.clone(); - let (refreshed, scope_outcome) = refresh_claude_credentials_with( - &original, - &scope, - move |template| { - let raw = fs::read_to_string(&reload_path) - .map_err(|error| format!("reload Claude test credentials: {error}"))?; - let mut credentials = - parse_claude_credentials_data(&raw, ClaudeCredentialSource::File)?; - credentials.scope_slot = template.scope_slot.clone(); - Ok(credentials) - }, - move |refresh_token| async move { - assert_eq!(refresh_token, "claude-old-refresh"); - Ok(response) - }, - move |credentials| save_claude_credentials_to_file(credentials, &save_path), - checkpoint_at(None), - ) + async fn claude_file_refresh_rejects_concurrent_target_change_without_touching_b() { + const B_BYTES: &[u8] = br#"{ + "claudeAiOauth": { + "accessToken": "account-b-access", + "refreshToken": "account-b-refresh", + "expiresAt": 4102444800000 + }, + "sibling": {"writer": "b", "revision": 2} +} +"#; + let (scope, path, original, old_scope, _, _) = + setup_claude_refresh("claude-file-target-race"); + let reload_path = path.clone(); + let request_path = path.clone(); + let save_path = path.clone(); + let save_failed = std::rc::Rc::new(std::cell::Cell::new(false)); + let observed_save_failure = std::rc::Rc::clone(&save_failed); + + let (refreshed, scope_outcome, cache_binding) = refresh_claude_credentials_with( + &original, + &scope, + move |template| { + let raw = fs::read_to_string(&reload_path) + .map_err(|error| format!("reload Claude test credentials: {error}"))?; + let mut credentials = + parse_claude_credentials_data(&raw, ClaudeCredentialSource::File)?; + credentials.scope_slot = template.scope_slot.clone(); + Ok(credentials) + }, + move |refresh_token, _attempt_binding| async move { + assert_eq!(refresh_token, "claude-old-refresh"); + fs::write(&request_path, B_BYTES).unwrap(); + Ok(ClaudeRefreshResponse { + access_token: "claude-new-access".to_string(), + refresh_token: Some("claude-new-refresh".to_string()), + expires_in: 3_600, + }) + }, + move |credentials| { + let result = save_claude_credentials_to_file(credentials, &save_path); + observed_save_failure.set(result.is_err()); + result + }, + checkpoint_at(None), + ) + .await + .unwrap(); + + assert_eq!(refreshed.access_token, "claude-new-access"); + assert_eq!(scope_outcome, old_scope); + assert_eq!(cache_binding, None); + assert!(save_failed.get()); + assert_eq!(fs::read(&path).unwrap(), B_BYTES); + scope.cleanup(); + } + + #[tokio::test] + async fn claude_file_refresh_preserves_concurrent_top_level_sibling() { + const CURRENT_WITH_NEW_SIBLING: &str = r#"{ + "claudeAiOauth": { + "accessToken": "claude-old-access", + "refreshToken": "claude-old-refresh", + "expiresAt": 0 + }, + "sibling": {"writer": "claude-cli", "revision": 2} + }"#; + let (scope, path, original, old_scope, _, _) = + setup_claude_refresh("claude-file-sibling-race"); + let reload_path = path.clone(); + let request_path = path.clone(); + let save_path = path.clone(); + + let (refreshed, scope_outcome, cache_binding) = refresh_claude_credentials_with( + &original, + &scope, + move |template| { + let raw = fs::read_to_string(&reload_path) + .map_err(|error| format!("reload Claude test credentials: {error}"))?; + let mut credentials = + parse_claude_credentials_data(&raw, ClaudeCredentialSource::File)?; + credentials.scope_slot = template.scope_slot.clone(); + Ok(credentials) + }, + move |refresh_token, _attempt_binding| async move { + assert_eq!(refresh_token, "claude-old-refresh"); + fs::write(&request_path, CURRENT_WITH_NEW_SIBLING).unwrap(); + Ok(ClaudeRefreshResponse { + access_token: "claude-new-access".to_string(), + refresh_token: Some("claude-new-refresh".to_string()), + expires_in: 3_600, + }) + }, + move |credentials| save_claude_credentials_to_file(credentials, &save_path), + checkpoint_at(None), + ) + .await + .unwrap(); + + assert_eq!(refreshed.access_token, "claude-new-access"); + assert_eq!(scope_outcome, old_scope); + assert_eq!( + cache_binding, + Some(ProviderCacheBinding::primary(old_scope.clone())) + ); + let stored: Value = serde_json::from_str(&fs::read_to_string(&path).unwrap()).unwrap(); + assert_eq!(stored["claudeAiOauth"]["accessToken"], "claude-new-access"); + assert_eq!( + stored["claudeAiOauth"]["refreshToken"], + "claude-new-refresh" + ); + assert_eq!(stored["sibling"]["writer"], "claude-cli"); + assert_eq!(stored["sibling"]["revision"], 2); + scope.cleanup(); + } + + #[tokio::test] + async fn claude_refresh_invalid_new_refresh_preserves_old_marker_and_store() { + for (tag, refresh_value) in [ + ("claude-invalid-refresh-empty", serde_json::json!("")), + ( + "claude-invalid-refresh-non-string", + serde_json::json!({ "unexpected": true }), + ), + ] { + let (scope, path, original, old_scope, _, location) = setup_claude_refresh(tag); + let response: ClaudeRefreshResponse = serde_json::from_value(serde_json::json!({ + "access_token": "claude-new-access", + "refresh_token": refresh_value, + "expires_in": 3600 + })) + .unwrap(); + let reload_path = path.clone(); + let save_path = path.clone(); + let (refreshed, scope_outcome, cache_binding) = refresh_claude_credentials_with( + &original, + &scope, + move |template| { + let raw = fs::read_to_string(&reload_path) + .map_err(|error| format!("reload Claude test credentials: {error}"))?; + let mut credentials = + parse_claude_credentials_data(&raw, ClaudeCredentialSource::File)?; + credentials.scope_slot = template.scope_slot.clone(); + Ok(credentials) + }, + move |refresh_token, _attempt_binding| async move { + assert_eq!(refresh_token, "claude-old-refresh"); + Ok(response) + }, + move |credentials| save_claude_credentials_to_file(credentials, &save_path), + checkpoint_at(None), + ) .await .unwrap(); - assert_eq!(refreshed.access_token, "claude-new-access", "{tag}"); + assert_eq!(refreshed.access_token, "claude-new-access"); assert_eq!( refreshed.refresh_token.as_deref(), - Some("claude-old-refresh"), - "{tag}" + Some("claude-old-refresh") + ); + assert_eq!(scope_outcome, old_scope); + assert_eq!( + cache_binding, + Some(ProviderCacheBinding::primary(old_scope.clone())) ); - assert_eq!(scope_outcome.unwrap(), old_scope, "{tag}"); assert_eq!( scope .resolve_current("claude-login-file", &location, b"claude-old-refresh") .unwrap(), - old_scope, - "{tag}" + old_scope ); - let stored = stored_claude_credentials(&path); - assert_eq!(stored.access_token, "claude-new-access", "{tag}"); assert_eq!( - stored.refresh_token.as_deref(), - Some("claude-old-refresh"), - "{tag}" + stored_claude_refresh_token(&path).as_deref(), + Some("claude-old-refresh") + ); + let stored: Value = serde_json::from_str(&fs::read_to_string(&path).unwrap()).unwrap(); + assert_eq!( + stored["claudeAiOauth"]["refreshToken"], + Value::String("claude-old-refresh".to_string()) ); scope.cleanup(); } } #[tokio::test] - async fn claude_refresh_transfer_and_crash_boundaries_use_production_sequence() { + async fn claude_refresh_blank_access_token_is_terminal_before_metadata_or_save() { + let (scope, path, original, old_scope, before, location) = + setup_claude_refresh("claude-blank-access"); + let store_before = fs::read(&path).unwrap(); + let reload_path = path.clone(); + let save_calls = std::cell::Cell::new(0); + + let failure = refresh_claude_credentials_with( + &original, + &scope, + move |template| { + let raw = fs::read_to_string(&reload_path) + .map_err(|error| format!("reload Claude test credentials: {error}"))?; + let mut credentials = + parse_claude_credentials_data(&raw, ClaudeCredentialSource::File)?; + credentials.scope_slot = template.scope_slot.clone(); + Ok(credentials) + }, + |refresh_token, _attempt_binding| async move { + assert_eq!(refresh_token, "claude-old-refresh"); + Ok(ClaudeRefreshResponse { + access_token: " \t\n ".to_string(), + refresh_token: Some("claude-new-refresh".to_string()), + expires_in: 3_600, + }) + }, + |_| { + save_calls.set(save_calls.get() + 1); + Ok(()) + }, + checkpoint_at(None), + ) + .await + .unwrap_err(); + + assert!(matches!( + failure, + ProviderFetchFailure::Terminal { ref display } + if display == "Claude OAuth refresh response has no access token." + )); + assert_eq!(save_calls.get(), 0); + assert_eq!(scope.metadata_bytes(), before); + assert_eq!(fs::read(&path).unwrap(), store_before); + assert_eq!( + scope + .resolve_current("claude-login-file", &location, b"claude-old-refresh") + .unwrap(), + old_scope + ); + scope.cleanup(); + } + + #[tokio::test] + async fn claude_refresh_crash_boundaries_and_scope_gate_use_production_sequence() { for boundary in [ RefreshCheckpoint::Reloaded, RefreshCheckpoint::NetworkReturned, @@ -6683,14 +7989,15 @@ mod tests { ] { let (scope, path, original, old_scope, before, location) = setup_claude_refresh("claude-crash"); + let failure = run_claude_refresh(&scope, &path, &original, Some(boundary)) + .await + .unwrap_err(); + assert!(matches!( + failure, + ProviderFetchFailure::Terminal { ref display } if display == "injected crash" + )); assert_eq!( - run_claude_refresh(&scope, &path, &original, Some(boundary)) - .await - .unwrap_err(), - "injected crash" - ); - assert_eq!( - stored_claude_credentials(&path).refresh_token.as_deref(), + stored_claude_refresh_token(&path).as_deref(), Some(if boundary == RefreshCheckpoint::CredentialsPersisted { "claude-new-refresh" } else { @@ -6723,14 +8030,13 @@ mod tests { let (scope, path, original, old_scope, before, location) = setup_claude_refresh("claude-metadata-fail"); scope.fail_metadata_save(); - let (refreshed, scope_outcome) = run_claude_refresh(&scope, &path, &original, None) + let failure = run_claude_refresh(&scope, &path, &original, None) .await - .unwrap(); - assert_eq!(refreshed.access_token, "claude-new-access"); - assert_eq!(scope_outcome, Err(AccountScopeError::MetadataWrite)); + .unwrap_err(); + assert!(matches!(failure, ProviderFetchFailure::Terminal { .. })); assert_eq!(scope.metadata_bytes(), before); assert_eq!( - stored_claude_credentials(&path).refresh_token.as_deref(), + stored_claude_refresh_token(&path).as_deref(), Some("claude-old-refresh") ); assert_eq!( @@ -6741,12 +8047,10 @@ mod tests { ); scope.cleanup(); - let (scope, path, original, _old_scope, before, _) = - setup_claude_refresh("claude-metadata-fail-unchanged"); - scope.fail_metadata_save(); + let (scope, path, original, old_scope, _, location) = + setup_claude_refresh("claude-save-fail"); let reload_path = path.clone(); - let save_path = path.clone(); - let (refreshed, scope_outcome) = refresh_claude_credentials_with( + let (refreshed, scope_outcome, cache_binding) = refresh_claude_credentials_with( &original, &scope, move |template| { @@ -6757,68 +8061,844 @@ mod tests { credentials.scope_slot = template.scope_slot.clone(); Ok(credentials) }, - |refresh_token| async move { - assert_eq!(refresh_token, "claude-old-refresh"); - Ok(ClaudeRefreshResponse { - access_token: "claude-new-access".to_string(), - refresh_token: None, - expires_in: 3_600, - }) - }, - move |credentials| save_claude_credentials_to_file(credentials, &save_path), + claude_test_response, + |_| Err("injected save failure".to_string()), checkpoint_at(None), ) .await .unwrap(); - assert_eq!(scope_outcome, Err(AccountScopeError::MetadataWrite)); - assert_eq!(scope.metadata_bytes(), before); + assert_eq!(refreshed.access_token, "claude-new-access"); + assert_eq!(scope_outcome, old_scope); + assert_eq!(cache_binding, None); assert_eq!( - refreshed.refresh_token.as_deref(), + stored_claude_refresh_token(&path).as_deref(), Some("claude-old-refresh") ); - let persisted = stored_claude_credentials(&path); - assert_eq!(persisted.access_token, "claude-new-access"); assert_eq!( - persisted.refresh_token.as_deref(), - Some("claude-old-refresh") + scope + .resolve_current("claude-login-file", &location, b"claude-new-refresh") + .unwrap(), + old_scope ); scope.cleanup(); let (scope, path, original, old_scope, _, location) = setup_claude_refresh("claude-success"); - let (_, scope_outcome) = run_claude_refresh(&scope, &path, &original, None) + let (_, scope_outcome, cache_binding) = run_claude_refresh(&scope, &path, &original, None) .await .unwrap(); - assert_eq!(scope_outcome.unwrap(), old_scope); + assert_eq!(scope_outcome, old_scope); + assert_eq!( + cache_binding, + Some(ProviderCacheBinding::primary(old_scope.clone())) + ); assert_eq!( scope .resolve_current("claude-login-file", &location, b"claude-new-refresh") .unwrap(), old_scope ); - assert_eq!( - stored_claude_credentials(&path).refresh_token.as_deref(), - Some("claude-new-refresh") - ); + scope.cleanup(); + } + + #[tokio::test] + async fn claude_refresh_transient_uses_lock_reloaded_binding_not_outer_binding() { + let (scope, path, mut original, inner_scope, _, location) = + setup_claude_refresh("claude-lock-binding"); + original.refresh_token = Some("outer-refresh-a".to_string()); + let outer_scope = scope + .resolve_current("claude-login-file", &location, b"outer-refresh-a") + .unwrap(); + assert_ne!(outer_scope, inner_scope); + let expected = ProviderCacheBinding::primary(inner_scope); + let request_expected = expected.clone(); + let reload_path = path.clone(); + + let failure = refresh_claude_credentials_with( + &original, + &scope, + move |template| { + let raw = fs::read_to_string(&reload_path) + .map_err(|error| format!("reload Claude test credentials: {error}"))?; + let mut credentials = + parse_claude_credentials_data(&raw, ClaudeCredentialSource::File)?; + credentials.scope_slot = template.scope_slot.clone(); + Ok(credentials) + }, + move |refresh_token, attempt_binding| async move { + assert_eq!(refresh_token, "claude-old-refresh"); + assert_eq!(attempt_binding, request_expected); + Err(ProviderFetchFailure::transient( + "Claude OAuth refresh failed. Retrying automatically.", + Some(attempt_binding), + SafeTransportDiagnostic::from_facts(TransportErrorFacts::synthetic( + true, + false, + TransportPhase::Request, + None, + )), + )) + }, + |_| Ok(()), + checkpoint_at(None), + ) + .await + .unwrap_err(); + + match failure { + ProviderFetchFailure::Transient { + attempt_binding, .. + } => assert_eq!(attempt_binding, Some(expected)), + ProviderFetchFailure::Terminal { .. } => panic!("timeout must remain transient"), + } scope.cleanup(); } #[test] - fn refreshes_or_expires_cached_windows() { - let base = Utc.timestamp_opt(1_700_000_000, 0).single().unwrap(); - let window = - unified_ratelimit_window("Session", Some(0.2), Some(1_700_000_000 + 3600), base) - .unwrap(); + fn stage4_codex_and_claude_matrix_assigns_semantic_keys() { + let now = Utc.timestamp_opt(1_700_000_000, 0).single().unwrap(); + let rate_limit = CodexRateLimit { + primary_window: Some(CodexWindow { + used_percent: 8.0, + reset_at: now.timestamp() + 18_000, + limit_window_seconds: 18_000, + }), + secondary_window: Some(CodexWindow { + used_percent: 35.0, + reset_at: now.timestamp() + 604_800, + limit_window_seconds: 604_800, + }), + }; + let codex = codex_windows(Some(&rate_limit), None, now); + assert_eq!( + codex[0].pace_status.window_key.as_deref(), + Some("main.session.v1") + ); + assert_eq!( + codex[1].pace_status.window_key.as_deref(), + Some("main.weekly.v1") + ); - // 30 min later, still before the reset: reset_text recomputed to the - // shorter countdown (not the frozen original). - let later = base + chrono::Duration::seconds(1800); - let refreshed = refresh_cached_windows(std::slice::from_ref(&window), later).unwrap(); - assert_eq!(refreshed.len(), 1); - assert!(refreshed[0].reset_text.as_deref().unwrap().contains("30m")); + let claude = ClaudeUsageResponse { + five_hour: Some(ClaudeWindow { + utilization: Some(10.0), + resets_at: Some("2026-07-18T00:00:00Z".to_string()), + }), + seven_day: Some(ClaudeWindow { + utilization: Some(20.0), + resets_at: Some("2026-07-19T00:00:00Z".to_string()), + }), + ..Default::default() + }; + let claude = claude_windows(&claude, now); + assert_eq!( + claude[0].pace_status.window_key.as_deref(), + Some("session.v1") + ); + assert_eq!( + claude[1].pace_status.window_key.as_deref(), + Some("weekly.v1") + ); + assert_eq!(claude[0].window_minutes_for_test(), Some(300)); + assert_eq!(claude[1].window_minutes_for_test(), Some(10_080)); + } - // Past the reset: stale -> expire (None) so the caller re-probes. - let after = base + chrono::Duration::seconds(3700); - assert!(refresh_cached_windows(std::slice::from_ref(&window), after).is_none()); + #[test] + fn stage4_duplicate_snapshot_rows_are_removed_before_history_and_wire() { + let scope = TestRefreshScope::new("stage4", "duplicate-rows"); + let account_scope = scope + .resolve_current("fixture", "duplicate", b"duplicate-marker") + .unwrap(); + let now = 1_700_000_000; + let reset = Utc.timestamp_opt(now + 86_400, 0).single().unwrap(); + let make_window = |label: &str, card_id: &str, window_key: &str, used: f64| { + UsageWindow::from_provider_used_percent( + label.to_string(), + used, + Some(reset), + Utc.timestamp_opt(now, 0).single().unwrap(), + ) + .with_identity( + card_id, + Some(window_key.to_string()), + None, + Some(DurationEvidence::contract(86_400)), + ) + }; + let mut snapshot = AgentUsageSnapshot { + client_id: "fixture".to_string(), + source: "fixture".to_string(), + updated_at: String::new(), + identity: None, + account_scope: Ok(account_scope), + windows: vec![ + make_window("First", "shared-card.v1", "first.v1", 10.0), + make_window("Duplicate key", "second-card.v1", "first.v1", 20.0), + make_window("Duplicate card", "shared-card.v1", "third.v1", 30.0), + ], + credits: None, + error: None, + transport_diagnostic: None, + }; + + enrich_snapshot_with(&mut snapshot, now, |active, observations, _| { + assert_eq!(active.len(), 1); + assert_eq!(observations.len(), 1); + assert_eq!(active[0].window_key, "first.v1"); + assert_eq!(observations[0].used_percent, 10.0); + Ok(vec![Ok((HistoryOutcome::LearningDuration, None, 0))]) + }); + + assert_eq!(snapshot.windows.len(), 1); + assert_eq!(snapshot.windows[0].label_for_test(), "First"); + let wire = serde_json::to_value(&snapshot).unwrap(); + let rows = wire["windows"].as_array().unwrap(); + assert_eq!(rows.len(), 1); + assert_eq!(rows[0]["cardId"], "shared-card.v1"); + assert_eq!(rows[0]["paceStatus"]["windowKey"], "first.v1"); + scope.cleanup(); + } + + #[test] + fn stage4_chained_identity_collisions_keep_only_actual_uniques() { + let scope = TestRefreshScope::new("stage4", "chained-collisions"); + let account_scope = scope + .resolve_current("fixture", "chained", b"chained-marker") + .unwrap(); + let now = 1_700_000_000; + let reset = Utc.timestamp_opt(now + 86_400, 0).single().unwrap(); + let make_window = |label: &str, card_id: &str, window_key: &str, used: f64| { + UsageWindow::from_provider_used_percent( + label.to_string(), + used, + Some(reset), + Utc.timestamp_opt(now, 0).single().unwrap(), + ) + .with_identity( + card_id, + Some(window_key.to_string()), + None, + Some(DurationEvidence::contract(86_400)), + ) + }; + let mut snapshot = AgentUsageSnapshot { + client_id: "fixture".to_string(), + source: "fixture".to_string(), + updated_at: String::new(), + identity: None, + account_scope: Ok(account_scope), + windows: vec![ + make_window("A/X", "a.v1", "x.v1", 10.0), + make_window("A/Y", "a.v1", "y.v1", 20.0), + make_window("C/Y", "c.v1", "y.v1", 30.0), + make_window("B/X", "b.v1", "x.v1", 40.0), + make_window("B/Z", "b.v1", "z.v1", 50.0), + ], + credits: None, + error: None, + transport_diagnostic: None, + }; + + enrich_snapshot_with(&mut snapshot, now, |active, observations, _| { + assert_eq!(active.len(), 3); + assert_eq!(observations.len(), 3); + assert_eq!(active[0].window_key, "x.v1"); + assert_eq!(active[1].window_key, "y.v1"); + assert_eq!(active[2].window_key, "z.v1"); + assert_eq!( + observations + .iter() + .map(|observation| observation.used_percent) + .collect::>(), + vec![10.0, 30.0, 50.0] + ); + Ok(vec![ + Ok((HistoryOutcome::LearningDuration, None, 0)), + Ok((HistoryOutcome::LearningDuration, None, 0)), + Ok((HistoryOutcome::LearningDuration, None, 0)), + ]) + }); + + assert_eq!(snapshot.windows.len(), 3); + assert_eq!( + snapshot + .windows + .iter() + .map(UsageWindow::label_for_test) + .collect::>(), + vec!["A/X", "C/Y", "B/Z"] + ); + let wire = serde_json::to_value(&snapshot).unwrap(); + let rows = wire["windows"].as_array().unwrap(); + assert_eq!( + rows.iter() + .map(|row| row["cardId"].as_str().unwrap()) + .collect::>(), + vec!["a.v1", "c.v1", "b.v1"] + ); + assert_eq!( + rows.iter() + .map(|row| row["paceStatus"]["windowKey"].as_str().unwrap()) + .collect::>(), + vec!["x.v1", "y.v1", "z.v1"] + ); + scope.cleanup(); + } + + #[test] + fn stage4_batch_maps_results_once_without_network() { + let scope = TestRefreshScope::new("stage4", "batch-map"); + let account_scope = scope + .resolve_current("fixture", "batch", b"batch-marker") + .unwrap(); + let now = 1_700_000_000; + let reset = Utc.timestamp_opt(now + 86_400, 0).single().unwrap(); + let mut snapshot = AgentUsageSnapshot { + client_id: "fixture".to_string(), + source: "fixture".to_string(), + updated_at: String::new(), + identity: None, + account_scope: Ok(account_scope), + windows: vec![ + UsageWindow::from_provider_used_percent( + "First".to_string(), + 20.0, + Some(reset), + Utc.timestamp_opt(now, 0).single().unwrap(), + ) + .with_identity( + "first.v1", + Some("first.v1".to_string()), + None, + Some(DurationEvidence::contract(86_400)), + ), + UsageWindow::from_provider_used_percent( + "Second".to_string(), + 40.0, + Some(reset), + Utc.timestamp_opt(now, 0).single().unwrap(), + ) + .with_identity( + "second.v1", + Some("second.v1".to_string()), + None, + Some(DurationEvidence::contract(86_400)), + ), + ], + credits: None, + error: None, + transport_diagnostic: None, + }; + let calls = std::cell::Cell::new(0); + enrich_snapshot_with(&mut snapshot, now, |active, observations, _| { + calls.set(calls.get() + 1); + assert_eq!(active.len(), 2); + assert_eq!(observations.len(), 2); + assert_eq!(active[0].window_key, "first.v1"); + assert_eq!(active[1].window_key, "second.v1"); + Ok(vec![ + Ok((HistoryOutcome::LearningDuration, None, 0)), + Ok(( + HistoryOutcome::Ready { + duration_seconds: 86_400, + source: DurationSource::Contract, + sampled: true, + }, + Some(HistoricalPace { + expected_percent: 42.0, + eta_seconds: Some(900.0), + will_last_to_reset: false, + run_out_probability: Some(0.25), + }), + 4, + )), + ]) + }); + assert_eq!( + calls.get(), + 1, + "one snapshot means one batch and no new request" + ); + assert_eq!( + snapshot.windows[0].pace_status.state, + PaceState::LearningDuration + ); + assert_eq!(snapshot.windows[1].pace_status.state, PaceState::Available); + assert_eq!(snapshot.windows[1].window_minutes_for_test(), Some(1_440)); + let wire = serde_json::to_value(&snapshot).unwrap(); + assert_eq!(wire["windows"][0]["paceStatus"]["completeCycles"], 0); + assert_eq!(wire["windows"][1]["paceStatus"]["completeCycles"], 4); + scope.cleanup(); + } + + #[test] + fn stage4_learning_history_uses_batch_complete_cycles() { + let scope = TestRefreshScope::new("stage4", "learning-history"); + let account_scope = scope + .resolve_current("fixture", "learning", b"learning-marker") + .unwrap(); + let now = 1_700_000_000; + let reset = Utc.timestamp_opt(now + 86_400, 0).single().unwrap(); + let mut snapshot = AgentUsageSnapshot { + client_id: "fixture".to_string(), + source: "fixture".to_string(), + updated_at: String::new(), + identity: None, + account_scope: Ok(account_scope), + windows: vec![UsageWindow::from_provider_used_percent( + "Weekly".to_string(), + 20.0, + Some(reset), + Utc.timestamp_opt(now, 0).single().unwrap(), + ) + .with_identity( + "weekly.v1", + Some("weekly.v1".to_string()), + None, + Some(DurationEvidence::contract(86_400)), + )], + credits: None, + error: None, + transport_diagnostic: None, + }; + + enrich_snapshot_with(&mut snapshot, now, |_, _, _| { + Ok(vec![Ok(( + HistoryOutcome::Ready { + duration_seconds: 86_400, + source: DurationSource::Contract, + sampled: true, + }, + None, + 2, + ))]) + }); + + assert_eq!( + snapshot.windows[0].pace_status.state, + PaceState::LearningHistory + ); + assert_eq!(snapshot.windows[0].pace_status.complete_cycles, 2); + let wire = serde_json::to_value(&snapshot).unwrap(); + assert_eq!(wire["windows"][0]["paceStatus"]["completeCycles"], 2); + scope.cleanup(); + } + + #[test] + fn stage4_incoherent_historical_result_is_typed_unavailable() { + let scope = TestRefreshScope::new("stage4", "incoherent-history"); + let account_scope = scope + .resolve_current("fixture", "incoherent", b"incoherent-marker") + .unwrap(); + let now = 1_700_000_000; + let reset = Utc.timestamp_opt(now + 86_400, 0).single().unwrap(); + let mut snapshot = AgentUsageSnapshot { + client_id: "fixture".to_string(), + source: "fixture".to_string(), + updated_at: String::new(), + identity: None, + account_scope: Ok(account_scope), + windows: vec![UsageWindow::from_provider_used_percent( + "Weekly".to_string(), + 20.0, + Some(reset), + Utc.timestamp_opt(now, 0).single().unwrap(), + ) + .with_identity( + "weekly.v1", + Some("weekly.v1".to_string()), + None, + Some(DurationEvidence::contract(86_400)), + )], + credits: None, + error: None, + transport_diagnostic: None, + }; + + enrich_snapshot_with(&mut snapshot, now, |_, _, _| { + Ok(vec![Ok(( + HistoryOutcome::Ready { + duration_seconds: 86_400, + source: DurationSource::Contract, + sampled: true, + }, + Some(HistoricalPace { + expected_percent: 42.0, + eta_seconds: Some(900.0), + will_last_to_reset: true, + run_out_probability: Some(0.25), + }), + 4, + ))]) + }); + + assert_eq!( + snapshot.windows[0].pace_status.state, + PaceState::Unavailable + ); + assert_eq!( + snapshot.windows[0].pace_status.reason.as_deref(), + Some("history") + ); + let wire = serde_json::to_value(&snapshot).unwrap(); + assert_eq!(wire["windows"][0]["paceStatus"]["state"], "unavailable"); + assert!(wire["windows"][0].get("historicalPace").is_none()); + scope.cleanup(); + } + + #[test] + fn stage4_historical_eta_and_will_last_are_exactly_coherent() { + let now = Utc.timestamp_opt(1_700_000_000, 0).single().unwrap(); + let reset = now + chrono::Duration::days(1); + let base = + UsageWindow::from_provider_used_percent("Daily".to_string(), 30.0, Some(reset), now) + .with_identity( + "daily.v1", + Some("daily.v1".to_string()), + None, + Some(DurationEvidence::contract(86_400)), + ); + let cases = [ + ( + "will-last", + HistoricalPace { + expected_percent: 42.0, + eta_seconds: None, + will_last_to_reset: true, + run_out_probability: Some(0.1), + }, + true, + ), + ( + "will-run-out", + HistoricalPace { + expected_percent: 42.0, + eta_seconds: Some(900.0), + will_last_to_reset: false, + run_out_probability: Some(0.25), + }, + true, + ), + ( + "will-last-with-eta", + HistoricalPace { + expected_percent: 42.0, + eta_seconds: Some(900.0), + will_last_to_reset: true, + run_out_probability: Some(0.25), + }, + false, + ), + ( + "will-run-out-without-eta", + HistoricalPace { + expected_percent: 42.0, + eta_seconds: None, + will_last_to_reset: false, + run_out_probability: Some(0.25), + }, + false, + ), + ]; + + for (label, pace, expected) in cases { + assert_eq!(historical_pace_is_coherent(&pace), expected, "{label}"); + let mut window = base.clone(); + window.pace_status.state = PaceState::Available; + window.historical_pace = Some(historical_pace_payload(pace)); + assert_eq!(serde_json::to_value(&window).is_ok(), expected, "{label}"); + } + } + + #[test] + fn stage4_scope_error_is_sticky_and_skips_history() { + let now = Utc.timestamp_opt(1_700_000_000, 0).single().unwrap(); + let mut snapshot = AgentUsageSnapshot { + client_id: "fixture".to_string(), + source: "fixture".to_string(), + updated_at: String::new(), + identity: None, + account_scope: Err(AccountScopeError::MetadataWrite), + windows: vec![ + UsageWindow::from_provider_used_percent( + "Session".to_string(), + 20.0, + Some(now + chrono::Duration::hours(5)), + now, + ) + .with_identity( + "session.v1", + Some("session.v1".to_string()), + None, + Some(DurationEvidence::contract(300 * 60)), + ), + UsageWindow::from_provider_used_percent( + "Unknown".to_string(), + 30.0, + Some(now + chrono::Duration::hours(5)), + now, + ) + .with_identity("row.unknown.v1", None, None, None), + ], + credits: None, + error: None, + transport_diagnostic: None, + }; + let calls = std::cell::Cell::new(0); + enrich_snapshot_with(&mut snapshot, now.timestamp(), |_, _, _| { + calls.set(calls.get() + 1); + Ok(Vec::new()) + }); + assert_eq!(calls.get(), 0); + assert_eq!( + snapshot.windows[0].pace_status.reason.as_deref(), + Some("accountScope") + ); + assert_eq!( + snapshot.windows[0].pace_status.state, + PaceState::Unavailable + ); + assert_eq!( + snapshot.windows[1].pace_status.reason.as_deref(), + Some("windowIdentity") + ); + assert!(snapshot.windows[1].pace_status.window_key.is_none()); + assert!(serde_json::to_value(&snapshot).is_ok()); + } + + #[test] + fn stage4_wire_rejects_internal_nested_drift_and_preserves_observed_learning() { + let now = Utc.timestamp_opt(1_700_000_000, 0).single().unwrap(); + let reset = now + chrono::Duration::days(1); + let base = + UsageWindow::from_provider_used_percent("Daily".to_string(), 30.0, Some(reset), now) + .with_identity( + "daily.v1", + Some("daily.v1".to_string()), + None, + Some(DurationEvidence::contract(86_400)), + ); + + let mut key_drift = base.clone(); + key_drift.pace_status.window_key = Some("other.v1".to_string()); + assert!(serde_json::to_value(&key_drift).is_err()); + + let mut duration_drift = base.clone(); + duration_drift.pace_status.duration_seconds = Some(3_600); + assert!(serde_json::to_value(&duration_drift).is_err()); + + let mut source_drift = base.clone(); + source_drift.pace_status.duration_source = Some(DurationSource::Provider); + assert!(serde_json::to_value(&source_drift).is_err()); + + let mut minutes_drift = base.clone(); + minutes_drift.window_minutes = Some(1); + assert!(serde_json::to_value(&minutes_drift).is_err()); + + let mut learning = + UsageWindow::from_provider_used_percent("Learning".to_string(), 30.0, Some(reset), now) + .with_identity("learning.v1", Some("learning.v1".to_string()), None, None); + learning.duration_source = Some(DurationSource::Observed); + learning.pace_status.duration_source = Some(DurationSource::Observed); + let wire = serde_json::to_value(&learning).unwrap(); + assert_eq!(wire["paceStatus"]["state"], "learningDuration"); + assert_eq!(wire["paceStatus"]["durationSource"], "observed"); + assert!(wire["paceStatus"].get("durationSeconds").is_none()); + } + + #[test] + fn stage4_wire_rejects_available_without_historical_pace() { + let now = Utc.timestamp_opt(1_700_000_000, 0).single().unwrap(); + let mut window = UsageWindow::from_provider_used_percent( + "Weekly".to_string(), + 30.0, + Some(now + chrono::Duration::days(7)), + now, + ) + .with_identity( + "weekly.v1", + Some("weekly.v1".to_string()), + None, + Some(DurationEvidence::contract(7 * 24 * 60 * 60)), + ); + window.pace_status.state = PaceState::Available; + window.historical_pace = None; + assert!(serde_json::to_value(&window).is_err()); + } + + #[test] + fn provider_quota_pace_v3_fixture_locks_production_serializer() { + fn window( + card_id: &str, + label: &str, + used_percent: f64, + resets_at: Option<&str>, + window_key: Option<&str>, + state: PaceState, + duration_seconds: Option, + duration_source: Option, + complete_cycles: usize, + reason: Option<&str>, + historical_pace: Option, + ) -> UsageWindow { + UsageWindow { + card_id: card_id.to_string(), + label: label.to_string(), + used_percent, + remaining_percent: 100.0 - used_percent, + resets_at: resets_at.map(|value| value.to_string()), + reset_text: None, + window_minutes: duration_seconds.map(|seconds| seconds / 60), + window_key: window_key.map(|value| value.to_string()), + duration_seconds, + duration_source, + provider_duration: None, + contract_duration: None, + pace_status: PaceStatusPayload { + state, + window_key: window_key.map(|value| value.to_string()), + duration_seconds, + duration_source, + complete_cycles, + reason: reason.map(|value| value.to_string()), + }, + historical_pace, + } + } + + let payload = AgentUsagePayload { + generated_at: "2026-07-10T12:00:00.000Z".to_string(), + publication_generation: 1, + agents: vec![AgentUsageSnapshot { + client_id: "provider-fixture.invalid".to_string(), + source: "fixture.invalid".to_string(), + updated_at: "2026-07-10T12:00:00.000Z".to_string(), + identity: None, + account_scope: Err(AccountScopeError::NoTrustedEvidence), + windows: vec![ + window( + "ahead.invalid", + "Ahead quota", + 72.0, + Some("2026-07-10T15:00:00Z"), + Some("quota.ahead.invalid"), + PaceState::Available, + Some(18_000), + Some(DurationSource::Provider), + 5, + None, + Some(HistoricalPacePayload { + expected_used_percent: 32.0, + eta_seconds: Some(3_600.0), + will_last_to_reset: false, + run_out_probability: Some(0.75), + }), + ), + window( + "behind.invalid", + "Behind quota", + 28.0, + Some("2026-07-15T12:00:00Z"), + Some("quota.behind.invalid"), + PaceState::Available, + Some(604_800), + Some(DurationSource::Contract), + 7, + None, + Some(HistoricalPacePayload { + expected_used_percent: 56.0, + eta_seconds: None, + will_last_to_reset: true, + run_out_probability: Some(0.2), + }), + ), + window( + "learning-history.invalid", + "Learning history", + 40.0, + Some("2026-07-10T15:00:00Z"), + Some("quota.learning-history.invalid"), + PaceState::LearningHistory, + Some(18_000), + Some(DurationSource::Provider), + 2, + None, + None, + ), + window( + "learning-duration.invalid", + "Learning duration", + 40.0, + Some("2026-07-10T15:00:00Z"), + Some("quota.learning-duration.invalid"), + PaceState::LearningDuration, + None, + Some(DurationSource::Observed), + 0, + None, + None, + ), + window( + "missing-reset.invalid", + "Missing reset", + 50.0, + None, + Some("quota.missing-reset.invalid"), + PaceState::Unavailable, + None, + None, + 0, + Some("missingReset"), + None, + ), + window( + "shared-first.invalid", + "Shared label", + 10.0, + Some("2026-07-10T15:00:00Z"), + Some("quota.shared-first.invalid"), + PaceState::LearningHistory, + Some(18_000), + Some(DurationSource::Provider), + 2, + None, + None, + ), + window( + "shared-second.invalid", + "Shared label", + 20.0, + Some("2026-07-10T15:00:00Z"), + Some("quota.shared-second.invalid"), + PaceState::LearningHistory, + Some(18_000), + Some(DurationSource::Provider), + 2, + None, + None, + ), + ], + credits: None, + error: None, + transport_diagnostic: None, + }], + opencode_subscriptions: Vec::new(), + }; + + let fixture_path = Path::new(env!("CARGO_MANIFEST_DIR")) + .join("../../Fixtures/CrossCheck/provider-quota-pace-v3.json"); + let fixture: Value = serde_json::from_str( + &fs::read_to_string(&fixture_path) + .unwrap_or_else(|error| panic!("read {}: {error}", fixture_path.display())), + ) + .unwrap_or_else(|error| panic!("decode {}: {error}", fixture_path.display())); + assert_eq!(fixture["schemaVersion"], 3); + let mut serialized = serde_json::to_value(payload).unwrap(); + assert_eq!(serialized["publicationGeneration"], 1); + serialized + .as_object_mut() + .expect("payload serializes as an object") + .remove("publicationGeneration"); + assert_eq!(fixture["payload"], serialized); } } diff --git a/crates/tb_core_ffi/src/lib.rs b/crates/tb_core_ffi/src/lib.rs index c7e68de..83c69b7 100644 --- a/crates/tb_core_ffi/src/lib.rs +++ b/crates/tb_core_ffi/src/lib.rs @@ -122,6 +122,39 @@ static RUNTIME: LazyLock = LazyLock::new(|| { .expect("build tokio runtime for tb_core_ffi") }); +#[derive(Default)] +struct PublicationGate { + generation: u64, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +struct PublicationGenerationExhausted; + +static AGENT_USAGE_PUBLICATION_GATE: LazyLock> = + LazyLock::new(|| Mutex::new(PublicationGate::default())); + +/// Serialize the complete publication path and assign its order at gate entry; +/// this gate is the sole generation source, not a timestamp or caller ordering. +/// Exhaustion fails closed instead of publishing a duplicate generation. +fn with_publication_gate( + gate: &Mutex, + body: impl FnOnce(u64) -> T, +) -> Result { + let mut state = gate.lock().unwrap_or_else(|poisoned| poisoned.into_inner()); + let generation = state + .generation + .checked_add(1) + .ok_or(PublicationGenerationExhausted)?; + state.generation = generation; + Ok(body(generation)) +} + +fn with_agent_usage_publication_gate( + body: impl FnOnce(u64) -> T, +) -> Result { + with_publication_gate(&AGENT_USAGE_PUBLICATION_GATE, body) +} + /// Cap rayon's global thread pool to 2 workers. tokscale-core uses rayon for /// parallel log parsing (55+ par_iter sites); the default pool size is num_cpus /// which is fine for a one-shot CLI but ruinous for a resident menu-bar daemon: @@ -137,9 +170,9 @@ static RAYON_INIT: LazyLock<()> = LazyLock::new(|| { /// year → (computed-at, source token, mapped graph payload). Same role as /// the Tauri AppState cache, plus a change token: when the cache entry ages -/// past the oneshot window but `latest_source_mtime_ms` still matches the -/// token, the entry is re-stamped and served — an idle machine never pays -/// for a full re-aggregation just because time passed. +/// past the oneshot window but the topology-sensitive token still matches, +/// the entry is re-stamped and served — an idle machine never pays for a full +/// re-aggregation just because time passed. type GraphCacheEntry = (Instant, u64, serde_json::Value); static GRAPH_CACHE: LazyLock>> = LazyLock::new(|| Mutex::new(HashMap::new())); @@ -181,10 +214,10 @@ fn envelope(result: Result) -> *mut c_char { /// rest of the menu-bar app running rather than aborting the whole process. The /// default panic hook still prints the panic location to stderr before we catch. /// -/// `AssertUnwindSafe` is sound here. State shared across calls is the three -/// std::sync Mutex statics (GRAPH_CACHE / TAIL_TICK / CLAUDE_USAGE_GATE), each -/// recovered from poison on the next lock via `into_inner()`, plus the live -/// tail's parking_lot Mutexes, which never poison and release cleanly on unwind. +/// `AssertUnwindSafe` is sound here. Process-wide std::sync Mutex state is +/// updated only while locked, and each lock helper recovers poison on the next +/// call via `into_inner()`. The live tail's parking_lot Mutexes never poison and +/// release cleanly on unwind. /// A caught panic can leave a cache entry stale or a tail tick un-run, never /// torn: the next call re-derives the graph, and `tail_tick_if_stale` clears its /// in-flight flag without stamping on a tick panic so the tail re-parses next. @@ -242,9 +275,9 @@ unsafe fn clients_from(clients: *const c_char) -> Result>, St fn graph_cached(year: &str, max_age: Duration) -> Option { // Read the entry and release the lock before any filesystem I/O — never hold - // GRAPH_CACHE across the mtime stat sweep below (mirrors graph_compute, which - // probes outside the lock too), so concurrent tb_graph callers don't queue - // behind one another's stat. + // GRAPH_CACHE across the source-state probe below (mirrors graph_compute, + // which probes outside the lock too), so concurrent tb_graph callers don't + // queue behind one another's stat sweep. let (fresh_enough, token, data) = { let cache = GRAPH_CACHE.lock().unwrap_or_else(|p| p.into_inner()); let (at, token, data) = cache.get(year)?; @@ -253,13 +286,14 @@ fn graph_cached(year: &str, max_age: Duration) -> Option { if fresh_enough { return Some(data); } - // Aged out — but if no source file changed since the compute, the graph + // Aged out — but if no source state changed since the compute, the graph // cannot have changed either. Probe with the lock released, then re-acquire // briefly to re-stamp so the next calls inside the oneshot window skip the // probe entirely. A lost re-stamp (entry evicted/replaced meanwhile) just // degrades to the next call re-probing — benign. let context = LocalSourceContext::current(); - let fresh = tokscale_core::latest_source_mtime_ms(&context.parse_options(None, None)).ok()?; + let fresh = + tokscale_core::local_source_change_token(&context.parse_options(None, None)).ok()?; if fresh == token { let mut cache = GRAPH_CACHE.lock().unwrap_or_else(|p| p.into_inner()); if let Some(entry) = cache.get_mut(year) { @@ -271,13 +305,13 @@ fn graph_cached(year: &str, max_age: Duration) -> Option { } fn graph_compute(year: &str) -> Result { - // Probe before parsing: a write that lands mid-compute moves the mtime - // past this token, so the next aged-out read recomputes rather than - // serving a graph that missed it. Keep the same context for both paths so - // the mtime token and report scan observe identical source roots. + // Probe before parsing: a source write or topology change that lands + // mid-compute changes the token, so the next aged-out read recomputes + // rather than serving a graph that missed it. Keep the same context for + // both paths so the probe and report scan observe identical source roots. let context = LocalSourceContext::current(); - let token = tokscale_core::latest_source_mtime_ms(&context.parse_options(None, None)) - .unwrap_or(0); + let token = + tokscale_core::local_source_change_token(&context.parse_options(None, None)).unwrap_or(0); let data = usage_graph::run(&context, year)?; GRAPH_CACHE .lock() @@ -462,19 +496,30 @@ pub extern "C" fn tb_tokens_per_min() -> *mut c_char { /// codex/claude/antigravity/copilot/grok, fetched concurrently. Network-bound — /// call from a background thread. Per-provider failures land in each /// snapshot's `error` field; the call itself only fails on serialization. +/// The publication gate assigns `publicationGeneration` and serializes the +/// provider run, JSON/envelope construction, and pointer creation. The gate is +/// released before this extern function returns, so Swift still needs its own +/// generation guard for caller return/apply order. #[no_mangle] pub extern "C" fn tb_agent_usage() -> *mut c_char { - guarded("tb_agent_usage", || { - // No outer timeout on purpose: each provider carries its own 30s - // per-request reqwest timeout (which covers connect, so nothing hangs - // unbounded), and they run concurrently via tokio::join!. A single outer - // ceiling would instead collapse the whole payload to one error — losing - // the providers that already succeeded — and could cut off the legitimate - // expired-token path (sequential refresh + fetch, up to ~60s). - let payload = RUNTIME.block_on(agent_usage::run()); - envelope( - serde_json::to_value(payload).map_err(|e| format!("serialize agent usage: {}", e)), - ) + with_agent_usage_publication_gate(|generation| { + guarded("tb_agent_usage", || { + // No outer timeout on purpose: each provider carries its own 30s + // per-request reqwest timeout (which covers connect, so nothing hangs + // unbounded), and they run concurrently via tokio::join!. A single outer + // ceiling would instead collapse the whole payload to one error — losing + // the providers that already succeeded — and could cut off the legitimate + // expired-token path (sequential refresh + fetch, up to ~60s). + let payload = RUNTIME.block_on(agent_usage::run(generation)); + envelope( + serde_json::to_value(payload).map_err(|e| format!("serialize agent usage: {}", e)), + ) + }) + }) + .unwrap_or_else(|_| { + envelope(Err( + "agent usage publication generation exhausted".to_string() + )) }) } @@ -575,6 +620,203 @@ mod tests { assert!(s.contains("tb_test panicked: boom"), "got: {s}"); } + #[test] + fn publication_gate_serializes_pointer_publication() { + use std::sync::{mpsc, Arc, Barrier}; + use std::thread; + + let gate = Arc::new(Mutex::new(PublicationGate::default())); + let (events_tx, events_rx) = mpsc::channel(); + let (release_a_tx, release_a_rx) = mpsc::channel(); + let (b_called_tx, b_called_rx) = mpsc::channel(); + let a_ready = Arc::new(Barrier::new(2)); + + let a_gate = Arc::clone(&gate); + let a_ready_thread = Arc::clone(&a_ready); + let a_events = events_tx.clone(); + let a = thread::spawn(move || { + with_publication_gate(a_gate.as_ref(), |_| { + a_events.send("A run").unwrap(); + a_ready_thread.wait(); + release_a_rx.recv().unwrap(); + + let pointer = CString::new("A").unwrap().into_raw(); + a_events.send("A publish").unwrap(); + unsafe { tb_free(pointer) }; + }) + .unwrap(); + }); + a_ready.wait(); + + let b_gate = Arc::clone(&gate); + let b_events = events_tx.clone(); + let b = thread::spawn(move || { + b_called_tx.send(()).unwrap(); + with_publication_gate(b_gate.as_ref(), |_| { + b_events.send("B run").unwrap(); + let pointer = CString::new("B").unwrap().into_raw(); + b_events.send("B publish").unwrap(); + unsafe { tb_free(pointer) }; + }) + .unwrap(); + }); + b_called_rx.recv().unwrap(); + release_a_tx.send(()).unwrap(); + + a.join().unwrap(); + b.join().unwrap(); + drop(events_tx); + + let events: Vec<_> = events_rx.iter().collect(); + assert_eq!(events, vec!["A run", "A publish", "B run", "B publish"]); + } + + #[test] + fn publication_generation_handles_return_order_reversal() { + use std::sync::{mpsc, Arc}; + use std::thread; + + let gate = Arc::new(Mutex::new(PublicationGate::default())); + let (a_gate_returned_tx, a_gate_returned_rx) = mpsc::channel(); + let (release_a_tx, release_a_rx) = mpsc::channel(); + let (returns_tx, returns_rx) = mpsc::channel(); + + let a_gate = Arc::clone(&gate); + let a_returns = returns_tx.clone(); + let a = thread::spawn(move || { + let (generation, pointer) = with_publication_gate(a_gate.as_ref(), |generation| { + let pointer = CString::new(format!( + "{{\"run\":\"A\",\"generation\":{generation},\"result\":\"success\"}}" + )) + .unwrap() + .into_raw(); + (generation, pointer) + }) + .unwrap(); + // The gate is released here. Pause before the wrapper records the + // return so B can publish and return first. + a_gate_returned_tx.send(generation).unwrap(); + release_a_rx.recv().unwrap(); + let payload = unsafe { take(pointer) }; + a_returns.send(("A returned", generation, payload)).unwrap(); + }); + + assert_eq!(a_gate_returned_rx.recv().unwrap(), 1); + + let b_gate = Arc::clone(&gate); + let b_returns = returns_tx.clone(); + let b = thread::spawn(move || { + let (generation, pointer) = with_publication_gate(b_gate.as_ref(), |generation| { + let pointer = CString::new(format!( + "{{\"run\":\"B\",\"generation\":{generation},\"result\":\"terminal\"}}" + )) + .unwrap() + .into_raw(); + (generation, pointer) + }) + .unwrap(); + let payload = unsafe { take(pointer) }; + b_returns.send(("B returned", generation, payload)).unwrap(); + }); + + let b_return = returns_rx.recv().unwrap(); + assert_eq!(b_return.0, "B returned"); + assert_eq!(b_return.1, 2); + assert_eq!( + b_return.2, + r#"{"run":"B","generation":2,"result":"terminal"}"# + ); + + release_a_tx.send(()).unwrap(); + let a_return = returns_rx.recv().unwrap(); + assert_eq!(a_return.0, "A returned"); + assert_eq!(a_return.1, 1); + assert_eq!( + a_return.2, + r#"{"run":"A","generation":1,"result":"success"}"# + ); + + a.join().unwrap(); + b.join().unwrap(); + } + + #[test] + fn publication_generation_exhaustion_fails_closed() { + let gate = Mutex::new(PublicationGate { + generation: u64::MAX - 1, + }); + assert_eq!( + with_publication_gate(&gate, |generation| generation).unwrap(), + u64::MAX + ); + + let mut body_called = false; + let exhausted = with_publication_gate(&gate, |generation| { + body_called = true; + generation + }); + assert_eq!(exhausted, Err(PublicationGenerationExhausted)); + assert!( + !body_called, + "exhaustion must not publish a duplicate generation" + ); + } + + #[test] + fn publication_gate_keeps_panic_envelope_before_next_run() { + use std::sync::{mpsc, Arc, Barrier}; + use std::thread; + + let gate = Arc::new(Mutex::new(PublicationGate::default())); + let (events_tx, events_rx) = mpsc::channel(); + let (release_a_tx, release_a_rx) = mpsc::channel(); + let (b_called_tx, b_called_rx) = mpsc::channel(); + let a_ready = Arc::new(Barrier::new(2)); + + let a_gate = Arc::clone(&gate); + let a_ready_thread = Arc::clone(&a_ready); + let a_events = events_tx.clone(); + let a = thread::spawn(move || { + with_publication_gate(a_gate.as_ref(), |_| { + let pointer = guarded("tb_test", || { + a_events.send("A run").unwrap(); + a_ready_thread.wait(); + release_a_rx.recv().unwrap(); + panic!("boom"); + }); + a_events.send("A panic publish").unwrap(); + unsafe { tb_free(pointer) }; + }) + .unwrap(); + }); + a_ready.wait(); + + let b_gate = Arc::clone(&gate); + let b_events = events_tx.clone(); + let b = thread::spawn(move || { + b_called_tx.send(()).unwrap(); + with_publication_gate(b_gate.as_ref(), |_| { + b_events.send("B run").unwrap(); + let pointer = CString::new("B").unwrap().into_raw(); + b_events.send("B publish").unwrap(); + unsafe { tb_free(pointer) }; + }) + .unwrap(); + }); + b_called_rx.recv().unwrap(); + release_a_tx.send(()).unwrap(); + + a.join().unwrap(); + b.join().unwrap(); + drop(events_tx); + + let events: Vec<_> = events_rx.iter().collect(); + assert_eq!( + events, + vec!["A run", "A panic publish", "B run", "B publish"] + ); + } + #[test] fn tick_guard_clears_in_flight_without_stamping_on_panic() { // Simulates a panic during TAILER.tick(): the guard, dropped mid-unwind, diff --git a/crates/tb_core_ffi/src/model_report.rs b/crates/tb_core_ffi/src/model_report.rs index f0ff0b9..e130a8c 100644 --- a/crates/tb_core_ffi/src/model_report.rs +++ b/crates/tb_core_ffi/src/model_report.rs @@ -47,10 +47,7 @@ struct ModelReportData { } /// Build the per-model report for `year` (empty string = all time). -pub(crate) fn run( - context: &crate::LocalSourceContext, - year: &str, -) -> Result { +pub(crate) fn run(context: &crate::LocalSourceContext, year: &str) -> Result { let year = normalize_year(year)?; let options = context.report_options(year, None); diff --git a/crates/tb_core_ffi/src/opencode_integrations.rs b/crates/tb_core_ffi/src/opencode_integrations.rs index 5cb391d..c56120a 100644 --- a/crates/tb_core_ffi/src/opencode_integrations.rs +++ b/crates/tb_core_ffi/src/opencode_integrations.rs @@ -2,7 +2,7 @@ //! //! opencode can sign in to providers via OAuth (a shared subscription, e.g. //! "Sign in with ChatGPT" = the Codex/ChatGPT plan) or via API keys (metered). -//! Its `~/.local/share/opencode/auth.json` records each provider with a `type`. +//! Its `auth.json` under the XDG data root records each provider with a `type`. //! We surface the `type: "oauth"` providers so the user can see which agent //! subscriptions opencode also draws on (its usage counts against those plans). @@ -59,7 +59,21 @@ fn subscription_label(provider: &str) -> String { } fn auth_path() -> Option { - crate::user_home_dir().map(|home| home.join(".local/share/opencode/auth.json")) + auth_path_from( + std::env::var("XDG_DATA_HOME"), + crate::user_home_dir().as_deref(), + ) +} + +fn auth_path_from( + xdg_data_home: Result, + user_home: Option<&std::path::Path>, +) -> Option { + let root = match xdg_data_home { + Ok(root) if !root.is_empty() => root, + Ok(_) | Err(_) => format!("{}/.local/share", user_home?.to_string_lossy()), + }; + Some(PathBuf::from(format!("{root}/opencode/auth.json"))) } pub(crate) struct GitHubCopilotCredential { @@ -69,155 +83,373 @@ pub(crate) struct GitHubCopilotCredential { pub(crate) canonical_location: String, } -/// The durable GitHub OAuth credential opencode stored for its Copilot login. -/// The trimmed refresh token is preferred; a trimmed access token is the fallback. -pub(crate) fn github_copilot_credential() -> Option { - let path = auth_path()?; - let raw = std::fs::read_to_string(&path).ok()?; - let json = serde_json::from_str::(&raw).ok()?; - github_copilot_credential_from(&path, &json) +pub(crate) enum GitHubCopilotCredentialLoad { + Absent, + Present(GitHubCopilotCredential), + Terminal(String), } -pub(crate) fn github_copilot_credential_from( +/// Load only the durable GitHub OAuth credential used by Copilot. A missing +/// file/entry or explicit logout is `Absent`; storage, syntax, canonicalization, +/// or a present malformed OAuth entry is `Terminal` and must not be hidden as a +/// signed-out state. +pub(crate) fn github_copilot_credential() -> GitHubCopilotCredentialLoad { + let Some(path) = auth_path() else { + return GitHubCopilotCredentialLoad::Terminal( + "OpenCode auth location could not be resolved.".to_string(), + ); + }; + github_copilot_credential_at(&path) +} + +fn github_copilot_credential_at(path: &std::path::Path) -> GitHubCopilotCredentialLoad { + let raw = match std::fs::read_to_string(path) { + Ok(raw) => raw, + Err(error) if error.kind() == std::io::ErrorKind::NotFound => { + return GitHubCopilotCredentialLoad::Absent; + } + Err(_) => { + return GitHubCopilotCredentialLoad::Terminal( + "OpenCode auth file could not be read.".to_string(), + ); + } + }; + let json = match serde_json::from_str::(&raw) { + Ok(json) => json, + Err(_) => { + return GitHubCopilotCredentialLoad::Terminal( + "OpenCode auth file could not be decoded.".to_string(), + ); + } + }; + github_copilot_credential_from(path, &json) +} + +fn github_copilot_credential_from( path: &std::path::Path, json: &serde_json::Value, -) -> Option { - let entry = json.get("github-copilot")?; - if entry.get("type").and_then(|t| t.as_str()) != Some("oauth") { - return None; +) -> GitHubCopilotCredentialLoad { + let Some(entry) = json.get("github-copilot") else { + return GitHubCopilotCredentialLoad::Absent; + }; + let Some(entry) = entry.as_object() else { + return GitHubCopilotCredentialLoad::Terminal( + "OpenCode Copilot OAuth entry is malformed.".to_string(), + ); + }; + match entry.get("type") { + Some(serde_json::Value::String(kind)) if kind.eq_ignore_ascii_case("oauth") => {} + Some(serde_json::Value::String(_)) => return GitHubCopilotCredentialLoad::Absent, + _ => { + return GitHubCopilotCredentialLoad::Terminal( + "OpenCode Copilot OAuth entry is malformed.".to_string(), + ); + } } - let token = ["refresh", "access"] - .into_iter() - .filter_map(|key| entry.get(key).and_then(serde_json::Value::as_str)) - .map(str::trim) - .find(|token| !token.is_empty())? - .to_string(); - Some(GitHubCopilotCredential { + + let mut token = None; + let mut saw_token_field = false; + for key in ["refresh", "access"] { + match entry.get(key) { + None => {} + Some(serde_json::Value::Null) => saw_token_field = true, + Some(serde_json::Value::String(value)) if value.trim().is_empty() => { + saw_token_field = true; + } + Some(serde_json::Value::String(value)) if token.is_none() => { + saw_token_field = true; + token = Some(value.trim().to_string()); + } + Some(serde_json::Value::String(_)) => saw_token_field = true, + Some(_) => { + return GitHubCopilotCredentialLoad::Terminal( + "OpenCode Copilot OAuth entry is malformed.".to_string(), + ); + } + } + } + let Some(token) = token else { + return if saw_token_field { + GitHubCopilotCredentialLoad::Absent + } else { + GitHubCopilotCredentialLoad::Terminal( + "OpenCode Copilot OAuth entry is malformed.".to_string(), + ) + }; + }; + let canonical_location = + match crate::agent_account_scope::canonical_file_location(path, Some("github-copilot")) { + Ok(location) => location, + Err(_) => { + return GitHubCopilotCredentialLoad::Terminal( + "OpenCode Copilot auth location could not be verified.".to_string(), + ); + } + }; + GitHubCopilotCredentialLoad::Present(GitHubCopilotCredential { request_token: token.clone(), marker: token.into_bytes(), semantic_source: "opencode-auth-json", - canonical_location: crate::agent_account_scope::canonical_file_location( - path, - Some("github-copilot"), - ) - .ok()?, + canonical_location, }) } #[cfg(test)] mod tests { use super::*; - use serde_json::json; - use std::sync::atomic::{AtomicU64, Ordering}; - static TEMP_COUNTER: AtomicU64 = AtomicU64::new(0); + #[test] + fn labels_oauth_providers_only() { + assert_eq!(subscription_label("openai"), "Codex"); + assert_eq!(subscription_label("github-copilot"), "Copilot"); + assert_eq!(subscription_label("anthropic"), "Claude"); + assert_eq!( + subscription_label("minimax-coding-plan"), + "Minimax-coding-plan" + ); + } - fn temp_auth_path(tag: &str) -> PathBuf { - let root = std::env::temp_dir().join(format!( - "tb-opencode-{tag}-{}-{}", - std::process::id(), - TEMP_COUNTER.fetch_add(1, Ordering::Relaxed) - )); - std::fs::create_dir_all(&root).unwrap(); - root.join("auth.json") + #[test] + fn auth_path_uses_nonempty_xdg_data_root() { + assert_eq!( + auth_path_from(Ok("configured-xdg-data".to_string()), None), + Some(PathBuf::from("configured-xdg-data/opencode/auth.json")) + ); } - fn fixture(path: &std::path::Path, json: &serde_json::Value) -> Option { - std::fs::write(path, serde_json::to_vec(json).unwrap()).unwrap(); - github_copilot_credential_from(path, json) + #[test] + fn auth_path_preserves_whitespace_only_xdg_data_root() { + assert_eq!( + auth_path_from(Ok(" ".to_string()), None), + Some(PathBuf::from(" /opencode/auth.json")) + ); } #[test] - fn labels_oauth_providers_only() { - assert_eq!(subscription_label("openai"), "Codex"); - assert_eq!(subscription_label("github-copilot"), "Copilot"); - assert_eq!(subscription_label("anthropic"), "Claude"); - assert_eq!(subscription_label("minimax-coding-plan"), "Minimax-coding-plan"); + fn auth_path_falls_back_for_empty_xdg_data_root() { + let home = PathBuf::from("resolved-home"); + assert_eq!( + auth_path_from(Ok(String::new()), Some(&home)), + Some(PathBuf::from( + "resolved-home/.local/share/opencode/auth.json" + )) + ); } #[test] - fn copilot_credential_uses_trimmed_refresh_then_access() { - let path = temp_auth_path("token-precedence"); + fn auth_path_falls_back_on_environment_errors() { + let home = PathBuf::from("resolved-home"); + let fallback = Some(PathBuf::from( + "resolved-home/.local/share/opencode/auth.json", + )); + for error in [ + std::env::VarError::NotPresent, + std::env::VarError::NotUnicode(std::ffi::OsString::new()), + ] { + assert_eq!(auth_path_from(Err(error), Some(&home)), fallback); + } + } + + #[test] + fn copilot_lineage_marker_uses_first_valid_refresh_or_access() { + let path = std::env::temp_dir().join("fixture-opencode-auth.json"); let cases = [ ( - json!({"github-copilot": { - "type": "oauth", "refresh": " refresh-token ", "access": "access-token" - }}), - Some("refresh-token"), + "refresh preferred", + serde_json::json!({ + "github-copilot": { + "type": "oauth", + "refresh": " refresh-marker ", + "access": "access-marker" + } + }), + Some("refresh-marker"), ), ( - json!({"github-copilot": { - "type": "oauth", "refresh": " ", "access": " access-token " - }}), - Some("access-token"), + "refresh missing", + serde_json::json!({ + "github-copilot": { "type": "oauth", "access": " access-marker " } + }), + Some("access-marker"), ), ( - json!({"github-copilot": { - "type": "oauth", "refresh": null, "access": "\t\n" - }}), + "refresh null", + serde_json::json!({ + "github-copilot": { + "type": "oauth", + "refresh": null, + "access": "access-marker" + } + }), + Some("access-marker"), + ), + ( + "refresh empty", + serde_json::json!({ + "github-copilot": { + "type": "oauth", + "refresh": "", + "access": "access-marker" + } + }), + Some("access-marker"), + ), + ( + "refresh whitespace", + serde_json::json!({ + "github-copilot": { + "type": "oauth", + "refresh": " \t\n ", + "access": "access-marker" + } + }), + Some("access-marker"), + ), + ( + "refresh non-string", + serde_json::json!({ + "github-copilot": { + "type": "oauth", + "refresh": { "unexpected": true }, + "access": "access-marker" + } + }), + None, + ), + ( + "both invalid", + serde_json::json!({ + "github-copilot": { + "type": "oauth", + "refresh": false, + "access": " " + } + }), None, ), ]; - for (json, expected) in cases { - let credential = fixture(&path, &json); + + for (label, json, expected) in cases { + let credential = github_copilot_credential_from(&path, &json); match expected { - Some(expected) => { - let credential = credential.unwrap(); - assert_eq!(credential.request_token, expected); - assert_eq!(credential.marker, expected.as_bytes()); - } - None => assert!(credential.is_none()), + Some(expected) => match credential { + GitHubCopilotCredentialLoad::Present(credential) => { + assert_eq!(credential.request_token, expected, "{label}"); + assert_eq!(credential.marker, expected.as_bytes(), "{label}"); + } + GitHubCopilotCredentialLoad::Absent => panic!("{label}: unexpectedly absent"), + GitHubCopilotCredentialLoad::Terminal(display) => { + panic!("{label}: unexpectedly terminal: {display}") + } + }, + None => assert!( + matches!(credential, GitHubCopilotCredentialLoad::Terminal(_)), + "{label}" + ), } } - let _ = std::fs::remove_dir_all(path.parent().unwrap()); } #[test] - fn copilot_credential_requires_exact_oauth_entry_and_scopes_location() { - let path = temp_auth_path("exact-entry"); - let credential = fixture( - &path, - &json!({ - "github-copilot": {"type": "oauth", "refresh": " exact-token "}, - "copilot": {"type": "oauth", "refresh": "sibling-token"}, - "GitHub-Copilot": {"type": "oauth", "refresh": "case-token"}, - "github-copilot-enterprise": {"type": "oauth", "refresh": "lookalike-token"}, - "github": {"type": "oauth", "refresh": "foreign-token"} - }), - ) - .unwrap(); - assert_eq!(credential.request_token, "exact-token"); - assert_eq!(credential.semantic_source, "opencode-auth-json"); - assert_eq!( - credential.canonical_location, - crate::agent_account_scope::canonical_file_location(&path, Some("github-copilot")) - .unwrap() - ); - assert_ne!( - credential.canonical_location, - crate::agent_account_scope::canonical_file_location(&path, None).unwrap() - ); - assert_ne!( - credential.canonical_location, - crate::agent_account_scope::canonical_file_location(&path, Some("copilot")).unwrap() - ); - assert!(fixture( - &path, - &json!({ - "github-copilot": {"type": "api", "refresh": "wrong-type"}, - "copilot": {"type": "oauth", "refresh": "sibling-token"}, - "GitHub-Copilot": {"type": "oauth", "refresh": "case-token"} - }) - ) - .is_none()); - let other_path = temp_auth_path("other-location"); - let other = fixture( - &other_path, - &json!({"github-copilot": {"type": "oauth", "refresh": "exact-token"}}), - ) - .unwrap(); - assert_ne!(credential.canonical_location, other.canonical_location); - let _ = std::fs::remove_dir_all(path.parent().unwrap()); - let _ = std::fs::remove_dir_all(other_path.parent().unwrap()); + fn copilot_loader_distinguishes_absent_from_terminal_evidence() { + let path = std::env::temp_dir().join("fixture-opencode-auth.json"); + for (label, json) in [ + ("missing entry", serde_json::json!({})), + ( + "non-oauth entry", + serde_json::json!({ "github-copilot": { "type": "api" } }), + ), + ( + "logged out entry", + serde_json::json!({ + "github-copilot": { + "type": "oauth", + "refresh": null, + "access": " " + } + }), + ), + ] { + assert!( + matches!( + github_copilot_credential_from(&path, &json), + GitHubCopilotCredentialLoad::Absent + ), + "{label}" + ); + } + + for (label, json) in [ + ( + "non-object entry", + serde_json::json!({ "github-copilot": "oauth" }), + ), + ( + "oauth entry without token fields", + serde_json::json!({ "github-copilot": { "type": "oauth" } }), + ), + ( + "missing type", + serde_json::json!({ "github-copilot": { "access": "token" } }), + ), + ( + "malformed token sibling", + serde_json::json!({ + "github-copilot": { + "type": "oauth", + "refresh": false, + "access": "token" + } + }), + ), + ] { + assert!( + matches!( + github_copilot_credential_from(&path, &json), + GitHubCopilotCredentialLoad::Terminal(_) + ), + "{label}" + ); + } + } + + #[test] + fn copilot_file_loader_treats_missing_as_absent_and_io_or_json_as_terminal() { + let root = + std::env::temp_dir().join(format!("tokenbar-copilot-loader-{}", std::process::id())); + let _ = std::fs::remove_dir_all(&root); + std::fs::create_dir_all(&root).unwrap(); + + assert!(matches!( + github_copilot_credential_at(&root.join("missing.json")), + GitHubCopilotCredentialLoad::Absent + )); + assert!(matches!( + github_copilot_credential_at(&root), + GitHubCopilotCredentialLoad::Terminal(_) + )); + + let invalid = root.join("invalid.json"); + std::fs::write(&invalid, "not json").unwrap(); + assert!(matches!( + github_copilot_credential_at(&invalid), + GitHubCopilotCredentialLoad::Terminal(_) + )); + std::fs::remove_dir_all(root).unwrap(); + } + + #[cfg(unix)] + #[test] + fn copilot_loader_rejects_non_utf8_canonical_location() { + use std::os::unix::ffi::OsStringExt; + + let path = std::path::PathBuf::from(std::ffi::OsString::from_vec(vec![0xff])); + let json = serde_json::json!({ + "github-copilot": { "type": "oauth", "access": "token" } + }); + assert!(matches!( + github_copilot_credential_from(&path, &json), + GitHubCopilotCredentialLoad::Terminal(_) + )); } } diff --git a/crates/tb_core_ffi/src/usage_graph.rs b/crates/tb_core_ffi/src/usage_graph.rs index 46576b3..7bd356a 100644 --- a/crates/tb_core_ffi/src/usage_graph.rs +++ b/crates/tb_core_ffi/src/usage_graph.rs @@ -103,10 +103,7 @@ struct TokenContributionData { /// the async `generate_local_graph_report`. That entry point uses cached /// pricing with a graceful offline fallback, and `PricingService` is a process /// -wide `OnceCell`, so the network fetch happens at most once per launch. -pub(crate) fn run( - context: &crate::LocalSourceContext, - year: &str, -) -> Result { +pub(crate) fn run(context: &crate::LocalSourceContext, year: &str) -> Result { let year = normalize_year(year)?; let options = context.report_options(year, None); diff --git a/crates/tb_core_ffi/src/usage_tail.rs b/crates/tb_core_ffi/src/usage_tail.rs index 125f0d8..e02bea1 100644 --- a/crates/tb_core_ffi/src/usage_tail.rs +++ b/crates/tb_core_ffi/src/usage_tail.rs @@ -26,8 +26,10 @@ pub struct UsageEvent { pub model: String, pub input: i64, pub output: i64, + pub reasoning: i64, pub cache_read: i64, pub cache_write: i64, + pub message_count: i32, } impl UsageEvent { @@ -38,6 +40,7 @@ impl UsageEvent { // pattern). self.input .saturating_add(self.output) + .saturating_add(self.reasoning) .saturating_add(self.cache_read) .saturating_add(self.cache_write) } @@ -55,8 +58,8 @@ pub struct TraceBucket { pub struct UsageTailer { events: Mutex>, - /// `latest_source_mtime_ms` token from the last parse; when it hasn't - /// moved, the event window is still correct (rate queries re-filter by + /// Topology-sensitive source token from the last parse; when it hasn't + /// changed, the event window is still correct (rate queries re-filter by /// timestamp on read) and the tick skips the parse entirely. last_source_token: Mutex>, } @@ -91,7 +94,7 @@ impl UsageTailer { // No source changed since the last parse → the window is already // correct; skip the parse. Probe failure falls through to a parse. - let token = tokscale_core::latest_source_mtime_ms(&options).ok(); + let token = tokscale_core::local_source_change_token(&options).ok(); if token.is_some() && *self.last_source_token.lock() == token { return self.events.lock().len(); } @@ -119,8 +122,10 @@ impl UsageTailer { model: m.model_id, input: m.input, output: m.output, + reasoning: m.reasoning, cache_read: m.cache_read, cache_write: m.cache_write, + message_count: m.message_count, } }) .collect(); @@ -180,7 +185,7 @@ impl UsageTailer { let slot = groups.entry(key).or_insert((0, 0)); // saturating_add: same cross-event overflow class as window_total. slot.0 = slot.0.saturating_add(e.total()); - slot.1 += 1; + slot.1 = slot.1.saturating_add(e.message_count.max(0) as u32); } let window_min = (window_secs as f32 / 60.0).max(1.0 / 60.0); let mut out: Vec = groups @@ -222,8 +227,10 @@ mod tests { model: "gemini-3-pro".to_string(), input: i64::MAX, output: 0, + reasoning: 0, cache_read: i64::MAX, cache_write: 0, + message_count: 1, } } @@ -249,6 +256,44 @@ mod tests { assert!(rate.is_finite(), "rate_in_window must not produce NaN/inf"); } + #[test] + fn trace_counts_messages_without_dropping_zero_message_loop_tokens() { + let tailer = UsageTailer::new(); + let now = now_ms(); + *tailer.events.lock() = vec![ + UsageEvent { + ts_ms: now, + client: "grok".to_string(), + agent: "grok".to_string(), + model: "grok-build".to_string(), + input: 10, + output: 2, + reasoning: 0, + cache_read: 3, + cache_write: 0, + message_count: 1, + }, + UsageEvent { + ts_ms: now, + client: "grok".to_string(), + agent: "grok".to_string(), + model: "grok-build".to_string(), + input: 0, + output: 0, + reasoning: 30, + cache_read: 0, + cache_write: 0, + message_count: 0, + }, + ]; + + let buckets = tailer.trace(3600); + assert_eq!(buckets.len(), 1); + assert_eq!(buckets[0].tokens, 45); + assert_eq!(buckets[0].messages, 1); + assert_eq!(tailer.window_total(3600), 45); + } + #[test] fn trace_saturates_across_overlarge_events_in_same_bucket() { let tailer = UsageTailer::new(); diff --git a/crosscheck/README.md b/crosscheck/README.md index 08ba94c..2832f3a 100644 --- a/crosscheck/README.md +++ b/crosscheck/README.md @@ -13,10 +13,11 @@ diagnostic only because serializer formatting and line endings differ by host. | Surface | Canonical source | |---|---| | Legacy `usage-pace.json` and `format.json` | macOS commit `2ed256ee` | -| Provider quota pace v3 | final macOS tree `1e00e7b769a1b0d20b2de077b0a3aab52a24c088` | +| Provider quota pace v3 | final merged Native PR #102 commit `4dfed5ffab26e2707a8cd82ee99864520b75892b` | -Run the Swift side from a clean archive or worktree at the exact provider-v3 -SHA. Do not use an unrelated dirty macOS checkout as the reference. +Run the Swift side from a clean archive or worktree at the exact tested +provider-v3 merged SHA. Do not use an unrelated dirty macOS checkout as the +reference. Canonical fixture fingerprints: @@ -165,8 +166,8 @@ projection differences. ## Running -Set `TOKENBAR_MAC_CANONICAL` to a clean archive or worktree at -`1e00e7b769a1b0d20b2de077b0a3aab52a24c088`. +Set `TOKENBAR_MAC_CANONICAL` to a clean archive or worktree at final merged +Native PR #102 commit `4dfed5ffab26e2707a8cd82ee99864520b75892b`. `Package.swift` links `target/release/libtb_core_ffi.a` by a path relative to the canonical macOS repo root. Build the Rust static library first, and run the @@ -175,7 +176,7 @@ unrelated checkout's `target/` artifact. ```bash export TOKENBAR_WINDOWS="$(pwd)" -export TOKENBAR_MAC_CANONICAL=/private/tmp/tokenbar-mac-1e00e7b-stage4c +export TOKENBAR_MAC_CANONICAL="${TMPDIR:-/tmp}/tokenbar-mac-4dfed5ff" ( cd "$TOKENBAR_MAC_CANONICAL" diff --git a/include/ctb.h b/include/ctb.h index a6b9ad8..562536e 100644 --- a/include/ctb.h +++ b/include/ctb.h @@ -9,8 +9,43 @@ // Envelope: every entry point except tb_probe returns // {"ok":true,"data":} on success // {"ok":false,"err":"..."} on failure -// The data shapes mirror the Tauri frontend contract (TokenBar-tokcat -// src/lib/types.ts and src/lib/agentUsage.ts) field-for-field. +// Payload fields use the Tauri frontend's camelCase contract. In particular, +// AgentUsagePayload is `{generatedAt, publicationGeneration?, agents, +// opencodeSubscriptions}` (the subscription array is omitted when empty). +// `publicationGeneration` is an additive optional Rust `u64` JSON integer for +// generated payloads; demo/legacy payloads omit it. +// Rust assigns it with a checked increment at process-wide publication-gate +// entry before the complete provider run; exhaustion returns an outer error +// rather than repeating a generation. The gate orders generations and pointer +// creation, but is released before `tb_agent_usage` returns and therefore does +// not promise C return order. Swift's shared MainActor publication coordinator +// rejects a lower generation for dashboard, Settings, tray, and snapshot +// consumers. AgentUsage snapshots may additionally carry the additive optional +// camelCase field +// `transportDiagnostic?: {category, status?, osCode?}`. `error` remains the +// user-visible provider status and may coexist with last-good windows; +// `transportDiagnostic` is the only provider failure detail permitted in the +// public Unified Log. Its `category` is limited to timeout/dns/tls/ +// connectionRefused/connectionReset/connect/request/responseBody/rateLimited/ +// serverError. `rateLimited` accepts only status 429; `serverError` accepts only +// 500...599. HTTP categories do not carry `osCode`; non-HTTP categories do not +// carry `status`. `osCode`, when present, is a 32-bit OS error integer. Neither +// field carries token, header, +// body, URL/query, email, account ID, credential path, or free-form cause data. +// Other report payloads retain their existing camelCase shapes from the Tauri +// contract. Adding these fields does not change C function signatures, ownership, +// ABI, or any other wire fields. Each v3 quota window uses +// `{cardId, label, usedPercent, remainingPercent, resetsAt, resetText, +// windowMinutes, paceStatus, historicalPace}`. `paceStatus` is required and +// carries `{state, windowKey, durationSeconds, durationSource, completeCycles, +// reason}`; positive durationSeconds is the pace calculation source of truth, +// while windowMinutes is compatibility output derived by integer division. +// historicalPace is present only for `available` and carries one coherent Rust +// result: `{expectedUsedPercent, etaSeconds, willLastToReset, +// runOutProbability}`. A legacy payload missing the entire paceStatus key is +// not eligible for an implicit Linear fallback. ETA/risk remain optional inside +// an available historical result. Other report payloads retain their existing +// camelCase shapes from the Tauri contract. // tb_probe keeps its Phase 0 shape: {"ok":true,"messages":N} / {"ok":false,...}. // // `year` parameters may be NULL or "" for the all-time view, otherwise a @@ -39,9 +74,7 @@ char *tb_usage_trace(int64_t window_secs); // Live rate: {"tokensPerMin": } (10-minute-window average). char *tb_tokens_per_min(void); -// OAuth quota cards (AgentUsagePayload) for codex/claude/grok/antigravity/copilot. -// Each quota window has a stable cardId and typed paceStatus; historicalPace is -// present only when the Rust core has a coherent historical projection. +// OAuth quota cards (AgentUsagePayload) for codex/claude/antigravity/copilot/grok. // Network-bound; per-provider failures are reported inside each snapshot. char *tb_agent_usage(void); diff --git a/src/TokenBar.App/DashboardModel.cs b/src/TokenBar.App/DashboardModel.cs index f12fb1d..9bc5c63 100644 --- a/src/TokenBar.App/DashboardModel.cs +++ b/src/TokenBar.App/DashboardModel.cs @@ -519,7 +519,9 @@ private void RefreshQuota() { try { - var quota = TryFetch(() => TbCore.AgentUsage(), "agentUsage"); + var quota = TryFetch( + () => AgentUsageFetchCoordinator.Shared.FetchAsync().GetAwaiter().GetResult(), + "agentUsage"); if (quota is not null) { _latestQuota = quota; diff --git a/src/TokenBar.App/TrayFeed.cs b/src/TokenBar.App/TrayFeed.cs index 2269ffd..063646a 100644 --- a/src/TokenBar.App/TrayFeed.cs +++ b/src/TokenBar.App/TrayFeed.cs @@ -170,7 +170,9 @@ private void RefreshSlow() _lastFullRefresh = DateTimeOffset.Now; } - var quota = TryFetch(() => TbCore.AgentUsage(), "tray quota"); + var quota = TryFetch( + () => AgentUsageFetchCoordinator.Shared.FetchAsync().GetAwaiter().GetResult(), + "tray quota"); _ = _dispatcher.TryEnqueue(() => { if (_disposed) diff --git a/src/TokenBar.Core.Tests/AgentUsageFetchCoordinatorTests.cs b/src/TokenBar.Core.Tests/AgentUsageFetchCoordinatorTests.cs new file mode 100644 index 0000000..702e39e --- /dev/null +++ b/src/TokenBar.Core.Tests/AgentUsageFetchCoordinatorTests.cs @@ -0,0 +1,56 @@ +using TokenBar.Core; +using TokenBar.Interop; + +namespace TokenBar.Core.Tests; + +public class AgentUsageFetchCoordinatorTests +{ + [Fact] + public async Task ConcurrentCallersShareOneFetchAndNextCallStartsFresh() + { + var started = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously); + var release = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously); + var calls = 0; + var payload = new AgentUsagePayload("now", []); + var coordinator = new AgentUsageFetchCoordinator(() => + { + Interlocked.Increment(ref calls); + started.TrySetResult(); + release.Task.GetAwaiter().GetResult(); + return payload; + }); + + var first = coordinator.FetchAsync(); + await started.Task; + var second = coordinator.FetchAsync(); + + Assert.Same(first, second); + release.SetResult(); + Assert.Same(payload, await first); + Assert.Same(payload, await second); + Assert.Equal(1, Volatile.Read(ref calls)); + + Assert.Same(payload, await coordinator.FetchAsync()); + Assert.Equal(2, Volatile.Read(ref calls)); + } + + [Fact] + public async Task FailedFetchIsNotCached() + { + var calls = 0; + var payload = new AgentUsagePayload("now", []); + var coordinator = new AgentUsageFetchCoordinator(() => + { + if (Interlocked.Increment(ref calls) == 1) + { + throw new InvalidOperationException("synthetic failure"); + } + + return payload; + }); + + await Assert.ThrowsAsync(() => coordinator.FetchAsync()); + Assert.Same(payload, await coordinator.FetchAsync()); + Assert.Equal(2, Volatile.Read(ref calls)); + } +} diff --git a/src/TokenBar.Core.Tests/DtoDecodeTests.cs b/src/TokenBar.Core.Tests/DtoDecodeTests.cs index 23375d5..26097bf 100644 --- a/src/TokenBar.Core.Tests/DtoDecodeTests.cs +++ b/src/TokenBar.Core.Tests/DtoDecodeTests.cs @@ -117,6 +117,19 @@ private static AgentUsagePayload DecodeAgentUsagePayload( Web) ?? throw new InvalidOperationException("agent usage payload decoded to null"); } + private static AgentUsageSnapshot DecodeTransportDiagnostic(string? diagnosticJson = null) + { + var diagnostic = diagnosticJson is null + ? "" + : $$""","transportDiagnostic":{{diagnosticJson}}"""; + var payload = TbCore.DecodeEnvelope( + $$$"""{"ok":true,"data":{"generatedAt":"now","agents":[{"clientId":"codex","source":"oauth","updatedAt":"now","windows":[{"label":"Weekly","usedPercent":40,"remainingPercent":60}]{{{diagnostic}}}}]}}"""); + Assert.Equal("now", payload.GeneratedAt); + var snapshot = Assert.Single(payload.Agents); + Assert.Single(snapshot.Windows); + return snapshot; + } + private static void AssertRejects(string json) => Assert.Throws(() => DecodeWindow(json)); @@ -184,6 +197,85 @@ public void AgentUsageDecodesLegacyFieldsWithoutTrustingScalars() Assert.Null(p.OpencodeSubscriptions); } + [Fact] + public void AgentUsagePayloadIgnoresUnknownPublicationGeneration() + { + var payload = TbCore.DecodeEnvelope( + """{"ok":true,"data":{"generatedAt":"now","publicationGeneration":42,"agents":[]}}"""); + Assert.Equal("now", payload.GeneratedAt); + Assert.Empty(payload.Agents); + } + + [Fact] + public void AgentUsageTransportDiagnosticDecodesHttpAndOsFacts() + { + var http = Assert.IsType( + DecodeTransportDiagnostic( + """{"category":"rateLimited","status":429}""").TransportDiagnostic); + Assert.Equal(new AgentUsageTransportDiagnostic("rateLimited", Status: 429), http); + + var os = Assert.IsType( + DecodeTransportDiagnostic( + """{"category":"connectionRefused","osCode":10061}""").TransportDiagnostic); + Assert.Equal(new AgentUsageTransportDiagnostic("connectionRefused", OsCode: 10061), os); + } + + [Fact] + public void AgentUsageTransportDiagnosticOmittedNullOrNonObjectDecodesAsNull() + { + foreach (var diagnosticJson in new string?[] { null, "null", "[]" }) + { + Assert.Null(DecodeTransportDiagnostic(diagnosticJson).TransportDiagnostic); + } + } + + [Fact] + public void AgentUsageTransportDiagnosticCategoryTypeMismatchPreservesOtherFields() + { + var diagnostic = Assert.IsType( + DecodeTransportDiagnostic( + """{"category":42,"status":429,"osCode":10061}""").TransportDiagnostic); + Assert.Equal(new AgentUsageTransportDiagnostic(Status: 429, OsCode: 10061), diagnostic); + } + + [Theory] + [InlineData("\"429\"")] + [InlineData("9223372036854775808")] + public void AgentUsageTransportDiagnosticMalformedStatusPreservesCategoryAndOsCode( + string statusJson) + { + var diagnostic = Assert.IsType( + DecodeTransportDiagnostic( + $$"""{"category":"connectionRefused","status":{{statusJson}},"osCode":10061}""") + .TransportDiagnostic); + Assert.Equal( + new AgentUsageTransportDiagnostic("connectionRefused", OsCode: 10061), + diagnostic); + } + + [Fact] + public void AgentUsageTransportDiagnosticPreservesUnknownCategory() + { + var diagnostic = Assert.IsType( + DecodeTransportDiagnostic( + """{"category":"futureTransport","status":599}""").TransportDiagnostic); + Assert.Equal(new AgentUsageTransportDiagnostic("futureTransport", Status: 599), diagnostic); + } + + [Fact] + public void AgentUsageTransportDiagnosticDoesNotRetainUnknownSensitiveMember() + { + var snapshot = DecodeTransportDiagnostic( + """{"category":"timeout","authorization":"Bearer token-secret"}"""); + Assert.Equal( + new AgentUsageTransportDiagnostic("timeout"), + Assert.IsType(snapshot.TransportDiagnostic)); + + var serialized = JsonSerializer.Serialize(snapshot, Web); + Assert.False(serialized.Contains("authorization", StringComparison.OrdinalIgnoreCase)); + Assert.False(serialized.Contains("token-secret", StringComparison.Ordinal)); + } + [Fact] public void AgentUsagePayloadDecodesStringSubscriptions() { diff --git a/src/TokenBar.Core/AgentUsageFetchCoordinator.cs b/src/TokenBar.Core/AgentUsageFetchCoordinator.cs new file mode 100644 index 0000000..13a4972 --- /dev/null +++ b/src/TokenBar.Core/AgentUsageFetchCoordinator.cs @@ -0,0 +1,41 @@ +using TokenBar.Interop; + +namespace TokenBar.Core; + +/// Shares one blocking agent-usage fetch across concurrent UI callers. +public sealed class AgentUsageFetchCoordinator(Func fetch) +{ + private readonly object _gate = new(); + private Task? _inFlight; + + public static AgentUsageFetchCoordinator Shared { get; } = new(TbCore.AgentUsage); + + public Task FetchAsync() + { + lock (_gate) + { + if (_inFlight is { } current) + { + return current; + } + + var next = Task.Run(fetch); + _inFlight = next; + _ = next.ContinueWith( + completed => + { + lock (_gate) + { + if (ReferenceEquals(_inFlight, completed)) + { + _inFlight = null; + } + } + }, + CancellationToken.None, + TaskContinuationOptions.ExecuteSynchronously, + TaskScheduler.Default); + return next; + } + } +} diff --git a/src/TokenBar.Interop/AgentUsage.cs b/src/TokenBar.Interop/AgentUsage.cs index 2afc59a..f830876 100644 --- a/src/TokenBar.Interop/AgentUsage.cs +++ b/src/TokenBar.Interop/AgentUsage.cs @@ -647,6 +647,65 @@ public sealed record CreditsSnapshot( bool Unlimited, double? Remaining = null); +public sealed record AgentUsageTransportDiagnostic( + string? Category = null, + long? Status = null, + long? OsCode = null); + +public sealed class AgentUsageTransportDiagnosticJsonConverter : + JsonConverter +{ + public override bool HandleNull => true; + + public override AgentUsageTransportDiagnostic? Read( + ref Utf8JsonReader reader, + Type typeToConvert, + JsonSerializerOptions options) + { + using var document = JsonDocument.ParseValue(ref reader); + var root = document.RootElement; + if (root.ValueKind != JsonValueKind.Object) + { + return null; + } + + return new AgentUsageTransportDiagnostic( + Category: OptionalString(root, "category"), + Status: OptionalInt64(root, "status"), + OsCode: OptionalInt64(root, "osCode")); + } + + public override void Write( + Utf8JsonWriter writer, + AgentUsageTransportDiagnostic? value, + JsonSerializerOptions options) + { + if (value is null) + { + writer.WriteNullValue(); + return; + } + + writer.WriteStartObject(); + if (value.Category is not null) writer.WriteString("category", value.Category); + if (value.Status is { } status) writer.WriteNumber("status", status); + if (value.OsCode is { } osCode) writer.WriteNumber("osCode", osCode); + writer.WriteEndObject(); + } + + private static string? OptionalString(JsonElement obj, string name) => + obj.TryGetProperty(name, out var value) && value.ValueKind == JsonValueKind.String + ? value.GetString() + : null; + + private static long? OptionalInt64(JsonElement obj, string name) => + obj.TryGetProperty(name, out var value) && + value.ValueKind == JsonValueKind.Number && + value.TryGetInt64(out var number) + ? number + : null; +} + public sealed record AgentUsageSnapshot( string ClientId, string Source, @@ -654,7 +713,9 @@ public sealed record AgentUsageSnapshot( IReadOnlyList Windows, AgentIdentity? Identity = null, CreditsSnapshot? Credits = null, - string? Error = null) : IJsonOnDeserialized + string? Error = null, + [property: JsonConverter(typeof(AgentUsageTransportDiagnosticJsonConverter))] + AgentUsageTransportDiagnostic? TransportDiagnostic = null) : IJsonOnDeserialized { void IJsonOnDeserialized.OnDeserialized() { diff --git a/vendor/README.md b/vendor/README.md index e815149..c2d47d5 100644 --- a/vendor/README.md +++ b/vendor/README.md @@ -16,6 +16,210 @@ This file remains the exact vendor ledger. The selective-port method and streami > which is stale. It sits between upstream v3.0.0 and v3.1.x, plus the local > patches below, plus the cherry-picked upstream commits listed next. +## Current selective-alignment checkpoint + +M21 merged in PR [#71](https://github.com/Nanako0129/TokenBar/pull/71) at [`471a7f239f0270b4ebfaed04894335c506d588d3`](https://github.com/Nanako0129/TokenBar/commit/471a7f239f0270b4ebfaed04894335c506d588d3), preserving the shared materialized, shipping streaming, count, and report paths for Kimi Code, Junie, and OpenCodeReview. M22 then reclassified five Zcode-bearing rows without changing runtime code, producing `74/8/0/15/13/1`. + +M25 merged in PR [#75](https://github.com/Nanako0129/TokenBar/pull/75) and moved reloadable grouping-only model aliases (`9a5aeb65`) to `ALREADY_VENDORED`, producing `75/7/0/15/13/1`. M25 itself required no bump from the then-active schema 31; the later Grok `turn_completed.usage` correction in PR #77 advanced the current monolithic cache to schema 32. + +M22 PR [#72](https://github.com/Nanako0129/TokenBar/pull/72) is closed and unmerged (`merged=false`), so none of its implementation is part of main. Its head is [`c41b864b2b867bd84b69e82ff391b0197177775e`](https://github.com/Nanako0129/TokenBar/commit/c41b864b2b867bd84b69e82ff391b0197177775e), spanning 17 commits and 12 files (`+4392/-52`), with production drift about 2.1–2.5x upstream and total drift about 3.1–3.5x, including 15 review-driven fixes and about 520 lines of custom cross-store matching. Under the fidelity rule, M22 is `DEFER` until upstream converges; do not treat its code as landed or assume `ClientId::Zcode` / `COUNT=34`. + +M23 was split after PR #74 (`e274f2ad`) reached about 3.1x Copilot production drift and 14 review rounds without converging; PR #74 is closed unmerged as fidelity evidence. M23-H merged in PR [#82](https://github.com/Nanako0129/TokenBar/pull/82) at [`1a8ee0c6`](https://github.com/Nanako0129/TokenBar/commit/1a8ee0c62429472589308c2836c026d6e5432243) with only the faithful Hermes Windows discovery residual (`c1aef5e9`). M23-D then merged in PR [#83](https://github.com/Nanako0129/TokenBar/pull/83) at [`f99d9274`](https://github.com/Nanako0129/TokenBar/commit/f99d9274fcfdfc8fc228e21e52808584f822385f), rebuilding Copilot Desktop from `f6f7eced + 0b454e60` with one pre-aggregation OTEL-session authority selector. M23-V keeps VS Code `chatSessions` (`074619f7`) in `DEFER` until upstream fixes ObjectMutationLog replay. The merged checkpoint is `78/3/0/16/13/1`, with cache schema 32 unchanged. + +The Copilot duplicate-span contribution reached upstream in issue [#938](https://github.com/junhoyeo/tokscale/issues/938) and merged PR [#939](https://github.com/junhoyeo/tokscale/pull/939) at merge commit [`1652852f`](https://github.com/junhoyeo/tokscale/commit/1652852f0f35d14fe4ff87f6cc32cc9d02a2ddd8). Post-merge verification found deterministic agent-attribution and partial-timestamp gaps, tracked in issue [#942](https://github.com/junhoyeo/tokscale/issues/942) and follow-up PR [#943](https://github.com/junhoyeo/tokscale/pull/943). PR #943 head `9c399cc5` is maintainer-ready: current-head Codex is clean, all CI checks pass, no review threads remain unresolved, the full upstream format/clippy/test/release-build gates pass, and fresh verification confirms parser-version 7 invalidation plus the agent/timing regressions. These upstream-only commits are outside the fixed 111-row audit range and do not change this vendor tree, its ledger counts, or local cache schema 32. + +M24 PR [#86](https://github.com/Nanako0129/TokenBar/pull/86) is closed unmerged as fidelity evidence; its [close rationale](https://github.com/Nanako0129/TokenBar/pull/86#issuecomment-5047332480) records the public stop decision. Current-head Codex found four independent security/state-machine defects, and the first local fix round was rejected by fresh security verification because process-local Warp source/revocation state could race with the shared cross-process singleton app cache: one process handling a 401 or Disconnect for scope B could delete another process's valid scope-A cache. Holding the refresh lock through remote I/O violates the approved network boundary; releasing it exposes the cache-ownership conflict. The second consecutive review round therefore triggered the M24 stop condition. No review fix was pushed, `63a44d7c` moved from `TAKE` to `DEFER`, and the resulting checkpoint was `78/2/0/17/13/1`. + +The approved replacement delivery graph removes M24 from the cache dependency chain without reviving Warp: `M26-A → M26-B → M19-BP → M19-BQ → PT0 → M19-B0 → M19-B1 → D2`. M26-A merged in PR [#90](https://github.com/Nanako0129/TokenBar/pull/90) at [`95c819c7`](https://github.com/Nanako0129/TokenBar/commit/95c819c7cf6532be7276b64386490b1b03a0c1ae); M26-B merged in PR [#91](https://github.com/Nanako0129/TokenBar/pull/91) at [`cc52c3b9`](https://github.com/Nanako0129/TokenBar/commit/cc52c3b9e01836cbf1c7ab7c77b4dd8b5e2df7b2), moving `cd07bf78` to `DEFER` because its Devin residual remains excluded. M19-BP, M19-BQ, and PT0 then merged in PRs [#92](https://github.com/Nanako0129/TokenBar/pull/92), [#93](https://github.com/Nanako0129/TokenBar/pull/93), and [#94](https://github.com/Nanako0129/TokenBar/pull/94) without changing the ledger. M19-B0 merged in PR [#99](https://github.com/Nanako0129/TokenBar/pull/99) at [`f820b06f`](https://github.com/Nanako0129/TokenBar/commit/f820b06fd99a53cada8495338bd7d58898525a7b), followed by the Native review fixes in PR [#101](https://github.com/Nanako0129/TokenBar/pull/101) at [`7abd205a`](https://github.com/Nanako0129/TokenBar/commit/7abd205ab692950c9b4574a14da72f01733d1935). M19-B1 is validating the exact Windows sync; Warp remains `parse_local: false`. + +The immutable audited set is the 111 hashes produced in a clean upstream clone: + +```bash +git rev-list --reverse \ + 0c820a5d406241d85dc6c7fc65ac5a1ee026ccfd..366ce64395594abf111e0409581d91016561b25a \ + -- crates/tokscale-core +``` + +The classification union has no duplicates and no symmetric difference from that range: + +| Classification | Count | +|---|---:| +| `ALREADY_VENDORED` | 79 | +| `TAKE` | 0 | +| `ADAPT_FOR_STREAMING` | 0 | +| `DEFER` | 18 | +| `SKIP` | 13 | +| `SUPERSEDED` | 1 | +| **Total** | **111** | + +### Exact 111-commit classification + +
+ALREADY_VENDORED — 79 + +```text +6dfd79f5 d9f2a9b7 44055841 1a305f0f 5c1fe659 7500b303 8493048f 2d90f41d +d4a3bd32 1492b962 b43dc5f8 4101711b 28aec200 aebe4ea8 5017eefb 0ce3d73f +3a68cf52 a75533e6 70fd5249 cbbd0dff 81d721fd 7f48257a 0549e2ed 686f3cf2 +783bbb8d 5ff7bf44 c2156fea cb6ebf61 4bc2aa03 afa65ed6 235230ac bc06d4ee +979b7015 7403dafa d5f2c6c4 59421da9 31deb7e6 b8156e64 dcb053e0 23cf62e0 +4cbc2f6b 0f84d174 da5e06d2 1752636f b7277d49 85669602 b49cec19 3587f745 +d50da475 24e3771c e5cfbae2 b64e4f14 72bf6667 46e01977 31bfd167 09344531 +163ec570 a2f7cef5 a0929482 366ce643 405ded4a 315549b4 6899ea03 b59979c5 +9155018c 18cd13cc a87f0ab6 959cce84 6c804711 839ce378 052f43de 633ea946 +77948d9d 302d39c3 9a5aeb65 c1aef5e9 f6f7eced 0b454e60 ae36db5c +``` + +
+ +
+TAKE — 0 + +```text +(empty) +``` + +
+ +
+ADAPT_FOR_STREAMING — 0 + +```text +(empty) +``` + +
+ +
+DEFER — 18 + +```text +18c7e87f db88138b 1c91cb34 6a1535d1 90d28ec0 +20f6d4dd b9b7d09f 0097ba7e ed64e77b 34cfbb50 +640e97b9 f7a124da ed6f8b95 65f8f3e2 b64d861e 074619f7 63a44d7c cd07bf78 +``` + +
+ +
+SKIP — 13 + +```text +b2b8c1fc 7ddfa748 b48af31e e644f966 010acd85 46f8fff9 c634d1a5 +471ad5a2 9b85b671 64f92fe9 8256280a cd394af2 d01db0a7 +``` + +
+ +
+SUPERSEDED — 1 + +```text +88b32ac8 +``` + +
+ +### Selected work + +M20 moved `366ce643` to `ALREADY_VENDORED`; M15-B moved `405ded4a` and `315549b4`; M16 moved `6899ea03`, `b59979c5`, `9155018c`, and `18cd13cc` to `ALREADY_VENDORED` while `34cfbb50` moved to `DEFER`; M19-A moved `a87f0ab6` after taking only its Windows atomic-replacement hunk; M17 used a non-main source and left the audited counts unchanged; M18 moved `959cce84` and `6c804711`; M21 moved `839ce378`, `052f43de`, `633ea946`, `77948d9d`, and `302d39c3` to `ALREADY_VENDORED`. M22 is closed unmerged, so no implementation row moved to `ALREADY_VENDORED`; the product decision instead reclassified its five Zcode-bearing rows from `TAKE` to `DEFER`. M25 moved `9a5aeb65` to `ALREADY_VENDORED`. M23-H moved `c1aef5e9` to `ALREADY_VENDORED` and reclassified `074619f7` to `DEFER`; merged M23-D moved `f6f7eced` and `0b454e60` to `ALREADY_VENDORED`. M24 PR #86 is closed unmerged, so `63a44d7c` moves from `TAKE` to `DEFER` and no Warp runtime code enters main. M26-A merged in PR #90 at `95c819c7`, moving `ae36db5c` to `ALREADY_VENDORED` after selectively adopting the format-1 shard engine and adapting both Native cache lanes. M26-B takes only generic format-2 metadata and Claude cached-parent recovery from `cd07bf78`; its Devin residual remains excluded, so `cd07bf78` moves from `TAKE` to `DEFER`. Mixed `b64d861e` remains one `DEFER` row because only its Kiro, Jcode, Junie, and OpenCodeReview hunks are vendored while Zcode and Devin remain excluded. + +| Milestone | Selected scope | Audited-range commits | +|---|---|---| +| M20 — merged in PR #65 | OpenCode v2 SQLite | `366ce643` | +| M15-B — merged in PR #66 | Kiro structured sessions | `405ded4a 315549b4` + `b64d861e` Kiro hunk | +| M16 — merged in PR #67 | Codex, Claude, Copilot, Jcode, provider, and Antigravity correctness | `6899ea03 b59979c5 9155018c 18cd13cc` + `34cfbb50` provider hunks + `b64d861e` Jcode hunk | +| M17 — merged in PR #69 | Grok unified-log precedence across all cache/report lanes | Non-main `ed798642` | +| M21 — merged in PR #71 | Kimi Code, Junie, and OpenCodeReview | `839ce378 052f43de 633ea946 77948d9d 302d39c3` + `b64d861e` Junie/OpenCodeReview hunks | +| M22 — PR #72 closed unmerged / DEFER | Zcode legacy and v2; no implementation landed | `640e97b9 f7a124da ed6f8b95 65f8f3e2` + `b64d861e` Zcode hunk remain deferred | +| M23-H — merged in PR #82 | Hermes Windows default/profile discovery | `c1aef5e9` | +| M23-D — merged in PR #83 | Copilot Desktop token source with OTEL whole-session authority | `f6f7eced 0b454e60` | +| M23-V — DEFER | Copilot VS Code `chatSessions`; upstream ObjectMutationLog replay is not yet trustworthy | `074619f7` | +| M18 — merged in PR #70 | Sakana/Fugu pricing and the full routed-pricing pipeline | `959cce84 6c804711` | +| M25 — merged in PR #75 | Reloadable configurable model aliases | `9a5aeb65` | +| M24 — PR #86 closed unmerged / DEFER | Warp producer and local reporting; process-local source ownership conflicts with the shared cross-process singleton cache | `63a44d7c` remains unvendored and deferred | +| M19-A — merged in PR #68 | Windows atomic replacement retry in the canonical Native source | `a87f0ab6` Windows hunk | +| M26-A — merged in PR #90 at `95c819c7` | Identity-aware 256-shard source-message cache with per-parser versions, bounded generic/SQLite sampling, Codex prefix state, concurrent-writer merge, and legacy-monolith preservation; Warp stays disabled | `ae36db5c: TAKE → ALREADY_VENDORED` | +| M26-B — based on `43fa8ad6` | Format-2 generic related-file path/existence metadata and Claude cached-parent recovery; Pi/Devin hunks excluded | `cd07bf78: TAKE → DEFER` because the Devin residual remains excluded | + +The selected non-main semantic sources stay outside the 111-row ledger: Grok unified-log `ed798642` (#849) for M17, request-level long-context pricing `548dc124` (#862) and routed prefix/suffix composition `6ea27ca1` (#846) for M18. Warp producer commit `d1cd03c2` (#636) predates the audited anchor and is only a semantic source for M24. + +### Deferred and skipped scope + +| Decision | Scope | Current audited rows | +|---|---|---| +| `DEFER` | Command Code | `18c7e87f db88138b` | +| `DEFER` | CodeBuddy / WorkBuddy | `1c91cb34 6a1535d1 90d28ec0 20f6d4dd b9b7d09f` | +| `DEFER` | Zcode legacy / v2 | `640e97b9 f7a124da ed6f8b95 65f8f3e2` plus the Zcode hunk of mixed `b64d861e`; PR #72 closed unmerged after exceeding the fidelity threshold | +| `DEFER` | Copilot VS Code `chatSessions` | `074619f7`; PR #74 reached about 3.1x Copilot production drift and still mis-replays ObjectMutationLog `kind:2`, so reassess only after upstream format semantics converge | +| `DEFER` | Warp producer and local reporting | `63a44d7c`; PR #86 closed unmerged after two consecutive review rounds exposed new security failure classes, culminating in a process-local state versus cross-process singleton-cache ownership conflict | +| `DEFER` | Devin CLI / Desktop | `0097ba7e ed64e77b` plus the Devin hunk of mixed `b64d861e`; M26-B leaves the Devin residual of `cd07bf78` excluded | +| `DEFER` | 9Router | `34cfbb50`; M16 takes only its provider hardening, leaving the bridge and 9Router product integration excluded | +| `SKIP` | Sakana subscription billing-console scrape | `c634d1a5` (#745); Fugu model pricing is selected separately in M18 | + +### Mixed-commit bookkeeping + +| Commit | Current state | Reassessment boundary | +|---|---|---| +| `34cfbb50` | `DEFER`; M16 vendors the provider hardening, while the remaining bridge and scanner hunks are 9Router-only | Remains `DEFER` | +| `b64d861e` | Kiro, Jcode, Junie, and OpenCodeReview start-anchor hunks are vendored through M21; Zcode did not merge | `DEFER`; its remaining Zcode and Devin hunks both wait for upstream convergence | +| `c1aef5e9` | `ALREADY_VENDORED`; M11 carried the macOS/profile scope and M23-H adds the remaining Windows home candidates without widening explicit `HERMES_HOME` | Remains `ALREADY_VENDORED` | +| `63a44d7c` | `DEFER`; the active-day aggregator hunk was already vendored earlier, but PR #86's Warp enablement/local producer did not merge | Reassess only after a coherent cross-process source/cache ownership contract exists | +| `ae36db5c` | `ALREADY_VENDORED`; M26-A adopts the format-1 shard architecture and retains the previously vendored Claude/Droid/Kimi/Kiro dependency seams | Per-client parser versions are append-only after this checkpoint | +| `cd07bf78` | Generic format-2 cache metadata and Claude cached-parent recovery are vendored; the Devin residual remains excluded, so the row is `DEFER` | Reassess the excluded Devin scope only after upstream convergence | + +### Execution order and ledger transitions + +```mermaid +flowchart TD + T[M15-T canonical docs/ledger] --> O[M20 OpenCode v2] + O --> K[M15-B Kiro structured] + K --> C[M16 existing-parser correctness] + C --> G[M17 Grok unified] + G --> N1[M21 Kimi Code + Junie + OpenCodeReview] + N1 --> H[M23-H Hermes Windows] + H --> D[M23-D Copilot Desktop] + C --> P[M18 Sakana + long-context + routed pricing] + P --> A[M25 reloadable model aliases] + A --> W[M24 fidelity stop — PR #86 / DEFER] + T --> F[M19-A Windows atomic retry] + D --> SA[M26-A format-1 shards] + F --> SA + SA --> SB[M26-B format-2 metadata] + SB --> BP[M19-BP shared roots] + BP --> BQ[M19-BQ Claude probe] + BQ --> PT0[PT0 provider transport] + PT0 --> W0[M19-B0 Native secure storage] + W0 --> W1[M19-B1 Windows final sync] + W1 --> D2[D2 final docs checkpoint] +``` + +| Milestone | Ledger transition or current state | +|---|---| +| M20 | `366ce643: TAKE → ALREADY_VENDORED` | +| M15-B | `405ded4a 315549b4: TAKE → ALREADY_VENDORED`; `b64d861e` remains `TAKE` | +| M16 | `6899ea03 b59979c5 9155018c 18cd13cc: TAKE → ALREADY_VENDORED`; `34cfbb50: TAKE → DEFER`; `b64d861e` remains `TAKE` | +| M17 | Non-main source; counts unchanged | +| M21 | `839ce378 052f43de 633ea946 77948d9d 302d39c3: TAKE → ALREADY_VENDORED`; `b64d861e` remains `TAKE` | +| M22 | Closed unmerged / DEFER; `640e97b9 f7a124da ed6f8b95 65f8f3e2 b64d861e: TAKE → DEFER` under the fidelity threshold | +| M23-H | `c1aef5e9: TAKE → ALREADY_VENDORED`; `074619f7: TAKE → DEFER` after the PR #74 fidelity stop | +| M23-D | PR #83 merged `f6f7eced 0b454e60: TAKE → ALREADY_VENDORED` | +| M23-V | No runtime branch; `074619f7` remains `DEFER` until upstream ObjectMutationLog semantics converge | +| M18 | `959cce84 6c804711: TAKE → ALREADY_VENDORED`; non-main sources do not change counts | +| M25 | `9a5aeb65: TAKE → ALREADY_VENDORED` | +| M24 | PR #86 closed unmerged; `63a44d7c: TAKE → DEFER` after the security fidelity stop | +| M19-A | `a87f0ab6: TAKE → ALREADY_VENDORED`; the TUI signal hunk remains irrelevant to TokenBar | +| M26-A | PR #90 merged at `95c819c7`; `ae36db5c: TAKE → ALREADY_VENDORED`; format-1 identity-aware shards become active while M24 remains excluded | +| M26-B | PR #91 merged at `cc52c3b9`; generic format-2 metadata and Claude cached-parent recovery landed, while `cd07bf78: TAKE → DEFER` leaves the Devin residual and Pi excluded | +| M19-BP / M19-BQ / PT0 | PRs #92／#93／#94 merged shared roots, the Claude version probe, and provider transport without changing counts | +| M19-B0 / M19-B1 | M19-B0 merged in PR #99 with Native review fixes in PR #101; M19-B1 validates the exact Native → Windows shared-core sync and keeps Windows shared-tree local patches empty; counts unchanged | + +The current checkpoint is `ALREADY_VENDORED 79`, `TAKE 0`, `ADAPT_FOR_STREAMING 0`, `DEFER 18`, `SKIP 13`, and `SUPERSEDED 1`, total 111. M24 remains deferred and outside the cache dependency graph; M26-A, M26-B, M19-BP, M19-BQ, PT0, and M19-B0 are merged, so the approved remaining sequence is M19-B1, then D2. Every later transition must start from this actual ledger, regenerate all six sets, and rerun duplicate and symmetric-difference checks. + +M26-A made shard format 1 active at `source-message-cache-v2//shard-XX.bin`; M26-B makes format 2 the active cache contract by persisting related-file paths and existence state and reusing cached Claude parent candidates only after the primary fingerprint matches. Existing format-1 shards are locally stale and cold-rebuild under format 2. The former schema-32 `source-message-cache.bin` remains inert provenance: the shard reader does not load, migrate, rewrite, or delete it. Per-client parser versions remain independent of this global storage-format bump. + +Public issue #45 is the designated remote inventory. M25 merged in PR #75; M22 PR #72 is closed unmerged; M23-H merged in PR #82 at `1a8ee0c6`; M23-D merged in PR #83 at `f99d9274`; M26-A merged in PR #90 at `95c819c7`; M26-B merged in PR #91 at `cc52c3b9`; M19-BP, M19-BQ, and PT0 merged in PRs #92／#93／#94; M19-B0 and its Native review fixes merged in PRs #99／#101; PR #74 is closed unmerged at `e274f2ad`; and M24 PR #86 is closed unmerged as the latest fidelity evidence. The private Project tracks executable milestones only; it does not duplicate the 111 commit rows. + ## Cherry-picked upstream commits (ahead of baseline) Specific later upstream fixes pulled in selectively (we do **not** re-vendor @@ -31,7 +235,7 @@ next sync should treat these as already-present and not re-apply them. | `#649` (44055841), `#651` (2d90f41d), `#681` (b43dc5f8), `#719` (3a68cf52) | Codex fork-replay correctness: skip replayed parent `token_count` rows in fork/subagent logs, scope the dedup key to the fork parent so sibling replays collapse, and keep user-fork own-turns after a repeated child `session_meta`. Fixes codex token over-counting (nested-parent + sibling forks) and under-counting (user forks). `CACHE_SCHEMA_VERSION` bumped to 19 to reparse stale entries. `lib.rs` change is the ported upstream integration tests (plus a streaming-path regression); production logic is in `codex.rs`. `#646` (turn detection) was already vendored. | `src/sessions/codex.rs`, `src/message_cache.rs`, `src/lib.rs` (tests) | | `#d4a3bd32` (Cline) | New client: Cline (the project Roo Code / Kilo Code forked from). VS Code globalStorage `ui_messages.json`; parser is a one-line wrapper over `roocode::parse_roo_kilo_file`. Lane added to OUR `scan_messages_streaming` (not upstream's materialized path). `ClientId::Cline = 25`. | `src/sessions/cline.rs`, `src/clients.rs`, `src/scanner.rs`, `src/lib.rs` | | `#70fd5249` (#713, Antigravity CLI) | New client: Antigravity CLI (terminal agent, distinct from the IDE `antigravity`). Reads `~/.gemini/antigravity-cli/conversations/*.db` SQLite directly via a dependency-free protobuf wire-reader (reverse-engineered fields). `ClientId::AntigravityCli = 26` (**upstream numbers it 31** — intervening clients gjc/grok/jcode/micode/commandcode were not yet vendored at the time). Streaming lane parses + prices + dedups on responseId **namespaced by session** (upstream has no cross-file gate; a responseId is unique only within a conversation DB, so the gate is conversation-scoped to keep our numbers identical to upstream while staying collision-proof). Adds the `gemini-3-flash-a` pricing alias. | `src/sessions/antigravity_cli.rs`, `src/clients.rs`, `src/scanner.rs`, `src/pricing/aliases.rs`, `src/lib.rs` | -| Grok Build (upstream `sessions/grok.rs`) | New client: Grok Build. Reads `$GROK_HOME/sessions/*/*/updates.jsonl` (default `~/.grok`). **Upstream behavior:** cumulative `totalTokens` deltas become input tokens; every counter rewind is treated as streaming jitter and skipped; a sibling `signals.json` then reconciles any remaining compaction undercount difference-based (`extra = signals_total - updates_total`, dropped when `<= 0`). `ClientId::Grok = 30` (**upstream numbers it 27**). Ports `SourceFingerprint::from_grok_path` (updates + signals; **locally extended to also fingerprint the `summary.json`/`events.jsonl` metadata siblings** — see the grok fingerprint row under Local patches). Scanner gains an `updates.jsonl` exact-name arm. Streaming `simple_lane!` + materialized path wired; every metadata sibling (`signals.json`/`summary.json`/`events.jsonl`) is included in both the live-tail change probe and the per-session mtime pruning decision. **The port is bit-faithful except for the compaction counter-epoch handling — a local divergence, see the grok row under Local patches below.** | `src/sessions/grok.rs`, `src/clients.rs`, `src/scanner.rs`, `src/message_cache.rs`, `src/lib.rs` | +| Grok Build (upstream `sessions/grok.rs`; M17 non-main `ed798642`) | Reads legacy `$GROK_HOME/sessions/*/*/updates.jsonl` plus only exact top-level `$GROK_HOME/logs/unified.jsonl` sources. `scanner.extraScanPaths.grok` and `TOKSCALE_EXTRA_DIRS` remain generic roots: direct Grok-home and exact `sessions` roots add their candidate, while recursively discovered legacy updates infer the owning Grok home for nested/ancestor roots; canonical file dedup collapses overlaps without recursively admitting archive or backup unified copies. Legacy cumulative `totalTokens` retains TokenBar's compaction counter-epoch handling and `signals.json` reconciliation; legacy fingerprints and mtime/pruning probes also include `summary.json` and `events.jsonl`. The unified log emits non-overlapping input/cache/output/reasoning buckets with session-scoped model precedence and a self-contained fingerprint; pidless session model events remain session-authoritative, PID-scoped model fallbacks expire on the explicit `AuthManager::new` process-start record before PID reuse, cached prompt counts above prompt totals clamp to the prompt total without dropping completion/reasoning usage, token-complete replay identity preserves distinct same-base inferences, and a missing loop index counts as loop one while explicit later loops do not increment message/turn counts. Materialized, shipping streaming, and count paths cache raw rows per source, collect the required Grok authority cohort, then call one pure selector before any report fold: a unified-covered session suppresses matching legacy rows after receiving unambiguous legacy model and workspace metadata, while legacy-only sessions remain. Both materialized and streaming lanes price selected rows after model carry-over. Shipping cache hits refresh timestamp-derived fields before filters, live trace and both materialized/streaming sessionization folds carry parser message counts without dropping later-loop tokens or timestamps, modified-after pruning preserves the authority/cohort relation, and removal of the unified file restores cached legacy rows. `ClientId::Grok = 30` (**upstream numbers it 27**); schema remains 31 because this is a new source rather than changed legacy parser output. | `src/sessions/grok.rs`, `src/clients.rs`, `src/scanner.rs`, `src/message_cache.rs`, `src/lib.rs`, `../crates/tb_core_ffi/src/usage_tail.rs` | | `#0ce3d73f` (#718, jcode) | New client: jcode (`~/.jcode/sessions/session_*.json`, JSONL with authoritative per-message `token_usage`). `ClientId::Jcode = 27` (**upstream numbers it 28**). Ports `SourceFingerprint::from_jcode_path` (snapshot + sibling `.journal.jsonl`, since jcode appends turns to the journal between snapshot rewrites) and **generalizes the `simple_lane!` macro to accept a custom fingerprint fn** so the streaming + materialized cached lanes invalidate on a journal-only write. Scanner gains a `session_*.json` glob (distinct from `session-*.json`). | `src/sessions/jcode.rs`, `src/message_cache.rs`, `src/clients.rs`, `src/scanner.rs`, `src/lib.rs` | | `#a75533e6` (micode / MiMo Code) | New client: MiMo Code (Xiaomi), WAL-mode SQLite at `$XDG_DATA_HOME/micode/*.db`, authoritative per-message tokens/cost/agent. `ClientId::MiMoCode = 28` (**upstream numbers it 30**). Reached via the generic `*.db` glob with the same WAL treatment as Antigravity CLI/Hermes/Zed: streaming `simple_lane!` fingerprints via `from_sqlite_path`, exempt from mtime pruning, `-wal`-probed in `latest_source_mtime_ms`. Authoritative cost survives because MiMo models are unpriced (`apply_pricing` only overwrites a non-zero calculated cost) — mirrors upstream. | `src/sessions/micode.rs`, `src/clients.rs`, `src/lib.rs` | | `#1492b962` + `#690` (gjc / gajae-code) | New client: gjc, JSONL at `$GJC_CODING_AGENT_DIR/sessions/*.jsonl` (default `~/.gjc/agent`) with an embedded `usage.cost.total` (USD) per assistant message. `ClientId::Gjc = 29` (**upstream numbers it 26**). Message-level `should_keep_deduped_message` dedup collapses depth-1/depth-2 replays. The original numeric `cost <= 0.0` lane guard has since been superseded by the typed `CostSource` policy in the M8 `#791` row below. | `src/sessions/gjc.rs`, `src/clients.rs`, `src/lib.rs` | @@ -58,11 +262,123 @@ next sync should treat these as already-present and not re-apply them. | `#837` (`4cbc2f6b`) + `#856` (`ae36db5c`, Claude parent-session dependency only) — **M10-D** | Claude workflow transcripts under `subagents/workflows//agent-*.jsonl` are discovered and counted, while `journal.jsonl` orchestration files are excluded. Deep workflow sidechains recover Tier-2 parent-session agent attribution. **TokenBar adaptation:** `parent_session_paths_for_cache` retains nested-first and flat-fallback candidates, bounded by an 8 MiB probe, so missing or changed parent transcripts invalidate the monolithic source cache; `CACHE_SCHEMA_VERSION` bumps **26→27** because schema-26 entries may replay usage-shaped workflow journals or stale/generic deep Tier-2 attribution. #856's shard cache, `CacheIdentity`, `parser_version` architecture, and unrelated client changes are deliberately excluded. Cowork discovery, cc-mirror #659 variant/provider/dedup behavior, and synthetic Claude-row drop #30 remain intact. | `src/sessions/claudecode.rs`, `src/scanner.rs`, `src/message_cache.rs` | | `#834` (`23cf62e0`, issue #829) — **M10-E** | Production baseline is upstream #834's final Copilot resolver: collect the full `parentSpanId` hierarchy and choose the first root `invoke_agent` span with an agent id, so nested sub-agent invokes exported first cannot hijack the trace fallback. The compact #821 reversed child-before-parent regression and its semantic are retained, without importing #821's intermediate `agent_from_invoke_agent` field. **TokenBar hardening over upstream #834, reported in issue [#879](https://github.com/junhoyeo/tokscale/issues/879) and merged via PR [#880](https://github.com/junhoyeo/tokscale/pull/880) (upstream commit `20d9096a68a40d4a4e83581b0e0dd308aadc5ab7`; GitHub merge commit `b7277d49a14ae905c17195be214d632e365b3ca6`):** first, collect every `traceId` + `spanId` + `parentSpanId` edge before the `attributes` gate, so attribute-less intermediate task/tool spans remain in the parent hierarchy; attributes are only used to identify `invoke_agent` and read agent ids. Second, scope `parent_of` and invoke-agent membership by `(traceId, spanId)` because OTel span ids are only unique within a trace, preventing reused ids in another trace from cross-contaminating the hierarchy. The local monolithic cache schema remains **27→28**; schema 28 is merged on TokenBar main, and both hardenings landed together in M10-E, so neither required an additional bump beyond 28. Preserves #723 cache-token aliases and #724/#751 per-record attribution. No scanner, fingerprint, FFI, Swift, or unrelated cache architecture changes. | `src/sessions/copilot.rs`, `src/message_cache.rs` | | `#833` (`85669602`) — **M10-F regression lock** | Codex scans both `~/.codex/sessions/` and `~/.codex/archived_sessions/`, and a transcript briefly present under both roots is counted once by content-derived session identity. The production scanner and dedup behavior were already present on the M10-F baseline, so this is deliberately a baseline-pass regression and parser-documentation lock rather than a production fix. **TokenBar adaptation:** one hermetic fixture proves live-only, archive-only, and duplicated shared sessions total exactly 3 messages / 150 input / 15 output across the materialized API, shipping streaming model report, and `parse_local_clients` count path. No parser output changed and cache schema remains 28. | `src/sessions/codex.rs`, `src/lib.rs` (tests) | -| `#848` (`c1aef5e9`) — **M11** | Hermes root homes now discover immediate `profiles//state.db` databases even when the default `state.db` is absent, while a profile-scoped `HERMES_HOME=/profiles/` remains isolated to its own database and never widens into siblings or the default profile. **TokenBar adaptation:** derive the Hermes home from the existing env-aware resolved `state.db`, preserve profile isolation through symlinked aliases by checking the canonical home as well as its lexical path, retain configured extra paths and deterministic default-first dedup, and reuse the plural materialized/streaming/count, WAL-probe, and fail-open pruning seams. Windows `%LOCALAPPDATA%` / `AppData/Local` discovery is intentionally omitted for the macOS-only product. No parser output changed and cache schema remains 28. | `src/scanner.rs`, `src/sessions/hermes.rs`, `src/lib.rs` (tests) | +| `#848` (`c1aef5e9`) — **M11 + M23-H** | Hermes root homes discover immediate `profiles//state.db` databases even when the default DB is absent, while profile-scoped `HERMES_HOME=/profiles/` remains isolated from siblings/defaults, including through symlink aliases. M23-H completes the selected Windows residual: when `HERMES_HOME` is unset, discovery checks native `%LOCALAPPDATA%/hermes` and supplied-home `AppData/Local/hermes`; an explicit root or profile remains authoritative and never widens. The existing plural materialized/streaming/count, physical-path dedup, WAL probe, and fail-open pruning seams are reused unchanged. No parser output changed; M11 kept its then-active schema 28 and M23-H keeps the current schema 32. | `src/scanner.rs`, `src/sessions/hermes.rs`, `src/lib.rs` (tests) | +| `#800` (`f6f7eced`) + follow-up `0b454e60` — **M23-D** | Adds the fixed `~/.copilot/data.db` token source, reads only token-bearing `sessions` rows, normalizes input-inclusive cache reads through the existing OTEL helper, parses SQLite fractional timestamps, and enriches model/workspace from sorted `session-state/*/events.jsonl` dependencies. **TokenBar adaptation:** `sessions.agent` is preserved as raw attribution; AIU-only rows, provider-reported cost, custom roots, and VS Code stores remain excluded. One raw `UnifiedMessage` selector gives OTEL whole-session authority before pricing, client/date filters, sessionization, and every report fold across materialized, shipping streaming, and count consumers. The Desktop fingerprint, latest-mtime probe, topology token, and fail-open pruning share the DB/WAL/event dependency set; a fresh Desktop source retains the complete OTEL suppressor cohort so incremental authority equals a full scan. Raw sources cache independently and the new source identity keeps schema 32. Hermetic fixtures cover agent attribution, AIU exclusion, date-window authority, cold/warm event and WAL invalidation, dependency-probe fail-open behavior, modified-after suppression, three-consumer projection, and graph/model/monthly/hourly/Agents parity. Parser production is 285 lines versus 261 lines in the final selected upstream `0b454e60` parser (+9.2%), within the M23 fidelity stop. | `src/sessions/copilot_desktop.rs`, `src/sessions/copilot.rs`, `src/sessions/mod.rs`, `src/scanner.rs`, `src/message_cache.rs`, `src/lib.rs` | | `#741` (`b49cec19`) — **M12 streaming-freshness completion** | RooCode, KiloCode, and Cline already used the specialized `from_roo_path` fingerprint in both materialized and streaming cache lanes, but TokenBar's live-tail change token and `modified_after` pruning still observed only `ui_messages.json`. A rewrite of the parser-read `api_conversation_history.json` could therefore leave model, agent, and pricing stale or prune an active task before fingerprint comparison. **TokenBar adaptation:** one parser-owned related-path helper now drives parsing, fingerprinting, latest-mtime probing, and sibling-aware fail-open pruning for all three wrappers. Hermetic regressions cover an absent-then-created sibling, history-only warm-cache refresh with materialized/streaming parity, all three change-token lanes, fresh-history pruning retention, and related-file stat failure. The fingerprint layout and parser output are unchanged, so cache schema remains 28. | `src/sessions/roocode.rs`, `src/message_cache.rs`, `src/lib.rs` | | `#856` (`ae36db5c`) — **M13 selective parser-dependency completion** | Existing Droid settings snapshots can read a same-session fallback JSONL for the model, legacy Kimi wire logs read the shared `~/.kimi/config.json`, and existing Kiro CLI session headers read a same-stem JSONL for estimated tokens, prompt time, and duration. Those dependencies were parser-only: both cache lanes, the live-tail token, and pruning still observed just the scanned primary. **TokenBar adaptation:** parser-owned path helpers now drive specialized monolithic fingerprints, materialized and shipping streaming lanes, latest-mtime probing, and sibling-aware fail-open pruning. Hermetic regressions cover absent-then-created and rewritten dependencies, independently warmed materialized/streaming caches, all three change-token lanes, fresh-dependency retention, missing-dependency pruning, and stat failure. The upstream shard cache, `CacheIdentity`, `parser_version`, Kimi Code, Kiro IDE/globalStorage/`.chat`/Windows sources, and unrelated client changes remain excluded. Related-file fingerprint storage already exists, so cache schema remains 28. | `src/sessions/droid.rs`, `src/sessions/kimi.rs`, `src/sessions/kiro.rs`, `src/message_cache.rs`, `src/lib.rs` | | `#878` (`3587f745`) — **M14** | Codex token snapshots now attribute `duration_ms` to non-overlapping intervals from the previous accepted snapshot rather than repeatedly measuring from the turn start. Invalid, equal, backward, replayed, duplicate, regressive, and zero-token rows do not advance the cursor; a new `turn_context` resets it. **TokenBar adaptation:** the cursor is serialized in the existing incremental state, the local monolithic cache schema bumps **28→29** so unchanged sources cannot replay overlapping durations or resume without the cursor, and hermetic regressions cover parser edge cases, incremental append parity, same-fingerprint schema-28 rebuild, materialized warm-cache parity, and shipping streaming cold/warm performance (`7,000 ms / 170 timed tokens / 3 samples`). Upstream's shard-cache `parser_version` bump is deliberately excluded; the existing Rust → FFI → Swift `ModelPerformance` shape is unchanged. | `src/sessions/codex.rs`, `src/message_cache.rs`, `src/lib.rs` (tests), `tests/fixtures/codex_duration_timing.jsonl` | +## M15-A Kiro IDE globalStorage completion + +M15-A is the macOS-only Kiro IDE globalStorage cohort, completed without changing `ClientId::Kiro = 22`, public client id `kiro`, CLI JSON/JSONL or SQLite behavior, attribution, dedup, model fallback, or `CACHE_SCHEMA_VERSION = 29`. + +The upstream semantic chain is preserved from the reviewed public tokscale commits: [#715](https://github.com/junhoyeo/tokscale/pull/715) / [`d50da475`](https://github.com/junhoyeo/tokscale/commit/d50da475) (initial globalStorage support), [#752](https://github.com/junhoyeo/tokscale/pull/752) / [`24e3771c`](https://github.com/junhoyeo/tokscale/commit/24e3771c0ca805a5e9983ecf562c483adcf8196a) (alias-safe snapshot traversal and model-key alignment), [#796](https://github.com/junhoyeo/tokscale/pull/796) / [`e5cfbae2`](https://github.com/junhoyeo/tokscale/commit/e5cfbae2483fd5c753f00cf2b130fb9c5343253e) (recursive model lookup that skips Kiro placeholders), [#799](https://github.com/junhoyeo/tokscale/pull/799) / [`b64e4f14`](https://github.com/junhoyeo/tokscale/commit/b64e4f14d667f03038ee9e7982477e33a575cb8d) (successful execution parsing), [`72bf6667`](https://github.com/junhoyeo/tokscale/commit/72bf6667) (execution workspace/model/duration/timestamp attribution), [`46e01977`](https://github.com/junhoyeo/tokscale/commit/46e01977) (execution precedence over session snapshots), [`31bfd167`](https://github.com/junhoyeo/tokscale/commit/31bfd167) (`.chat` executionId linkage and real chat roles), and [#814](https://github.com/junhoyeo/tokscale/pull/814) / [`09344531`](https://github.com/junhoyeo/tokscale/commit/09344531c6bb61a89e19a171d7139ca74bd2f265) (workspace-session `promptLogs` and `input.data.messages`). The earlier provider/time hardening from #760 remains in the existing Kiro parser row above. + +TokenBar's local M15-A adaptation adds only the macOS roots `~/Library/Application Support/Kiro/User/globalStorage/kiro.kiroagent` and `~/Library/Application Support/kiro/User/globalStorage/kiro.kiroagent`. The special `kiro-globalstorage` scanner follows the known storage layouts instead of accepting every matching file recursively: legacy snapshots are direct `/*.chat` files, execution records are `.json` or extensionless files one store directory deeper, and workspace-session JSON lives under `workspace-sessions//`. Deeper mirrored-project trees, root-level JSON/extensionless artifacts, `index.sqlite`, and unrelated extensions are excluded; generic role/content traversal is additionally limited to `.chat`, while JSON and extensionless sources must match an execution shape and workspace-session parsing is gated to the exact `workspace-sessions//*.json` subtree. Snapshot role/text traversal is recursive with alias-value deduplication and distinct-subtree counting. Explicit nested `tool` or unknown role/type markers clear an inherited user/assistant role, while neutral content types such as `text` retain it; this local hardening prevents injected tool payloads from inflating estimated usage. Model lookup is recursive and skips the Kiro placeholders `agent`, `auto`, and `qdev`, falling back to `auto`; this is model selection, not suppression of usage messages. Successful execution records support both input shapes, both output shapes, timestamp variants, duration, and execution precedence over snapshots; numeric-string epochs use the same seconds-versus-milliseconds discrimination as JSON numbers. Status-bearing non-`.chat` records with `context.messages` remain valid when `actions` is absent, while legacy `.chat` snapshots are not reclassified by `status` alone. Workspace-session paths attribute the nested workspace component rather than the literal `workspace-sessions` bucket. Source paths stay attached through the batch merge, so only execution messages parsed from globalStorage or `.chat` sources can seed suppression and only IDE-cohort messages can be suppression targets; exact dedup keys are likewise scoped to the IDE or CLI cohort, and successful execution identities are additionally scoped by workspace, so arbitrary CLI session ids or repeated execution IDs in another workspace cannot suppress, impersonate, collide with, or drop globalStorage usage. Workspace-session records retain their global session-id suppression semantics. + +Because precedence crosses files, TokenBar caches raw parser output per source and runs batch suppression only after every Kiro source has been collected. The materialized, streaming, count, and model/monthly/hourly/Agents report paths all consume the same post-suppression dedup gate, so mirrored `.chat` and `.json` snapshots with the same identity remain exact-once on every lane. The cache remains raw so an execution rewrite, failure, or removal can restore a cached snapshot; mtime probing uses each primary globalStorage file, while `modified_after` retains the complete IDE cohort whenever any IDE source passes the threshold so an older authoritative execution cannot be pruned before suppressing a newer snapshot. Kiro CLI sources remain independently prunable and stat failures remain fail-open. M15-A itself added no Linux or Windows globalStorage roots and did not include the structured `sess_*` / `session.json` / `messages.jsonl` cohort; M15-B below adds that cohort as a separate source with distinct identities. Upstream shard-cache, `parser_version`, and CLI-TUI-only changes are also excluded. + +## M20 OpenCode v2 SQLite completion + +M20 selectively ports upstream [#920](https://github.com/junhoyeo/tokscale/pull/920) / [`366ce643`](https://github.com/junhoyeo/tokscale/commit/366ce64395594abf111e0409581d91016561b25a). OpenCode v2 assistant rows are read from `session_message` only when `type = 'assistant'`; model identity falls back from v1 top-level `modelID` / `providerID` to v2 `model.id` / `model.providerID`, with top-level values retaining precedence. The v1 `message` table and v2 table feed one accumulator: fork copies with the same non-cost payload fingerprint and embedded id collapse, rows with different embedded ids remain distinct even when every fingerprint field collides, and row ids remain the fallback dedup key. Cost is deliberately excluded from that payload fingerprint; when cost-only copies merge, the first provider-reported cost upgrades an estimate without changing logical identity. When an overlapping v1 row has no embedded id, the merged row retains its `message.id` fallback as an alternate source key; if neither copy has an embedded id, the v1 key remains primary and the v2 row key becomes the alias. Provider authority expands across connected primary/alias groups before any streaming emission, while full payload identity still prevents unrelated embedded ids with coincidentally equal tokens from collapsing. TokenBar's shared post-parser identity includes timestamp, duration, model/provider, tokens, and agent, so incompatible SQLite rows that reuse one embedded id remain distinct in materialized, shipping streaming, and count lanes. Legacy JSON authority matches an exact deferred payload through any primary/alias source key before applying source-level suppression and replaces only that one SQLite identity, while cost-only estimated/provider-reported copies retain source precedence. Legacy JSON files still require an explicit `role = "assistant"`; making the SQLite role optional does not broaden file parsing. Workspace attribution, negative-token clamping, provider canonicalization, and provider-reported `CostSource` behavior retain TokenBar's existing hardening. + +The scanner already accepts `opencode-next.db` through the existing `opencode-.db` rule, and all materialized, shipping streaming, count, model, monthly, hourly, and Agents paths already call the shared SQLite parser, so no scanner or registry change is required. Cache schema advances **29→30** for a real non-empty stale case: a hybrid database can have a schema-29 entry containing only its v1 rows, with the same SQLite/WAL-aware fingerprint that the new parser would otherwise reuse without reading v2 rows. A hermetic fixture writes exactly that schema-29 v1-only entry, proves the database fingerprint is unchanged, verifies rejection and rebuild to v1+v2, then confirms warm-cache, streaming, count, and report parity. Pure v2, v1/v2 overlap, same-id forks, distinct-id collisions, role filtering, workspace, and clamping remain covered in parser-local fixtures. The shard-cache architecture and per-client parser versions remain deferred to M26. + +## M15-B Kiro structured-session completion + +M15-B selectively ports upstream [#836](https://github.com/junhoyeo/tokscale/pull/836) / [`405ded4a`](https://github.com/junhoyeo/tokscale/commit/405ded4a529013ad6bc668ba519bbe9eb68e8ad4), [#847](https://github.com/junhoyeo/tokscale/pull/847) / [`315549b4`](https://github.com/junhoyeo/tokscale/commit/315549b4bbaa53638672ff5977c0d6cb77e43793), and only the Kiro start-anchor hunk of mixed commit [`b64d861e`](https://github.com/junhoyeo/tokscale/commit/b64d861e04824984fef2638df5db8df1d672cb2b). The scanner discovers only `~/.kiro/sessions//sess_*/session.json` anchors; sibling `messages.jsonl` files are parser dependencies, not independent sessions. Structured `payload.type` records emit one message per completed turn: `contextUsage.usagePercentage` estimates input against a 200,000-token window, assistant text and tool-call arguments estimate output, `usage_summary.elapsedTime` supplies duration, and a missing prompt timestamp back-calculates the start anchor from `turn_end - elapsedTime`. Older flat role/content JSONL remains an aggregated fallback. Session metadata supplies model and workspace when present; only this new IDE source falls back to model `auto`, while existing CLI and SQLite sources retain their schema-30 `unknown` fallback. + +The existing `kiro_related_messages_path()` seam now identifies same-stem CLI JSONL or IDE sibling `messages.jsonl`, so one source of truth drives specialized fingerprints, materialized and shipping streaming cache loads, live latest-mtime detection, and sibling-aware fail-open pruning. Hermetic coverage proves absent-then-created and sibling-only rewritten files invalidate the specialized fingerprint without changing the primary fingerprint, warm cache rebuilds from one to two turns, `modified_after` retains a source whose sibling is newer, and materialized, streaming, count, model, monthly, hourly, and Agents results agree. M15-A globalStorage and M15-B structured messages keep distinct dedup identities, so coincident text does not suppress either cohort. This is a new discovery source rather than a changed existing parser output, so `CACHE_SCHEMA_VERSION` remains **30**; shard cache and per-client parser versions remain deferred to M26. + +## M16 existing-parser correctness completion + +M16 selectively ports upstream [#896](https://github.com/junhoyeo/tokscale/pull/896) / [`6899ea03`](https://github.com/junhoyeo/tokscale/commit/6899ea03640b6de47bb3108c409be6fea7227bed), [#892](https://github.com/junhoyeo/tokscale/pull/892) / [`b59979c5`](https://github.com/junhoyeo/tokscale/commit/b59979c571d972418916df35df02f75f802d6b0b), [#890](https://github.com/junhoyeo/tokscale/pull/890) / [`9155018c`](https://github.com/junhoyeo/tokscale/commit/9155018c683fec78acf6c0ff9e77793ba02e34dc), [#891](https://github.com/junhoyeo/tokscale/pull/891) / [`18cd13cc`](https://github.com/junhoyeo/tokscale/commit/18cd13cc921f42a309cab5a57b0331e1aebbda4c), only the provider-hardening hunks of [#887](https://github.com/junhoyeo/tokscale/pull/887) / mixed [`34cfbb50`](https://github.com/junhoyeo/tokscale/commit/34cfbb50796c5e084c55b595385f6fd3b48e802d), and only the Jcode hunk of [#898](https://github.com/junhoyeo/tokscale/pull/898) / mixed [`b64d861e`](https://github.com/junhoyeo/tokscale/commit/b64d861e04824984fef2638df5db8df1d672cb2b). Every 9Router bridge, scanner, registry, filter, pricing, script, and service hunk remains excluded; `34cfbb50` therefore becomes `DEFER` rather than `ALREADY_VENDORED`. + +Codex leniently decodes numeric `task_started.started_at`, keeps legacy UUID-v4 ancestor replays behind the child boundary until a child-local task start is causally ordered, and start-anchors token snapshots at `turn_context` or a resumed human `user_message`. Claude bare files under `~/.claude/transcripts/` no longer synthesize char-estimated tool-result tokens when no project or cc-mirror identity exists, but explicit tool-result counts remain authoritative; normal project transcripts retain estimation. Claude request timestamps now remain at the activity start while duplicate chunks retain per-field token maxima and the longest observed duration. Copilot prefers OTEL `startTime`, back-calculates end-only records when duration exists, and collapses repeated span identities before every report fold with per-bucket maxima, earliest start, maximum duration, and missing-agent recovery. Jcode assistant completion timestamps with positive `tool_duration_ms` are back-calculated to the turn start while preserving the recorded duration. + +Provider identity is corrected before pricing and aggregation: delimited Kimi models infer `moonshotai`, MiMo infers `xiaomi`, GLM infers `zai`, and missing or empty GJC/Pi providers infer from the model before falling back to the client. Current Antigravity IDE placeholders and CLI response-model IDs resolve through single-hop machine aliases, including distinct Low/Medium tiers and the verified Gemini 3.5 Flash High mappings; hermetic synthetic-catalog tests prove representative aliases reach priced entries without changing forced-source or routed-pricing precedence. + +Because these are existing-source parser-output changes, TokenBar advances its monolithic `CACHE_SCHEMA_VERSION` **30→31** instead of importing upstream parser-version shards before M26. The integration fixture writes a real schema-30 Jcode entry with the unchanged specialized fingerprint and old completion timestamp, proves schema 31 rejects it, rebuilds the start-anchored message, and verifies cold/warm materialized, shipping streaming, count, model, monthly, hourly, and Agents parity. Parser-local fixtures cover each selected old-fail/new-pass behavior; no scanner, client registry, FFI, Swift, or C ABI change is required. + +## M17 Grok unified-log precedence completion + +M17 selectively ports non-main upstream [#849](https://github.com/junhoyeo/tokscale/pull/849) / [`ed798642`](https://github.com/junhoyeo/tokscale/commit/ed7986425efb8bad8df62e16269e30c15a82801a) without replacing TokenBar's hardened legacy Grok parser. The scanner adds only exact top-level `$GROK_HOME/logs/unified.jsonl` sources beside existing per-session `updates.jsonl`; generic `scanner.extraScanPaths` and `TOKSCALE_EXTRA_DIRS` roots may name a Grok home, its `sessions` directory, or an ancestor/nested root whose discovered updates identify the owning home, while canonical file dedup collapses overlaps and recursive archive or backup unified copies remain excluded. Unified `shell.turn.inference_done` records split prompt/cache and completion/reasoning totals into non-overlapping buckets, clamp cached prompt overlap to the prompt total rather than dropping usable output, count a missing `loop_index` as loop one while explicit later loops add no user message/turn, deduplicate exact replays by the complete token identity rather than the coarse session/timestamp/PID/loop tuple, preserve distinct same-base inferences, and prefer session-scoped model events—including events without `pid`—over PID-only catalog fallbacks. An explicit `AuthManager::new` process-start record expires both PID-only and PID-plus-session model authority for that PID, so an OS-reused PID fails closed to unknown or legacy carry-over until the new process emits model authority. + +One pure `prefer_unified_log_messages()` selector defines authority before aggregation: it first discovers every session represented by a `grok-unified:` row, carries an unambiguous non-placeholder legacy model plus unambiguous legacy workspace metadata onto matching unified rows that lack those values, then drops legacy rows only for those covered sessions. Conflicting legacy model or workspace values fail closed instead of choosing by input order. Legacy-only sessions remain unchanged. Authority is intentionally session-scoped because legacy rows expose cumulative turn deltas without a stable inference identity; attempting a same-session timestamp merge would risk counting the two representations together. No downstream subtraction or arithmetic attempts to remove already-folded usage. Materialized, shipping streaming, and count lanes all collect the complete raw Grok cohort and call the selector once. Materialized rows are repriced after selection so a unified row that borrows its model does not retain the zero cost calculated for `grok-unknown`; the streaming lane follows the same order before client filtering, dedup, and the report sink. + +Legacy `updates.jsonl` retains the existing composite fingerprint and sibling-aware latest-mtime/pruning contract for `signals.json`, `summary.json`, and `events.jsonl`; the self-contained unified log uses the normal primary-file fingerprint and mtime. Modified-after pruning keeps an older unified authority file whenever a legacy source is fresh, and a fresh unified file retains the legacy cohort needed for workspace attribution. FFI graph and live-tail caches use a topology-sensitive source token rather than the maximum mtime alone, so creating, rewriting, or deleting a non-max unified source invalidates stale unified authority even when a newer legacy file masks its mtime. Shipping streaming refreshes timestamp-derived fields after raw cache selection and before date filtering, so a timezone change cannot leave a cache-hit Grok row in the old local day. Live trace events retain each parsed row's `message_count`: later tool loops still contribute every token to rate and trace totals while adding zero messages. A hermetic lifecycle fixture deliberately makes legacy newer than unified, omits the unified in-log model event, and proves legacy-only cold/warm behavior, absent→created unified authority, legacy model recovery, exact token/message preservation, materialized/streaming/count/model/monthly/hourly/Agents parity, and unified removal restoring cached legacy rows and the original source token. Focused regressions additionally poison a cached derived date, preserve two distinct inferences that share session/timestamp/PID/loop while collapsing an exact replay, fail closed on conflicting legacy models, verify cold/warm materialized pricing after model carry-over against the streaming lane, discover exact unified logs for primary/settings/environment roots without duplicates or archived copies, and keep a reasoning-only zero-message loop's tokens without inflating live trace messages. Fresh verification first refuted the trace claim because `UsageEvent` omitted the reasoning bucket; after carrying reasoning and parser message counts through the live-tail event, the production FFI reproduction reports all 11 tokens and one message and a new verifier returned `CONFIRMED`. Later Codex passes found that unified rows without an in-log model event lost available legacy model metadata, that materialized rows retained pre-selection unknown-model costs, and that configured Grok roots omitted their unified logs; unique non-placeholder model carry-over, post-selector repricing, exact configured-root discovery, and conflict fail-closed behavior now cover those paths. The following current-head pass found three adjacent cases: a session model event could omit `pid`, generic configured roots could name a Grok home or an ancestor rather than exact `sessions`, and cached prompt overlap could exceed the prompt total without invalidating usable completion/reasoning usage. Session-only model authority, update-derived Grok-home discovery, and prompt-bounded cache clamping now preserve those rows, with focused old-fail/new-pass fixtures. The next pass found that an omitted `loop_index` silently became zero messages, that a restarted process could inherit a prior process's PID-only model catalog entry after OS PID reuse, and that both sessionization folds coerced explicit zero-message tool loops back to one. Missing loop indexes now default to the counted first loop, an explicit `AuthManager::new` boundary clears only PID-scoped authority while session-only model authority survives, and both `sessionize()` and `SessionizeAccumulator::feed()` preserve explicit zero counts while still retaining each row's tokens and timestamp. Parser-local and sessionization fixtures cover bucket decomposition, missing and explicit loop message counts, process-restart PID reuse, zero-message token retention, model precedence, replay dedup, session-level fallback, workspace carry-over, exact top-level discovery, authority-cohort pruning, and input-order-independent selection. Fresh verification returned `CONFIRMED`. This was a new source with distinct fingerprints, not a changed existing-source output, so M17 kept the then-active monolithic cache at **31** and the 111-row ledger at `67/20/0/10/13/1`; current main is schema 32 after later PR #77. + +## M18 Sakana, long-context, and routed-pricing completion + +M18 selectively ports audited upstream [`959cce84`](https://github.com/junhoyeo/tokscale/commit/959cce84) and [`6c804711`](https://github.com/junhoyeo/tokscale/commit/6c804711), plus the verified request-level long-context rule from non-main [`548dc124`](https://github.com/junhoyeo/tokscale/commit/548dc124) and routed prefix/suffix composition from non-main [`6ea27ca1`](https://github.com/junhoyeo/tokscale/commit/6ea27ca1). Sakana `fugu-ultra` uses regular per-million rates of `$5` input, `$30` output, and `$0.50` cache-read; requests with verified Sakana identity and `input + cache_read > 272,000` use `$10`, `$45`, and `$1`. Verified LiteLLM GPT-5.4/GPT-5.5 identities use the same whole-request threshold behavior with their own catalog tiers. Exactly 272,000 remains regular, output plus reasoning follows the selected request tier, cache-write does not select the tier, and bare `fugu` derives Sakana provider identity while remaining intentionally unpriced. The Sakana subscription billing-console scrape remains excluded. + +The routed lookup preserves forced LiteLLM/OpenRouter isolation before custom pricing. Exact full raw IDs, complete custom keys, and parenthesized custom or built-in literals receive first refusal before tier stripping or suffix validation; invalid parenthesized suffixes fail closed after that exact miss. Provider-scoped `accounts//{models,routers}/...` paths never widen into a terminal-only custom or fuzzy cross-provider match. Bounded full-path suffixes precede terminal fallback, terminal custom lookup tries explicit raw and alias-normalized candidates, built-in direct/static aliases retain precedence, and every fallback result passes one Claude never-degrade guard. Provider ranking and cache-rate backfill remain unchanged. + +An early prepared-patch verifier reproduced a provider-scoped boundary bypass: `CustomPricing::lookup_with_key()` normalized a full Fireworks path to its terminal model before the fail-closed guard. Internal routed lookup now uses exact-only custom keys, while the public custom API retains its existing synthetic normalization and generic non-scoped terminal routes still work through explicit candidates. The final integrated verifier also reproduced case-sensitive forced-source matching (`OpenRouter` fell through to automatic lookup); the shared lookup boundary now lowercases the selector once so LiteLLM, OpenRouter, and Custom isolation remains case-insensitive through direct, exact, and fallback paths. GitHub review then found three adjacent provenance cases: provider-prefixed verified GPT-5.4/GPT-5.5 LiteLLM rows missed whole-request tiering, an unknown router path could let a Cursor model-part override beat the terminal upstream exact row, and an explicit non-Sakana provider could inherit the Sakana override. Whole-request identity now recognizes the exact terminal GPT key, only truly unknown router prefixes defer model-part overrides until terminal upstream resolution, and Sakana pricing requires every explicit path segment and provider hint to agree with Sakana. Post-review fresh verification extended the router case to suffix normalization (`cx/gpt-5.3-codex-high`) and the provider case to both a nested model path (`sakana/openrouter/fugu-ultra`) and a composite provider hint (`sakana/openrouter`): full-path suffix candidates remain exact-only for unknown routers before terminal suffix resolution can choose the upstream exact row, while any non-Sakana routing or hint tag now makes the Fugu override fail closed. Explicit `unknown`, empty, or whitespace-only hints retain that explicitness for the Sakana gate instead of becoming an absent hint; the cached lookup path bypasses cache for those invalid hints so an unhinted Sakana result cannot be reused. Hermetic pricing fixtures cover the 272,000/272,001 boundary, cache-read threshold selection, output/reasoning tiering, cache-write exclusion, bare `fugu`, non-verified marginal behavior, exact/custom/alias/parenthesized/path/terminal routes, mixed-case forced sources, provider-scoped isolation, ranking/backfill, and every Claude never-degrade path. Pricing remains post-cache retrieval, so raw model identity and source-message cache identity are unchanged; monolithic schema remains **31**. M18 moves `959cce84` and `6c804711` to `ALREADY_VENDORED`, producing the duplicate-free exact classification `69/18/0/10/13/1`. + +## M21 Kimi Code, Junie, and OpenCodeReview completion + +M21 selectively ports Kimi Code from upstream [`839ce378`](https://github.com/junhoyeo/tokscale/commit/839ce378) and [`052f43de`](https://github.com/junhoyeo/tokscale/commit/052f43de), Junie from [`633ea946`](https://github.com/junhoyeo/tokscale/commit/633ea946) and [`77948d9d`](https://github.com/junhoyeo/tokscale/commit/77948d9d), OpenCodeReview from [`302d39c3`](https://github.com/junhoyeo/tokscale/commit/302d39c3), and only the Junie/OpenCodeReview start-anchor hunks of mixed [`b64d861e`](https://github.com/junhoyeo/tokscale/commit/b64d861e). Kimi Code reuses public client ID `kimi`, is selected by the `sessions///agents//wire.jsonl` topology, counts only explicit turn-scoped `usage.record` rows, and collapses exact replays without merging distinct turns that happen to use the same token counts. The default `~/.kimi-code` root and explicit `KIMI_CODE_HOME` root coexist with legacy `~/.kimi`; environment overrides are ignored when `use_env_roots` is false. Legacy Kimi keeps its config-dependent model lookup, while self-contained Kimi Code files exclude that unrelated sibling from fingerprint, mtime, pruning, and change-token semantics. + +Junie reads `events.jsonl`, preserves each finite non-negative provider-reported cost as authoritative, marks only the first valid usage in the next response event as a turn start while consuming prompt ownership even when that response has no `modelUsage` array or no countable usage, and start-anchors explicit response end timestamps by duration. OpenCodeReview reads session JSONL, carries normalized workspace metadata from `session_start`, and likewise start-anchors duration while retaining the recorded end timestamp in replay identity. Append-only `ClientId::Junie = 31` and `ClientId::OpenCodeReview = 32` raise `COUNT` to 33 without changing existing IDs; Swift uses `Kimi`, `Junie`, and `OpenCodeReview` display metadata and retains the existing initial-letter fallback instead of adding unverified logo assets. Materialized, shipping streaming, count, model, monthly, hourly, and Agents paths share parser selection, per-client dedup, raw source caching, and post-cache pricing. Hermetic fixtures prove Kimi legacy/Code coexistence, in-file and cross-agent replay collapse with distinct-turn retention, overflow-safe preservation of extreme nonzero token buckets, Junie authoritative cost through cold/warm rewrite, OpenCodeReview removal without stale cache replay, count/report parity, and unchanged totals across every lane. These were new source identities rather than changed existing-source serialization, so M21 kept the then-active monolithic cache at **31**; current main is schema 32 after later PR #77. M21 moves five audited rows to `ALREADY_VENDORED`; M22 is closed unmerged, and its five selected Zcode rows are `DEFER`, producing the duplicate-free exact classification `74/8/0/15/13/1`. Mixed `b64d861e` is `DEFER` because its Zcode hunk did not land. + +## M25 reloadable grouping model aliases + +M25 selectively ports upstream [#850](https://github.com/junhoyeo/tokscale/pull/850) / [`9a5aeb65`](https://github.com/junhoyeo/tokscale/commit/9a5aeb65): a config-driven `{alias → canonical}` map that folds model-name variants for **local report grouping only**. The fold is the terminal step of `normalize_model_for_grouping`; `canonical_model_id` is the alias-free syntactic path used by graph `ClientContribution` keys and any future submit/export/persist surface. Pricing continues to resolve the raw message `model_id` (and the static `pricing/aliases.rs` machine-id table remains a separate layer). TokenBar adaptation vs upstream's load-once `OnceLock`: the process-wide map is **reloadable** via `set_model_aliases` / `clear_model_aliases`, bumps `model_alias_generation`, and fires every `register_usage_data_invalidation_hook` so usage-data consumers refresh without a process restart. Excluded: tokscale-cli TUI/settings/tests, any M25 message-cache schema bump (it kept the then-active **31**; current main is 32 after later PR #77), Copilot Desktop/VS Code/Hermes discovery (M23), and Swift/FFI settings wiring (core API is ready; call from settings later). Hermetic tests prove grouping folds while pricing and canonical identity stay on the raw path, and that reload/clear bump generation and fire the invalidation seam. Ledger transition: `9a5aeb65: TAKE → ALREADY_VENDORED`, post-M25-only classification `75/7/0/15/13/1`. + +## M24 Warp fidelity stop + +M24 PR [#86](https://github.com/Nanako0129/TokenBar/pull/86) proposed a Rust-owned normalized Warp source from either an explicit process-memory bearer or one user-selected exact `usage.json`. Its current head `7dcb6985` included bounded fixed-endpoint GraphQL transport, opaque installation-key HMAC identities, mutually exclusive app/external modes, MAC-verified scoped app cache, generic source-message-cache bypass, C ABI/Swift settings wiring, and materialized/streaming/count/report parity. Hermetic and full runtime gates passed, but those green checks did not establish a coherent multi-process cache owner. + +Current-head Codex found four new failure classes: an inactive 401 could purge an unrelated account cache, an old refresh failure could stale or clear a newer source, a failed Disconnect could forget the only purge-retry state, and a stale bearer connect could replace a newer source. The first local fix round added generation guards, matching-scope inactive purge, retryable Disconnect cleanup, and remote I/O outside the refresh/storage transaction. That round was not pushed because fresh security verification reproduced a further production interleaving: Warp source and revoked-scope state are process-local, but `warp-usage-v1.json` is one cross-process singleton. While process P1 waits on remote I/O for scope B, process P2 can write scope A; P1 can then reacquire the provider lock and delete A during same-scope 401 handling or Disconnect. Keeping the lock across remote I/O would avoid that interleaving but violate the approved network boundary. + +This second consecutive review round triggers the fidelity stop rather than another local patch. PR #86 is closed unmerged, its rejected fix remains unpushed, and the remote branch is retained as evidence. Main has no M24 runtime, schema stays 32, and `63a44d7c` moves from `TAKE` to `DEFER`, producing `78/2/0/17/13/1`. The approved replacement plan then removes M24 from the cache dependency graph without reviving Warp. + +## M26-A identity-aware shard cache + +M26-A selectively ports upstream [#856](https://github.com/junhoyeo/tokscale/pull/856) / [`ae36db5c`](https://github.com/junhoyeo/tokscale/commit/ae36db5cd1f23e24b93768f9c2f169ef825072ef) as a cache-engine boundary replacement. Active storage is format 1 under `source-message-cache-v2//shard-XX.bin`: 256 shards, a 256 MiB per-shard decode/write limit, identity keys containing parser namespace plus path, and payload envelopes that validate format/namespace/parser version before decoding entries. Initial parser versions are Codex 4, Jcode 4, Copilot 3, synthetic 1, and every other current Native client 1. Generic and SQLite sources use at most five 4 KiB sample windows; Codex retains full prefix hashing and incremental state. The selected WalkDir fast path reuses each entry's known file type instead of repeating `stat()` for regular files while preserving followed metadata for file symlinks. Dirty/deleted keys are bucketed once, affected shards are reread under the existing cross-process lock, stale deletes compare their old fingerprint before removal, and durable writes reuse `fs_atomic::replace_file`. + +The former schema-32 `source-message-cache.bin` is deliberately inert: M26-A never reads, migrates, rewrites, or deletes it, and the byte/metadata sentinel verifies a cold shard build leaves it untouched. `HASH_MEMO` and monolith-only `STORE_MEMO` are retired rather than layered over bounded sampling or shard persistence. Cache entries remain raw per-source parser output; cohort authority, dedup, pricing, filters, and reports still run after cache retrieval. Warp stays `parse_local: false`, and no M24 producer, credential, account scope, network, settings, or ABI code is introduced. + +| Parser identity | Preserved dependencies and authority | Format-1 evidence | Format-2 result | +|---|---|---|---| +| Claude | `.meta.json`, cc-mirror variant, nested/flat parent transcripts | Sidecar/parent invalidation, materialized and streaming cache paths | Exact paths/existence persist; cached parent paths reuse after a primary match | +| Copilot | OTEL files plus Desktop DB, optional WAL, dynamic `session-state/*/events.jsonl`; OTEL whole-session authority stays post-cache | WAL/event invalidation and unreadable-dependency fail-open | Dynamic related-set discovery, exact path/cardinality checks, and fail-open behavior remain intact | +| Jcode | Snapshot plus `.journal.jsonl` | Journal-only mutation invalidates warm cache | Absent journal path persists and creation/removal invalidates | +| Droid | Settings snapshot plus fallback `session.jsonl` | Fallback-only mutation invalidates warm cache | Absent fallback path persists and creation/removal invalidates | +| Kimi | Legacy `config.json`; Kimi Code remains self-contained | Legacy config mutations invalidate; nearby Kimi Code config does not | Exact legacy path/existence persists without adding a Code dependency | +| Kiro | CLI same-stem JSONL and IDE `messages.jsonl`; suppression remains post-cache | Both sibling forms invalidate independently | Exact candidate path/existence persists | +| Roo / Kilo / Cline | `api_conversation_history.json` | History-only model/agent mutation invalidates | Exact history path/existence persists | +| Grok | Legacy `signals.json`, `summary.json`, `events.jsonl`; unified log remains self-contained and precedence stays post-cache | All legacy siblings invalidate independently | Exact sibling paths/existence persist | +| Codex | `consumed_offset`, newline boundary, prefix hash, parse state | Full-prefix reuse, append resume, mismatch cold parse | No generic related-file expansion | +| Generic / SQLite | Primary samples; SQLite `-wal` | Same-size restored-mtime mutation in a sampled window invalidates without full hashing | Related path/existence metadata persists where applicable | + +The ledger transition is `ae36db5c: TAKE → ALREADY_VENDORED`, producing `79/1/0/17/13/1` at the M26-A checkpoint. M26-A merged in PR #90 at `95c819c7`; its actual child base for M26-B is `43fa8ad6`. + +## M26-B format-2 related-file metadata + +M26-B selectively ports the generic cache hunks from upstream `cd07bf78` without importing its Pi or Devin runtime work. `CACHE_FORMAT_VERSION` advances from 1 to 2 because each serialized `RelatedFileFingerprint` now retains the exact `CachedPath` and an `exists` bit. Missing dependencies are persisted as zero-metadata records, so later creation invalidates a warm entry; any non-`NotFound` metadata, read, or bounded-sampling failure prevents a warm hit and forces a cold parse instead of treating an unreadable dependency as absent. Warm validation separates the primary sample check from exact related-file count, suffix, path, existence transition, metadata, and bounded sample checks. Claude reuses cached parent-session candidates only when its primary fingerprint matches, avoiding repeated sidechain probing while still invalidating on parent creation or change. Both the materialized compatibility lane and the shipping streaming reports use this status-aware checker. + +The active cache is format 2. Existing format-1 shards are locally stale and rebuild cold; the legacy schema-32 `source-message-cache.bin` monolith remains unread, unmodified, and undeleted. Copilot Desktop dynamically discovers `session-state/*/events.jsonl`; a malformed non-directory entry returns a non-`NotFound` dependency error so existing callers cold-parse and fail open instead of reusing an incomplete fingerprint. Native Grok/Kiro/Droid/Kimi/Jcode/Roo/SQLite helpers, M26-A shard merge/durability/pruning semantics, Codex prefix state, bounded generic/SQLite samples, and all post-cache authority, deduplication, pricing, filtering, reports, FFI, and Swift behavior remain unchanged. The Devin residual of `cd07bf78` remains excluded, so the audited row moves `TAKE → DEFER` and the exact classification is `79/0/0/18/13/1`, total 111. + +## M19-A Windows atomic-replacement completion + +M19-A selectively ports only the `fs_atomic.rs` hunk from upstream [#906](https://github.com/junhoyeo/tokscale/pull/906) / [`a87f0ab6`](https://github.com/junhoyeo/tokscale/commit/a87f0ab6). On Windows, `MoveFileExW(MOVEFILE_REPLACE_EXISTING | MOVEFILE_WRITE_THROUGH)` retries only `ERROR_ACCESS_DENIED` (5) and `ERROR_SHARING_VIOLATION` (32), stops after five total attempts, and sleeps for 10, 20, 30, then 40 milliseconds between failures. Other errors return after one attempt. The non-Windows path remains the existing single `std::fs::rename` call. + +TokenBar factors the retry loop into an injected replacement/sleep helper compiled on Windows and in tests, so deterministic non-Windows fixtures cover both retryable codes, eventual success, persistent failure at attempt five, immediate non-transient failure, and the exact attempt/backoff sequence without pretending to exercise `MoveFileExW` on macOS. The upstream TUI signal and background-screen changes are excluded. No serialized parser output, cache identity, layout, FFI, Swift, or C ABI changes, so M19-A kept the then-active monolithic cache at 31; current main is schema 32 after later PR #77. Any future approved shard-cache plan can reuse this bounded replacement behavior rather than inventing another Windows retry path. + +## Recovered Windows downstream commits + +The Windows port started from TokenBar [`2ed256ee`](https://github.com/Nanako0129/TokenBar/commit/2ed256eea7f6761e85198e3bc584e08a8d3d8ac1), then accumulated and validated the vendor-only sequence recorded in the [issue #45 handoff](https://github.com/Nanako0129/TokenBar/issues/45#issuecomment-5002865759). These commits are now recovered into this repository so it remains the canonical source for the next Windows sync; the detailed Windows build, hostile-environment, profile-manifest, and FFI evidence stays in that handoff rather than being duplicated here. + +| Commits | What | Files | Cache / upstream status | +|---|---|---|---| +| `e5200634` → `db2a96a3` → `15f418ee` → `e807f333` → `0979cdb0` → `26b892a6` → `6ac77c03` → `e91fb2c6` → `fbecb99c` → `3c8bfc52` → `aec5bd88` | Makes core tests hermetic with panic-safe environment/current-directory guards, serial coordination, isolated cache/XDG roots, platform-safe JSON/path fixtures, and explicit-home parser/scanner fixtures while retaining dedicated positive environment cases. Production fixes release the temporary cache writer before Windows atomic replacement, reopen the final cache read/write for the durability sync, derive Windows explicit-home config and Zed paths from supplied `/AppData/Roaming/tokscale` and `/AppData/Local/Zed/threads/threads.db` while limiting process known-folder fallbacks to env-aware scans, and compare extra-path warnings against the supplied scan home. | `src/clients.rs`, `src/lib.rs`, `src/message_cache.rs`, `src/pricing/cache.rs`, `src/scanner.rs`, `src/sessions/claudecode.rs`, `src/sessions/opencode.rs` | No serialized parser output, cache layout, FFI, or public API change; `CACHE_SCHEMA_VERSION` remained 29 at that Windows-recovery checkpoint; M20 later advances the Native source to schema 30. The upstream-applicable portions are merged: cache PR [#914](https://github.com/junhoyeo/tokscale/pull/914) (`163ec570`) maps the handle ordering to upstream's shard writer, explicit-home PR [#916](https://github.com/junhoyeo/tokscale/pull/916) (`a2f7cef5`) matches the supplied-home and warning semantics including positive `AppData` roots, and core hermetic PR [#917](https://github.com/junhoyeo/tokscale/pull/917) (`a0929482`) carries the same fixture-isolation failure class on upstream's current test layout. The remaining structural differences are TokenBar-specific cache and test adaptations, not missing upstream behavior. M26-A later replaced the monolithic writer with the active shard writer, so M19-B1-R2 reapplies the same upstream-approved Windows handle ordering to that new persistence seam without a schema change. | + +M19-B1-R2 also closes the other Windows portability gaps exposed by the exact downstream sync. The format-2 shard writer now drops its temporary `BufWriter` before `fs_atomic::replace_file` and reopens the installed shard read/write for the durability sync; Copilot Desktop rejects non-directory `session-state` entries before probing `events.jsonl`; the cc-mirror JSON fixture uses `serde_json`; and the Kiro scanner fixture compares native `PathBuf` values. Hosted Windows test fixtures pin `TOKSCALE_CONFIG_DIR` to each sandbox because Windows Known Folders ignore HOME/XDG, open mtime targets with writable handles, and drop manually serialized stale-shard writers before readers or replacement. In the same Native canonical repair, `tb_core_ffi` restores cfg-gated PowerShell/CIM Antigravity process and PID-bound listener discovery. These changes affect `src/lib.rs`, `src/message_cache.rs`, `src/sessions/copilot_desktop.rs`, `src/scanner.rs`, and `crates/tb_core_ffi/src/agent_antigravity.rs`; they do not change parser output, parser versions, cache identity, `CACHE_FORMAT_VERSION = 2`, FFI/C ABI, or the inert schema-32 monolith. Windows consumes them through the next byte-identical Native sync rather than a downstream shared-tree patch. + +The PR #102 current-head review adds two shared-tree corrections before that sync. Codex OAuth credential write-back now commits its staged `auth.json` through the same `fs_atomic::replace_file` primitive, so Windows receives the bounded 5/32 retry policy while the existing unique sibling temp, create-new semantics, file sync, Unix directory sync, live-token revalidation, JSON sibling preservation, receipt/lineage ordering, and rollback boundary stay unchanged; Windows-only fixtures cover transient success and persistent-lock failure without losing the original file. Copilot duplicate OTEL spans now merge the earliest start with the latest known endpoint instead of pairing the earliest start with the largest independent duration. Because `duration_ms` is serialized parser output, the Copilot parser version advances **3→4** and a same-fingerprint stale-v3 shard regression proves the old duration is rejected and rebuilt. Hosted Windows shipping-cache fixtures now build related-file fingerprints, seed, and query with the scanner-returned path spelling, preserving the production cache key and related-path identity when explicit-home resolution uses `/` but `TempDir::join` constructs the same physical path with `\`. `CACHE_FORMAT_VERSION` remains 2, every other parser namespace stays warm, FFI/C ABI is unchanged, and the schema-32 monolith remains inert. + ## Upstream fixes reported but not yet vendored Bugs we found *while* vendoring, reported upstream, and that were fixed there — @@ -77,18 +393,37 @@ these are upstream fixes a sync will *gain*.) | Patch | Files | Status upstream | |---|---|---| -| PR #2 (perf): `HASH_MEMO` + `STORE_MEMO` process-level memos; `LocalParseOptions.modified_after` mtime pruning; `latest_source_mtime_ms()` change probe | `src/message_cache.rs`, `src/lib.rs` | not yet forwarded to junhoyeo/tokscale | +| PR #2 (perf, partially retired by M26-A): `LocalParseOptions.modified_after` mtime pruning and `latest_source_mtime_ms()` change probing remain active. M26-A removes `HASH_MEMO` because path/size/mtime memoization can bypass bounded sample validation, and removes monolith-only `STORE_MEMO` because shard dirty tracking owns persistence | `src/message_cache.rs`, `src/lib.rs` | Remaining mtime/pruning seam not yet forwarded; retired memos must not be re-applied | | PR #3 (perf): streaming per-file aggregation replaces materialize-then-aggregate for the graph/model/monthly/hourly reports — `StreamingAggregator` + `SessionizeAccumulator` folded by `scan_messages_streaming` in one cache-aware pass (no full-history `Vec`). Each client lane owns its dedup set (follow-up `0752e35`: prevents cross-client `dedup_key` collisions). | `src/aggregator.rs`, `src/lib.rs`, `src/sessionize.rs`, `tests/streaming_snapshot.rs` | not yet forwarded to junhoyeo/tokscale | | #6 (fix): the **agents report** now folds over `scan_messages_streaming` too — new `get_agents_report` (mirrors `get_model_report`, `resolve_report_clients` + a single streaming pass into `AgentAccumulator`), so it shares the one deduped/per-client-gated/priced stream as every other report (resolves the issue #6 divergence: agents no longer over-counts copilot/codebuff/kimi/cursor/warp/… duplicate `dedup_key`s, and scans the same client set). `parse_local_unified_messages` survives as public API only (footgun-documented, no in-repo callers). `crates/tb_core_ffi/src/agents_report.rs` is now a thin mapper like `model_report.rs` (no longer byte-identical to the archived Tauri original — accepted). | `src/lib.rs`, `crates/tb_core_ffi/src/agents_report.rs` | not yet forwarded to junhoyeo/tokscale | | #35/#36 (fix): **two-level client filter for the hourly & agents reports.** TokenBar passes the user's displayed client slice into `ReportOptions.clients` so shared hour/agent buckets carry only the selected clients' totals (a Swift membership filter can't split a mixed fold). But `scan_messages_streaming` selects scanner *lanes* from that list, and a `cc-mirror/` id (produced during Claude-lane parsing, #659) is not a lane — requesting it alone would scan nothing. New `split_report_client_filter` / `report_message_client_passes` helpers split the request into (a) lanes to scan (each `cc-mirror/*` → its producing `claude` lane) and (b) an EXACT client-id set the fold keeps; `get_agents_report` **and** `get_hourly_report` call the split + gate `msg_filter` on it. Unlike the scan's built-in `retain_for_requested_clients`, requesting `claude` here excludes the distinct `cc-mirror/*` variants (they are their own client ids in the graph/model/daily payloads); the `synthetic` special-case is preserved. `get_hourly_report` was previously byte-upstream — this makes it a local patch; **re-apply on any re-vendor of `lib.rs`.** | `src/lib.rs`, `crates/tb_core_ffi/src/lib.rs`, `crates/tb_core_ffi/src/{hourly,agents}_report.rs` | not yet forwarded to junhoyeo/tokscale | -| #5 (feat): discover Claude desktop "Cowork" (local-agent-mode) transcripts. `discover_cowork_project_roots()` recurses `~/Library/Application Support/Claude/local-agent-mode-sessions/**/.claude/projects` and feeds the roots into `built_in_extra_scan_paths_for` as `ClientId::Claude`. Returns the per-session `projects` roots only, so the sibling `audit.jsonl` (a mirror of the same `usage` records) is never scanned — scanning it would double-count. | `src/scanner.rs` | not yet forwarded to junhoyeo/tokscale | +| #5 (feat): discover Claude desktop "Cowork" (local-agent-mode) transcripts. `discover_cowork_project_roots()` recurses `~/Library/Application Support/Claude/local-agent-mode-sessions/**/.claude/projects` and feeds the roots into `built_in_extra_scan_paths_for` as `ClientId::Claude`. Returns the per-session `projects` roots only, so the sibling `audit.jsonl` (a mirror of the same `usage` records) is never scanned — scanning it would double-count. | `src/scanner.rs` | Covered by closed-unmerged upstream PR #708 together with `` filtering; revive that PR rather than opening a duplicate. | | pricing (fix): **cache-rate backfill** — `choose_best_source_result` is wrapped so the chosen pricing source has any missing cache read/write rates grafted from the runner-up source (`backfill_cache_costs` + `prefer_litellm_over_openrouter`). Without it, a provider-hint that selects an entry lacking cache rates (e.g. an OpenRouter row) bills cache reads at $0 (fable-5 showed $11.50 instead of $45). Upstream `#658`/`#707` do **not** subsume this: `has_any_usable_pricing` is an `.any()` gate, so a row that prices everything *except* cache still passes through unfilled — the root cause stays ours to fix. | `src/pricing/lookup.rs` | not yet forwarded to junhoyeo/tokscale | | pricing (perf): **in-memory auto-refresh** — `PRICING_SERVICE` is a `RwLock>` with `IN_MEMORY_TTL = 3600s` (was a never-refreshing `OnceCell`), so prices re-read the file cache / network roughly hourly instead of being frozen for the process lifetime; adds `pricing_cached_at()` (Models card "Prices updated …"). Spans `mod.rs` + the additive `litellm::cached_at` / `cache::cache_timestamp` helpers. | `src/pricing/mod.rs`, `src/pricing/litellm.rs`, `src/pricing/cache.rs` | not yet forwarded to junhoyeo/tokscale | -| streaming `simple_lane!` fingerprint-fn arm (M4 enabler): the local streaming-aggregation `simple_lane!` macro (part of PR #3 — it does **not** exist upstream) gained a 3-arg arm taking a custom fingerprint fn; the 2-arg arm delegates to it with `SourceFingerprint::from_path` (no behaviour change for the existing callers). jcode passes `from_jcode_path` (snapshot + `.journal.jsonl`) and micode passes `from_sqlite_path`. **Must be re-applied on any re-vendor of `lib.rs`** — dropping it reverts jcode/micode to a snapshot-only fingerprint and reintroduces stale (missing-turn) data after a sibling-only write. | `src/lib.rs` | local (the macro is local; not forwardable as-is) | +| streaming `simple_lane!` fingerprint-fn arm (M4 enabler, adapted for M26-A): the local streaming-aggregation macro (part of PR #3 — it does **not** exist upstream) retains its custom fingerprint arm and identity-aware shard lookup/writeback. Default sources use bounded sample fingerprints; Jcode, Droid, Kimi, Roo/Kilo/Cline, and Micode retain specialized dependency functions. **Must be re-applied on any re-vendor of `lib.rs`** — dropping it reverts sibling-aware streaming invalidation or loses parser namespace isolation. | `src/lib.rs` | local streaming adaptation over upstream's format-1 cache contract | | streaming `simple_lane!` cost-guarded arm + `reprice_lane_message` (M7 — #742 Part 2): the `simple_lane!` macro gained a 4-arg arm `(client, parse, fingerprint, guard_authoritative_cost: bool)`; the 2-/3-arg arms delegate with `false` (unconditional reprice — unchanged behaviour). The body reprices via `reprice_lane_message`, which for `guard=true` only reprices when the message's embedded cost is `<= 0.0`. **micode** passes `true` so MiMo Code's authoritative per-message cost is never overwritten by a recomputed `tokens*rate` when a MiMo model resolves to a price (today a no-op — MiMo models are unpriced — but future-proof). The M8 `CostSource` port intentionally does **not** mark Micode provider-reported: this lane-specific guard remains the Micode source of truth. The micode materialized path (dead code, public API only) is guarded the same way (loader called with `pricing: None`, then reprice-if-`cost<=0.0`). Upstream's #742 guard is in its materialized lane (dead code for us), so this is the local equivalent. **Must be re-applied on any re-vendor of `lib.rs`** — dropping it reverts micode to unconditional repricing (clobbers the authoritative cost once MiMo is priced). | `src/lib.rs` | local (streaming lane + cost-guard are local; upstream guards its own materialized path) | -| mux dedup key workspace-scoped (fix, over `#760`): upstream #760 keys mux's cross-file dedup on `mux::`. The index is unstable across re-parses, and `mux::0` collides cross-workspace — so in our per-client streaming seen set two `.mux/sessions//session-usage.json` files with the same model drop one, a net regression vs the pre-#760 `None` key that kept everything. We key on `mux::` (workspace = the session file's parent-dir name): stable across re-parses and unique per workspace. **Re-apply on any re-vendor of `mux.rs`** (do not restore the upstream index key). | `src/sessions/mux.rs` | Reported upstream — issue #816 / PR #817, reframed as defensive hardening: upstream does not consume mux's `dedup_key` (bare `extend`, no `_seen` filter), so the collision is inert there; only our streaming seen set consumes the key, hence the drop we saw. Keep this local key until upstream merges. | -| drop synthetic Claude placeholder rows (fix, PR #30 by @starburst3190): claudecode parsing discards assistant turns whose model is the literal `` placeholder. Claude Code fabricates these locally (cancelled requests, injected continuations); they never hit a real model and carry zero tokens/cost (verified against real-world data: every sampled `` record carried an all-zero `usage`), so they only surface as a phantom zero-token model row. Guarded at all four message-emitting sites (`parse_claude_file_with_cache_and_home`, `extract_claude_tool_result_message`, `extract_claude_headless_message`, `finalize_headless_state`); the streaming lane inherits it because it calls `parse_claude_file_with_cache_and_home`. `CACHE_SCHEMA_VERSION` bumped 22→23 so sessions cached with the phantom row reparse. **Re-apply on any re-vendor of `claudecode.rs`.** | `src/sessions/claudecode.rs`, `src/message_cache.rs` | Net-new local hardening — upstream still emits synthetic rows (its parser test expects `messages.len() == 5`, counting the `` row as an `unknown`-provider entry). Not yet forwarded; candidate for an upstream issue/PR. | +| mux dedup key workspace-scoped (fix, over `#760`): upstream #760 keys mux's cross-file dedup on `mux::`. The index is unstable across re-parses, and `mux::0` collides cross-workspace — so in our per-client streaming seen set two `.mux/sessions//session-usage.json` files with the same model drop one, a net regression vs the pre-#760 `None` key that kept everything. We key on `mux::` (workspace = the session file's parent-dir name): stable across re-parses and unique per workspace. **Re-apply on any re-vendor of `mux.rs`** (do not restore the upstream index key). | `src/sessions/mux.rs` | Already fixed upstream via PR #817; retain the local key when selectively re-vendoring because upstream does not consume this dedup key in its materialized path. | +| drop synthetic Claude placeholder rows (fix, PR #30 by @starburst3190): claudecode parsing discards assistant turns whose model is the literal `` placeholder. Claude Code fabricates these locally (cancelled requests, injected continuations); they never hit a real model and carry zero tokens/cost (verified against real-world data: every sampled `` record carried an all-zero `usage`), so they only surface as a phantom zero-token model row. Guarded at all four message-emitting sites (`parse_claude_file_with_cache_and_home`, `extract_claude_tool_result_message`, `extract_claude_headless_message`, `finalize_headless_state`); the streaming lane inherits it because it calls `parse_claude_file_with_cache_and_home`. `CACHE_SCHEMA_VERSION` bumped 22→23 so sessions cached with the phantom row reparse. **Re-apply on any re-vendor of `claudecode.rs`.** | `src/sessions/claudecode.rs`, `src/message_cache.rs` | Covered by closed-unmerged upstream PR #708; revive that existing PR rather than opening a duplicate. | | grok compaction counter epochs (fix, over the faithful port): upstream's grok parser treats **every** cumulative-counter rewind as streaming jitter and `continue`s past it, relying solely on `signals.json` to recover a compacted session's lost tokens. When `signals.json` is absent or lags the live session, the post-compaction tokens are silently dropped. Our port adds in-stream counter epochs: a rewind that both drops `>= 32_000` tokens (`COMPACTION_MIN_DROP_TOKENS`) AND at least halves the counter (`is_compaction_reset`) opens a new epoch (`start_new_counter_epoch`), banking the prior epoch's peak into `completed_epoch_tokens` so the turn's total accumulates across compactions instead of resetting. Verified: a `180000 -> 40000 -> 500000` signals-absent stream yields **680000** with the delta vs **500000** upstream (the pre-compaction 180000 is lost). It does **not** double-count against the signals reconciliation, which is difference-based and idempotent (`extra = signals_total - updates_total`, skipped when `<= 0`): with signals present covering the epochs the two mechanisms produce identical totals (regression test `compaction_epoch_survives_without_signals_reconciliation` asserts both the signals-absent 680000 and the with-signals idempotence). **Must be re-applied on any re-vendor of `grok.rs`** — a wholesale re-vendor would silently drop it and regress the signals-absent undercount. | `src/sessions/grok.rs` | Net-new local divergence — upstream has no counter-epoch handling. Candidate to upstream (fixes a real undercount). Not yet forwarded. | -| grok fingerprint spans every metadata sibling (fix, over the faithful port): `read_metadata` reads the session model id (and timestamp) from three siblings of `updates.jsonl` — `signals.json`, `summary.json` (`current_model_id`/`model_id` + `updated_at`/`created_at`), and `events.jsonl` (`model_id`/`session_id`/`ts`) — but upstream's `SourceFingerprint::from_grok_path` fingerprints only `signals.json`. So when `summary.json`/`events.jsonl` is the sole carrier of the real model id and arrives or changes *after* `updates.jsonl` was cached, the stale parse (fallback `grok`/unknown model → wrong pricing + grouping) is served until an unrelated file changes. Our `from_grok_path` now folds all three siblings into the fingerprint (shared `message_cache::GROK_METADATA_SIBLINGS` const; regression tests `from_grok_path_invalidates_on_summary_only_change` / `_on_events_only_change`). **The `lib.rs` mtime change probes must watch the same set or the extended fingerprint is unreachable:** `latest_source_mtime_ms` (live-tail change token) and `grok_source_mtime_ms` (per-session mtime pruning) now max across all four files via the same const — a `summary.json`- or `events.jsonl`-only write otherwise leaves the change token/prune decision unchanged, so `graph_cached` serves the old graph and the live tail skips the re-parse (probe tests `test_latest_source_mtime_ms_probes_grok_summary`/`_events`, prune tests `test_modified_after_prunes_grok_keeps_session_with_fresh_summary`/`_events`). Same class of gap as our reported #741 (roo history sibling) / #822 (fold). **Must be re-applied on any re-vendor of `message_cache.rs` or `lib.rs`.** | `src/message_cache.rs`, `src/lib.rs` | Net-new local divergence — upstream `from_grok_path` fingerprints only `signals.json` despite reading all three siblings, and its mtime probes match. Candidate to report upstream (same class as #741/#822). Not yet forwarded. | +| grok `turn_completed.usage` primary path (fix, over the faithful port): modern Grok Build logs expose authoritative per-turn API usage on `sessionUpdate: turn_completed` → `params.update.usage` (`inputTokens` / `outputTokens` / `reasoningTokens` / `cachedReadTokens` / `modelUsage` / `costUsdTicks`). Upstream (and our prior port) still treat cumulative `params._meta.totalTokens` deltas as spend — but that field is **context occupancy**, so multi-turn sessions undercount by ~25× on real data (long sessions ~50×). Primary path emits usage rows (cache netted out of `input` so TokenBar's additive total does not double-count; provider cost = `costUsdTicks / 1e9` when present). Legacy context-delta + compaction epoch + `signals.json` reconciliation remain for logs with no usage objects; an open turn after the last completed usage still contributes a live context partial. Signals reconciliation is **skipped** on the usage path (signals measure context, not cumulative spend). `CACHE_SCHEMA_VERSION` **31→32** so schema-31 context-only rows rebuild. Plan: [`docs/knowledge/plans/grok-turn-completed-usage.md`](../docs/knowledge/plans/grok-turn-completed-usage.md). **Must be re-applied on any re-vendor of `grok.rs` / `message_cache.rs`.** | `src/sessions/grok.rs`, `src/message_cache.rs` | Net-new local divergence — upstream still uses context deltas only. Candidate to upstream. Not yet forwarded. | +| grok fingerprint spans every metadata sibling (fix, over the faithful port): `read_metadata` reads the session model id (and timestamp) from three siblings of `updates.jsonl` — `signals.json`, `summary.json` (`current_model_id`/`model_id` + `updated_at`/`created_at`), and `events.jsonl` (`model_id`/`session_id`/`ts`) — but upstream's `SourceFingerprint::from_grok_path` fingerprints only `signals.json`. So when `summary.json`/`events.jsonl` is the sole carrier of the real model id and arrives or changes *after* `updates.jsonl` was cached, the stale parse (fallback `grok`/unknown model → wrong pricing + grouping) is served until an unrelated file changes. Our `from_grok_path` now folds all three siblings into the fingerprint (shared `message_cache::GROK_METADATA_SIBLINGS` const; regression tests `from_grok_path_invalidates_on_summary_only_change` / `_on_events_only_change`). **The `lib.rs` mtime change probes must watch the same set or the extended fingerprint is unreachable:** `latest_source_mtime_ms` (live-tail change token) and `grok_source_mtime_ms` (per-session mtime pruning) now max across all four files via the same const — a `summary.json`- or `events.jsonl`-only write otherwise leaves the change token/prune decision unchanged, so `graph_cached` serves the old graph and the live tail skips the re-parse (probe tests `test_latest_source_mtime_ms_probes_grok_summary`/`_events`, prune tests `test_modified_after_prunes_grok_keeps_session_with_fresh_summary`/`_events`). Same class of gap as our reported #741 (roo history sibling) / #822 (fold). **Must be re-applied on any re-vendor of `message_cache.rs` or `lib.rs`.** | `src/message_cache.rs`, `src/lib.rs` | Already fixed upstream via PR #856; retain the sibling sweep when selectively re-vendoring until the matching upstream tree is present. | | vendor-local sibling sweep alongside #823 (fix): #823 saturated the 6 report/session per-field token folds, but a fresh-context verifier pass found 2 more bare `+=` folds in the same failure class that #823 does not cover — `AgentAccumulator::add` (the streaming agents report's own per-field token fold, parallel to but distinct from `aggregate_model_usage_entries`) and `aggregator.rs`'s yearly-totals fold (`entry.tokens += c.totals.tokens`, where a contribution's total can already be saturated to `i64::MAX` by `TokenBreakdown::total`). Both switched to `saturating_add`. A subsequent external-review (Codex) pass caught one more in the same file: `calculate_summary`'s cross-day `total_tokens` sum (`contributions.iter().map(|c| c.totals.tokens).sum()`) had the identical overflow class and is now a saturating fold too. **Not part of upstream #823**, kept separate here since none of these sites exist in the same form upstream (upstream has no `AgentAccumulator`, and its yearly/summary-totals equivalents were not covered by PR #823 either). | `src/lib.rs`, `src/aggregator.rs` | local-only (not upstream; nothing to forward) | +| Kimi Code empty-root parity (fix, M19-B1 review): `KIMI_CODE_HOME` is accepted only when its original value contains at least one non-whitespace character. Missing, exactly empty, and whitespace-only values fall back to the supplied `/.kimi-code` root instead of scanning `/sessions` or a whitespace-named relative root; every other non-empty value is preserved byte-for-byte. The existing `use_env_roots=false` path still ignores the process environment. This changes discovery only, not parser output, cache serialization, or public APIs; discovered source paths remain cache identities, so active format 2 and the inert schema-32 monolith are unchanged. **Re-apply on any re-vendor of `scanner.rs` until upstream carries the same fallback.** | `src/scanner.rs` | Reported upstream as [#963](https://github.com/junhoyeo/tokscale/issues/963); narrow PR not yet sent | +| OpenCode/XDG empty-root parity (fix, PR #97): OpenCode's pinned `xdg-basedir` treats exactly empty `XDG_DATA_HOME` as unset, while whitespace-only and every other non-empty value remain configured roots. `PathRoot::XdgData` now matches that contract, and OpenCode legacy JSON, SQLite, plus synthetic Octofriend discovery all reuse the shared resolver instead of maintaining separate environment parsing. The FFI `auth.json` resolver mirrors the same selection so subscription labels and Copilot quota credentials stay on the scanned OpenCode installation. Empty-root fixtures cover both legacy messages and `opencode.db`. No parser output, serialized cache layout, or public API changed; discovered source paths remain cache keys, so the current schema 32 and active shard format 2 remain valid. **Re-apply on any re-vendor of `clients.rs` or `scanner.rs` until upstream carries the same fallback.** | `src/clients.rs`, `src/scanner.rs` | Candidate for a narrow upstream fix; not yet forwarded | | reqwest TLS: `native-tls-vendored` (0.12) → 0.13 `rustls` (build): vendored OpenSSL needs Perl+NASM on Windows builds (CI runners ship them; local dev boxes often don't) and is the slowest unit of a clean build. reqwest 0.13's `rustls` feature verifies via **rustls-platform-verifier** (Security.framework on macOS, SChannel on Windows), so scoped Keychain trust constraints, corporate MITM roots, and AIA/AuthRoot chain fetching all keep native semantics — deliberately NOT `rustls-tls-native-roots`, which flattens per-host/per-app trust into universal WebPKI roots and drops AIA chasing (Codex review finding). Matches `tb_core_ffi`'s reqwest line exactly (one reqwest build in the workspace) and zeroes the Windows repo's local-patch table so its re-syncs become a plain rsync. **Re-apply on any re-vendor of `Cargo.toml`.** | `Cargo.toml` | local build-infra divergence — upstream ships `native-tls-vendored`; nothing to forward | + +## Current-upstream reportability audit + +This is a deduplicated inventory for the current upstream target. It separates reportable upstream gaps from already reported, fixed, confirmed, local-only, and evidence-gated items; it does not imply that every local patch should be upstreamed. + +| Category | Items | +|---|---| +| `REPORTABLE_NOW` | Pricing cache-rate backfill (highest priority; linked upstream [#57](https://github.com/junhoyeo/tokscale/issues/57) / [#56](https://github.com/junhoyeo/tokscale/pull/56)); Grok counter epochs (linked merged [#798](https://github.com/junhoyeo/tokscale/pull/798)); Kimi exact replay (issue-first); Kimi extreme-token overflow (narrow PR); OpenCodeReview `u64` clamp (narrow PR); OpenCode logical identity/authority (issue-first); OpenCode/XDG empty-root parity (narrow PR); long-lived in-memory pricing TTL (optional, linked [#333](https://github.com/junhoyeo/tokscale/pull/333)). | +| `REVIVE_EXISTING` | Cowork and `` share closed-unmerged upstream PR [#708](https://github.com/junhoyeo/tokscale/pull/708); do not open a duplicate. | +| `ALREADY_REPORTED_OPEN` | Copilot duplicate-span follow-up issue [#942](https://github.com/junhoyeo/tokscale/issues/942) and maintainer-ready PR [#943](https://github.com/junhoyeo/tokscale/pull/943), after the base fix merged through [#938](https://github.com/junhoyeo/tokscale/issues/938) / [#939](https://github.com/junhoyeo/tokscale/pull/939); Copilot Desktop lifetime-total time attribution [#962](https://github.com/junhoyeo/tokscale/issues/962); Kimi Code empty-root parity [#963](https://github.com/junhoyeo/tokscale/issues/963); Grok unified [#849](https://github.com/junhoyeo/tokscale/issues/849); routed prefix+suffix [#846](https://github.com/junhoyeo/tokscale/issues/846); request-level long-context [#862](https://github.com/junhoyeo/tokscale/issues/862). | +| `NEEDS_CONFIRMATION` | Kiro nested `sess_*` topology and CLI/SQLite `auto` versus `unknown` fallback; obtain format and compatibility evidence before calling either a confirmed bug. | +| `ALREADY_FIXED` | mux [#817](https://github.com/junhoyeo/tokscale/pull/817); Jcode [#819](https://github.com/junhoyeo/tokscale/pull/819); Copilot [#880](https://github.com/junhoyeo/tokscale/pull/880); folds [#823](https://github.com/junhoyeo/tokscale/pull/823); Grok metadata sibling fingerprint [#856](https://github.com/junhoyeo/tokscale/pull/856); M19-A `fs_atomic` current upstream. The mux and Grok fingerprint local-patch rows above are stale as reportability statuses and are retained only as re-vendor instructions. | +| `TOKENBAR_LOCAL` | `HASH`/`STORE` memo; streaming aggregator; Agents/hourly filters; `simple_lane` arms; extra sibling folds; live-tail/FFI/cache/report seams; M18 local pricing policy; reqwest TLS. | +| `DEFER-EVIDENCE` | PR #72 Zcode cross-store/parser/schema/provider/scanner and GLM fail-closed scope; wait for an upstream Zcode update and re-audit rather than sending the current drift upstream. | + +The M22 implementation is not main history: PR #72 is closed with `merged=false`, and its 17-commit, 12-file head remains evidence for the fidelity decision only. The rule is to preserve necessary TokenBar streaming/FFI seams, but stop when core parser/authority needs continuous systematic repair, custom algorithm size exceeds upstream scope, or review keeps exposing new failure classes; defer until upstream converges instead of continuing by sunk cost. diff --git a/vendor/tokscale-core/SYNC.md b/vendor/tokscale-core/SYNC.md index c8bb71b..976e684 100644 --- a/vendor/tokscale-core/SYNC.md +++ b/vendor/tokscale-core/SYNC.md @@ -1,60 +1,65 @@ -# Vendor sync provenance - -This repo's `crates/tb_core_ffi/` and `vendor/tokscale-core/` are **copies** from -the macOS repo (`Nanako0129/TokenBar`, local `~/side-project/TokenBar-Native`), -which is the **single sync source** for the shared Rust core. Upstream -(junhoyeo/tokscale) syncs land in the macOS repo first, then get re-copied here. +# Vendored tok-scale sync | Field | Value | |---|---| -| Source repo | `Nanako0129/TokenBar` (macOS) | -| Copied at commit | `2ed256ee` (= v1.4.0 `fe19eebc` + the two backports below; landed on macOS main 2026-07-16 via ff-merge, SHA preserved) | -| Copied on | 2026-07-15 | -| Upstream milestone state | M6 + M7 + M9 backports merged; adds Grok Build + Hermes clients, cost-provenance contract, client-selection filter for hourly/agents (ctb.h signature change) | -| Cache schema version | 29 (`vendor/tokscale-core/src/message_cache.rs`) | - -## Local patches (Windows repo only) - -Patches that exist here but NOT yet in the macOS repo. Every local vendor drift -must be recorded here before delivery. Platform-neutral or cfg-gated fixes -should be upstreamed to the macOS sync source before the next copy; otherwise -the listed commits must be deliberately reapplied after syncing. - -| Commit | Patch | Files | Upstreamed to macOS repo? | -|---|---|---|---| -| `e520063` | Hermetic Windows test portability: panic-safe env/cwd guards, serial coordination, platform-safe JSON/path fixtures, isolated cache/XDG roots, and non-mutating Windows legacy-path coverage | `src/clients.rs`, `src/lib.rs`, `src/message_cache.rs`, `src/pricing/cache.rs`, `src/scanner.rs`, `src/sessions/claudecode.rs`, `src/sessions/opencode.rs` | No — Windows repo only (2026-07-17) | -| `db2a96a` | Release the temp writer before Windows atomic replacement and reopen the final cache read/write for the durability sync | `src/message_cache.rs` | No — Windows repo only (2026-07-17) | -| `15f418e` | Route pure scanner fixtures through `use_env_roots=false` so host `TOKSCALE_EXTRA_DIRS` and other roots cannot change fixed-count tests | `src/scanner.rs` | No — Windows repo only (2026-07-17) | -| `e807f33` | Gate the Windows Zed known-folder fallback on `use_env_roots` | `src/scanner.rs` | No — Windows repo only (2026-07-17) | -| `0979cdb` | Make Windows `PathRoot::Config` use `%APPDATA%` only for environment-aware scans; explicit-home scans use the supplied home | `src/clients.rs` | No — Windows repo only (2026-07-17) | -| `26b892a` | Compare extra-path warnings with the scanner's supplied home instead of querying the process user's platform known folder | `src/scanner.rs` | No — Windows repo only (2026-07-17) | -| `6ac77c0` | Isolate env-aware scanner fixtures from unrelated host `TOKSCALE_EXTRA_DIRS` while retaining the dedicated positive extra-dir case | `src/scanner.rs` | No — Windows repo only (2026-07-17) | -| `e91fb2c` | Route Claude parser unit fixtures through the explicit no-home entry so they cannot enumerate or read the process user's `.cc-mirror` variants | `src/sessions/claudecode.rs` | No — Windows repo only (2026-07-17) | -| `fbecb99` | Isolate scanner-settings and headless fixtures from unrelated host client roots while retaining dedicated env-override coverage | `src/scanner.rs` | No — Windows repo only (2026-07-17) | -| `3c8bfc5` | Make the Grok registration fixture assert the supplied-home fallback instead of inheriting host `GROK_HOME` | `src/clients.rs` | No — Windows repo only (2026-07-17) | -| `aec5bd8` | Route materialized parser fixtures through the explicit-home strategy so unrelated host extra directories and client roots cannot change fixed outputs | `src/lib.rs` | No — Windows repo only (2026-07-17) | - -Earlier divergences were upstreamed to the macOS repo on 2026-07-15 -(branch `fix-vendor-rustls-tls`, the sync source above), so `crates/` + -`vendor/` were byte-identical at copied commit `2ed256ee` before the local -patches listed above: - -- reqwest TLS: now 0.13 `rustls` (rustls-platform-verifier — native trust - semantics on both OSes: Security.framework / SChannel, keeps scoped trust - + AIA chasing; upgraded from our original `rustls-tls-native-roots` after - a Codex review finding) — macOS commit `220cf8ab`. -- `crate-type = ["cdylib", "staticlib"]` for tb_core_ffi (P/Invoke needs the - cdylib; was an unrecorded local divergence since Phase 0) — macOS commit - `2ed256ee`. - -Note: the macOS repo's own local-patch table vs junhoyeo/tokscale lives in -`vendor/README.md` (copied along) — that provenance chain still applies. - -## How to re-sync from the macOS repo +| Source repo | [Nanako0129/TokenBar](https://github.com/Nanako0129/TokenBar) | +| Copied commit | [`4dfed5ffab26e2707a8cd82ee99864520b75892b`](https://github.com/Nanako0129/TokenBar/commit/4dfed5ffab26e2707a8cd82ee99864520b75892b) | +| Copied on | 2026-07-27 | + +`4dfed5ffab26e2707a8cd82ee99864520b75892b` is the final merged Native PR +#102 commit on `main`. In addition to the Antigravity discovery, format-2 writer +lifecycle, malformed Copilot Desktop entry, and fixture portability +corrections, this source retries Codex credential replacement, merges duplicate +Copilot span endpoints, advances only the Copilot parser identity to 4, and +makes shipping cache fixtures build fingerprints, seed, and query with the +scanner-returned path spelling. The post-review merge also preserves valid +Copilot quota rows when the optional reset field has a non-string type and +emits DNS/TLS transport categories from typed error sources without exposing +free-form error text. This is the canonical post-merge M19-B checkpoint. + +The active cache is format 2 at `source-message-cache-v2`; format-1 shards are +stale and rebuild cold under format 2. The legacy schema-32 monolith +`source-message-cache.bin` remains inert, unread, unmodified, and undeleted. + +The Windows shared-tree local patch table is **none**. The Native-only +`vendor/AGENTS.md` repository adapter is intentionally excluded because this +repository keeps local agent guides untracked; it is not runtime source. Every +tracked runtime vendor file is byte-identical to Native except this Windows-only +`vendor/tokscale-core/SYNC.md` provenance record. The Rust serializer-lock +fixture at `Fixtures/CrossCheck/provider-quota-pace-v3.json` and the Windows +cross-check copy are byte-identical to the same Native fixture. + +## Sync procedure + +Run from a clean Windows checkout with `TOKENBAR_NATIVE` pointing to a clean +Native checkout. Preserve this Windows-only file across the exact replacement, +then update its commit, date, and verification fields before committing: ```bash -cd ~/side-project/TokenBar-Native && git archive -- crates vendor \ - | tar -x -C ~/side-project/TokenBar-Windows -# then: restore this SYNC.md, update the table above, re-apply any local -# patches listed here, run scripts/check.sh + full smoke on Windows. +: "${TOKENBAR_NATIVE:?set TOKENBAR_NATIVE to a clean Native checkout}" +source_commit=4dfed5ffab26e2707a8cd82ee99864520b75892b +stage="$(mktemp -d)" +sync_record="$(mktemp)" +trap 'rm -rf "$stage"; rm -f "$sync_record"' EXIT + +cp vendor/tokscale-core/SYNC.md "$sync_record" +git -C "$TOKENBAR_NATIVE" archive \ + "$source_commit" crates vendor Sources/CTB/include/ctb.h \ + Fixtures/CrossCheck/provider-quota-pace-v3.json \ + | tar -x -C "$stage" + +rm -rf crates vendor +cp -a "$stage/crates" crates +cp -a "$stage/vendor" vendor +rm -f vendor/AGENTS.md +cp "$stage/Sources/CTB/include/ctb.h" include/ctb.h +mkdir -p Fixtures/CrossCheck +cp "$stage/Fixtures/CrossCheck/provider-quota-pace-v3.json" \ + Fixtures/CrossCheck/provider-quota-pace-v3.json +cp "$stage/Fixtures/CrossCheck/provider-quota-pace-v3.json" \ + crosscheck/fixtures/provider-quota-pace-v3.json +cp "$sync_record" vendor/tokscale-core/SYNC.md ``` + +Do not write private local paths or credentials. This record does not claim that +real ARM64 runtime validation is complete. diff --git a/vendor/tokscale-core/src/aggregator.rs b/vendor/tokscale-core/src/aggregator.rs index 7e41540..fcd862b 100644 --- a/vendor/tokscale-core/src/aggregator.rs +++ b/vendor/tokscale-core/src/aggregator.rs @@ -281,18 +281,22 @@ impl DayAccumulator { .reasoning .saturating_add(msg.tokens.reasoning); - // Update client contribution + // Update client contribution. + // Canonical (alias-free) id: this contribution is serialized into the + // graph/export-facing payload, so a machine-local modelAliases config + // must not rewrite the model identity that leaves the machine. Local + // report grouping uses `normalize_model_for_grouping` instead. let key = format!( "{}:{}", msg.client, - crate::normalize_model_for_grouping(&msg.model_id) + crate::canonical_model_id(&msg.model_id) ); let client_entry = self .clients .entry(key) .or_insert_with(|| ClientContribution { client: msg.client.clone(), - model_id: crate::normalize_model_for_grouping(&msg.model_id), + model_id: crate::canonical_model_id(&msg.model_id), provider_id: msg.provider_id.clone(), tokens: TokenBreakdown::default(), cost: 0.0, @@ -517,7 +521,9 @@ impl SessionAccumulator { .saturating_add(msg.tokens.reasoning); // Track tightest (client, provider, model) by cost contribution. - let normalized_model = crate::normalize_model_for_grouping(&msg.model_id); + // Canonical (alias-free) id — feeds the graph/export-facing payload, so + // machine-local aliases must not rewrite it (see `add_message`). + let normalized_model = crate::canonical_model_id(&msg.model_id); let key = format!("{}:{}:{}", msg.client, msg.provider_id, normalized_model); let client_entry = self .clients @@ -905,6 +911,7 @@ mod tests { message_count: 1, agent: None, dedup_key: None, + dedup_aliases: Vec::new(), is_turn_start: false, } } @@ -1575,6 +1582,7 @@ mod tests { message_count: 1, agent: None, dedup_key: None, + dedup_aliases: Vec::new(), is_turn_start: false, duration_ms: None, } @@ -1864,6 +1872,7 @@ mod tests { message_count: 1, agent: None, dedup_key: dedup_key.map(|s| s.to_string()), + dedup_aliases: Vec::new(), is_turn_start: false, } } diff --git a/vendor/tokscale-core/src/clients.rs b/vendor/tokscale-core/src/clients.rs index c0d1509..c486e81 100644 --- a/vendor/tokscale-core/src/clients.rs +++ b/vendor/tokscale-core/src/clients.rs @@ -15,8 +15,10 @@ impl PathRoot { PathRoot::Home => home_dir.to_string(), PathRoot::XdgData => { if use_env_roots { - std::env::var("XDG_DATA_HOME") - .unwrap_or_else(|_| format!("{}/.local/share", home_dir)) + match std::env::var("XDG_DATA_HOME") { + Ok(root) if !root.is_empty() => root, + Ok(_) | Err(_) => format!("{}/.local/share", home_dir), + } } else { format!("{}/.local/share", home_dir) } @@ -42,6 +44,14 @@ impl PathRoot { } } + #[cfg(target_os = "windows")] + if !use_env_roots { + return std::path::Path::new(home_dir) + .join("AppData/Roaming/tokscale") + .to_string_lossy() + .into_owned(); + } + format!("{home_dir}/.config/tokscale") } PathRoot::EnvVar { @@ -492,6 +502,24 @@ define_clients!( headless: false, parse_local: true, submit_default: true + }, + Junie = 31 => { + id: "junie", + root: PathRoot::Home, + relative: ".junie/sessions", + pattern: "events.jsonl", + headless: false, + parse_local: true, + submit_default: true + }, + OpenCodeReview = 32 => { + id: "opencodereview", + root: PathRoot::Home, + relative: ".opencodereview/sessions", + pattern: "*.jsonl", + headless: false, + parse_local: true, + submit_default: true } ); @@ -581,7 +609,27 @@ mod tests { #[test] fn test_client_id_count() { - assert_eq!(ClientId::COUNT, 31); + assert_eq!(ClientId::COUNT, 33); + } + + #[test] + fn test_m21_clients_registered() { + let junie = ClientId::Junie.data(); + assert_eq!(junie.resolve_path("/tmp/home"), "/tmp/home/.junie/sessions"); + assert_eq!(junie.pattern, "events.jsonl"); + assert!(junie.parse_local); + assert!(junie.submit_default); + assert!(!junie.headless); + + let review = ClientId::OpenCodeReview.data(); + assert_eq!( + review.resolve_path("/tmp/home"), + "/tmp/home/.opencodereview/sessions" + ); + assert_eq!(review.pattern, "*.jsonl"); + assert!(review.parse_local); + assert!(review.submit_default); + assert!(!review.headless); } #[test] @@ -622,6 +670,26 @@ mod tests { assert_eq!(resolved, "/tmp/xdg-data-home"); } + #[test] + #[serial] + fn test_path_root_xdg_data_preserves_whitespace_when_set() { + let mut _env = EnvGuard::capture(&["XDG_DATA_HOME"]); + _env.set("XDG_DATA_HOME", " "); + + let resolved = PathRoot::XdgData.resolve("/tmp/home"); + assert_eq!(resolved, " "); + } + + #[test] + #[serial] + fn test_path_root_xdg_data_falls_back_when_empty() { + let mut _env = EnvGuard::capture(&["XDG_DATA_HOME"]); + _env.set("XDG_DATA_HOME", ""); + + let resolved = PathRoot::XdgData.resolve("/tmp/home"); + assert_eq!(resolved, "/tmp/home/.local/share"); + } + #[test] #[serial] fn test_path_root_xdg_data_falls_back_when_unset() { @@ -697,7 +765,15 @@ mod tests { _env.set("XDG_CONFIG_HOME", "/tmp/xdg-config-home"); let resolved = PathRoot::Config.resolve_with_env_strategy("/tmp/home", false); - assert_eq!(resolved, "/tmp/home/.config/tokscale"); + let expected = if cfg!(target_os = "windows") { + std::path::Path::new("/tmp/home") + .join("AppData/Roaming/tokscale") + .to_string_lossy() + .into_owned() + } else { + "/tmp/home/.config/tokscale".to_string() + }; + assert_eq!(resolved, expected); } #[test] diff --git a/vendor/tokscale-core/src/fs_atomic.rs b/vendor/tokscale-core/src/fs_atomic.rs index 876c2a0..3478f0c 100644 --- a/vendor/tokscale-core/src/fs_atomic.rs +++ b/vendor/tokscale-core/src/fs_atomic.rs @@ -13,6 +13,41 @@ pub fn replace_file(tmp_path: &Path, final_path: &Path) -> io::Result<()> { } } +#[cfg(any(target_os = "windows", test))] +fn retry_atomic_replace(mut replace: F, mut sleep: S) -> io::Result<()> +where + F: FnMut() -> io::Result<()>, + S: FnMut(std::time::Duration), +{ + // MoveFileExW replacing an existing file is a well-known source of + // transient ERROR_ACCESS_DENIED (5) / ERROR_SHARING_VIOLATION (32) on + // Windows: antivirus, indexing, and cloud-sync agents routinely hold a + // brief scan handle open on a just-written file. Retry a handful of + // times with a short backoff before giving up, rather than surfacing a + // one-shot failure for what is usually a few-millisecond lock. + const ERROR_ACCESS_DENIED: i32 = 5; + const ERROR_SHARING_VIOLATION: i32 = 32; + const MAX_ATTEMPTS: u32 = 5; + + for attempt in 1..=MAX_ATTEMPTS { + match replace() { + Ok(()) => return Ok(()), + Err(error) => { + let is_retryable = matches!( + error.raw_os_error(), + Some(ERROR_ACCESS_DENIED) | Some(ERROR_SHARING_VIOLATION) + ); + if !is_retryable || attempt == MAX_ATTEMPTS { + return Err(error); + } + sleep(std::time::Duration::from_millis(10 * attempt as u64)); + } + } + } + + unreachable!("loop always returns on its final attempt") +} + #[cfg(target_os = "windows")] fn windows_replace_file(tmp_path: &Path, final_path: &Path) -> io::Result<()> { use std::ffi::OsStr; @@ -38,16 +73,152 @@ fn windows_replace_file(tmp_path: &Path, final_path: &Path) -> io::Result<()> { let existing = encode(tmp_path); let new = encode(final_path); - let result = unsafe { - MoveFileExW( - existing.as_ptr(), - new.as_ptr(), - MOVEFILE_REPLACE_EXISTING | MOVEFILE_WRITE_THROUGH, - ) - }; - if result == 0 { - Err(io::Error::last_os_error()) - } else { - Ok(()) + retry_atomic_replace( + || { + let result = unsafe { + MoveFileExW( + existing.as_ptr(), + new.as_ptr(), + MOVEFILE_REPLACE_EXISTING | MOVEFILE_WRITE_THROUGH, + ) + }; + if result == 0 { + Err(io::Error::last_os_error()) + } else { + Ok(()) + } + }, + std::thread::sleep, + ) +} + +#[cfg(test)] +mod tests { + use super::retry_atomic_replace; + use std::io; + use std::time::Duration; + + fn os_error(code: i32) -> io::Error { + io::Error::from_raw_os_error(code) + } + + #[test] + fn retries_access_denied_then_succeeds() { + let mut attempts = 0; + let mut delays = Vec::new(); + let result = retry_atomic_replace( + || { + attempts += 1; + if attempts == 1 { + Err(os_error(5)) + } else { + Ok(()) + } + }, + |delay| delays.push(delay), + ); + + assert!(result.is_ok()); + assert_eq!(attempts, 2); + assert_eq!(delays.as_slice(), &[Duration::from_millis(10)]); + } + + #[test] + fn retries_sharing_violation_then_succeeds() { + let mut attempts = 0; + let mut delays = Vec::new(); + let result = retry_atomic_replace( + || { + attempts += 1; + if attempts == 1 { + Err(os_error(32)) + } else { + Ok(()) + } + }, + |delay| delays.push(delay), + ); + + assert!(result.is_ok()); + assert_eq!(attempts, 2); + assert_eq!(delays.as_slice(), &[Duration::from_millis(10)]); + } + + #[test] + fn persistent_retryable_error_stops_after_five_attempts() { + let mut attempts = 0; + let mut delays = Vec::new(); + let result = retry_atomic_replace( + || { + attempts += 1; + Err(os_error(5)) + }, + |delay| delays.push(delay), + ); + + let error = result.expect_err("persistent retryable errors must fail"); + assert_eq!(error.raw_os_error(), Some(5)); + assert_eq!(attempts, 5); + assert_eq!( + delays.as_slice(), + &[ + Duration::from_millis(10), + Duration::from_millis(20), + Duration::from_millis(30), + Duration::from_millis(40), + ] + ); + } + + #[test] + fn non_transient_error_attempts_once() { + let mut attempts = 0; + let mut delays = Vec::new(); + let result = retry_atomic_replace( + || { + attempts += 1; + Err(os_error(2)) + }, + |delay| delays.push(delay), + ); + + let error = result.expect_err("non-transient errors must return immediately"); + assert_eq!(error.raw_os_error(), Some(2)); + assert_eq!(attempts, 1); + assert!(delays.is_empty()); + } + + #[test] + fn retry_attempts_and_backoff_follow_exact_sequence() { + let outcomes = [ + Err(os_error(5)), + Err(os_error(32)), + Err(os_error(5)), + Err(os_error(32)), + Ok(()), + ]; + let mut outcomes = outcomes.into_iter(); + let mut attempts = Vec::new(); + let mut delays = Vec::new(); + let result = retry_atomic_replace( + || { + let attempt = attempts.len() + 1; + attempts.push(attempt); + outcomes.next().expect("test provides five outcomes") + }, + |delay| delays.push(delay), + ); + + assert!(result.is_ok()); + assert_eq!(attempts, [1, 2, 3, 4, 5]); + assert_eq!( + delays.as_slice(), + &[ + Duration::from_millis(10), + Duration::from_millis(20), + Duration::from_millis(30), + Duration::from_millis(40), + ] + ); } } diff --git a/vendor/tokscale-core/src/lib.rs b/vendor/tokscale-core/src/lib.rs index b16084e..464f1d0 100644 --- a/vendor/tokscale-core/src/lib.rs +++ b/vendor/tokscale-core/src/lib.rs @@ -5,6 +5,7 @@ mod cc_mirror; pub mod clients; pub mod fs_atomic; mod message_cache; +pub mod model_alias; mod parser; pub mod paths; pub mod pricing; @@ -15,11 +16,16 @@ pub mod sessions; pub use aggregator::*; pub use clients::{ClientCounts, ClientDef, ClientId, PathRoot}; +pub use model_alias::{ + clear_model_aliases, model_alias_generation, model_aliases, + register_usage_data_invalidation_hook, set_model_aliases, snapshot_grouping_aliases, + GroupingAliasSnapshot, ModelAliasMap, +}; pub use parser::*; pub use scanner::*; pub use sessionize::{ - compute_daily_active_time, compute_time_metrics, sessionize, SessionizeAccumulator, - SessionInterval, TimeMetrics, DEFAULT_IDLE_GAP_MS, + compute_daily_active_time, compute_time_metrics, sessionize, SessionInterval, + SessionizeAccumulator, TimeMetrics, DEFAULT_IDLE_GAP_MS, }; pub use sessions::{CostSource, UnifiedMessage}; @@ -55,7 +61,45 @@ pub(crate) fn strip_parenthesized_reasoning_tier(model_id: &str) -> Option<&str> Some(base_model) } +/// Canonical model identity — the model id that leaves the machine. +/// +/// This is [`normalize_syntactic`] with **no alias folding**: purely structural +/// canonicalization (lowercase, strip a `(reasoning-tier)` suffix, strip a +/// trailing `-YYYYMMDD` date, rewrite `.`→`-` inside claude version numbers, and +/// fold an `anthropic/claude-…` prefix). It never consults the user's +/// machine-local model aliases. +/// +/// Every path that submits, uploads, exports as raw data, or persists a model id +/// MUST use this, not [`normalize_model_for_grouping`]. A machine-local alias +/// config must never rewrite the model identity persisted server-side, or usage +/// history would fragment and fork across a user's devices. Graph +/// `ClientContribution` keys also use this so export/raw identity stays stable. +pub fn canonical_model_id(model_id: &str) -> String { + normalize_syntactic(model_id) +} + +/// Local display/grouping model name: [`canonical_model_id`] plus the user's +/// configured model-alias fold. Every local report-grouping surface — the models +/// report, every `GroupBy`, monthly, and hourly — routes through this so name +/// variants fold uniformly for presentation. +/// +/// The alias fold is **presentation only** and must never reach the +/// submit/upload/export/persist path (those use [`canonical_model_id`]), pricing +/// (which resolves the raw message `model_id`), or the message-cache key space. +/// An empty/unset alias config makes this identical to [`canonical_model_id`]. pub fn normalize_model_for_grouping(model_id: &str) -> String { + model_alias::apply_global(normalize_syntactic(model_id)) +} + +/// Structural-only model-name normalization: lowercase, strip a +/// `(reasoning-tier)` suffix, strip a trailing `-YYYYMMDD` date, rewrite `.`→`-` +/// inside claude version numbers, and fold an `anthropic/claude-…` prefix. +/// +/// This is the syntactic half of [`normalize_model_for_grouping`] / +/// [`canonical_model_id`]. It is also used by [`model_alias`] to normalize +/// configured alias keys and values into the same space, so a configured alias +/// matches its model regardless of case, dated suffix, or `.`-vs-`-` spelling. +pub(crate) fn normalize_syntactic(model_id: &str) -> String { let mut name = model_id.to_lowercase(); if let Some(base_model) = strip_parenthesized_reasoning_tier(&name) { @@ -124,6 +168,22 @@ fn retain_for_requested_clients( && sessions::synthetic::matches_synthetic_filter(client, model_id, provider_id)) } +fn prefer_copilot_otel_messages( + mut otel_messages: Vec, + desktop_messages: Vec, +) -> Vec { + let otel_sessions: HashSet = otel_messages + .iter() + .map(|message| message.session_id.clone()) + .collect(); + otel_messages.extend( + desktop_messages + .into_iter() + .filter(|message| !otel_sessions.contains(message.session_id.as_str())), + ); + otel_messages +} + #[derive(Debug, Clone, Default, PartialEq, serde::Serialize)] pub enum GroupBy { Model, @@ -306,11 +366,12 @@ pub struct LocalParseOptions { /// Persistent scanner config loaded from `~/.config/tokscale/settings.json`. /// Defaults to empty when callers don't care about user-configured paths. pub scanner_settings: scanner::ScannerSettings, - /// Skip parsing file-backed session logs whose mtime (unix ms) is older - /// than this. Lets high-frequency callers (live tails) avoid re-parsing - /// an entire history when they only need recent messages — callers align - /// it with `since`. Database-backed sources (SQLite) are always parsed: - /// WAL writes may not touch the main db file's mtime. + /// Skip parsing session sources whose newest relevant mtime (unix ms) is + /// older than this. Lets high-frequency callers (live tails) avoid + /// re-parsing an entire history when they only need recent messages — + /// callers align it with `since`. SQLite lanes are either retained or use + /// source-specific DB/WAL/dependency probes because WAL writes may not touch + /// the main database mtime. pub modified_after: Option, } @@ -501,6 +562,14 @@ pub fn get_home_dir_string(home_dir_option: &Option) -> Result Vec { + if sessions::kimi::is_kimi_code_path(path) { + sessions::kimi::parse_kimi_code_file(path) + } else { + sessions::kimi::parse_kimi_file(path) + } +} + #[allow(dead_code)] fn parse_all_messages_with_pricing( home_dir: &str, @@ -516,6 +585,155 @@ fn parse_all_messages_with_pricing( ) } +#[derive(Debug)] +struct CodexRawCacheOutcome { + messages: Vec, + fallback_timestamp_indices: Vec, + cache_entry: Option, + invalidate_cache: bool, +} + +fn build_codex_cache_entry( + path: &Path, + fingerprint: message_cache::SourceFingerprint, + raw_messages: Vec, + consumed_offset: u64, + state: sessions::codex::CodexParseState, + fallback_timestamp_indices: Vec, +) -> Option { + if fingerprint.size != consumed_offset { + return None; + } + + let codex_incremental = message_cache::build_codex_incremental_cache_with_prefix_hash( + path, + consumed_offset, + state, + fingerprint.content_hash, + )?; + + Some(message_cache::CachedSourceEntry::new( + message_cache::CacheIdentity::for_client(ClientId::Codex), + path, + fingerprint, + raw_messages, + fallback_timestamp_indices, + Some(codex_incremental), + )) +} + +fn parse_full_codex_raw_source(path: &Path) -> CodexRawCacheOutcome { + let parsed = sessions::codex::parse_codex_file_incremental( + path, + 0, + sessions::codex::CodexParseState::default(), + ); + if !parsed.parse_succeeded || parsed.unresolved_model_events { + return CodexRawCacheOutcome { + messages: parsed.messages, + fallback_timestamp_indices: parsed.fallback_timestamp_indices, + cache_entry: None, + invalidate_cache: false, + }; + } + + let cache_entry = message_cache::SourceFingerprint::from_path(path).and_then(|fingerprint| { + build_codex_cache_entry( + path, + fingerprint, + parsed.messages.clone(), + parsed.consumed_offset, + parsed.state, + parsed.fallback_timestamp_indices.clone(), + ) + }); + + CodexRawCacheOutcome { + messages: parsed.messages, + fallback_timestamp_indices: parsed.fallback_timestamp_indices, + cache_entry, + invalidate_cache: false, + } +} + +fn load_or_parse_codex_raw_source( + path: &Path, + source_cache: &message_cache::SourceMessageCache, +) -> CodexRawCacheOutcome { + let identity = message_cache::CacheIdentity::for_client(ClientId::Codex); + let Some(cached) = source_cache.get(identity, path) else { + return parse_full_codex_raw_source(path); + }; + let Some(fingerprint_status) = + message_cache::SourceFingerprint::check_path(path, Some(&cached.fingerprint)) + else { + return parse_full_codex_raw_source(path); + }; + let fingerprint = match fingerprint_status { + message_cache::FingerprintStatus::Unchanged => cached.fingerprint.clone(), + message_cache::FingerprintStatus::Changed(fingerprint) => fingerprint, + }; + let reparse_from_start = |invalidate_cache: bool| { + let mut outcome = parse_full_codex_raw_source(path); + outcome.invalidate_cache = invalidate_cache && outcome.cache_entry.is_none(); + outcome + }; + + if cached.fingerprint == fingerprint { + if message_cache::codex_cache_entry_matches_fingerprint(cached, &fingerprint) { + return CodexRawCacheOutcome { + messages: cached.messages.clone(), + fallback_timestamp_indices: cached.fallback_timestamp_indices.clone(), + cache_entry: None, + invalidate_cache: false, + }; + } + return reparse_from_start(true); + } + + if let Some(codex_incremental) = cached.codex_incremental.as_ref() { + if fingerprint.size > codex_incremental.consumed_offset + && message_cache::codex_prefix_matches(path, codex_incremental) + { + let parsed = sessions::codex::parse_codex_file_incremental( + path, + codex_incremental.consumed_offset, + codex_incremental.state.clone(), + ); + if parsed.parse_succeeded && !parsed.unresolved_model_events { + let mut messages = cached.messages.clone(); + let mut fallback_timestamp_indices = cached.fallback_timestamp_indices.clone(); + let existing_len = messages.len(); + fallback_timestamp_indices.extend( + parsed + .fallback_timestamp_indices + .iter() + .map(|index| existing_len + index), + ); + messages.extend(parsed.messages); + let Some(cache_entry) = build_codex_cache_entry( + path, + fingerprint, + messages.clone(), + parsed.consumed_offset, + parsed.state, + fallback_timestamp_indices.clone(), + ) else { + return reparse_from_start(true); + }; + return CodexRawCacheOutcome { + messages, + fallback_timestamp_indices, + cache_entry: Some(cache_entry), + invalidate_cache: false, + }; + } + } + } + + reparse_from_start(true) +} + // All report consumers (graph/model/monthly/hourly/agents) now fold over // scan_messages_streaming. The materialized path below survives only behind the // public `parse_local_unified_messages` (no in-repo callers — see its footgun @@ -553,55 +771,6 @@ fn parse_all_messages_with_pricing_with_env_strategy( messages } - fn parse_full_log_source( - path: &Path, - pricing: Option<&pricing::PricingService>, - is_headless: bool, - ) -> CachedParseOutcome { - let fallback_timestamp = sessions::utils::file_modified_timestamp_ms(path); - let parsed = sessions::codex::parse_codex_file_incremental( - path, - 0, - sessions::codex::CodexParseState::default(), - ); - let messages = finalize_codex_messages( - parsed.messages.clone(), - pricing, - is_headless, - &parsed.fallback_timestamp_indices, - fallback_timestamp, - ); - if !parsed.parse_succeeded { - return CachedParseOutcome { - messages, - cache_entry: None, - invalidate_cache: false, - }; - } - - if parsed.unresolved_model_events { - return CachedParseOutcome { - messages, - cache_entry: None, - invalidate_cache: false, - }; - } - - let cache_entry = build_codex_cache_entry( - path, - parsed.messages, - parsed.consumed_offset, - parsed.state, - parsed.fallback_timestamp_indices, - ); - - CachedParseOutcome { - messages, - cache_entry, - invalidate_cache: false, - } - } - fn finalize_codex_messages( mut messages: Vec, pricing: Option<&pricing::PricingService>, @@ -621,32 +790,9 @@ fn parse_all_messages_with_pricing_with_env_strategy( messages } - fn build_codex_cache_entry( - path: &Path, - raw_messages: Vec, - consumed_offset: u64, - state: sessions::codex::CodexParseState, - fallback_timestamp_indices: Vec, - ) -> Option { - let fingerprint = message_cache::SourceFingerprint::from_path(path)?; - if fingerprint.size != consumed_offset { - return None; - } - - let codex_incremental = - message_cache::build_codex_incremental_cache(path, consumed_offset, state)?; - - Some(message_cache::CachedSourceEntry::new( - path, - fingerprint, - raw_messages, - fallback_timestamp_indices, - Some(codex_incremental), - )) - } - fn load_or_parse_source_with_fingerprint_and_policy( path: &Path, + identity: message_cache::CacheIdentity, source_cache: &message_cache::SourceMessageCache, pricing: Option<&pricing::PricingService>, fingerprint_from_path: FingerprintFn, @@ -654,9 +800,15 @@ fn parse_all_messages_with_pricing_with_env_strategy( ) -> CachedParseOutcome where F: Fn(&Path) -> (Vec, bool), - FingerprintFn: Fn(&Path) -> Option, + FingerprintFn: Fn( + &Path, + Option<&message_cache::SourceFingerprint>, + ) -> Option, { - let Some(fingerprint) = fingerprint_from_path(path) else { + let cached = source_cache.get(identity, path); + let Some(fingerprint_status) = + fingerprint_from_path(path, cached.map(|entry| &entry.fingerprint)) + else { let (mut messages, _) = parse(path); apply_pricing_to_messages(&mut messages, pricing); return CachedParseOutcome { @@ -666,7 +818,24 @@ fn parse_all_messages_with_pricing_with_env_strategy( }; }; - if let Some(cached) = source_cache.get(path) { + let fingerprint = match fingerprint_status { + message_cache::FingerprintStatus::Unchanged => { + let Some(cached) = cached else { + unreachable!("an uncached source always builds a complete fingerprint") + }; + if !cached.messages.is_empty() { + return CachedParseOutcome { + messages: cached_messages(cached, pricing), + cache_entry: None, + invalidate_cache: false, + }; + } + cached.fingerprint.clone() + } + message_cache::FingerprintStatus::Changed(fingerprint) => fingerprint, + }; + + if let Some(cached) = cached { if cached.fingerprint == fingerprint && !cached.messages.is_empty() { return CachedParseOutcome { messages: cached_messages(cached, pricing), @@ -681,6 +850,7 @@ fn parse_all_messages_with_pricing_with_env_strategy( None } else { Some(message_cache::CachedSourceEntry::new( + identity, path, fingerprint, messages.clone(), @@ -699,6 +869,7 @@ fn parse_all_messages_with_pricing_with_env_strategy( fn load_or_parse_source_with_fingerprint( path: &Path, + identity: message_cache::CacheIdentity, source_cache: &message_cache::SourceMessageCache, pricing: Option<&pricing::PricingService>, fingerprint_from_path: FingerprintFn, @@ -706,10 +877,14 @@ fn parse_all_messages_with_pricing_with_env_strategy( ) -> CachedParseOutcome where F: Fn(&Path) -> Vec, - FingerprintFn: Fn(&Path) -> Option, + FingerprintFn: Fn( + &Path, + Option<&message_cache::SourceFingerprint>, + ) -> Option, { load_or_parse_source_with_fingerprint_and_policy( path, + identity, source_cache, pricing, fingerprint_from_path, @@ -719,6 +894,7 @@ fn parse_all_messages_with_pricing_with_env_strategy( fn load_or_parse_source( path: &Path, + identity: message_cache::CacheIdentity, source_cache: &message_cache::SourceMessageCache, pricing: Option<&pricing::PricingService>, parse: F, @@ -728,15 +904,17 @@ fn parse_all_messages_with_pricing_with_env_strategy( { load_or_parse_source_with_fingerprint( path, + identity, source_cache, pricing, - message_cache::SourceFingerprint::from_path, + message_cache::SourceFingerprint::check_path_samples_only, parse, ) } fn load_or_parse_sqlite_source( path: &Path, + identity: message_cache::CacheIdentity, source_cache: &message_cache::SourceMessageCache, pricing: Option<&pricing::PricingService>, parse: F, @@ -746,9 +924,10 @@ fn parse_all_messages_with_pricing_with_env_strategy( { load_or_parse_source_with_fingerprint( path, + identity, source_cache, pricing, - message_cache::SourceFingerprint::from_sqlite_path, + message_cache::SourceFingerprint::check_sqlite_path, parse, ) } @@ -759,88 +938,19 @@ fn parse_all_messages_with_pricing_with_env_strategy( pricing: Option<&pricing::PricingService>, headless_roots: &[PathBuf], ) -> CachedParseOutcome { - let is_headless = is_headless_path(path, headless_roots); - let Some(fingerprint) = message_cache::SourceFingerprint::from_path(path) else { - return parse_full_log_source(path, pricing, is_headless); - }; - let fallback_timestamp = sessions::utils::file_modified_timestamp_ms(path); - - if let Some(cached) = source_cache.get(path) { - let reparse_from_start = |invalidate_cache: bool| { - let mut outcome = parse_full_log_source(path, pricing, is_headless); - outcome.invalidate_cache = invalidate_cache && outcome.cache_entry.is_none(); - outcome - }; - - if cached.fingerprint == fingerprint { - if message_cache::codex_cache_entry_matches_fingerprint(cached, &fingerprint) { - return CachedParseOutcome { - messages: finalize_codex_messages( - cached.messages.clone(), - pricing, - is_headless, - &cached.fallback_timestamp_indices, - fallback_timestamp, - ), - cache_entry: None, - invalidate_cache: false, - }; - } - - return reparse_from_start(true); - } - - if let Some(codex_incremental) = cached.codex_incremental.as_ref() { - if fingerprint.size > codex_incremental.consumed_offset - && message_cache::codex_prefix_matches(path, codex_incremental) - { - let parsed = sessions::codex::parse_codex_file_incremental( - path, - codex_incremental.consumed_offset, - codex_incremental.state.clone(), - ); - if parsed.parse_succeeded && !parsed.unresolved_model_events { - let mut raw_messages = cached.messages.clone(); - let mut fallback_timestamp_indices = - cached.fallback_timestamp_indices.clone(); - let existing_len = raw_messages.len(); - fallback_timestamp_indices.extend( - parsed - .fallback_timestamp_indices - .iter() - .map(|index| existing_len + index), - ); - raw_messages.extend(parsed.messages.clone()); - let cache_entry = build_codex_cache_entry( - path, - raw_messages.clone(), - parsed.consumed_offset, - parsed.state, - fallback_timestamp_indices.clone(), - ); - let Some(cache_entry) = cache_entry else { - return reparse_from_start(true); - }; - let messages = finalize_codex_messages( - raw_messages, - pricing, - is_headless, - &fallback_timestamp_indices, - fallback_timestamp, - ); - return CachedParseOutcome { - messages, - cache_entry: Some(cache_entry), - invalidate_cache: false, - }; - } - } - } - - return reparse_from_start(true); + let raw = load_or_parse_codex_raw_source(path, source_cache); + let messages = finalize_codex_messages( + raw.messages, + pricing, + is_headless_path(path, headless_roots), + &raw.fallback_timestamp_indices, + sessions::utils::file_modified_timestamp_ms(path), + ); + CachedParseOutcome { + messages, + cache_entry: raw.cache_entry, + invalidate_cache: raw.invalidate_cache, } - - parse_full_log_source(path, pricing, is_headless) } let scan_result = scanner::scan_all_clients_with_scanner_settings( @@ -862,46 +972,64 @@ fn parse_all_messages_with_pricing_with_env_strategy( .opencode_dbs .iter() .map(|db_path| { - load_or_parse_sqlite_source(db_path, &source_cache, pricing, |path| { - sessions::opencode::parse_opencode_sqlite(path) - }) + load_or_parse_sqlite_source( + db_path, + message_cache::CacheIdentity::for_client(ClientId::OpenCode), + &source_cache, + pricing, + sessions::opencode::parse_opencode_sqlite, + ) }) .collect(); let opencode_json_outcomes: Vec = scan_result .get(ClientId::OpenCode) .par_iter() .map(|path| { - load_or_parse_source(path, &source_cache, pricing, |path| { - sessions::opencode::parse_opencode_file(path) - .into_iter() - .collect() - }) + load_or_parse_source( + path, + message_cache::CacheIdentity::for_client(ClientId::OpenCode), + &source_cache, + pricing, + |path| { + sessions::opencode::parse_opencode_file(path) + .into_iter() + .collect() + }, + ) }) .collect(); - let opencode_authoritative: HashSet = opencode_sqlite_outcomes - .iter() - .chain(opencode_json_outcomes.iter()) - .flat_map(|outcome| outcome.messages.iter()) - .filter(|message| message.cost_source == CostSource::ProviderReported) - .filter_map(|message| message.dedup_key.clone()) - .collect(); - let mut opencode_seen: HashSet = HashSet::new(); + let opencode_authoritative = opencode_authoritative_sources( + opencode_sqlite_outcomes + .iter() + .chain(opencode_json_outcomes.iter()) + .flat_map(|outcome| outcome.messages.iter()) + .map(opencode_identity_group), + ); + let mut opencode_selection = OpenCodeSelection::new(opencode_authoritative); - for outcome in opencode_sqlite_outcomes - .into_iter() - .chain(opencode_json_outcomes) - { - all_messages.extend(outcome.messages.into_iter().filter(|message| { - message.dedup_key.as_ref().is_none_or(|key| { - (!opencode_authoritative.contains(key) - || message.cost_source == CostSource::ProviderReported) - && opencode_seen.insert(key.clone()) - }) - })); + for outcome in opencode_sqlite_outcomes { + all_messages.extend( + outcome + .messages + .into_iter() + .filter_map(|message| opencode_selection.select_sqlite(message)), + ); + if let Some(entry) = outcome.cache_entry { + source_cache.insert(entry); + } + } + for outcome in opencode_json_outcomes { + all_messages.extend( + outcome + .messages + .into_iter() + .filter_map(|message| opencode_selection.select_json(message, true)), + ); if let Some(entry) = outcome.cache_entry { source_cache.insert(entry); } } + all_messages.extend(opencode_selection.finish()); let claude_home = PathBuf::from(home_dir); let claude_outcomes: Vec = scan_result @@ -910,11 +1038,13 @@ fn parse_all_messages_with_pricing_with_env_strategy( .map(|path| { load_or_parse_source_with_fingerprint( path, + message_cache::CacheIdentity::for_client(ClientId::Claude), &source_cache, pricing, - |path| { - message_cache::SourceFingerprint::from_claude_code_path_with_home( + |path, cached| { + message_cache::SourceFingerprint::check_claude_code_path_with_home_samples_only( path, + cached, Some(&claude_home), ) }, @@ -962,7 +1092,10 @@ fn parse_all_messages_with_pricing_with_env_strategy( if let Some(entry) = outcome.cache_entry { source_cache.insert(entry); } else if outcome.invalidate_cache { - source_cache.remove(&path); + source_cache.remove( + message_cache::CacheIdentity::for_client(ClientId::Codex), + &path, + ); } } @@ -970,17 +1103,45 @@ fn parse_all_messages_with_pricing_with_env_strategy( .get(ClientId::Copilot) .par_iter() .map(|path| { - load_or_parse_source(path, &source_cache, pricing, |path| { - sessions::copilot::parse_copilot_file(path) - }) + load_or_parse_source( + path, + message_cache::CacheIdentity::for_client(ClientId::Copilot), + &source_cache, + None, + sessions::copilot::parse_copilot_file, + ) }) .collect(); + let mut copilot_otel_messages = Vec::new(); for outcome in copilot_outcomes { - all_messages.extend(outcome.messages); + copilot_otel_messages.extend(outcome.messages); + if let Some(entry) = outcome.cache_entry { + source_cache.insert(entry); + } + } + let mut copilot_desktop_messages = Vec::new(); + if let Some(db_path) = &scan_result.copilot_desktop_db { + let outcome = load_or_parse_source_with_fingerprint( + db_path, + message_cache::CacheIdentity::for_client(ClientId::Copilot), + &source_cache, + None, + message_cache::SourceFingerprint::check_copilot_desktop_path_samples_only, + sessions::copilot_desktop::parse_copilot_desktop_db, + ); + copilot_desktop_messages = outcome.messages; if let Some(entry) = outcome.cache_entry { source_cache.insert(entry); } } + all_messages.extend( + prefer_copilot_otel_messages(copilot_otel_messages, copilot_desktop_messages) + .into_iter() + .map(|mut message| { + apply_pricing_if_available(&mut message, pricing); + message + }), + ); let gemini_outcomes: Vec<(PathBuf, CachedParseOutcome)> = scan_result .get(ClientId::Gemini) @@ -988,9 +1149,10 @@ fn parse_all_messages_with_pricing_with_env_strategy( .map(|path| { let outcome = load_or_parse_source_with_fingerprint_and_policy( path, + message_cache::CacheIdentity::for_client(ClientId::Gemini), &source_cache, pricing, - message_cache::SourceFingerprint::from_path, + message_cache::SourceFingerprint::check_path_samples_only, |path| { let parsed = sessions::gemini::parse_gemini_file_with_cache_status(path); (parsed.messages, parsed.cacheable) @@ -1004,7 +1166,10 @@ fn parse_all_messages_with_pricing_with_env_strategy( if let Some(entry) = outcome.cache_entry { source_cache.insert(entry); } else if outcome.invalidate_cache { - source_cache.remove(&path); + source_cache.remove( + message_cache::CacheIdentity::for_client(ClientId::Gemini), + &path, + ); } } @@ -1012,9 +1177,13 @@ fn parse_all_messages_with_pricing_with_env_strategy( .get(ClientId::Cursor) .par_iter() .map(|path| { - load_or_parse_source(path, &source_cache, pricing, |path| { - sessions::cursor::parse_cursor_file(path) - }) + load_or_parse_source( + path, + message_cache::CacheIdentity::for_client(ClientId::Cursor), + &source_cache, + pricing, + sessions::cursor::parse_cursor_file, + ) }) .collect(); for outcome in cursor_outcomes { @@ -1028,9 +1197,13 @@ fn parse_all_messages_with_pricing_with_env_strategy( .get(ClientId::Warp) .par_iter() .map(|path| { - load_or_parse_source(path, &source_cache, pricing, |path| { - sessions::warp::parse_warp_file(path) - }) + load_or_parse_source( + path, + message_cache::CacheIdentity::for_client(ClientId::Warp), + &source_cache, + pricing, + sessions::warp::parse_warp_file, + ) }) .collect(); for outcome in warp_outcomes { @@ -1044,9 +1217,13 @@ fn parse_all_messages_with_pricing_with_env_strategy( .get(ClientId::Amp) .par_iter() .map(|path| { - load_or_parse_source(path, &source_cache, pricing, |path| { - sessions::amp::parse_amp_file(path) - }) + load_or_parse_source( + path, + message_cache::CacheIdentity::for_client(ClientId::Amp), + &source_cache, + pricing, + sessions::amp::parse_amp_file, + ) }) .collect(); for outcome in amp_outcomes { @@ -1060,9 +1237,13 @@ fn parse_all_messages_with_pricing_with_env_strategy( .get(ClientId::Codebuff) .par_iter() .map(|path| { - load_or_parse_source(path, &source_cache, pricing, |path| { - sessions::codebuff::parse_codebuff_file(path) - }) + load_or_parse_source( + path, + message_cache::CacheIdentity::for_client(ClientId::Codebuff), + &source_cache, + pricing, + sessions::codebuff::parse_codebuff_file, + ) }) .collect(); for outcome in codebuff_outcomes { @@ -1078,9 +1259,10 @@ fn parse_all_messages_with_pricing_with_env_strategy( .map(|path| { load_or_parse_source_with_fingerprint( path, + message_cache::CacheIdentity::for_client(ClientId::Droid), &source_cache, pricing, - message_cache::SourceFingerprint::from_droid_path, + message_cache::SourceFingerprint::check_droid_path_samples_only, sessions::droid::parse_droid_file, ) }) @@ -1096,9 +1278,13 @@ fn parse_all_messages_with_pricing_with_env_strategy( .get(ClientId::OpenClaw) .par_iter() .map(|path| { - load_or_parse_source(path, &source_cache, pricing, |path| { - sessions::openclaw::parse_openclaw_transcript(path) - }) + load_or_parse_source( + path, + message_cache::CacheIdentity::for_client(ClientId::OpenClaw), + &source_cache, + pricing, + sessions::openclaw::parse_openclaw_transcript, + ) }) .collect(); for outcome in openclaw_outcomes { @@ -1112,9 +1298,13 @@ fn parse_all_messages_with_pricing_with_env_strategy( .get(ClientId::Pi) .par_iter() .map(|path| { - load_or_parse_source(path, &source_cache, pricing, |path| { - sessions::pi::parse_pi_file(path) - }) + load_or_parse_source( + path, + message_cache::CacheIdentity::for_client(ClientId::Pi), + &source_cache, + pricing, + sessions::pi::parse_pi_file, + ) }) .collect(); for outcome in pi_outcomes { @@ -1124,34 +1314,104 @@ fn parse_all_messages_with_pricing_with_env_strategy( } } - let kimi_outcomes: Vec = scan_result + let kimi_outcomes: Vec<(bool, CachedParseOutcome)> = scan_result .get(ClientId::Kimi) .par_iter() .map(|path| { - load_or_parse_source_with_fingerprint( - path, - &source_cache, - pricing, - message_cache::SourceFingerprint::from_kimi_path, - sessions::kimi::parse_kimi_file, + ( + sessions::kimi::is_kimi_code_path(path), + load_or_parse_source_with_fingerprint( + path, + message_cache::CacheIdentity::for_client(ClientId::Kimi), + &source_cache, + pricing, + message_cache::SourceFingerprint::check_kimi_path_samples_only, + parse_kimi_source, + ), ) }) .collect(); - for outcome in kimi_outcomes { - all_messages.extend(outcome.messages); + let mut kimi_code_seen: HashSet = HashSet::new(); + for (is_kimi_code, outcome) in kimi_outcomes { + if is_kimi_code { + all_messages.extend( + outcome + .messages + .into_iter() + .filter(|message| should_keep_deduped_message(&mut kimi_code_seen, message)), + ); + } else { + all_messages.extend(outcome.messages); + } if let Some(entry) = outcome.cache_entry { source_cache.insert(entry); } } - // Parse Qwen files - let qwen_outcomes: Vec = scan_result + let junie_outcomes: Vec = scan_result + .get(ClientId::Junie) + .par_iter() + .map(|path| { + load_or_parse_source( + path, + message_cache::CacheIdentity::for_client(ClientId::Junie), + &source_cache, + pricing, + sessions::junie::parse_junie_file, + ) + }) + .collect(); + let mut junie_seen: HashSet = HashSet::new(); + for outcome in junie_outcomes { + all_messages.extend( + outcome + .messages + .into_iter() + .filter(|message| should_keep_deduped_message(&mut junie_seen, message)), + ); + if let Some(entry) = outcome.cache_entry { + source_cache.insert(entry); + } + } + + let opencodereview_outcomes: Vec = scan_result + .get(ClientId::OpenCodeReview) + .par_iter() + .map(|path| { + load_or_parse_source( + path, + message_cache::CacheIdentity::for_client(ClientId::OpenCodeReview), + &source_cache, + pricing, + sessions::opencodereview::parse_opencodereview_file, + ) + }) + .collect(); + let mut opencodereview_seen: HashSet = HashSet::new(); + for outcome in opencodereview_outcomes { + all_messages.extend( + outcome + .messages + .into_iter() + .filter(|message| should_keep_deduped_message(&mut opencodereview_seen, message)), + ); + if let Some(entry) = outcome.cache_entry { + source_cache.insert(entry); + } + } + + // Parse Qwen files + let qwen_outcomes: Vec = scan_result .get(ClientId::Qwen) .par_iter() .map(|path| { - load_or_parse_source(path, &source_cache, pricing, |path| { - sessions::qwen::parse_qwen_file(path) - }) + load_or_parse_source( + path, + message_cache::CacheIdentity::for_client(ClientId::Qwen), + &source_cache, + pricing, + sessions::qwen::parse_qwen_file, + ) }) .collect(); for outcome in qwen_outcomes { @@ -1170,9 +1430,10 @@ fn parse_all_messages_with_pricing_with_env_strategy( // a history-only rewrite invalidates the cache (#741). load_or_parse_source_with_fingerprint( path, + message_cache::CacheIdentity::for_client(ClientId::RooCode), &source_cache, pricing, - message_cache::SourceFingerprint::from_roo_path, + message_cache::SourceFingerprint::check_roo_path_samples_only, sessions::roocode::parse_roocode_file, ) }) @@ -1190,9 +1451,10 @@ fn parse_all_messages_with_pricing_with_env_strategy( .map(|path| { load_or_parse_source_with_fingerprint( path, + message_cache::CacheIdentity::for_client(ClientId::KiloCode), &source_cache, pricing, - message_cache::SourceFingerprint::from_roo_path, + message_cache::SourceFingerprint::check_roo_path_samples_only, sessions::kilocode::parse_kilocode_file, ) }) @@ -1210,9 +1472,10 @@ fn parse_all_messages_with_pricing_with_env_strategy( .map(|path| { load_or_parse_source_with_fingerprint( path, + message_cache::CacheIdentity::for_client(ClientId::Cline), &source_cache, pricing, - message_cache::SourceFingerprint::from_roo_path, + message_cache::SourceFingerprint::check_roo_path_samples_only, sessions::cline::parse_cline_file, ) }) @@ -1230,9 +1493,10 @@ fn parse_all_messages_with_pricing_with_env_strategy( .map(|path| { load_or_parse_source_with_fingerprint( path, + message_cache::CacheIdentity::for_client(ClientId::Jcode), &source_cache, pricing, - message_cache::SourceFingerprint::from_jcode_path, + message_cache::SourceFingerprint::check_jcode_path_samples_only, sessions::jcode::parse_jcode_file, ) }) @@ -1261,9 +1525,13 @@ fn parse_all_messages_with_pricing_with_env_strategy( .get(ClientId::MiMoCode) .par_iter() .map(|path| { - load_or_parse_sqlite_source(path, &source_cache, None, |path| { - sessions::micode::parse_micode_sqlite(path) - }) + load_or_parse_sqlite_source( + path, + message_cache::CacheIdentity::for_client(ClientId::MiMoCode), + &source_cache, + None, + sessions::micode::parse_micode_sqlite, + ) }) .collect(); let mut micode_seen: HashSet = HashSet::new(); @@ -1308,37 +1576,52 @@ fn parse_all_messages_with_pricing_with_env_strategy( .filter(|message| should_keep_deduped_message(&mut gjc_seen, message)), ); - // Grok Build: updates.jsonl + metadata-sibling fingerprint (signals/summary/ - // events; a compaction rollup or late model id must invalidate the cache). - // Cumulative totalTokens deltas land as - // input tokens; pricing treats them as input at the resolved xai rates. + // Grok Build has two representations of the same sessions: legacy + // per-session updates and a global per-inference unified log. Collect both + // raw sets without pricing and apply unified-over-legacy precedence exactly + // once before pricing or any report fold; a downstream aggregate cannot + // subtract covered legacy usage safely. let grok_outcomes: Vec = scan_result .get(ClientId::Grok) .par_iter() .map(|path| { load_or_parse_source_with_fingerprint( path, + message_cache::CacheIdentity::for_client(ClientId::Grok), &source_cache, - pricing, - message_cache::SourceFingerprint::from_grok_path, - sessions::grok::parse_grok_updates_file, + None, + message_cache::SourceFingerprint::check_grok_path_samples_only, + sessions::grok::parse_grok_file, ) }) .collect(); + let mut grok_messages = Vec::new(); for outcome in grok_outcomes { - all_messages.extend(outcome.messages); + grok_messages.extend(outcome.messages); if let Some(entry) = outcome.cache_entry { source_cache.insert(entry); } } + all_messages.extend( + sessions::grok::prefer_unified_log_messages(grok_messages) + .into_iter() + .map(|mut message| { + apply_pricing_if_available(&mut message, pricing); + message + }), + ); let mux_outcomes: Vec = scan_result .get(ClientId::Mux) .par_iter() .map(|path| { - load_or_parse_source(path, &source_cache, pricing, |path| { - sessions::mux::parse_mux_file(path) - }) + load_or_parse_source( + path, + message_cache::CacheIdentity::for_client(ClientId::Mux), + &source_cache, + pricing, + sessions::mux::parse_mux_file, + ) }) .collect(); for outcome in mux_outcomes { @@ -1382,34 +1665,50 @@ fn parse_all_messages_with_pricing_with_env_strategy( } for db_path in scan_result.zed_db_paths() { - let outcome = load_or_parse_sqlite_source(&db_path, &source_cache, pricing, |path| { - sessions::zed::parse_zed_sqlite(path) - }); + let outcome = load_or_parse_sqlite_source( + &db_path, + message_cache::CacheIdentity::for_client(ClientId::Zed), + &source_cache, + pricing, + sessions::zed::parse_zed_sqlite, + ); all_messages.extend(outcome.messages); if let Some(entry) = outcome.cache_entry { source_cache.insert(entry); } } - let kiro_outcomes: Vec = scan_result + // Kiro globalStorage has a precedence relation between self-contained + // snapshots and execution records. Cache each source's raw parser output, + // then suppress only after every file has been collected. The suppression + // result must never be written back into the per-source cache. + let kiro_outcomes: Vec<(PathBuf, CachedParseOutcome)> = scan_result .get(ClientId::Kiro) .par_iter() .map(|path| { - load_or_parse_source_with_fingerprint( - path, - &source_cache, - pricing, - message_cache::SourceFingerprint::from_kiro_path, - sessions::kiro::parse_kiro_file, + ( + path.clone(), + load_or_parse_source_with_fingerprint( + path, + message_cache::CacheIdentity::for_client(ClientId::Kiro), + &source_cache, + None, + message_cache::SourceFingerprint::check_kiro_path_samples_only, + sessions::kiro::parse_kiro_file, + ), ) }) .collect(); - for outcome in kiro_outcomes { - all_messages.extend(outcome.messages); + let mut kiro_sources = Vec::new(); + for (path, outcome) in kiro_outcomes { + kiro_sources.push((path, outcome.messages)); if let Some(entry) = outcome.cache_entry { source_cache.insert(entry); } } + let mut kiro_messages = sessions::kiro::merge_kiro_source_messages(kiro_sources); + apply_pricing_to_messages(&mut kiro_messages, pricing); + all_messages.extend(kiro_messages); if let Some(db_path) = &scan_result.kiro_db { let kiro_db_messages: Vec = sessions::kiro::parse_kiro_sqlite(db_path) @@ -1476,9 +1775,13 @@ fn parse_all_messages_with_pricing_with_env_strategy( if include_synthetic { if let Some(db_path) = &scan_result.synthetic_db { - let outcome = load_or_parse_sqlite_source(db_path, &source_cache, pricing, |path| { - sessions::synthetic::parse_octofriend_sqlite(path) - }); + let outcome = load_or_parse_sqlite_source( + db_path, + message_cache::CacheIdentity::synthetic(), + &source_cache, + pricing, + sessions::synthetic::parse_octofriend_sqlite, + ); all_messages.extend(outcome.messages); if let Some(entry) = outcome.cache_entry { source_cache.insert(entry); @@ -1588,9 +1891,12 @@ fn aggregate_model_usage_entries( group_by: &GroupBy, ) -> Vec { let mut model_map: HashMap = HashMap::new(); + // One alias snapshot for the whole fold so a mid-fold set_model_aliases + // cannot split messages across two grouping configs. + let aliases = model_alias::snapshot_grouping_aliases(); for msg in messages { - let normalized = normalize_model_for_grouping(&msg.model_id); + let normalized = aliases.fold(normalize_syntactic(&msg.model_id)); let (workspace_group_key, workspace_key, workspace_label) = workspace_bucket(&msg); let key = match group_by { GroupBy::Model => normalized.clone(), @@ -1813,91 +2119,326 @@ fn dedup_gate_passes(key: &str, seen: &mut HashSet) -> bool { true } +/// Cross-store OpenCode source identity. A migrated SQLite message can carry +/// both an embedded id and a v1 row/file fallback, so callers compare every +/// primary and alternate key at the same creation timestamp. +#[derive(Debug, Clone, PartialEq, Eq, Hash)] +struct OpenCodeSourceIdentity { + key: String, + timestamp: i64, +} + +impl OpenCodeSourceIdentity { + fn all_from_message(message: &UnifiedMessage) -> Vec { + let mut identities = Vec::new(); + let mut seen = HashSet::new(); + for key in message + .dedup_key + .iter() + .chain(message.dedup_aliases.iter()) + { + if !key.is_empty() && seen.insert(key.clone()) { + identities.push(Self { + key: key.clone(), + timestamp: message.timestamp, + }); + } + } + identities + } +} + +/// Logical OpenCode payload identity used after parser-local fork collapse. +/// +/// Source keys are indexed separately because one migrated request can have +/// both an embedded id and a row/file fallback. Session and workspace remain +/// excluded because true fork copies move between both. Cost is excluded so a +/// provider-reported copy can still replace an estimated one. +#[derive(Debug, Clone, PartialEq, Eq, Hash)] +struct OpenCodePayloadIdentity { + timestamp: i64, + duration_ms: Option, + model_id: String, + provider_id: String, + input: i64, + output: i64, + cache_read: i64, + cache_write: i64, + reasoning: i64, + agent: Option, +} + +impl OpenCodePayloadIdentity { + fn from_message(message: &UnifiedMessage) -> Self { + Self { + timestamp: message.timestamp, + duration_ms: message.duration_ms, + model_id: message.model_id.clone(), + provider_id: message.provider_id.clone(), + input: message.tokens.input, + output: message.tokens.output, + cache_read: message.tokens.cache_read, + cache_write: message.tokens.cache_write, + reasoning: message.tokens.reasoning, + agent: message.agent.clone(), + } + } +} + +fn opencode_identity_group( + message: &UnifiedMessage, +) -> (bool, Vec) { + ( + message.cost_source == CostSource::ProviderReported, + OpenCodeSourceIdentity::all_from_message(message), + ) +} + +fn opencode_authoritative_sources( + groups: impl IntoIterator)>, +) -> HashSet { + let groups: Vec<_> = groups.into_iter().collect(); + let mut authoritative = HashSet::new(); + for (is_provider_reported, sources) in &groups { + if *is_provider_reported { + authoritative.extend(sources.iter().cloned()); + } + } + loop { + let previous_len = authoritative.len(); + for (_, sources) in &groups { + if sources + .iter() + .any(|source| authoritative.contains(source)) + { + authoritative.extend(sources.iter().cloned()); + } + } + if authoritative.len() == previous_len { + return authoritative; + } + } +} + /// Applies one OpenCode source-priority snapshot without requiring sink -/// retraction. Only estimated SQLite messages hidden by that snapshot are -/// retained as fallbacks; JSON messages remain fully streaming. -struct OpenCodeStreamingSelection { - authoritative_snapshot: HashSet, - seen: HashSet, +/// retraction. SQLite fork copies deduplicate only when a source key (primary or +/// alias) and the full payload identity both match; distinct embedded ids remain +/// separate even when every payload field collides. +struct OpenCodeSelection { + authoritative_snapshot: HashSet, + seen_sqlite: HashMap>, + emitted_sources: HashSet, deferred_sqlite: Vec>, - deferred_by_key: HashMap, + deferred_sources: Vec>, + deferred_by_identity: + HashMap>>, + deferred_by_source: HashMap>, } -impl OpenCodeStreamingSelection { - fn new(authoritative_snapshot: HashSet) -> Self { +impl OpenCodeSelection { + fn new(authoritative_snapshot: HashSet) -> Self { Self { authoritative_snapshot, - seen: HashSet::new(), + seen_sqlite: HashMap::new(), + emitted_sources: HashSet::new(), deferred_sqlite: Vec::new(), - deferred_by_key: HashMap::new(), + deferred_sources: Vec::new(), + deferred_by_identity: HashMap::new(), + deferred_by_source: HashMap::new(), + } + } + + fn has_authority(&self, sources: &[OpenCodeSourceIdentity]) -> bool { + sources + .iter() + .any(|source| self.authoritative_snapshot.contains(source)) + } + + fn has_emitted(&self, sources: &[OpenCodeSourceIdentity]) -> bool { + sources + .iter() + .any(|source| self.emitted_sources.contains(source)) + } + + fn connected_deferred_sources( + &self, + sources: &[OpenCodeSourceIdentity], + ) -> Vec { + let mut connected: HashSet<_> = sources.iter().cloned().collect(); + let mut pending: Vec<_> = sources.to_vec(); + while let Some(source) = pending.pop() { + let Some(indices) = self.deferred_by_source.get(&source) else { + continue; + }; + for &index in indices { + for alias in &self.deferred_sources[index] { + if connected.insert(alias.clone()) { + pending.push(alias.clone()); + } + } + } + } + connected.into_iter().collect() + } + + fn mark_emitted(&mut self, sources: &[OpenCodeSourceIdentity]) { + let connected = self.connected_deferred_sources(sources); + self.emitted_sources.extend(connected); + } + + fn mark_sqlite_seen( + &mut self, + sources: &[OpenCodeSourceIdentity], + payload: &OpenCodePayloadIdentity, + ) -> bool { + let duplicate = sources.iter().any(|source| { + self.seen_sqlite + .get(source) + .is_some_and(|payloads| payloads.contains(payload)) + }); + for source in sources { + self.seen_sqlite + .entry(source.clone()) + .or_default() + .insert(payload.clone()); + } + duplicate + } + + fn index_deferred( + &mut self, + index: usize, + sources: &[OpenCodeSourceIdentity], + payload: &OpenCodePayloadIdentity, + ) { + for source in sources { + if !self.deferred_sources[index].contains(source) { + self.deferred_sources[index].push(source.clone()); + } + let identity_indices = self + .deferred_by_identity + .entry(source.clone()) + .or_default() + .entry(payload.clone()) + .or_default(); + if !identity_indices.contains(&index) { + identity_indices.push(index); + } + let source_indices = self.deferred_by_source.entry(source.clone()).or_default(); + if !source_indices.contains(&index) { + source_indices.push(index); + } } } + fn defer_sqlite( + &mut self, + sources: &[OpenCodeSourceIdentity], + payload: &OpenCodePayloadIdentity, + message: UnifiedMessage, + ) { + let index = self.deferred_sqlite.len(); + self.deferred_sqlite.push(Some(message)); + self.deferred_sources.push(Vec::new()); + self.index_deferred(index, sources, payload); + } + + fn first_active(&self, indices: &[usize]) -> Option { + indices + .iter() + .copied() + .find(|&index| self.deferred_sqlite[index].is_some()) + } + + fn find_deferred_identity( + &self, + sources: &[OpenCodeSourceIdentity], + payload: &OpenCodePayloadIdentity, + ) -> Option { + sources.iter().find_map(|source| { + self.deferred_by_identity + .get(source) + .and_then(|payloads| payloads.get(payload)) + .and_then(|indices| self.first_active(indices)) + }) + } + + fn find_deferred_source(&self, sources: &[OpenCodeSourceIdentity]) -> Option { + sources.iter().find_map(|source| { + self.deferred_by_source + .get(source) + .and_then(|indices| self.first_active(indices)) + }) + } + + fn take_deferred(&mut self, index: usize) -> Vec { + self.deferred_sqlite[index] = None; + self.deferred_sources[index].clone() + } + fn select_sqlite(&mut self, message: UnifiedMessage) -> Option { - let Some(key) = message.dedup_key.as_deref() else { + let sources = OpenCodeSourceIdentity::all_from_message(&message); + if sources.is_empty() { return Some(message); - }; - if self.seen.contains(key) { - return None; } - if let Some(&index) = self.deferred_by_key.get(key) { - if message.cost_source != CostSource::ProviderReported { + let payload = OpenCodePayloadIdentity::from_message(&message); + if self.mark_sqlite_seen(&sources, &payload) { + if let Some(index) = self.find_deferred_identity(&sources, &payload) { + self.index_deferred(index, &sources, &payload); + if message.cost_source == CostSource::ProviderReported { + let deferred_sources = self.take_deferred(index); + self.mark_emitted(&deferred_sources); + self.mark_emitted(&sources); + return Some(message); + } return None; } - self.deferred_by_key.remove(key); - self.deferred_sqlite[index] = None; - self.seen.insert(key.to_owned()); - return Some(message); + if self.has_emitted(&sources) { + self.mark_emitted(&sources); + } + return None; } - if self.authoritative_snapshot.contains(key) - && message.cost_source != CostSource::ProviderReported - { - let index = self.deferred_sqlite.len(); - self.deferred_by_key.insert(key.to_owned(), index); - self.deferred_sqlite.push(Some(message)); + if self.has_emitted(&sources) { + self.mark_emitted(&sources); + } + if self.has_authority(&sources) && message.cost_source != CostSource::ProviderReported { + self.defer_sqlite(&sources, &payload, message); return None; } - self.seen.insert(key.to_owned()); + self.mark_emitted(&sources); Some(message) } - fn select_json( - &mut self, - message: UnifiedMessage, - will_emit: bool, - ) -> Option { - let Some(key) = message.dedup_key.as_deref() else { + fn select_json(&mut self, message: UnifiedMessage, will_emit: bool) -> Option { + let sources = OpenCodeSourceIdentity::all_from_message(&message); + if sources.is_empty() { return will_emit.then_some(message); - }; - if self.seen.contains(key) { - return None; } - if self.authoritative_snapshot.contains(key) { + if self.has_authority(&sources) { if message.cost_source != CostSource::ProviderReported || !will_emit { return None; } - self.seen.insert(key.to_owned()); - if let Some(index) = self.deferred_by_key.remove(key) { - self.deferred_sqlite[index] = None; + let payload = OpenCodePayloadIdentity::from_message(&message); + let exact = self.find_deferred_identity(&sources, &payload); + if exact.is_none() && self.has_emitted(&sources) { + return None; + } + if let Some(index) = exact.or_else(|| self.find_deferred_source(&sources)) { + let deferred_sources = self.take_deferred(index); + self.mark_emitted(&deferred_sources); } + self.mark_emitted(&sources); return Some(message); } - self.seen.insert(key.to_owned()); + if self.has_emitted(&sources) { + return None; + } + self.mark_emitted(&sources); will_emit.then_some(message) } fn finish(self) -> impl Iterator { - let Self { - mut seen, - deferred_sqlite, - .. - } = self; - deferred_sqlite.into_iter().flatten().filter(move |message| { - message - .dedup_key - .as_ref() - .is_none_or(|key| seen.insert(key.clone())) - }) + self.deferred_sqlite.into_iter().flatten() } } @@ -1974,6 +2515,9 @@ pub async fn get_monthly_report(options: ReportOptions) -> Result = HashMap::new(); + // One alias snapshot for the whole fold so a mid-fold set_model_aliases + // cannot split messages across two grouping configs. + let aliases = model_alias::snapshot_grouping_aliases(); scan_messages_streaming( &home_dir, &clients, pricing.as_deref(), options.use_env_roots, &options.scanner_settings, @@ -1989,7 +2533,7 @@ pub async fn get_monthly_report(options: ReportOptions) -> Result Result = HashMap::new(); + // One alias snapshot for the whole fold so a mid-fold set_model_aliases + // cannot split messages across two grouping configs. + let aliases = model_alias::snapshot_grouping_aliases(); scan_messages_streaming( &home_dir, &clients, pricing.as_deref(), options.use_env_roots, &options.scanner_settings, @@ -2254,7 +2801,7 @@ pub async fn get_hourly_report(options: ReportOptions) -> Result = HashMap::new(); // ---- OpenCode SQLite + legacy JSON ---- - // The sink cannot retract an estimated duplicate, so pre-scan only the - // authoritative identities before streaming the lane in its existing order. - // Legacy JSON is parsed again below so this pass never retains message bodies. - let mut opencode_authoritative: HashSet = scan_result + // The sink cannot retract an estimated duplicate, so pre-scan source-key + // alias groups and expand provider authority across each connected group. + // Legacy JSON is parsed again below; this pass retains identities, not bodies. + let mut opencode_identity_groups: Vec<_> = scan_result .get(ClientId::OpenCode) .par_iter() .filter_map(|path| sessions::opencode::parse_opencode_file(path)) - .filter(|message| message.cost_source == CostSource::ProviderReported) - .filter_map(|message| message.dedup_key) + .map(|message| opencode_identity_group(&message)) .collect(); for db_path in &scan_result.opencode_dbs { - opencode_authoritative.extend( + opencode_identity_groups.extend( sessions::opencode::parse_opencode_sqlite(db_path) .into_iter() - .filter(|message| message.cost_source == CostSource::ProviderReported) - .filter_map(|message| message.dedup_key), + .map(|message| opencode_identity_group(&message)), ); } + let opencode_authoritative = opencode_authoritative_sources(opencode_identity_groups); - let mut opencode_selection = OpenCodeStreamingSelection::new(opencode_authoritative); + let mut opencode_selection = OpenCodeSelection::new(opencode_authoritative); for db_path in &scan_result.opencode_dbs { for mut message in sessions::opencode::parse_opencode_sqlite(db_path) { apply_pricing_if_available(&mut message, pricing); @@ -2413,8 +2959,31 @@ where let claude_home = PathBuf::from(home_dir); let mut claude_seen: HashSet = HashSet::new(); for path in scan_result.get(ClientId::Claude) { - let fp = message_cache::SourceFingerprint::from_claude_code_path_with_home(path, Some(&claude_home)); - let cache_hit = fp.as_ref().and_then(|fp| source_cache.get(path).filter(|c| &c.fingerprint == fp && !c.messages.is_empty())); + let identity = message_cache::CacheIdentity::for_client(ClientId::Claude); + let cached = source_cache.get(identity, path); + let fingerprint_status = + message_cache::SourceFingerprint::check_claude_code_path_with_home_samples_only( + path, + cached.map(|entry| &entry.fingerprint), + Some(&claude_home), + ); + let (cache_hit, fingerprint) = match fingerprint_status { + Some(message_cache::FingerprintStatus::Unchanged) => { + let cached = + cached.expect("an uncached Claude source always builds a complete fingerprint"); + if cached.messages.is_empty() { + (None, Some(cached.fingerprint.clone())) + } else { + (Some(cached), None) + } + } + Some(message_cache::FingerprintStatus::Changed(fingerprint)) => { + let cache_hit = cached + .filter(|entry| entry.fingerprint == fingerprint && !entry.messages.is_empty()); + (cache_hit, Some(fingerprint)) + } + None => (None, None), + }; if let Some(cached) = cache_hit { for msg in cached.messages.iter() { let mut m = msg.clone(); @@ -2426,6 +2995,18 @@ where } } else { let msgs = sessions::claudecode::parse_claude_file_with_home(path, Some(&claude_home)); + if !msgs.is_empty() { + if let Some(fingerprint) = fingerprint { + source_cache.insert(message_cache::CachedSourceEntry::new( + identity, + path, + fingerprint, + msgs.clone(), + Vec::new(), + None, + )); + } + } for mut m in msgs { apply_pricing_if_available(&mut m, pricing); if !passes_client(&m) { continue; } @@ -2435,41 +3016,34 @@ where } } - // ---- Codex JSONL (cache-aware, headless-aware) ---- + // ---- Codex JSONL (cache-aware, incremental, headless-aware) ---- let mut codex_seen: HashSet = HashSet::new(); for path in scan_result.get(ClientId::Codex) { - let fp = message_cache::SourceFingerprint::from_path(path); - let cache_hit = fp.as_ref().and_then(|fp| source_cache.get(path).filter(|c| &c.fingerprint == fp)); - if let Some(cached) = cache_hit { - let is_headless = is_headless_path(path, &headless_roots); - let fallback_ts = sessions::utils::file_modified_timestamp_ms(path); - let fti = &cached.fallback_timestamp_indices; - for (idx, msg) in cached.messages.iter().enumerate() { - let mut m = msg.clone(); - if fti.contains(&idx) { m.set_timestamp(fallback_ts); } else { m.refresh_derived_fields(); } - apply_pricing_if_available(&mut m, pricing); - apply_headless_agent(&mut m, is_headless); - if !passes_client(&m) { continue; } - let keep = m.dedup_key.as_ref().is_none_or(|k| k.is_empty() || dedup_gate_passes(k, &mut codex_seen)); - if keep && filter(&m) { sink(&m); } - } - } else { - let is_headless = is_headless_path(path, &headless_roots); - let fallback_ts = sessions::utils::file_modified_timestamp_ms(path); - let parsed = sessions::codex::parse_codex_file_incremental( - path, 0, sessions::codex::CodexParseState::default(), + let raw = load_or_parse_codex_raw_source(path, &source_cache); + if let Some(entry) = raw.cache_entry { + source_cache.insert(entry); + } else if raw.invalidate_cache { + source_cache.remove( + message_cache::CacheIdentity::for_client(ClientId::Codex), + path, ); - let mut msgs = parsed.messages; - for idx in &parsed.fallback_timestamp_indices { - if let Some(m) = msgs.get_mut(*idx) { m.set_timestamp(fallback_ts); } - } - for mut m in msgs { - apply_pricing_if_available(&mut m, pricing); - apply_headless_agent(&mut m, is_headless); - if !passes_client(&m) { continue; } - let keep = m.dedup_key.as_ref().is_none_or(|k| k.is_empty() || dedup_gate_passes(k, &mut codex_seen)); - if keep && filter(&m) { sink(&m); } + } + + let is_headless = is_headless_path(path, &headless_roots); + let fallback_timestamp = sessions::utils::file_modified_timestamp_ms(path); + for (index, mut message) in raw.messages.into_iter().enumerate() { + if raw.fallback_timestamp_indices.contains(&index) { + message.set_timestamp(fallback_timestamp); + } else { + message.refresh_derived_fields(); } + apply_pricing_if_available(&mut message, pricing); + apply_headless_agent(&mut message, is_headless); + if !passes_client(&message) { continue; } + let keep = message.dedup_key.as_ref().is_none_or(|key| { + key.is_empty() || dedup_gate_passes(key, &mut codex_seen) + }); + if keep && filter(&message) { sink(&message); } } } @@ -2484,7 +3058,7 @@ where simple_lane!( $client_id, $parse_fn, - message_cache::SourceFingerprint::from_path + message_cache::SourceFingerprint::from_path_samples_only ) }; // Custom fingerprint fn — for sources whose cache validity depends on a @@ -2508,7 +3082,9 @@ where for path in scan_result.get($client_id) { let fp = $fingerprint_fn(path); let cache_hit = fp.as_ref().and_then(|fp| { - source_cache.get(path).filter(|c| c.fingerprint == *fp && !c.messages.is_empty()) + source_cache + .get(message_cache::CacheIdentity::for_client($client_id), path) + .filter(|c| c.fingerprint == *fp && !c.messages.is_empty()) }); if let Some(cached) = cache_hit { for msg in cached.messages.iter() { @@ -2532,7 +3108,12 @@ where if !msgs.is_empty() { if let Some(fp) = $fingerprint_fn(path) { let entry = message_cache::CachedSourceEntry::new( - path, fp, msgs.clone(), Vec::new(), None, + message_cache::CacheIdentity::for_client($client_id), + path, + fp, + msgs.clone(), + Vec::new(), + None, ); source_cache.insert(entry); } @@ -2546,7 +3127,111 @@ where } }}; } - simple_lane!(ClientId::Copilot, sessions::copilot::parse_copilot_file); + // ---- Copilot OTEL + Desktop aggregate (whole-session OTEL authority) ---- + { + let otel_paths = scan_result.get(ClientId::Copilot); + let mut raw_by_path: Vec>> = + (0..otel_paths.len()).map(|_| None).collect(); + let mut miss_paths: Vec<(usize, &PathBuf)> = Vec::new(); + + for (index, path) in otel_paths.iter().enumerate() { + let fingerprint = message_cache::SourceFingerprint::from_path_samples_only(path); + let cache_hit = fingerprint.as_ref().and_then(|fingerprint| { + source_cache + .get( + message_cache::CacheIdentity::for_client(ClientId::Copilot), + path, + ) + .filter(|cached| { + cached.fingerprint == *fingerprint && !cached.messages.is_empty() + }) + }); + if let Some(cached) = cache_hit { + raw_by_path[index] = Some(cached.messages.clone()); + } else { + miss_paths.push((index, path)); + } + } + + let parsed_misses: Vec<(usize, &PathBuf, Vec)> = miss_paths + .par_iter() + .map(|(index, path)| (*index, *path, sessions::copilot::parse_copilot_file(path))) + .collect(); + for (index, path, messages) in parsed_misses { + if !messages.is_empty() { + if let Some(fingerprint) = + message_cache::SourceFingerprint::from_path_samples_only(path) + { + source_cache.insert(message_cache::CachedSourceEntry::new( + message_cache::CacheIdentity::for_client(ClientId::Copilot), + path, + fingerprint, + messages.clone(), + Vec::new(), + None, + )); + } + } + raw_by_path[index] = Some(messages); + } + + let otel_messages = raw_by_path + .into_iter() + .flatten() + .flatten() + .collect::>(); + let desktop_messages = scan_result + .copilot_desktop_db + .as_ref() + .map(|db_path| { + let fingerprint = + message_cache::SourceFingerprint::from_copilot_desktop_path(db_path); + if let Some(cached) = fingerprint.as_ref().and_then(|fingerprint| { + source_cache + .get( + message_cache::CacheIdentity::for_client(ClientId::Copilot), + db_path, + ) + .filter(|cached| { + cached.fingerprint == *fingerprint && !cached.messages.is_empty() + }) + }) { + return cached.messages.clone(); + } + + let messages = sessions::copilot_desktop::parse_copilot_desktop_db(db_path); + if !messages.is_empty() { + if let Some(fingerprint) = fingerprint { + source_cache.insert(message_cache::CachedSourceEntry::new( + message_cache::CacheIdentity::for_client(ClientId::Copilot), + db_path, + fingerprint, + messages.clone(), + Vec::new(), + None, + )); + } + } + messages + }) + .unwrap_or_default(); + + let mut seen_keys: HashSet = HashSet::new(); + for mut message in prefer_copilot_otel_messages(otel_messages, desktop_messages) { + message.refresh_derived_fields(); + reprice_lane_message(&mut message, pricing, false); + if !passes_client(&message) { + continue; + } + let keep = message + .dedup_key + .as_ref() + .is_none_or(|key| key.is_empty() || dedup_gate_passes(key, &mut seen_keys)); + if keep && filter(&message) { + sink(&message); + } + } + } simple_lane!(ClientId::Cursor, sessions::cursor::parse_cursor_file); simple_lane!(ClientId::Warp, sessions::warp::parse_warp_file); simple_lane!(ClientId::Amp, sessions::amp::parse_amp_file); @@ -2560,9 +3245,14 @@ where simple_lane!(ClientId::Pi, sessions::pi::parse_pi_file); simple_lane!( ClientId::Kimi, - sessions::kimi::parse_kimi_file, + parse_kimi_source, message_cache::SourceFingerprint::from_kimi_path ); + simple_lane!(ClientId::Junie, sessions::junie::parse_junie_file); + simple_lane!( + ClientId::OpenCodeReview, + sessions::opencodereview::parse_opencodereview_file + ); simple_lane!(ClientId::Qwen, sessions::qwen::parse_qwen_file); // roo family: fingerprint via from_roo_path so a history-only rewrite of the // sibling api_conversation_history.json (which parse_roo_kilo_file reads for @@ -2587,22 +3277,85 @@ where sessions::jcode::parse_jcode_file, message_cache::SourceFingerprint::from_jcode_path ); - // Grok Build: fingerprint updates.jsonl + every metadata sibling the parser - // reads (signals/summary/events) so a late compaction rollup or sibling-only - // model id invalidates the cache (parse reconciles signals totals). - simple_lane!( - ClientId::Grok, - sessions::grok::parse_grok_updates_file, - message_cache::SourceFingerprint::from_grok_path - ); - // micode is WAL-mode SQLite; fingerprint via from_sqlite_path so a `-wal` - // write invalidates the cache. MiMo Code records an authoritative per-message - // cost (usage.cost), so this lane is cost-guarded (`true`): apply_pricing - // only runs when the embedded cost is absent (`<= 0.0`), never overwriting a - // real embedded cost with a recomputed tokens*rate. Today MiMo models are - // absent from the pricing dataset so unconditional repricing would be a - // no-op, but the guard future-proofs against a priced provider routed - // through MiMo Code / the model being added to the dataset. (#742 Part 2 — + // ---- Grok legacy updates + unified log (batch precedence) ---- + // Cache each raw source independently, but do not emit cache hits early: the + // global unified log can suppress legacy rows from any session file. + { + let grok_paths = scan_result.get(ClientId::Grok); + let mut raw_by_path: Vec>> = + (0..grok_paths.len()).map(|_| None).collect(); + let mut miss_paths: Vec<(usize, &PathBuf)> = Vec::new(); + + for (index, path) in grok_paths.iter().enumerate() { + let fingerprint = message_cache::SourceFingerprint::from_grok_path(path); + let cache_hit = fingerprint.as_ref().and_then(|fingerprint| { + source_cache + .get( + message_cache::CacheIdentity::for_client(ClientId::Grok), + path, + ) + .filter(|cached| { + cached.fingerprint == *fingerprint && !cached.messages.is_empty() + }) + }); + if let Some(cached) = cache_hit { + raw_by_path[index] = Some(cached.messages.clone()); + } else { + miss_paths.push((index, path)); + } + } + + let parsed_misses: Vec<(usize, &PathBuf, Vec)> = miss_paths + .par_iter() + .map(|(index, path)| (*index, *path, sessions::grok::parse_grok_file(path))) + .collect(); + for (index, path, messages) in parsed_misses { + if !messages.is_empty() { + if let Some(fingerprint) = + message_cache::SourceFingerprint::from_grok_path(path) + { + source_cache.insert(message_cache::CachedSourceEntry::new( + message_cache::CacheIdentity::for_client(ClientId::Grok), + path, + fingerprint, + messages.clone(), + Vec::new(), + None, + )); + } + } + raw_by_path[index] = Some(messages); + } + + let raw_messages = raw_by_path + .into_iter() + .flatten() + .flatten() + .collect::>(); + let mut seen_keys: HashSet = HashSet::new(); + for mut message in sessions::grok::prefer_unified_log_messages(raw_messages) { + message.refresh_derived_fields(); + reprice_lane_message(&mut message, pricing, false); + if !passes_client(&message) { + continue; + } + let keep = message + .dedup_key + .as_ref() + .is_none_or(|key| key.is_empty() || dedup_gate_passes(key, &mut seen_keys)); + if keep && filter(&message) { + sink(&message); + } + } + } + // micode is WAL-mode SQLite; fingerprint via from_sqlite_path so a `-wal` + // write invalidates the cache. MiMo Code records an authoritative per-message + // cost (usage.cost), so this lane is cost-guarded (`true`): apply_pricing + // only runs when the embedded cost is absent (`<= 0.0`), never overwriting a + // real embedded cost with a recomputed tokens*rate. Today MiMo models are + // absent from the pricing dataset so unconditional repricing would be a + // no-op, but the guard future-proofs against a priced provider routed + // through MiMo Code / the model being added to the dataset. (#742 Part 2 — // upstream applies this guard in its materialized lane, which is dead code // for us; the streaming lane is where the app actually reprices.) simple_lane!( @@ -2612,11 +3365,82 @@ where true ); simple_lane!(ClientId::Mux, sessions::mux::parse_mux_file); - simple_lane!( - ClientId::Kiro, - sessions::kiro::parse_kiro_file, - message_cache::SourceFingerprint::from_kiro_path - ); + + // ---- Kiro globalStorage files (raw cache + batch suppression) ---- + // Snapshots and successful executions can describe the same conversation. + // Collect every raw source first so suppression runs before pricing, client + // gating, date/report filters, and the sink. Suppressed aggregates are never + // written to the per-source cache, allowing a later execution removal to + // restore the cached snapshot. + { + let kiro_paths = scan_result.get(ClientId::Kiro); + let mut raw_by_path: Vec>> = + (0..kiro_paths.len()).map(|_| None).collect(); + let mut miss_paths: Vec<(usize, &PathBuf)> = Vec::new(); + + for (index, path) in kiro_paths.iter().enumerate() { + let fingerprint = message_cache::SourceFingerprint::from_kiro_path(path); + let cache_hit = fingerprint.as_ref().and_then(|fingerprint| { + source_cache + .get( + message_cache::CacheIdentity::for_client(ClientId::Kiro), + path, + ) + .filter(|cached| { + cached.fingerprint == *fingerprint && !cached.messages.is_empty() + }) + }); + if let Some(cached) = cache_hit { + raw_by_path[index] = Some(cached.messages.clone()); + } else { + miss_paths.push((index, path)); + } + } + + let parsed_misses: Vec<(usize, &PathBuf, Vec)> = miss_paths + .par_iter() + .map(|(index, path)| (*index, *path, sessions::kiro::parse_kiro_file(path))) + .collect(); + for (index, path, messages) in parsed_misses { + if !messages.is_empty() { + if let Some(fingerprint) = message_cache::SourceFingerprint::from_kiro_path(path) { + source_cache.insert(message_cache::CachedSourceEntry::new( + message_cache::CacheIdentity::for_client(ClientId::Kiro), + path, + fingerprint, + messages.clone(), + Vec::new(), + None, + )); + } + } else { + // A changed source that now parses empty must not keep replaying + // a stale non-empty entry through a later same-fingerprint hit. + source_cache.remove( + message_cache::CacheIdentity::for_client(ClientId::Kiro), + path, + ); + } + raw_by_path[index] = Some(messages); + } + + let raw_sources = kiro_paths + .iter() + .cloned() + .zip(raw_by_path.into_iter().map(Option::unwrap_or_default)) + .collect(); + let messages = sessions::kiro::merge_kiro_source_messages(raw_sources); + for mut message in messages { + message.refresh_derived_fields(); + reprice_lane_message(&mut message, pricing, false); + if !passes_client(&message) { + continue; + } + if filter(&message) { + sink(&message); + } + } + } // ---- Gemini (cache-aware with invalidate_cache semantics) ---- // Uses load_or_parse_source_with_fingerprint_and_policy equivalent: @@ -2627,9 +3451,14 @@ where let mut seen_keys: HashSet = HashSet::new(); let mut gemini_miss_paths: Vec<&PathBuf> = Vec::new(); for path in scan_result.get(ClientId::Gemini) { - let fp = message_cache::SourceFingerprint::from_path(path); + let fp = message_cache::SourceFingerprint::from_path_samples_only(path); let cache_hit = fp.as_ref().and_then(|fp| { - source_cache.get(path).filter(|c| c.fingerprint == *fp && !c.messages.is_empty()) + source_cache + .get( + message_cache::CacheIdentity::for_client(ClientId::Gemini), + path, + ) + .filter(|c| c.fingerprint == *fp && !c.messages.is_empty()) }); if let Some(cached) = cache_hit { for msg in cached.messages.iter() { @@ -2646,18 +3475,31 @@ where } let gemini_parsed: Vec<(&PathBuf, sessions::gemini::GeminiParseResult)> = gemini_miss_paths .par_iter() - .map(|path| (*path, sessions::gemini::parse_gemini_file_with_cache_status(path))) + .map(|path| { + ( + *path, + sessions::gemini::parse_gemini_file_with_cache_status(path), + ) + }) .collect(); for (path, parsed) in gemini_parsed { if parsed.cacheable && !parsed.messages.is_empty() { - if let Some(fp) = message_cache::SourceFingerprint::from_path(path) { + if let Some(fp) = message_cache::SourceFingerprint::from_path_samples_only(path) { let entry = message_cache::CachedSourceEntry::new( - path, fp, parsed.messages.clone(), Vec::new(), None, + message_cache::CacheIdentity::for_client(ClientId::Gemini), + path, + fp, + parsed.messages.clone(), + Vec::new(), + None, ); source_cache.insert(entry); } } else if !parsed.cacheable { - source_cache.remove(path); + source_cache.remove( + message_cache::CacheIdentity::for_client(ClientId::Gemini), + path, + ); } for mut m in parsed.messages { apply_pricing_if_available(&mut m, pricing); @@ -2739,7 +3581,14 @@ where // ---- Zed SQLite (cache-aware reference-iterate) ---- for db_path in scan_result.zed_db_paths() { let fp = message_cache::SourceFingerprint::from_sqlite_path(&db_path); - let cache_hit = fp.as_ref().and_then(|fp| source_cache.get(&db_path).filter(|c| &c.fingerprint == fp)); + let cache_hit = fp.as_ref().and_then(|fp| { + source_cache + .get( + message_cache::CacheIdentity::for_client(ClientId::Zed), + &db_path, + ) + .filter(|c| &c.fingerprint == fp) + }); if let Some(cached) = cache_hit { for msg in cached.messages.iter() { let mut m = msg.clone(); @@ -2748,9 +3597,22 @@ where if passes_client(&m) && filter(&m) { sink(&m); } } } else { - for mut m in sessions::zed::parse_zed_sqlite(&db_path) { - apply_pricing_if_available(&mut m, pricing); - if passes_client(&m) && filter(&m) { sink(&m); } + let messages = sessions::zed::parse_zed_sqlite(&db_path); + if !messages.is_empty() { + if let Some(fingerprint) = fp { + source_cache.insert(message_cache::CachedSourceEntry::new( + message_cache::CacheIdentity::for_client(ClientId::Zed), + &db_path, + fingerprint, + messages.clone(), + Vec::new(), + None, + )); + } + } + for mut message in messages { + apply_pricing_if_available(&mut message, pricing); + if passes_client(&message) && filter(&message) { sink(&message); } } } } @@ -2804,9 +3666,13 @@ where && m.dedup_key.as_ref().is_some_and(|k| { existing.dedup_key.as_ref().is_none_or(|ek| k.as_str() > ek.as_str()) })); - if replace { *slot.get_mut() = m; } + if replace { + *slot.get_mut() = m; + } + } + std::collections::hash_map::Entry::Vacant(slot) => { + slot.insert(m); } - std::collections::hash_map::Entry::Vacant(slot) => { slot.insert(m); } } } } @@ -2814,7 +3680,11 @@ where // ---- Synthetic ---- if let Some(db_path) = scan_result.synthetic_db.as_ref().filter(|_| include_synthetic) { let fp = message_cache::SourceFingerprint::from_sqlite_path(db_path); - let cache_hit = fp.as_ref().and_then(|fp| source_cache.get(db_path).filter(|c| &c.fingerprint == fp)); + let cache_hit = fp.as_ref().and_then(|fp| { + source_cache + .get(message_cache::CacheIdentity::synthetic(), db_path) + .filter(|c| &c.fingerprint == fp) + }); if let Some(cached) = cache_hit { for msg in cached.messages.iter() { let mut m = msg.clone(); @@ -2824,10 +3694,26 @@ where if passes_client(&m) && filter(&m) { sink(&m); } } } else { - for mut m in sessions::synthetic::parse_octofriend_sqlite(db_path) { - apply_pricing_if_available(&mut m, pricing); - sessions::synthetic::normalize_synthetic_gateway_fields(&mut m.model_id, &mut m.provider_id); - if passes_client(&m) && filter(&m) { sink(&m); } + let messages = sessions::synthetic::parse_octofriend_sqlite(db_path); + if !messages.is_empty() { + if let Some(fingerprint) = fp { + source_cache.insert(message_cache::CachedSourceEntry::new( + message_cache::CacheIdentity::synthetic(), + db_path, + fingerprint, + messages.clone(), + Vec::new(), + None, + )); + } + } + for mut message in messages { + apply_pricing_if_available(&mut message, pricing); + sessions::synthetic::normalize_synthetic_gateway_fields( + &mut message.model_id, + &mut message.provider_id, + ); + if passes_client(&message) && filter(&message) { sink(&message); } } } } @@ -2840,7 +3726,6 @@ where source_cache.save_if_dirty(); } - async fn generate_graph_with_loaded_pricing( options: ReportOptions, pricing: Option<&pricing::PricingService>, @@ -2962,9 +3847,9 @@ pub fn build_graph_result_from_messages( messages: &[UnifiedMessage], since: Option<&str>, ) -> GraphResult { - let iter = messages.iter().filter(|msg| { - since.is_none_or(|s| msg.date.as_str() >= s) - }); + let iter = messages + .iter() + .filter(|msg| since.is_none_or(|s| msg.date.as_str() >= s)); let contributions = aggregator::fold_messages_iter(iter); aggregator::generate_graph_result(contributions, 0) } @@ -3115,21 +4000,28 @@ fn parse_local_unified_messages_resolved( ); Ok(filter_unified_messages(messages, &options)) } -/// Max mtime (unix ms) across every file the local scan would read — the -/// cheapest "did anything change" probe for callers that cache reports -/// derived from `parse_local_clients` / the unified-message parsers. -/// Database-backed sources contribute both the db file and its `-wal` -/// sidecar (WAL writes may leave the main db file's mtime untouched). -/// Stat failures contribute nothing, so a vanished file alone never -/// invalidates a caller's cache — its replacement or sibling will. +/// Max mtime (unix ms) across every file the local scan would read. This +/// remains the timestamp-shaped probe used by pruning and diagnostics; +/// in-process report caches that must observe source deletion use +/// `local_source_change_token` instead. Database-backed sources contribute both +/// the db file and its `-wal` sidecar. Stat failures contribute nothing, so a +/// vanished file alone does not change this max-mtime value. pub fn latest_source_mtime_ms(options: &LocalParseOptions) -> Result { + let scan_result = scan_local_sources(options)?; + Ok(latest_source_mtime_ms_from_scan(&scan_result)) +} + +fn scan_local_sources(options: &LocalParseOptions) -> Result { let (home_dir, clients) = resolve_local_parse_request(options)?; - let scan_result = scanner::scan_all_clients_with_scanner_settings( + Ok(scanner::scan_all_clients_with_scanner_settings( &home_dir, &clients, options.use_env_roots, &options.scanner_settings, - ); + )) +} + +fn latest_source_mtime_ms_from_scan(scan_result: &scanner::ScanResult) -> u64 { let mut latest: u64 = 0; for files in scan_result.files.iter() { for path in files { @@ -3165,6 +4057,9 @@ pub fn latest_source_mtime_ms(options: &LocalParseOptions) -> Result Result Result Result Result { + use std::hash::{Hash, Hasher}; + + let scan_result = scan_local_sources(options)?; + let mut paths: Vec = scan_result.files.iter().flatten().cloned().collect(); + + let mut dbs = scan_result.opencode_dbs.clone(); + let single_dbs = [ + &scan_result.synthetic_db, + &scan_result.kilo_db, + &scan_result.goose_db, + &scan_result.kiro_db, + ]; + dbs.extend(single_dbs.into_iter().flatten().cloned()); + dbs.extend(scan_result.hermes_db_paths()); + dbs.extend(scan_result.zed_db_paths()); + dbs.extend( + scan_result + .crush_dbs + .iter() + .map(|source| source.db_path.clone()), + ); + dbs.extend(scan_result.get(ClientId::AntigravityCli).iter().cloned()); + dbs.extend(scan_result.get(ClientId::MiMoCode).iter().cloned()); + for db in dbs { + paths.push(db.clone()); + let mut wal = db.into_os_string(); + wal.push("-wal"); + paths.push(PathBuf::from(wal)); + } + if let Some(db_path) = &scan_result.copilot_desktop_db { + probe_readable_dependency(db_path, false).map_err(|error| { + format!( + "Failed to read Copilot Desktop database {}: {error}", + db_path.display() + ) + })?; + paths.push(db_path.clone()); + let mut wal = db_path.clone().into_os_string(); + wal.push("-wal"); + let wal = PathBuf::from(wal); + probe_readable_dependency(&wal, true).map_err(|error| { + format!( + "Failed to read Copilot Desktop WAL {}: {error}", + wal.display() + ) + })?; + paths.push(wal); + paths.extend( + sessions::copilot_desktop::session_state_event_paths(db_path).map_err(|error| { + format!( + "Failed to probe Copilot Desktop session-state dependencies for {}: {error}", + db_path.display() + ) + })?, + ); + } + + for snapshot in scan_result.get(ClientId::Jcode) { + paths.push(message_cache::jcode_journal_path(snapshot)); + } + for source in scan_result.get(ClientId::Grok) { + if source.file_name().and_then(|name| name.to_str()) == Some("updates.jsonl") { + if let Some(parent) = source.parent() { + paths.extend( + message_cache::GROK_METADATA_SIBLINGS + .into_iter() + .map(|name| parent.join(name)), + ); + } + } + } + for client in [ClientId::RooCode, ClientId::KiloCode, ClientId::Cline] { + paths.extend( + scan_result + .get(client) + .iter() + .map(|path| sessions::roocode::history_path_for_ui_messages(path)), + ); + } + paths.extend( + scan_result + .get(ClientId::Droid) + .iter() + .filter_map(|path| sessions::droid::droid_jsonl_path(path)), + ); + paths.extend( + scan_result + .get(ClientId::Kimi) + .iter() + .filter_map(|path| sessions::kimi::kimi_config_path(path)), + ); + paths.extend( + scan_result + .get(ClientId::Kiro) + .iter() + .filter_map(|path| sessions::kiro::kiro_related_messages_path(path)), + ); + + paths.sort(); + paths.dedup(); + + let mut hasher = std::collections::hash_map::DefaultHasher::new(); + for path in paths { + path.hash(&mut hasher); + let state = std::fs::metadata(&path).ok().map(|metadata| { + let modified = metadata + .modified() + .ok() + .and_then(|value| value.duration_since(std::time::UNIX_EPOCH).ok()) + .map(|duration| (duration.as_secs(), duration.subsec_nanos())); + (metadata.len(), modified) + }); + state.hash(&mut hasher); + } + Ok(hasher.finish()) +} + +fn probe_readable_dependency(path: &Path, optional: bool) -> std::io::Result<()> { + match std::fs::File::open(path) { + Ok(_) => Ok(()), + Err(error) if optional && error.kind() == std::io::ErrorKind::NotFound => Ok(()), + Err(error) => Err(error), + } } /// File mtime as unix ms; `None` on any stat failure. @@ -3269,14 +4294,27 @@ fn kiro_source_mtime_ms(session_path: &Path) -> Option { ) } -/// Newest mtime for a Grok session's scanned updates file and every metadata -/// sibling the parser reads (`signals.json` / `summary.json` / `events.jsonl` — -/// kept in lockstep with the fingerprint via `GROK_METADATA_SIBLINGS`). `None` -/// keeps the source during pruning because over-parsing is safer than dropping a -/// session when the updates file or an existing sibling cannot be inspected. -fn grok_source_mtime_ms(updates_path: &Path) -> Option { - let mut latest = file_mtime_ms(updates_path)?; - let Some(parent) = updates_path.parent() else { +fn copilot_desktop_source_mtime_ms(db_path: &Path) -> Option { + probe_readable_dependency(db_path, false).ok()?; + let mut wal = db_path.to_path_buf().into_os_string(); + wal.push("-wal"); + let wal = PathBuf::from(wal); + probe_readable_dependency(&wal, true).ok()?; + let events = sessions::copilot_desktop::session_state_event_paths(db_path).ok()?; + source_with_related_mtime_ms(db_path, std::iter::once(wal).chain(events)) +} + +/// Newest mtime for a Grok source and every file that affects its parse. The +/// unified log is self-contained; legacy updates include metadata siblings kept +/// in lockstep with `SourceFingerprint::from_grok_path`. `None` keeps the source +/// during pruning because over-parsing is safer than silently dropping usage. +fn grok_source_mtime_ms(source_path: &Path) -> Option { + let mut latest = file_mtime_ms(source_path)?; + if source_path.file_name().and_then(|name| name.to_str()) == Some("unified.jsonl") { + return Some(latest); + } + + let Some(parent) = source_path.parent() else { return Some(latest); }; for name in message_cache::GROK_METADATA_SIBLINGS { @@ -3295,11 +4333,18 @@ fn grok_source_mtime_ms(updates_path: &Path) -> Option { /// (WAL writes may not bump the main `.db` mtime), while the jcode lane holds a /// `session_*.json` snapshot whose sibling `.journal.jsonl` is appended between /// snapshot rewrites. Those lanes remain exempt. Roo-family, Droid, legacy Kimi, -/// Kiro CLI, and Grok sources can still be bounded by folding every parser -/// dependency into their newest mtime. +/// Kiro file sources, and Grok sources can still be bounded by folding every +/// parser dependency into their newest mtime. /// Any stat failure keeps the file — over-parsing is safe, silently skipping is /// not. fn prune_scan_result_by_mtime(scan_result: &mut scanner::ScanResult, threshold_ms: u64) { + let copilot_desktop_fresh = scan_result.copilot_desktop_db.as_ref().is_some_and(|path| { + copilot_desktop_source_mtime_ms(path).is_none_or(|mtime| mtime >= threshold_ms) + }); + if !copilot_desktop_fresh { + scan_result.copilot_desktop_db = None; + } + // Lanes whose scanned file's mtime does not reflect a sibling write // (SQLite `-wal` or jcode's `.journal.jsonl`); kept in lockstep with the // sibling probes in `latest_source_mtime_ms`. @@ -3316,6 +4361,12 @@ fn prune_scan_result_by_mtime(scan_result: &mut scanner::ScanResult, threshold_m ClientId::Cline as usize, ]; for (lane, files) in scan_result.files.iter_mut().enumerate() { + if lane == ClientId::Copilot as usize && copilot_desktop_fresh { + // A fresh Desktop aggregate may overlap any older OTEL row from the + // same session. Keep the complete OTEL cohort so the shared selector + // cannot reopen Desktop usage that a full scan suppresses. + continue; + } if db_lanes.contains(&lane) { continue; } @@ -3337,14 +4388,45 @@ fn prune_scan_result_by_mtime(scan_result: &mut scanner::ScanResult, threshold_m continue; } if lane == ClientId::Kiro as usize { + // globalStorage precedence crosses files. If any IDE source changed, + // retain the complete IDE cohort so an older execution can still + // suppress a newer snapshot; CLI sources remain independently prunable. + let keep_global_storage_batch = files + .iter() + .filter(|path| sessions::kiro::is_kiro_global_storage_source(path)) + .any(|path| kiro_source_mtime_ms(path).is_none_or(|mtime| mtime >= threshold_ms)); files.retain(|path| { - kiro_source_mtime_ms(path).is_none_or(|mtime| mtime >= threshold_ms) + if sessions::kiro::is_kiro_global_storage_source(path) { + keep_global_storage_batch + } else { + kiro_source_mtime_ms(path).is_none_or(|mtime| mtime >= threshold_ms) + } }); continue; } if lane == ClientId::Grok as usize { + let is_unified = |path: &Path| { + path.file_name().and_then(|name| name.to_str()) == Some("unified.jsonl") + }; + let is_fresh = + |path: &Path| grok_source_mtime_ms(path).is_none_or(|mtime| mtime >= threshold_ms); + let unified_fresh = files.iter().any(|path| is_unified(path) && is_fresh(path)); + if unified_fresh { + // A fresh global authority file can cover any legacy session and + // also needs those rows for workspace attribution. + continue; + } + + let legacy_fresh = files.iter().any(|path| !is_unified(path) && is_fresh(path)); files.retain(|path| { - grok_source_mtime_ms(path).is_none_or(|mtime| mtime >= threshold_ms) + if is_unified(path) { + // An older authority file must remain available while a + // legacy session is fresh, or live-tail pruning can reopen + // rows that full reports correctly suppress. + legacy_fresh + } else { + is_fresh(path) + } }); continue; } @@ -3398,48 +4480,36 @@ pub fn parse_local_clients(options: LocalParseOptions) -> Result = HashSet::new(); - let mut count: i32 = 0; - - for db_path in &scan_result.opencode_dbs { - let sqlite_msgs: Vec<(String, ParsedMessage)> = - sessions::opencode::parse_opencode_sqlite(db_path) - .into_iter() - .filter_map(|msg| { - let key = msg.dedup_key.clone().unwrap_or_default(); - // Dedup across multiple channel-suffixed dbs: the - // same session can end up in both `opencode.db` and - // `opencode-.db` if the user switches - // channels mid-session. - if !key.is_empty() && !seen.insert(key.clone()) { - return None; - } - Some((key, unified_to_parsed(&msg))) - }) - .collect(); - count += sqlite_msgs.len() as i32; - for (_key, parsed) in sqlite_msgs { - messages.push(parsed); - } - } - - let json_msgs: Vec<(String, ParsedMessage)> = scan_result + let sqlite_messages: Vec = scan_result + .opencode_dbs + .iter() + .flat_map(|db_path| sessions::opencode::parse_opencode_sqlite(db_path)) + .collect(); + let json_messages: Vec = scan_result .get(ClientId::OpenCode) .par_iter() - .filter_map(|path| { - let msg = sessions::opencode::parse_opencode_file(path)?; - let key = msg.dedup_key.clone().unwrap_or_default(); - Some((key, unified_to_parsed(&msg))) - }) + .filter_map(|path| sessions::opencode::parse_opencode_file(path)) .collect(); - let deduped: Vec = json_msgs + let authoritative = opencode_authoritative_sources( + sqlite_messages + .iter() + .chain(json_messages.iter()) + .map(opencode_identity_group), + ); + let mut selection = OpenCodeSelection::new(authoritative); + let mut selected: Vec = sqlite_messages .into_iter() - .filter(|(key, _)| key.is_empty() || seen.insert(key.clone())) - .map(|(_, msg)| msg) + .filter_map(|message| selection.select_sqlite(message)) .collect(); - count += deduped.len() as i32; - messages.extend(deduped); + selected.extend( + json_messages + .into_iter() + .filter_map(|message| selection.select_json(message, true)), + ); + selected.extend(selection.finish()); + let count = selected.len() as i32; + messages.extend(selected.iter().map(unified_to_parsed)); count }; counts.set(ClientId::OpenCode, opencode_count); @@ -3498,16 +4568,21 @@ pub fn parse_local_clients(options: LocalParseOptions) -> Result = scan_result + let copilot_otel_messages: Vec = scan_result .get(ClientId::Copilot) .par_iter() - .flat_map(|path| { - sessions::copilot::parse_copilot_file(path) - .into_iter() - .map(|msg| unified_to_parsed(&msg)) - .collect::>() - }) + .flat_map(|path| sessions::copilot::parse_copilot_file(path)) .collect(); + let copilot_desktop_messages = scan_result + .copilot_desktop_db + .as_ref() + .map(|path| sessions::copilot_desktop::parse_copilot_desktop_db(path)) + .unwrap_or_default(); + let copilot_msgs: Vec = + prefer_copilot_otel_messages(copilot_otel_messages, copilot_desktop_messages) + .iter() + .map(unified_to_parsed) + .collect(); let copilot_count = copilot_msgs.len() as i32; counts.set(ClientId::Copilot, copilot_count); messages.extend(copilot_msgs); @@ -3596,21 +4671,63 @@ pub fn parse_local_clients(options: LocalParseOptions) -> Result = scan_result + let kimi_outcomes: Vec<(bool, Vec)> = scan_result .get(ClientId::Kimi) .par_iter() - .flat_map(|path| { - sessions::kimi::parse_kimi_file(path) + .map(|path| { + ( + sessions::kimi::is_kimi_code_path(path), + parse_kimi_source(path), + ) + }) + .collect(); + let mut kimi_code_seen: HashSet = HashSet::new(); + let kimi_msgs: Vec = kimi_outcomes + .into_iter() + .flat_map(|(is_kimi_code, messages)| { + messages .into_iter() - .map(|msg| unified_to_parsed(&msg)) + .filter(|message| { + !is_kimi_code || should_keep_deduped_message(&mut kimi_code_seen, message) + }) + .map(|message| unified_to_parsed(&message)) .collect::>() }) .collect(); - let kimi_count = kimi_msgs.len() as i32; + let kimi_count = summed_parsed_message_count(&kimi_msgs); counts.set(ClientId::Kimi, kimi_count); messages.extend(kimi_msgs); + let junie_msgs_raw: Vec = scan_result + .get(ClientId::Junie) + .par_iter() + .flat_map(|path| sessions::junie::parse_junie_file(path)) + .collect(); + let mut junie_seen: HashSet = HashSet::new(); + let junie_msgs: Vec = junie_msgs_raw + .into_iter() + .filter(|message| should_keep_deduped_message(&mut junie_seen, message)) + .map(|message| unified_to_parsed(&message)) + .collect(); + let junie_count = summed_parsed_message_count(&junie_msgs); + counts.set(ClientId::Junie, junie_count); + messages.extend(junie_msgs); + + let opencodereview_msgs_raw: Vec = scan_result + .get(ClientId::OpenCodeReview) + .par_iter() + .flat_map(|path| sessions::opencodereview::parse_opencodereview_file(path)) + .collect(); + let mut opencodereview_seen: HashSet = HashSet::new(); + let opencodereview_msgs: Vec = opencodereview_msgs_raw + .into_iter() + .filter(|message| should_keep_deduped_message(&mut opencodereview_seen, message)) + .map(|message| unified_to_parsed(&message)) + .collect(); + let opencodereview_count = summed_parsed_message_count(&opencodereview_msgs); + counts.set(ClientId::OpenCodeReview, opencodereview_count); + messages.extend(opencodereview_msgs); + // Parse Qwen JSONL files in parallel let qwen_msgs: Vec = scan_result .get(ClientId::Qwen) @@ -3718,16 +4835,16 @@ pub fn parse_local_clients(options: LocalParseOptions) -> Result = scan_result + let grok_messages: Vec = scan_result .get(ClientId::Grok) .par_iter() - .flat_map(|path| { - sessions::grok::parse_grok_updates_file(path) - .into_iter() - .map(|msg| unified_to_parsed(&msg)) - .collect::>() - }) + .flat_map(|path| sessions::grok::parse_grok_file(path)) .collect(); + let grok_msgs: Vec = + sessions::grok::prefer_unified_log_messages(grok_messages) + .into_iter() + .map(|message| unified_to_parsed(&message)) + .collect(); let grok_count = summed_parsed_message_count(&grok_msgs); counts.set(ClientId::Grok, grok_count); messages.extend(grok_msgs); @@ -3796,15 +4913,14 @@ pub fn parse_local_clients(options: LocalParseOptions) -> Result = scan_result + let kiro_sources: Vec<(PathBuf, Vec)> = scan_result .get(ClientId::Kiro) .par_iter() - .flat_map(|path| { - sessions::kiro::parse_kiro_file(path) - .into_iter() - .map(|msg| unified_to_parsed(&msg)) - .collect::>() - }) + .map(|path| (path.clone(), sessions::kiro::parse_kiro_file(path))) + .collect(); + let kiro_msgs: Vec = sessions::kiro::merge_kiro_source_messages(kiro_sources) + .iter() + .map(unified_to_parsed) .collect(); let kiro_count = summed_parsed_message_count(&kiro_msgs); counts.set(ClientId::Kiro, kiro_count); @@ -4040,6 +5156,7 @@ pub fn parsed_to_unified(msg: &ParsedMessage, cost: f64) -> UnifiedMessage { message_count: msg.message_count, agent: msg.agent.clone(), dedup_key: None, + dedup_aliases: Vec::new(), is_turn_start: false, } } @@ -4048,20 +5165,25 @@ pub fn parsed_to_unified(msg: &ParsedMessage, cost: f64) -> UnifiedMessage { mod tests { use super::{ agent_bucket_key, aggregate_model_usage_entries, apply_pricing_if_available, + canonical_model_id, clear_model_aliases, copilot_desktop_source_mtime_ms, dedupe_latest_trae_messages, fold_messages_streaming, get_agents_report, get_hourly_report, - get_model_report, get_monthly_report, latest_source_mtime_ms, message_cache, - normalize_model_for_grouping, parse_all_messages_with_pricing_with_env_strategy, - parse_local_clients, + get_model_report, get_monthly_report, latest_source_mtime_ms, local_source_change_token, + message_cache, model_alias_generation, normalize_model_for_grouping, normalize_syntactic, + opencode_authoritative_sources, opencode_identity_group, + parse_all_messages_with_pricing_with_env_strategy, parse_local_clients, parse_local_unified_messages, parsed_to_unified, pricing, prune_scan_result_by_mtime, - reprice_lane_message, retain_for_requested_clients, scan_messages_streaming, scanner, - select_local_parse_pricing, sessions, unified_to_parsed, AgentAccumulator, ClientId, - CostSource, GroupBy, LocalParseOptions, OpenCodeStreamingSelection, ReportOptions, - TokenBreakdown, UnifiedMessage, UNKNOWN_WORKSPACE_LABEL, + register_usage_data_invalidation_hook, reprice_lane_message, retain_for_requested_clients, + scan_messages_streaming, scanner, select_local_parse_pricing, sessions, set_model_aliases, + snapshot_grouping_aliases, unified_to_parsed, AgentAccumulator, ClientId, CostSource, + GroupBy, LocalParseOptions, ModelAliasMap, OpenCodeSelection, OpenCodeSourceIdentity, + ReportOptions, TokenBreakdown, UnifiedMessage, UNKNOWN_WORKSPACE_LABEL, }; - use std::collections::{HashMap, HashSet}; + use bincode::Options; + use std::collections::{BTreeMap, HashMap, HashSet}; use std::io::Write; use std::path::{Path, PathBuf}; use std::str::FromStr; + use std::sync::atomic::{AtomicUsize, Ordering}; use std::sync::Arc; struct EnvGuard(Vec<(&'static str, Option)>); @@ -4136,1032 +5258,2240 @@ mod tests { ) } - #[test] - #[serial_test::serial] - fn test_env_guard_restores_some_and_none_after_panic() { - const KEYS: [&str; 3] = ["HOME", "TOKSCALE_PRICING_CACHE_ONLY", "TOKSCALE_CONFIG_DIR"]; - let _original = EnvGuard::capture(&KEYS); + fn copilot_desktop_test_db(home: &Path) -> (PathBuf, rusqlite::Connection) { + let root = home.join(".copilot"); + std::fs::create_dir_all(&root).unwrap(); + let path = root.join("data.db"); + let connection = rusqlite::Connection::open(&path).unwrap(); + connection + .execute_batch( + r#" + CREATE TABLE sessions ( + id TEXT, + model TEXT, + total_input_tokens INTEGER, + total_output_tokens INTEGER, + total_cached_tokens INTEGER, + total_reasoning_tokens INTEGER, + total_nano_aiu INTEGER, + created_at TEXT, + agent TEXT + ); + "#, + ) + .unwrap(); + (path, connection) + } - unsafe { - std::env::set_var("HOME", "/tmp/tokscale-env-guard-home-before"); - std::env::remove_var("TOKSCALE_PRICING_CACHE_ONLY"); - std::env::set_var( - "TOKSCALE_CONFIG_DIR", - "/tmp/tokscale-env-guard-config-before", - ); - } - let first = std::panic::catch_unwind(|| { - let _guard = EnvGuard::set(&[ - ( - "HOME", - std::ffi::OsStr::new("/tmp/tokscale-env-guard-home-during"), - ), - ("TOKSCALE_PRICING_CACHE_ONLY", std::ffi::OsStr::new("1")), - ( - "TOKSCALE_CONFIG_DIR", - std::ffi::OsStr::new("/tmp/tokscale-env-guard-config-during"), - ), - ]); - panic!("exercise EnvGuard unwinding"); + #[allow(clippy::too_many_arguments)] + fn insert_copilot_desktop_session( + connection: &rusqlite::Connection, + session_id: &str, + model: &str, + input: i64, + output: i64, + cache_read: i64, + reasoning: i64, + created_at: &str, + agent: Option<&str>, + ) { + connection + .execute( + r#" + INSERT INTO sessions ( + id, model, total_input_tokens, total_output_tokens, + total_cached_tokens, total_reasoning_tokens, total_nano_aiu, + created_at, agent + ) VALUES (?1, ?2, ?3, ?4, ?5, ?6, 0, ?7, ?8) + "#, + rusqlite::params![ + session_id, model, input, output, cache_read, reasoning, created_at, agent + ], + ) + .unwrap(); + } + + fn write_copilot_otel_message( + home: &Path, + session_id: &str, + unix_seconds: i64, + input: i64, + output: i64, + agent: Option<&str>, + ) -> PathBuf { + let root = home.join(".copilot/otel"); + std::fs::create_dir_all(&root).unwrap(); + let path = root.join("copilot.jsonl"); + let record = serde_json::json!({ + "type": "span", + "traceId": format!("trace-{session_id}"), + "spanId": format!("span-{session_id}"), + "name": "chat gpt-5.4-mini", + "startTime": [unix_seconds, 0], + "attributes": { + "gen_ai.operation.name": "chat", + "gen_ai.response.model": "gpt-5.4-mini", + "gen_ai.conversation.id": session_id, + "gen_ai.usage.input_tokens": input, + "gen_ai.usage.output_tokens": output, + "gen_ai.agent.id": agent, + } }); - assert!(first.is_err()); - assert_eq!( - std::env::var_os("HOME"), - Some(std::ffi::OsString::from( - "/tmp/tokscale-env-guard-home-before" - )) + let mut file = std::fs::OpenOptions::new() + .create(true) + .append(true) + .open(&path) + .unwrap(); + writeln!(file, "{record}").unwrap(); + path + } + + fn write_copilot_desktop_events( + home: &Path, + session_id: &str, + model: &str, + workspace: &str, + ) -> PathBuf { + let root = home.join(".copilot/session-state").join(session_id); + std::fs::create_dir_all(&root).unwrap(); + let path = root.join("events.jsonl"); + std::fs::write( + &path, + format!( + "{{\"type\":\"session.start\",\"data\":{{\"context\":{{\"cwd\":\"{workspace}\"}}}}}}\n{{\"type\":\"session.model_change\",\"data\":{{\"newModel\":\"{model}\"}}}}\n" + ), + ) + .unwrap(); + path + } + + fn copilot_local_options(home: &Path) -> LocalParseOptions { + LocalParseOptions { + home_dir: Some(home.to_string_lossy().into_owned()), + use_env_roots: false, + clients: Some(vec!["copilot".to_string()]), + ..Default::default() + } + } + + fn copilot_report_options(home: &Path) -> ReportOptions { + ReportOptions { + home_dir: Some(home.to_string_lossy().into_owned()), + use_env_roots: false, + clients: Some(vec!["copilot".to_string()]), + ..Default::default() + } + } + + #[test] + fn copilot_desktop_selector_prefers_otel_per_session() { + let message = |session_id: &str, dedup_key: &str, input: i64| { + UnifiedMessage::new_with_dedup( + "copilot", + "gpt-5.4-mini", + "openai", + session_id, + 1_782_909_296_000, + TokenBreakdown { + input, + ..Default::default() + }, + 0.0, + Some(dedup_key.to_string()), + ) + }; + let selected = super::prefer_copilot_otel_messages( + vec![message("shared", "otel:shared", 10)], + vec![ + message("shared", "copilot-desktop:shared", 100), + message("desktop-only", "copilot-desktop:desktop-only", 20), + ], ); - assert_eq!(std::env::var_os("TOKSCALE_PRICING_CACHE_ONLY"), None); - assert_eq!( - std::env::var_os("TOKSCALE_CONFIG_DIR"), - Some(std::ffi::OsString::from( - "/tmp/tokscale-env-guard-config-before" - )) + + assert_eq!(selected.len(), 2); + assert!(selected + .iter() + .any(|message| message.dedup_key.as_deref() == Some("otel:shared"))); + assert!(selected.iter().any(|message| { + message.dedup_key.as_deref() == Some("copilot-desktop:desktop-only") + })); + assert!(!selected + .iter() + .any(|message| { message.dedup_key.as_deref() == Some("copilot-desktop:shared") })); + } + + #[test] + #[serial_test::serial] + fn copilot_desktop_authority_precedes_date_filters_in_all_consumers() { + let source_home = tempfile::TempDir::new().unwrap(); + let cache_home = tempfile::TempDir::new().unwrap(); + let _env = EnvGuard::set(&[ + ("HOME", cache_home.path().as_os_str()), + ("TOKSCALE_CONFIG_DIR", cache_home.path().as_os_str()), + ("TOKSCALE_PRICING_CACHE_ONLY", std::ffi::OsStr::new("1")), + ]); + let (_, connection) = copilot_desktop_test_db(source_home.path()); + insert_copilot_desktop_session( + &connection, + "shared", + "gpt-5.4-mini", + 100, + 10, + 0, + 0, + "2026-07-01T12:34:56Z", + Some("github.copilot.default"), + ); + drop(connection); + write_copilot_otel_message( + source_home.path(), + "shared", + 1_735_689_600, + 10, + 1, + Some("github.copilot.default"), ); - unsafe { - std::env::remove_var("HOME"); - std::env::set_var("TOKSCALE_PRICING_CACHE_ONLY", "before"); - std::env::remove_var("TOKSCALE_CONFIG_DIR"); - } - let second = std::panic::catch_unwind(|| { - let _guard = EnvGuard::set(&[ - ( - "HOME", - std::ffi::OsStr::new("/tmp/tokscale-env-guard-home-during"), - ), - ("TOKSCALE_PRICING_CACHE_ONLY", std::ffi::OsStr::new("1")), - ( - "TOKSCALE_CONFIG_DIR", - std::ffi::OsStr::new("/tmp/tokscale-env-guard-config-during"), - ), - ]); - panic!("exercise inverse EnvGuard unwinding"); - }); - assert!(second.is_err()); - assert_eq!(std::env::var_os("HOME"), None); - assert_eq!( - std::env::var_os("TOKSCALE_PRICING_CACHE_ONLY"), - Some(std::ffi::OsString::from("before")) + let mut options = copilot_local_options(source_home.path()); + options.since = Some("2026-07-01".to_string()); + options.until = Some("2026-07-01".to_string()); + let home = source_home.path().to_str().unwrap(); + let clients = ["copilot".to_string()]; + let materialized = + super::parse_local_unified_messages_resolved(options.clone(), home, &clients, None) + .unwrap(); + assert!(materialized.is_empty()); + + let mut streamed = Vec::new(); + scan_messages_streaming( + home, + &clients, + None, + false, + &scanner::ScannerSettings::default(), + &|message| message.date == "2026-07-01", + &mut |message| streamed.push(message.clone()), ); - assert_eq!(std::env::var_os("TOKSCALE_CONFIG_DIR"), None); + assert!(streamed.is_empty()); + + let counted = parse_local_clients(options).unwrap(); + assert!(counted.messages.is_empty()); + assert_eq!(counted.counts.get(ClientId::Copilot), 1); } #[test] #[serial_test::serial] - fn test_empty_reports_normalize_total_cost_to_positive_zero() { + fn copilot_desktop_cache_tracks_events_and_wal_updates() { let source_home = tempfile::TempDir::new().unwrap(); let cache_home = tempfile::TempDir::new().unwrap(); let _env = EnvGuard::set(&[ ("HOME", cache_home.path().as_os_str()), + ("TOKSCALE_CONFIG_DIR", cache_home.path().as_os_str()), ("TOKSCALE_PRICING_CACHE_ONLY", std::ffi::OsStr::new("1")), ]); + let (db_path, connection) = copilot_desktop_test_db(source_home.path()); + connection + .pragma_update(None, "journal_mode", "WAL") + .unwrap(); + insert_copilot_desktop_session( + &connection, + "session-1", + "auto", + 100, + 10, + 0, + 0, + "2026-07-01T12:34:56Z", + Some("github.copilot.default"), + ); + let home = source_home.path().to_str().unwrap(); + let clients = ["copilot".to_string()]; + let scan_result = scan_test_sources( + source_home.path(), + &clients, + &scanner::ScannerSettings::default(), + ); + let cache_db_path = scanner_spelling( + std::slice::from_ref( + scan_result + .copilot_desktop_db + .as_ref() + .expect("Copilot Desktop database must be discovered"), + ), + &db_path, + ); - let options = ReportOptions { - home_dir: Some(source_home.path().to_string_lossy().into_owned()), - use_env_roots: false, - clients: Some(vec!["opencode".to_string()]), - ..Default::default() - }; - let runtime = tokio::runtime::Builder::new_current_thread() - .enable_all() - .build() + let initial_fingerprint = + message_cache::SourceFingerprint::from_copilot_desktop_path(&db_path).unwrap(); + let initial_token = + local_source_change_token(&copilot_local_options(source_home.path())).unwrap(); + let cold = parse_all_messages_with_pricing(home, &clients, None); + let warm = parse_all_messages_with_pricing(home, &clients, None); + assert_eq!(cold, warm); + assert_eq!(warm.len(), 1); + assert_eq!(warm[0].model_id, "auto"); + + let events = write_copilot_desktop_events( + source_home.path(), + "session-1", + "claude-sonnet-4-5", + "/tmp/copilot-workspace", + ); + let event_fingerprint = + message_cache::SourceFingerprint::from_copilot_desktop_path(&db_path).unwrap(); + assert_ne!(event_fingerprint, initial_fingerprint); + assert_ne!( + local_source_change_token(&copilot_local_options(source_home.path())).unwrap(), + initial_token + ); + let event_refresh = parse_all_messages_with_pricing(home, &clients, None); + assert_eq!(event_refresh.len(), 1); + assert_eq!(event_refresh[0].model_id, "claude-sonnet-4-5"); + assert_eq!( + event_refresh[0].workspace_label.as_deref(), + Some("copilot-workspace") + ); + + let event_time = std::time::UNIX_EPOCH + std::time::Duration::from_secs(1_800_000_000); + std::fs::File::options() + .write(true) + .open(&events) + .unwrap() + .set_modified(event_time) .unwrap(); - let totals = [ - runtime - .block_on(get_model_report(options.clone())) - .unwrap() - .total_cost, - runtime - .block_on(get_monthly_report(options.clone())) - .unwrap() - .total_cost, - runtime - .block_on(get_hourly_report(options.clone())) - .unwrap() - .total_cost, - runtime - .block_on(get_agents_report(options)) - .unwrap() - .total_cost, - ]; + assert_eq!( + latest_source_mtime_ms(&copilot_local_options(source_home.path())).unwrap(), + 1_800_000_000_000 + ); - for total in totals { - assert_eq!(total.to_bits(), 0.0f64.to_bits()); - } + let before_wal = + message_cache::SourceFingerprint::from_copilot_desktop_path(&db_path).unwrap(); + insert_copilot_desktop_session( + &connection, + "session-2", + "gpt-5.4-mini", + 20, + 2, + 0, + 0, + "2026-07-01T13:00:00Z", + None, + ); + let after_wal = + message_cache::SourceFingerprint::from_copilot_desktop_path(&db_path).unwrap(); + assert_ne!(after_wal, before_wal); + let wal_refresh = parse_all_messages_with_pricing(home, &clients, None); + assert_eq!(wal_refresh.len(), 2); + + let cache = message_cache::SourceMessageCache::load(); + assert!(cache + .get( + message_cache::CacheIdentity::for_client(ClientId::Copilot), + &cache_db_path, + ) + .is_some_and(|entry| entry.messages.iter().all(|message| message.cost == 0.0))); } - fn make_opencode_selection_message(key: &str, cost: f64, source: CostSource) -> UnifiedMessage { - let mut message = UnifiedMessage::new_with_dedup( - "opencode", "gpt-4o", "openai", "oc-session", 1_733_011_200_000, - TokenBreakdown { input: 10, output: 5, cache_read: 0, cache_write: 0, reasoning: 0 }, - cost, Some(key.to_string()), + #[test] + fn copilot_desktop_modified_after_keeps_otel_suppressors_and_fails_open() { + let source_home = tempfile::TempDir::new().unwrap(); + let (db_path, connection) = copilot_desktop_test_db(source_home.path()); + drop(connection); + let otel_path = + write_copilot_otel_message(source_home.path(), "shared", 1_735_689_600, 10, 1, None); + let events = write_copilot_desktop_events( + source_home.path(), + "shared", + "gpt-5.4-mini", + "/tmp/workspace", ); - match source { - CostSource::ProviderReported => message.mark_provider_reported_cost(), - CostSource::Estimated => message.mark_estimated_cost(), - CostSource::Unknown => {} - } - message + let stale_time = std::time::UNIX_EPOCH + std::time::Duration::from_secs(1_700_000_000); + let fresh_time = std::time::UNIX_EPOCH + std::time::Duration::from_secs(1_700_086_400); + std::fs::File::options() + .write(true) + .open(&db_path) + .unwrap() + .set_modified(stale_time) + .unwrap(); + std::fs::File::options() + .write(true) + .open(&otel_path) + .unwrap() + .set_modified(stale_time) + .unwrap(); + std::fs::File::options() + .write(true) + .open(&events) + .unwrap() + .set_modified(fresh_time) + .unwrap(); + + let mut event_fresh = scanner::ScanResult { + copilot_desktop_db: Some(db_path.clone()), + ..Default::default() + }; + event_fresh + .get_mut(ClientId::Copilot) + .push(otel_path.clone()); + prune_scan_result_by_mtime(&mut event_fresh, 1_700_043_200_000); + assert_eq!(event_fresh.copilot_desktop_db.as_ref(), Some(&db_path)); + assert_eq!(event_fresh.get(ClientId::Copilot), &vec![otel_path.clone()]); + + std::fs::File::options() + .write(true) + .open(&events) + .unwrap() + .set_modified(stale_time) + .unwrap(); + let mut all_stale = scanner::ScanResult { + copilot_desktop_db: Some(db_path), + ..Default::default() + }; + all_stale.get_mut(ClientId::Copilot).push(otel_path); + prune_scan_result_by_mtime(&mut all_stale, 1_700_043_200_000); + assert!(all_stale.copilot_desktop_db.is_none()); + assert!(all_stale.get(ClientId::Copilot).is_empty()); + + let mut stat_failure = scanner::ScanResult { + copilot_desktop_db: Some(source_home.path().join("missing.db")), + ..Default::default() + }; + stat_failure + .get_mut(ClientId::Copilot) + .push(source_home.path().join("missing-otel.jsonl")); + prune_scan_result_by_mtime(&mut stat_failure, u64::MAX); + assert!(stat_failure.copilot_desktop_db.is_some()); + assert_eq!(stat_failure.get(ClientId::Copilot).len(), 1); } #[test] - fn test_opencode_streaming_selection_flushes_snapshot_fallback_on_json_drift() { - // A missing file and an invalid file both produce no second-pass message; - // a downgraded file produces an estimated message. All must flush SQLite. - for second_pass in [None, None, Some(CostSource::Estimated)] { - let key = "snapshot-authoritative"; - let mut selection = OpenCodeStreamingSelection::new(HashSet::from([key.to_string()])); - assert!(selection.select_sqlite(make_opencode_selection_message( - key, 0.25, CostSource::Estimated, - )).is_none()); - if let Some(source) = second_pass { - assert!(selection.select_json(make_opencode_selection_message( - key, 0.0, source, - ), true).is_none()); - } - let selected: Vec<_> = selection.finish().collect(); - assert_eq!(selected.len(), 1); - assert_eq!(selected[0].cost, 0.25); - assert_eq!(selected[0].cost_source, CostSource::Estimated); - } + fn copilot_desktop_dependency_probe_failure_fails_open() { + let source_home = tempfile::TempDir::new().unwrap(); + let (db_path, connection) = copilot_desktop_test_db(source_home.path()); + drop(connection); + let session_state = source_home.path().join(".copilot/session-state"); + std::fs::create_dir_all(&session_state).unwrap(); + std::fs::File::create(session_state.join("not-a-directory")).unwrap(); + let otel_path = + write_copilot_otel_message(source_home.path(), "shared", 1_735_689_600, 10, 1, None); + + assert!(message_cache::SourceFingerprint::from_copilot_desktop_path(&db_path).is_none()); + assert!(local_source_change_token(&copilot_local_options(source_home.path())).is_err()); + + let mut scan_result = scanner::ScanResult { + copilot_desktop_db: Some(db_path.clone()), + ..Default::default() + }; + scan_result + .get_mut(ClientId::Copilot) + .push(otel_path.clone()); + prune_scan_result_by_mtime(&mut scan_result, u64::MAX); + + assert_eq!(scan_result.copilot_desktop_db.as_ref(), Some(&db_path)); + assert_eq!(scan_result.get(ClientId::Copilot), &vec![otel_path]); } + #[cfg(unix)] #[test] - fn test_opencode_streaming_selection_replaces_deferred_sqlite_estimate() { - let key = "sqlite-authoritative-replacement"; - let mut selection = OpenCodeStreamingSelection::new(HashSet::from([key.to_string()])); - assert!(selection - .select_sqlite(make_opencode_selection_message( - key, 0.25, CostSource::Estimated, - )) - .is_none()); - - let selected = selection - .select_sqlite(make_opencode_selection_message( - key, 0.50, CostSource::ProviderReported, - )) - .expect("a later authoritative SQLite message must replace the fallback"); - assert_eq!(selected.cost, 0.50); - assert_eq!(selected.cost_source, CostSource::ProviderReported); - assert_eq!(selection.finish().count(), 0); - } + fn copilot_desktop_unreadable_event_bypasses_caches_and_pruning() { + use std::os::unix::fs::PermissionsExt; - #[test] - fn test_opencode_streaming_selection_keeps_first_sqlite_estimate() { - let key = "sqlite-estimated-first-wins"; - let mut selection = OpenCodeStreamingSelection::new(HashSet::from([key.to_string()])); - assert!(selection - .select_sqlite(make_opencode_selection_message( - key, 0.25, CostSource::Estimated, - )) - .is_none()); - assert!(selection - .select_sqlite(make_opencode_selection_message( - key, 0.50, CostSource::Estimated, - )) - .is_none()); + let source_home = tempfile::TempDir::new().unwrap(); + let (db_path, connection) = copilot_desktop_test_db(source_home.path()); + insert_copilot_desktop_session( + &connection, + "shared", + "auto", + 100, + 10, + 0, + 0, + "2026-07-01T12:34:56Z", + None, + ); + drop(connection); + let events = write_copilot_desktop_events( + source_home.path(), + "shared", + "gpt-5.4-mini", + "/tmp/workspace", + ); + let otel_path = + write_copilot_otel_message(source_home.path(), "shared", 1_735_689_600, 10, 1, None); + assert!(message_cache::SourceFingerprint::from_copilot_desktop_path(&db_path).is_some()); + assert!(local_source_change_token(&copilot_local_options(source_home.path())).is_ok()); - let selected: Vec<_> = selection.finish().collect(); - assert_eq!(selected.len(), 1); - assert_eq!(selected[0].cost, 0.25); - assert_eq!(selected[0].cost_source, CostSource::Estimated); - } + std::fs::set_permissions(&events, std::fs::Permissions::from_mode(0o000)).unwrap(); + assert!(message_cache::SourceFingerprint::from_copilot_desktop_path(&db_path).is_none()); + assert!(local_source_change_token(&copilot_local_options(source_home.path())).is_err()); - #[test] - fn test_opencode_streaming_selection_keeps_first_sqlite_authority() { - let key = "sqlite-authoritative-first-wins"; - let mut selection = OpenCodeStreamingSelection::new(HashSet::from([key.to_string()])); - let first = selection - .select_sqlite(make_opencode_selection_message( - key, 0.50, CostSource::ProviderReported, - )) - .expect("the first authoritative SQLite message must be selected"); - assert_eq!(first.cost, 0.50); - assert!(selection - .select_sqlite(make_opencode_selection_message( - key, 0.75, CostSource::ProviderReported, - )) - .is_none()); - assert_eq!(selection.finish().count(), 0); - } + let mut scan_result = scanner::ScanResult { + copilot_desktop_db: Some(db_path.clone()), + ..Default::default() + }; + scan_result + .get_mut(ClientId::Copilot) + .push(otel_path.clone()); + prune_scan_result_by_mtime(&mut scan_result, u64::MAX); + std::fs::set_permissions(&events, std::fs::Permissions::from_mode(0o600)).unwrap(); - #[test] - fn test_opencode_streaming_selection_keeps_fallback_until_json_is_emitted() { - let key = "filtered-authoritative"; - let mut selection = OpenCodeStreamingSelection::new(HashSet::from([key.to_string()])); - assert!(selection.select_sqlite(make_opencode_selection_message( - key, 0.25, CostSource::Estimated, - )).is_none()); - assert!(selection.select_json(make_opencode_selection_message( - key, 0.50, CostSource::ProviderReported, - ), false).is_none()); - let selected: Vec<_> = selection.finish().collect(); - assert_eq!(selected.len(), 1); - assert_eq!(selected[0].cost, 0.25); + assert_eq!(scan_result.copilot_desktop_db.as_ref(), Some(&db_path)); + assert_eq!(scan_result.get(ClientId::Copilot), &vec![otel_path]); } + #[cfg(unix)] #[test] - fn test_opencode_streaming_selection_does_not_double_count_new_authority() { - let key = "newly-authoritative"; - let mut selection = OpenCodeStreamingSelection::new(HashSet::new()); - assert!(selection.select_sqlite(make_opencode_selection_message( - key, 0.25, CostSource::Estimated, - )).is_some()); - assert!(selection.select_json(make_opencode_selection_message( - key, 0.50, CostSource::ProviderReported, - ), true).is_none()); - assert_eq!(selection.finish().count(), 0); - } + #[serial_test::serial] + fn copilot_desktop_unreadable_wal_never_reuses_db_only_cache() { + use std::os::unix::fs::PermissionsExt; - #[test] - fn test_opencode_streaming_selection_prefers_snapshot_authority() { - let key = "stable-authoritative"; - let mut selection = OpenCodeStreamingSelection::new(HashSet::from([key.to_string()])); - assert!(selection.select_sqlite(make_opencode_selection_message( - key, 0.25, CostSource::Estimated, - )).is_none()); - let json = selection.select_json(make_opencode_selection_message( - key, 0.50, CostSource::ProviderReported, - ), true).unwrap(); - assert_eq!(json.cost, 0.50); - assert_eq!(json.cost_source, CostSource::ProviderReported); - assert_eq!(selection.finish().count(), 0); - } + let source_home = tempfile::TempDir::new().unwrap(); + let cache_home = tempfile::TempDir::new().unwrap(); + let _env = EnvGuard::set(&[ + ("HOME", cache_home.path().as_os_str()), + ("TOKSCALE_CONFIG_DIR", cache_home.path().as_os_str()), + ("TOKSCALE_PRICING_CACHE_ONLY", std::ffi::OsStr::new("1")), + ]); + let (db_path, connection) = copilot_desktop_test_db(source_home.path()); + connection + .pragma_update(None, "journal_mode", "WAL") + .unwrap(); + insert_copilot_desktop_session( + &connection, + "one", + "gpt-5.4-mini", + 10, + 1, + 0, + 0, + "2026-07-01T12:34:56Z", + None, + ); + connection + .execute_batch("PRAGMA wal_checkpoint(TRUNCATE);") + .unwrap(); + drop(connection); + let mut wal = db_path.clone().into_os_string(); + wal.push("-wal"); + let wal = PathBuf::from(wal); + if wal.exists() { + std::fs::remove_file(&wal).unwrap(); + } - fn make_workspace_message( - client: &str, - model_id: &str, - provider_id: &str, - session_id: &str, - cost: f64, - workspace_key: Option<&str>, - workspace_label: Option<&str>, - ) -> UnifiedMessage { - let mut msg = UnifiedMessage::new( - client, - model_id, - provider_id, - session_id, - 1_733_011_200_000, + let home = source_home.path().to_str().unwrap(); + let clients = ["copilot".to_string()]; + assert_eq!( + parse_all_messages_with_pricing(home, &clients, None).len(), + 1 + ); + let db_only_fingerprint = + message_cache::SourceFingerprint::from_copilot_desktop_path(&db_path).unwrap(); + let sentinel = UnifiedMessage::new_with_dedup( + "copilot", + "cached-sentinel", + "github-copilot", + "cached-sentinel", + 1_782_909_296_000, TokenBreakdown { - input: 10, - output: 5, - cache_read: 0, - cache_write: 0, - reasoning: 0, + input: 999, + ..Default::default() }, - cost, + 0.0, + Some("copilot-desktop:cached-sentinel".to_string()), + ); + let mut cache = message_cache::SourceMessageCache::load(); + cache.insert(message_cache::CachedSourceEntry::new( + message_cache::CacheIdentity::for_client(ClientId::Copilot), + &db_path, + db_only_fingerprint, + vec![sentinel], + Vec::new(), + None, + )); + cache.save_if_dirty(); + + let writer = rusqlite::Connection::open(&db_path).unwrap(); + writer.pragma_update(None, "journal_mode", "WAL").unwrap(); + insert_copilot_desktop_session( + &writer, + "two", + "gpt-5.4-mini", + 20, + 2, + 0, + 0, + "2026-07-01T13:00:00Z", + None, ); - msg.set_workspace( - workspace_key.map(str::to_string), - workspace_label.map(str::to_string), + assert!(wal.is_file()); + assert!(message_cache::SourceFingerprint::from_copilot_desktop_path(&db_path).is_some()); + std::fs::set_permissions(&wal, std::fs::Permissions::from_mode(0o000)).unwrap(); + + assert!(message_cache::SourceFingerprint::from_copilot_desktop_path(&db_path).is_none()); + assert!(local_source_change_token(&copilot_local_options(source_home.path())).is_err()); + assert!(copilot_desktop_source_mtime_ms(&db_path).is_none()); + let materialized = parse_all_messages_with_pricing(home, &clients, None); + assert!(!materialized + .iter() + .any(|message| message.model_id == "cached-sentinel")); + let mut streamed = Vec::new(); + scan_messages_streaming( + home, + &clients, + None, + false, + &scanner::ScannerSettings::default(), + &|_| true, + &mut |message| streamed.push(message.clone()), ); - msg - } + assert!(!streamed + .iter() + .any(|message| message.model_id == "cached-sentinel")); - fn make_trae_message( - session_id: &str, - timestamp: i64, - dedup_key: Option<&str>, - cost: f64, - ) -> UnifiedMessage { - UnifiedMessage::new_with_dedup( - "trae", - "gpt-5.2", - "openai", - session_id, - timestamp, - TokenBreakdown { - input: 10, - output: 5, - cache_read: 0, - cache_write: 0, - reasoning: 0, - }, - cost, - dedup_key.map(str::to_string), - ) - } + std::fs::set_permissions(&wal, std::fs::Permissions::from_mode(0o600)).unwrap(); + drop(writer); + assert_eq!( + parse_all_messages_with_pricing(home, &clients, None).len(), + 2 + ); - #[allow(clippy::too_many_arguments)] - fn build_opencode_sqlite_payload( - created_ms: f64, - completed_ms: f64, - input: i64, - output: i64, - reasoning: i64, - cache_read: i64, - cache_write: i64, - cost: f64, - ) -> String { - format!( - r#"{{ - "role": "assistant", - "modelID": "claude-sonnet-4", - "providerID": "anthropic", - "cost": {cost}, - "tokens": {{ - "input": {input}, - "output": {output}, - "reasoning": {reasoning}, - "cache": {{ "read": {cache_read}, "write": {cache_write} }} - }}, - "time": {{ "created": {created_ms}, "completed": {completed_ms} }}, - "mode": "build" - }}"# - ) + std::fs::set_permissions(&db_path, std::fs::Permissions::from_mode(0o000)).unwrap(); + assert!(message_cache::SourceFingerprint::from_copilot_desktop_path(&db_path).is_none()); + assert!(local_source_change_token(&copilot_local_options(source_home.path())).is_err()); + assert!(copilot_desktop_source_mtime_ms(&db_path).is_none()); + std::fs::set_permissions(&db_path, std::fs::Permissions::from_mode(0o600)).unwrap(); } - fn create_opencode_sqlite_db(db_path: &std::path::Path) -> rusqlite::Connection { - let conn = rusqlite::Connection::open(db_path).unwrap(); - conn.execute_batch( - "CREATE TABLE message ( - id TEXT PRIMARY KEY, - session_id TEXT NOT NULL, - data TEXT NOT NULL - );", - ) - .unwrap(); - conn - } + #[test] + #[serial_test::serial] + fn copilot_desktop_materialized_streaming_count_and_report_parity() { + let source_home = tempfile::TempDir::new().unwrap(); + let cache_home = tempfile::TempDir::new().unwrap(); + let _env = EnvGuard::set(&[ + ("HOME", cache_home.path().as_os_str()), + ("TOKSCALE_CONFIG_DIR", cache_home.path().as_os_str()), + ("TOKSCALE_PRICING_CACHE_ONLY", std::ffi::OsStr::new("1")), + ]); + let (db_path, connection) = copilot_desktop_test_db(source_home.path()); + insert_copilot_desktop_session( + &connection, + "shared", + "gpt-5.4-mini", + 100, + 10, + 0, + 0, + "2026-07-01T12:34:56Z", + Some("github.copilot.default"), + ); + insert_copilot_desktop_session( + &connection, + "desktop-only", + "gpt-5.4-mini", + 40, + 10, + 0, + 0, + "2026-07-01T13:00:00Z", + Some("github.copilot.default"), + ); + drop(connection); + let otel_path = write_copilot_otel_message( + source_home.path(), + "shared", + 1_782_909_296, + 20, + 5, + Some("github.copilot.default"), + ); + let home = source_home.path().to_str().unwrap(); + let clients = ["copilot".to_string()]; + let scan_result = scan_test_sources( + source_home.path(), + &clients, + &scanner::ScannerSettings::default(), + ); + let cache_db_path = scanner_spelling( + std::slice::from_ref( + scan_result + .copilot_desktop_db + .as_ref() + .expect("Copilot Desktop database must be discovered"), + ), + &db_path, + ); + let cache_otel_path = scanner_spelling(scan_result.get(ClientId::Copilot), &otel_path); - fn create_hermes_sqlite_db(db_path: &std::path::Path) -> rusqlite::Connection { - let conn = rusqlite::Connection::open(db_path).unwrap(); - conn.execute_batch( - "CREATE TABLE sessions ( - id TEXT PRIMARY KEY, - source TEXT NOT NULL, - model TEXT, - started_at REAL NOT NULL, - message_count INTEGER DEFAULT 0, - input_tokens INTEGER DEFAULT 0, - output_tokens INTEGER DEFAULT 0, - cache_read_tokens INTEGER DEFAULT 0, - cache_write_tokens INTEGER DEFAULT 0, - reasoning_tokens INTEGER DEFAULT 0, - billing_provider TEXT, - estimated_cost_usd REAL, - actual_cost_usd REAL - );", - ) - .unwrap(); - conn - } - - fn create_zed_sqlite_db(db_path: &std::path::Path) -> rusqlite::Connection { - let conn = rusqlite::Connection::open(db_path).unwrap(); - conn.execute_batch( - "CREATE TABLE threads ( - id TEXT PRIMARY KEY, - summary TEXT NOT NULL, - updated_at TEXT NOT NULL, - data_type TEXT NOT NULL, - data BLOB NOT NULL - );", - ) - .unwrap(); - conn - } - - fn insert_zed_thread(conn: &rusqlite::Connection, id: &str, model: &str) { - let payload = format!( - r#"{{ - "version": "0.3.0", - "title": "Test thread", - "updated_at": "2026-05-01T12:30:00Z", - "request_token_usage": {{ - "turn-1": {{ - "input_tokens": 42, - "output_tokens": 7, - "cache_creation_input_tokens": 3, - "cache_read_input_tokens": 5 - }} - }}, - "model": {{ - "provider": "zed.dev", - "model": "{model}" - }}, - "imported": false - }}"# - ); - conn.execute( - "INSERT INTO threads (id, summary, updated_at, data_type, data) VALUES (?1, ?2, ?3, ?4, ?5)", - rusqlite::params![id, "Test thread", "2026-05-01T12:30:00Z", "json", payload.as_bytes()], - ) - .unwrap(); - } - - fn insert_hermes_session( - conn: &rusqlite::Connection, - id: &str, - model: &str, - message_count: i64, - input_tokens: i64, - output_tokens: i64, - actual_cost_usd: f64, - ) { - conn.execute( - "INSERT INTO sessions ( - id, source, model, started_at, message_count, - input_tokens, output_tokens, cache_read_tokens, cache_write_tokens, reasoning_tokens, - billing_provider, estimated_cost_usd, actual_cost_usd - ) VALUES (?1, 'cli', ?2, 1775001102.0, ?3, ?4, ?5, 0, 0, 0, 'anthropic', NULL, ?6)", - rusqlite::params![ - id, - model, - message_count, - input_tokens, - output_tokens, - actual_cost_usd - ], - ) - .unwrap(); - } - - #[test] - fn test_normalize_model_for_grouping() { - assert_eq!( - normalize_model_for_grouping("claude-opus-4-5-20251101"), - "claude-opus-4-5" - ); + let mut materialized = parse_all_messages_with_pricing(home, &clients, None); + let cold_materialized = materialized.clone(); + let mut warm_materialized = parse_all_messages_with_pricing(home, &clients, None); + materialized.sort_unstable_by(|left, right| left.session_id.cmp(&right.session_id)); + warm_materialized.sort_unstable_by(|left, right| left.session_id.cmp(&right.session_id)); + assert_eq!(materialized, warm_materialized); + assert_eq!(materialized.len(), 2); assert_eq!( - normalize_model_for_grouping("claude-sonnet-4-5-20250929"), - "claude-sonnet-4-5" + materialized + .iter() + .map(|message| message.tokens.input) + .sum::(), + 60 ); assert_eq!( - normalize_model_for_grouping("claude-sonnet-4-20250514"), - "claude-sonnet-4" + materialized + .iter() + .map(|message| message.tokens.output) + .sum::(), + 15 ); + assert!(!materialized + .iter() + .any(|message| { message.dedup_key.as_deref() == Some("copilot-desktop:shared") })); - assert_eq!( - normalize_model_for_grouping("claude-opus-4.5"), - "claude-opus-4-5" - ); - assert_eq!( - normalize_model_for_grouping("claude-sonnet-4.5"), - "claude-sonnet-4-5" - ); - assert_eq!( - normalize_model_for_grouping("claude-opus-4.6"), - "claude-opus-4-6" - ); - assert_eq!( - normalize_model_for_grouping("anthropic/claude-4-6-sonnet"), - "claude-sonnet-4-6" - ); - assert_eq!( - normalize_model_for_grouping("anthropic/claude-4-5-haiku"), - "claude-haiku-4-5" - ); - assert_eq!( - normalize_model_for_grouping("anthropic/claude-4-6-opus"), - "claude-opus-4-6" + let cache = message_cache::SourceMessageCache::load(); + assert!(cache + .get( + message_cache::CacheIdentity::for_client(ClientId::Copilot), + &cache_db_path, + ) + .is_some_and(|entry| entry.messages.len() == 2)); + assert!(cache + .get( + message_cache::CacheIdentity::for_client(ClientId::Copilot), + &cache_otel_path, + ) + .is_some_and(|entry| entry.messages.len() == 1)); + + let mut streamed = Vec::new(); + scan_messages_streaming( + home, + &clients, + None, + false, + &scanner::ScannerSettings::default(), + &|_| true, + &mut |message| streamed.push(message.clone()), ); + streamed.sort_unstable_by(|left, right| left.session_id.cmp(&right.session_id)); + assert_eq!(streamed, materialized); - assert_eq!(normalize_model_for_grouping("gpt-5.2"), "gpt-5.2"); - assert_eq!(normalize_model_for_grouping("gpt-5.4(xhigh)"), "gpt-5.4"); - assert_eq!(normalize_model_for_grouping("gpt-5.4(high)"), "gpt-5.4"); - assert_eq!(normalize_model_for_grouping("gpt-5.4(minimal)"), "gpt-5.4"); - assert_eq!(normalize_model_for_grouping("gpt-5.4(auto)"), "gpt-5.4"); - assert_eq!(normalize_model_for_grouping("gpt-5.4(none)"), "gpt-5.4"); + let counted = parse_local_clients(copilot_local_options(source_home.path())).unwrap(); + assert_eq!(counted.counts.get(ClientId::Copilot), 2); + assert_eq!(counted.messages.len(), 2); + for message in &counted.messages { + let source = materialized + .iter() + .find(|source| source.session_id == message.session_id) + .unwrap(); + assert_eq!(message.client, source.client); + assert_eq!(message.model_id, source.model_id); + assert_eq!(message.provider_id, source.provider_id); + assert_eq!(message.workspace_key, source.workspace_key); + assert_eq!(message.workspace_label, source.workspace_label); + assert_eq!(message.timestamp, source.timestamp); + assert_eq!(message.date, source.date); + assert_eq!(message.input, source.tokens.input); + assert_eq!(message.output, source.tokens.output); + assert_eq!(message.cache_read, source.tokens.cache_read); + assert_eq!(message.cache_write, source.tokens.cache_write); + assert_eq!(message.reasoning, source.tokens.reasoning); + assert_eq!(message.duration_ms, source.duration_ms); + assert_eq!(message.message_count, source.message_count); + assert_eq!(message.agent, source.agent); + } + + let runtime = tokio::runtime::Builder::new_current_thread() + .enable_all() + .build() + .unwrap(); + let options = copilot_report_options(source_home.path()); + let graph = runtime + .block_on(super::generate_graph_with_loaded_pricing( + options.clone(), + None, + )) + .unwrap(); + let model = runtime.block_on(get_model_report(options.clone())).unwrap(); + let monthly = runtime + .block_on(get_monthly_report(options.clone())) + .unwrap(); + let hourly = runtime + .block_on(get_hourly_report(options.clone())) + .unwrap(); + let agents = runtime.block_on(get_agents_report(options)).unwrap(); + let expected_tokens = cold_materialized + .iter() + .map(|message| message.tokens.total()) + .sum::(); + + assert_eq!(graph.summary.total_tokens, expected_tokens); assert_eq!( - normalize_model_for_grouping("gpt-5.4(weirdgarbage)"), - "gpt-5.4(weirdgarbage)" + graph + .contributions + .iter() + .map(|entry| entry.totals.messages) + .sum::(), + 2 ); + assert_eq!(model.total_messages, 2); + assert_eq!(model.total_input, 60); + assert_eq!(model.total_output, 15); assert_eq!( - normalize_model_for_grouping("claude-sonnet-4.5(high)"), - "claude-sonnet-4-5" + monthly + .entries + .iter() + .map(|entry| entry.message_count) + .sum::(), + 2 ); assert_eq!( - normalize_model_for_grouping("gemini-3-pro(auto)"), - "gemini-3-pro" + hourly + .entries + .iter() + .map(|entry| entry.message_count) + .sum::(), + 2 ); + assert_eq!(agents.total_messages, 2); assert_eq!( - normalize_model_for_grouping("gemini-2.5-pro"), - "gemini-2.5-pro" + agents.entries.iter().map(|entry| entry.input).sum::(), + 60 ); + } + #[test] + #[serial_test::serial] + fn test_env_guard_restores_some_and_none_after_panic() { + const KEYS: [&str; 3] = ["HOME", "TOKSCALE_PRICING_CACHE_ONLY", "TOKSCALE_CONFIG_DIR"]; + let _original = EnvGuard::capture(&KEYS); + + unsafe { + std::env::set_var("HOME", "/tmp/tokscale-env-guard-home-before"); + std::env::remove_var("TOKSCALE_PRICING_CACHE_ONLY"); + std::env::set_var( + "TOKSCALE_CONFIG_DIR", + "/tmp/tokscale-env-guard-config-before", + ); + } + let first = std::panic::catch_unwind(|| { + let _guard = EnvGuard::set(&[ + ( + "HOME", + std::ffi::OsStr::new("/tmp/tokscale-env-guard-home-during"), + ), + ("TOKSCALE_PRICING_CACHE_ONLY", std::ffi::OsStr::new("1")), + ( + "TOKSCALE_CONFIG_DIR", + std::ffi::OsStr::new("/tmp/tokscale-env-guard-config-during"), + ), + ]); + panic!("exercise EnvGuard unwinding"); + }); + assert!(first.is_err()); assert_eq!( - normalize_model_for_grouping("claude-opus-4-5-high"), - "claude-opus-4-5-high" - ); - assert_eq!( - normalize_model_for_grouping("claude-opus-4-5-thinking-high"), - "claude-opus-4-5-thinking-high" + std::env::var_os("HOME"), + Some(std::ffi::OsString::from( + "/tmp/tokscale-env-guard-home-before" + )) ); + assert_eq!(std::env::var_os("TOKSCALE_PRICING_CACHE_ONLY"), None); assert_eq!( - normalize_model_for_grouping("claude-sonnet-4-5-high"), - "claude-sonnet-4-5-high" + std::env::var_os("TOKSCALE_CONFIG_DIR"), + Some(std::ffi::OsString::from( + "/tmp/tokscale-env-guard-config-before" + )) ); + unsafe { + std::env::remove_var("HOME"); + std::env::set_var("TOKSCALE_PRICING_CACHE_ONLY", "before"); + std::env::remove_var("TOKSCALE_CONFIG_DIR"); + } + let second = std::panic::catch_unwind(|| { + let _guard = EnvGuard::set(&[ + ( + "HOME", + std::ffi::OsStr::new("/tmp/tokscale-env-guard-home-during"), + ), + ("TOKSCALE_PRICING_CACHE_ONLY", std::ffi::OsStr::new("1")), + ( + "TOKSCALE_CONFIG_DIR", + std::ffi::OsStr::new("/tmp/tokscale-env-guard-config-during"), + ), + ]); + panic!("exercise inverse EnvGuard unwinding"); + }); + assert!(second.is_err()); + assert_eq!(std::env::var_os("HOME"), None); assert_eq!( - normalize_model_for_grouping("claude-4-sonnet"), - "claude-4-sonnet" - ); - assert_eq!( - normalize_model_for_grouping("claude-4-opus-thinking"), - "claude-4-opus-thinking" + std::env::var_os("TOKSCALE_PRICING_CACHE_ONLY"), + Some(std::ffi::OsString::from("before")) ); - - assert_eq!(normalize_model_for_grouping("big-pickle"), "big-pickle"); - assert_eq!(normalize_model_for_grouping("grok-code"), "grok-code"); - - assert_eq!( - normalize_model_for_grouping("claude-opus-4.5-20251101"), - "claude-opus-4-5" - ); - } - - #[test] - fn test_group_by_from_str_valid_values() { - assert_eq!(GroupBy::from_str("model").unwrap(), GroupBy::Model); - assert_eq!( - GroupBy::from_str("client,model").unwrap(), - GroupBy::ClientModel - ); - assert_eq!( - GroupBy::from_str("client-model").unwrap(), - GroupBy::ClientModel - ); - assert_eq!( - GroupBy::from_str("client,provider,model").unwrap(), - GroupBy::ClientProviderModel - ); - assert_eq!( - GroupBy::from_str("client-provider-model").unwrap(), - GroupBy::ClientProviderModel - ); - assert_eq!( - GroupBy::from_str("workspace,model").unwrap(), - GroupBy::WorkspaceModel - ); - assert_eq!( - GroupBy::from_str("workspace-model").unwrap(), - GroupBy::WorkspaceModel - ); - assert_eq!(GroupBy::from_str("session").unwrap(), GroupBy::Session); - assert_eq!( - GroupBy::from_str("session,model").unwrap(), - GroupBy::Session - ); - assert_eq!( - GroupBy::from_str("session-model").unwrap(), - GroupBy::Session - ); - assert_eq!( - GroupBy::from_str("client,session").unwrap(), - GroupBy::ClientSession - ); - assert_eq!( - GroupBy::from_str("client,session,model").unwrap(), - GroupBy::ClientSession - ); - assert_eq!( - GroupBy::from_str("client-session-model").unwrap(), - GroupBy::ClientSession - ); - assert!(GroupBy::from_str("unknown").is_err()); - } + assert_eq!(std::env::var_os("TOKSCALE_CONFIG_DIR"), None); + } #[test] - fn test_group_by_default_is_client_model() { - assert_eq!(GroupBy::default(), GroupBy::ClientModel); - } + #[serial_test::serial] + fn test_empty_reports_normalize_total_cost_to_positive_zero() { + let source_home = tempfile::TempDir::new().unwrap(); + let cache_home = tempfile::TempDir::new().unwrap(); + let _env = EnvGuard::set(&[ + ("HOME", cache_home.path().as_os_str()), + ("TOKSCALE_PRICING_CACHE_ONLY", std::ffi::OsStr::new("1")), + ]); - #[test] - fn test_group_by_display_round_trips_with_from_str() { - let variants = [ - GroupBy::Model, - GroupBy::ClientModel, - GroupBy::ClientProviderModel, - GroupBy::WorkspaceModel, - GroupBy::Session, - GroupBy::ClientSession, + let options = ReportOptions { + home_dir: Some(source_home.path().to_string_lossy().into_owned()), + use_env_roots: false, + clients: Some(vec!["opencode".to_string()]), + ..Default::default() + }; + let runtime = tokio::runtime::Builder::new_current_thread() + .enable_all() + .build() + .unwrap(); + let totals = [ + runtime + .block_on(get_model_report(options.clone())) + .unwrap() + .total_cost, + runtime + .block_on(get_monthly_report(options.clone())) + .unwrap() + .total_cost, + runtime + .block_on(get_hourly_report(options.clone())) + .unwrap() + .total_cost, + runtime + .block_on(get_agents_report(options)) + .unwrap() + .total_cost, ]; - for variant in variants { - let rendered = variant.to_string(); - let parsed = GroupBy::from_str(&rendered).unwrap(); - assert_eq!(parsed, variant); + for total in totals { + assert_eq!(total.to_bits(), 0.0f64.to_bits()); } } - #[test] - fn test_group_by_from_str_whitespace_handling() { - assert_eq!( - GroupBy::from_str("client, model").unwrap(), - GroupBy::ClientModel - ); - assert_eq!(GroupBy::from_str(" model ").unwrap(), GroupBy::Model); - assert_eq!( - GroupBy::from_str("client , provider , model").unwrap(), - GroupBy::ClientProviderModel - ); - assert_eq!( - GroupBy::from_str("workspace, model").unwrap(), - GroupBy::WorkspaceModel + fn make_opencode_selection_message(key: &str, cost: f64, source: CostSource) -> UnifiedMessage { + let mut message = UnifiedMessage::new_with_dedup( + "opencode", "gpt-4o", "openai", "oc-session", 1_733_011_200_000, + TokenBreakdown { input: 10, output: 5, cache_read: 0, cache_write: 0, reasoning: 0 }, + cost, Some(key.to_string()), ); + match source { + CostSource::ProviderReported => message.mark_provider_reported_cost(), + CostSource::Estimated => message.mark_estimated_cost(), + CostSource::Unknown => {} + } + message } - #[test] - fn test_model_usage_performance_uses_only_timed_positive_token_messages() { - let mut timed = make_workspace_message( - "opencode", - "gpt-5.4", - "openai", - "session-1", - 0.0, - None, - None, - ); - timed.tokens = TokenBreakdown { - input: 100, - output: 50, - cache_read: 25, - cache_write: 0, - reasoning: 25, - }; - timed.duration_ms = Some(400); + fn opencode_authority_set(key: &str) -> HashSet { + let message = make_opencode_selection_message(key, 0.0, CostSource::ProviderReported); + OpenCodeSourceIdentity::all_from_message(&message) + .into_iter() + .collect() + } - let mut untimed = make_workspace_message( - "opencode", - "gpt-5.4", - "openai", - "session-2", - 0.0, - None, - None, - ); - untimed.tokens = TokenBreakdown { - input: 300, - output: 0, - cache_read: 0, - cache_write: 0, - reasoning: 0, - }; + #[test] + fn test_opencode_streaming_selection_flushes_snapshot_fallback_on_json_drift() { + // A missing file and an invalid file both produce no second-pass message; + // a downgraded file produces an estimated message. All must flush SQLite. + for second_pass in [None, None, Some(CostSource::Estimated)] { + let key = "snapshot-authoritative"; + let mut selection = OpenCodeSelection::new(opencode_authority_set(key)); + assert!(selection.select_sqlite(make_opencode_selection_message( + key, 0.25, CostSource::Estimated, + )).is_none()); + if let Some(source) = second_pass { + assert!(selection.select_json(make_opencode_selection_message( + key, 0.0, source, + ), true).is_none()); + } + let selected: Vec<_> = selection.finish().collect(); + assert_eq!(selected.len(), 1); + assert_eq!(selected[0].cost, 0.25); + assert_eq!(selected[0].cost_source, CostSource::Estimated); + } + } - let entries = aggregate_model_usage_entries(vec![timed, untimed], &GroupBy::ClientModel); + #[test] + fn test_opencode_streaming_selection_replaces_deferred_sqlite_estimate() { + let key = "sqlite-authoritative-replacement"; + let mut selection = OpenCodeSelection::new(opencode_authority_set(key)); + assert!(selection + .select_sqlite(make_opencode_selection_message( + key, 0.25, CostSource::Estimated, + )) + .is_none()); - assert_eq!(entries.len(), 1); - let performance = &entries[0].performance; - assert_eq!(performance.total_duration_ms, 400); - assert_eq!(performance.timed_tokens, 200); - assert_eq!(performance.sample_count, 1); - assert_eq!(performance.ms_per_1k_tokens, Some(2000.0)); - assert!((performance.token_coverage - 0.4).abs() < f64::EPSILON); + let selected = selection + .select_sqlite(make_opencode_selection_message( + key, 0.50, CostSource::ProviderReported, + )) + .expect("a later authoritative SQLite message must replace the fallback"); + assert_eq!(selected.cost, 0.50); + assert_eq!(selected.cost_source, CostSource::ProviderReported); + assert_eq!(selection.finish().count(), 0); } #[test] - fn test_model_usage_performance_is_null_without_duration_samples() { - let entries = aggregate_model_usage_entries( - vec![make_workspace_message( - "claude", - "claude-sonnet-4-5", - "anthropic", - "session-1", - 0.0, - None, - None, - )], - &GroupBy::ClientModel, - ); + fn test_opencode_streaming_selection_keeps_first_sqlite_estimate() { + let key = "sqlite-estimated-first-wins"; + let mut selection = OpenCodeSelection::new(opencode_authority_set(key)); + assert!(selection + .select_sqlite(make_opencode_selection_message( + key, 0.25, CostSource::Estimated, + )) + .is_none()); + assert!(selection + .select_sqlite(make_opencode_selection_message( + key, 0.50, CostSource::Estimated, + )) + .is_none()); - assert_eq!(entries.len(), 1); - assert_eq!(entries[0].performance.ms_per_1k_tokens, None); - assert_eq!(entries[0].performance.total_duration_ms, 0); - assert_eq!(entries[0].performance.timed_tokens, 0); - assert_eq!(entries[0].performance.token_coverage, 0.0); + let selected: Vec<_> = selection.finish().collect(); + assert_eq!(selected.len(), 1); + assert_eq!(selected[0].cost, 0.25); + assert_eq!(selected[0].cost_source, CostSource::Estimated); } #[test] - fn test_workspace_model_grouping_merges_same_workspace_and_model() { - let entries = aggregate_model_usage_entries( - vec![ - make_workspace_message( - "claude", - "claude-sonnet-4-5-20250929", - "anthropic", - "session-1", - 1.25, - Some("/repo-a"), - Some("repo-a"), - ), - make_workspace_message( - "qwen", - "claude-sonnet-4-5-20250929", - "anthropic", - "session-2", - 2.75, - Some("/repo-a"), - Some("repo-a"), - ), - ], - &GroupBy::WorkspaceModel, - ); + fn test_opencode_streaming_selection_keeps_first_sqlite_authority() { + let key = "sqlite-authoritative-first-wins"; + let mut selection = OpenCodeSelection::new(opencode_authority_set(key)); + let first = selection + .select_sqlite(make_opencode_selection_message( + key, 0.50, CostSource::ProviderReported, + )) + .expect("the first authoritative SQLite message must be selected"); + assert_eq!(first.cost, 0.50); + assert!(selection + .select_sqlite(make_opencode_selection_message( + key, 0.75, CostSource::ProviderReported, + )) + .is_none()); + assert_eq!(selection.finish().count(), 0); + } - assert_eq!(entries.len(), 1); - assert_eq!(entries[0].model, "claude-sonnet-4-5"); - assert_eq!(entries[0].workspace_key.as_deref(), Some("/repo-a")); - assert_eq!(entries[0].workspace_label.as_deref(), Some("repo-a")); - assert_eq!(entries[0].cost, 4.0); - assert_eq!(entries[0].message_count, 2); - assert_eq!(entries[0].merged_clients.as_deref(), Some("claude, qwen")); + #[test] + fn test_opencode_streaming_selection_keeps_fallback_until_json_is_emitted() { + let key = "filtered-authoritative"; + let mut selection = OpenCodeSelection::new(opencode_authority_set(key)); + assert!(selection.select_sqlite(make_opencode_selection_message( + key, 0.25, CostSource::Estimated, + )).is_none()); + assert!(selection.select_json(make_opencode_selection_message( + key, 0.50, CostSource::ProviderReported, + ), false).is_none()); + let selected: Vec<_> = selection.finish().collect(); + assert_eq!(selected.len(), 1); + assert_eq!(selected[0].cost, 0.25); } #[test] - fn test_model_grouping_merges_anthropic_prefixed_claude_variant_with_canonical_model() { - let entries = aggregate_model_usage_entries( - vec![ - make_workspace_message( - "claude", - "anthropic/claude-4-6-sonnet", - "anthropic", - "session-1", - 1.25, - Some("/repo-a"), - Some("repo-a"), - ), - make_workspace_message( - "claude", - "claude-sonnet-4-6", - "anthropic", - "session-2", - 2.75, - Some("/repo-b"), - Some("repo-b"), - ), - ], - &GroupBy::ClientModel, - ); + fn test_opencode_streaming_selection_does_not_double_count_new_authority() { + let key = "newly-authoritative"; + let mut selection = OpenCodeSelection::new(HashSet::new()); + assert!(selection.select_sqlite(make_opencode_selection_message( + key, 0.25, CostSource::Estimated, + )).is_some()); + assert!(selection.select_json(make_opencode_selection_message( + key, 0.50, CostSource::ProviderReported, + ), true).is_none()); + assert_eq!(selection.finish().count(), 0); + } - assert_eq!(entries.len(), 1); - assert_eq!(entries[0].model, "claude-sonnet-4-6"); - assert_eq!(entries[0].input, 20); - assert_eq!(entries[0].output, 10); - assert_eq!(entries[0].cost, 4.0); - assert_eq!(entries[0].message_count, 2); + #[test] + fn test_opencode_streaming_selection_keeps_incompatible_same_id_rows() { + let key = "reused-message-id"; + let mut selection = OpenCodeSelection::new(opencode_authority_set(key)); + let first = make_opencode_selection_message(key, 0.25, CostSource::Estimated); + let mut second = make_opencode_selection_message(key, 0.50, CostSource::Estimated); + second.tokens.input = 20; + let json = make_opencode_selection_message(key, 0.75, CostSource::ProviderReported); + + assert!(selection.select_sqlite(first).is_none()); + assert!(selection.select_sqlite(second).is_none()); + assert!(selection.select_json(json, true).is_some()); + let deferred: Vec<_> = selection.finish().collect(); + assert_eq!(deferred.len(), 1); + assert_eq!(deferred[0].tokens.input, 20); } #[test] - fn test_workspace_model_grouping_separates_different_workspaces() { - let entries = aggregate_model_usage_entries( - vec![ - make_workspace_message( - "claude", - "claude-sonnet-4-5-20250929", - "anthropic", - "session-1", - 1.0, - Some("/repo-a"), - Some("repo-a"), - ), - make_workspace_message( - "claude", - "claude-sonnet-4-5-20250929", - "anthropic", - "session-2", - 2.0, - Some("/repo-b"), - Some("repo-b"), - ), - ], - &GroupBy::WorkspaceModel, + fn test_opencode_streaming_selection_replaces_deferred_identity_after_other_emission() { + let key = "reused-authoritative-id"; + let mut selection = OpenCodeSelection::new(opencode_authority_set(key)); + let first = make_opencode_selection_message(key, 0.25, CostSource::ProviderReported); + let mut second = make_opencode_selection_message(key, 0.50, CostSource::Estimated); + second.tokens.input = 20; + let mut json = make_opencode_selection_message(key, 0.75, CostSource::ProviderReported); + json.tokens.input = 20; + + assert!(selection.select_sqlite(first).is_some()); + assert!(selection.select_sqlite(second).is_none()); + let selected = selection + .select_json(json, true) + .expect("JSON must replace its exact deferred identity"); + assert_eq!(selected.tokens.input, 20); + assert_eq!(selected.cost_source, CostSource::ProviderReported); + assert_eq!(selection.finish().count(), 0); + } + + #[test] + fn test_opencode_streaming_selection_expands_authority_through_aliases() { + let embedded = "shared-embedded"; + let fallback = "legacy-message"; + let pure = make_opencode_selection_message(embedded, 0.25, CostSource::Estimated); + let mut hybrid = make_opencode_selection_message(embedded, 0.50, CostSource::Estimated); + hybrid.dedup_aliases.push(fallback.to_string()); + let json = make_opencode_selection_message(fallback, 0.75, CostSource::ProviderReported); + let duplicate = + make_opencode_selection_message(embedded, 1.0, CostSource::ProviderReported); + let authoritative = opencode_authoritative_sources( + [&pure, &hybrid, &json] + .into_iter() + .map(opencode_identity_group), ); + let mut selection = OpenCodeSelection::new(authoritative); - assert_eq!(entries.len(), 2); - let labels: HashSet<_> = entries - .iter() - .map(|entry| entry.workspace_label.as_deref().unwrap()) - .collect(); - assert_eq!(labels, HashSet::from(["repo-a", "repo-b"])); + assert!(selection.select_sqlite(pure).is_none()); + assert!(selection.select_sqlite(hybrid).is_none()); + let selected = selection + .select_json(json, true) + .expect("JSON authority must replace the aliased SQLite identity"); + assert_eq!(selected.cost, 0.75); + assert!(selection.select_json(duplicate, true).is_none()); + assert_eq!(selection.finish().count(), 0); } #[test] - fn test_workspace_model_grouping_uses_unknown_bucket_without_workspace_metadata() { - let entries = aggregate_model_usage_entries( - vec![ - make_workspace_message( - "claude", - "claude-sonnet-4-5-20250929", - "anthropic", - "session-1", - 1.0, - None, - None, - ), - make_workspace_message( - "claude", - "claude-sonnet-4-5-20250929", - "anthropic", - "session-2", - "2.0".parse().unwrap(), - None, - None, - ), - ], - &GroupBy::WorkspaceModel, + fn test_opencode_streaming_selection_keeps_aliases_on_sqlite_replacement() { + let embedded = "shared-embedded"; + let fallback = "legacy-message"; + let mut hybrid = make_opencode_selection_message(embedded, 0.25, CostSource::Estimated); + hybrid.dedup_aliases.push(fallback.to_string()); + let sqlite = make_opencode_selection_message(embedded, 0.50, CostSource::ProviderReported); + let json = make_opencode_selection_message(fallback, 0.75, CostSource::ProviderReported); + let authoritative = opencode_authoritative_sources( + [&hybrid, &sqlite, &json] + .into_iter() + .map(opencode_identity_group), ); + let mut selection = OpenCodeSelection::new(authoritative); - assert_eq!(entries.len(), 1); - assert_eq!(entries[0].workspace_key, None); - assert_eq!( - entries[0].workspace_label.as_deref(), - Some(UNKNOWN_WORKSPACE_LABEL) + assert!(selection.select_sqlite(hybrid).is_none()); + assert!(selection.select_sqlite(sqlite).is_some()); + assert!(selection.select_json(json, true).is_none()); + assert_eq!(selection.finish().count(), 0); + } + + #[test] + fn test_opencode_streaming_selection_propagates_emitted_aliases_to_deferred_rows() { + let embedded = "shared-embedded"; + let fallback = "legacy-message"; + let first = make_opencode_selection_message(embedded, 0.25, CostSource::ProviderReported); + let mut deferred = make_opencode_selection_message(embedded, 0.50, CostSource::Estimated); + deferred.tokens.input = 20; + deferred.dedup_aliases.push(fallback.to_string()); + let mut json = make_opencode_selection_message(fallback, 0.75, CostSource::ProviderReported); + json.tokens.input = 30; + let authoritative = opencode_authoritative_sources( + [&first, &deferred, &json] + .into_iter() + .map(opencode_identity_group), + ); + let mut selection = OpenCodeSelection::new(authoritative); + + assert!(selection.select_sqlite(first).is_some()); + assert!(selection.select_sqlite(deferred).is_none()); + assert!(selection.select_json(json, true).is_none()); + let remaining: Vec<_> = selection.finish().collect(); + assert_eq!(remaining.len(), 1); + assert_eq!(remaining[0].tokens.input, 20); + } + + #[test] + fn test_opencode_streaming_selection_propagates_emitted_aliases_from_deferred_graph() { + let embedded = "shared-embedded"; + let fallback = "legacy-message"; + let mut deferred = make_opencode_selection_message(embedded, 0.25, CostSource::Estimated); + deferred.tokens.input = 20; + deferred.dedup_aliases.push(fallback.to_string()); + let sqlite = make_opencode_selection_message(embedded, 0.50, CostSource::ProviderReported); + let mut json = make_opencode_selection_message(fallback, 0.75, CostSource::ProviderReported); + json.tokens.input = 30; + let authoritative = opencode_authoritative_sources( + [&deferred, &sqlite, &json] + .into_iter() + .map(opencode_identity_group), ); - assert_eq!(entries[0].message_count, 2); - assert_eq!(entries[0].cost, 3.0); + let mut selection = OpenCodeSelection::new(authoritative); + + assert!(selection.select_sqlite(deferred).is_none()); + assert!(selection.select_sqlite(sqlite).is_some()); + assert!(selection.select_json(json, true).is_none()); + let remaining: Vec<_> = selection.finish().collect(); + assert_eq!(remaining.len(), 1); + assert_eq!(remaining[0].tokens.input, 20); } #[test] - fn test_parsed_round_trip_preserves_workspace_metadata() { - let mut unified = UnifiedMessage::new( - "qwen", - "qwen3.5-plus", - "qwen", - "session-1", - 1_742_390_400_000, + fn test_opencode_streaming_selection_prefers_snapshot_authority() { + let key = "stable-authoritative"; + let mut selection = OpenCodeSelection::new(opencode_authority_set(key)); + assert!(selection.select_sqlite(make_opencode_selection_message( + key, 0.25, CostSource::Estimated, + )).is_none()); + let json = selection.select_json(make_opencode_selection_message( + key, 0.50, CostSource::ProviderReported, + ), true).unwrap(); + assert_eq!(json.cost, 0.50); + assert_eq!(json.cost_source, CostSource::ProviderReported); + assert_eq!(selection.finish().count(), 0); + } + + fn make_workspace_message( + client: &str, + model_id: &str, + provider_id: &str, + session_id: &str, + cost: f64, + workspace_key: Option<&str>, + workspace_label: Option<&str>, + ) -> UnifiedMessage { + let mut msg = UnifiedMessage::new( + client, + model_id, + provider_id, + session_id, + 1_733_011_200_000, TokenBreakdown { input: 10, output: 5, - cache_read: 2, + cache_read: 0, cache_write: 0, - reasoning: 1, + reasoning: 0, }, - 1.25, - ); - unified.set_workspace( - Some("//server/share/demo-workspace".to_string()), - Some("demo-workspace".to_string()), - ); - unified.duration_ms = Some(2500); - - let parsed = unified_to_parsed(&unified); - let round_tripped = parsed_to_unified(&parsed, 2.5); - - assert_eq!( - round_tripped.workspace_key.as_deref(), - Some("//server/share/demo-workspace") + cost, ); - assert_eq!( - round_tripped.workspace_label.as_deref(), - Some("demo-workspace") + msg.set_workspace( + workspace_key.map(str::to_string), + workspace_label.map(str::to_string), ); - assert_eq!(round_tripped.cost, 2.5); - assert_eq!(round_tripped.duration_ms, Some(2500)); + msg } - #[test] - fn test_workspace_model_grouping_keeps_real_unknown_workspace_separate() { - let entries = aggregate_model_usage_entries( - vec![ - make_workspace_message( - "claude", - "claude-sonnet-4-5-20250929", - "anthropic", - "session-1", - 1.0, - Some("unknown-workspace"), - Some("unknown-workspace"), - ), - make_workspace_message( - "claude", - "claude-sonnet-4-5-20250929", - "anthropic", - "session-2", - 2.0, - None, - None, - ), - ], - &GroupBy::WorkspaceModel, - ); + fn make_trae_message( + session_id: &str, + timestamp: i64, + dedup_key: Option<&str>, + cost: f64, + ) -> UnifiedMessage { + UnifiedMessage::new_with_dedup( + "trae", + "gpt-5.2", + "openai", + session_id, + timestamp, + TokenBreakdown { + input: 10, + output: 5, + cache_read: 0, + cache_write: 0, + reasoning: 0, + }, + cost, + dedup_key.map(str::to_string), + ) + } - assert_eq!(entries.len(), 2); - assert!(entries.iter().any(|entry| { - entry.workspace_key.as_deref() == Some("unknown-workspace") - && entry.workspace_label.as_deref() == Some("unknown-workspace") - && (entry.cost - 1.0).abs() < f64::EPSILON - })); - assert!(entries.iter().any(|entry| { - entry.workspace_key.is_none() - && entry.workspace_label.as_deref() == Some(UNKNOWN_WORKSPACE_LABEL) - && (entry.cost - 2.0).abs() < f64::EPSILON - })); + #[allow(clippy::too_many_arguments)] + fn build_opencode_sqlite_payload( + created_ms: f64, + completed_ms: f64, + input: i64, + output: i64, + reasoning: i64, + cache_read: i64, + cache_write: i64, + cost: f64, + ) -> String { + format!( + r#"{{ + "role": "assistant", + "modelID": "claude-sonnet-4", + "providerID": "anthropic", + "cost": {cost}, + "tokens": {{ + "input": {input}, + "output": {output}, + "reasoning": {reasoning}, + "cache": {{ "read": {cache_read}, "write": {cache_write} }} + }}, + "time": {{ "created": {created_ms}, "completed": {completed_ms} }}, + "mode": "build" + }}"# + ) } - #[test] - fn test_session_grouping_merges_same_session_and_model() { - // Two messages with the same session_id + same model — should collapse - // into one row regardless of the client that produced them, because - // GroupBy::Session keys on (session_id, model) only. - let entries = aggregate_model_usage_entries( - vec![ - make_workspace_message( - "claude", - "claude-sonnet-4-5-20250929", - "anthropic", - "session-shared", - 1.25, - None, - None, - ), - make_workspace_message( - "amp", - "claude-sonnet-4-5-20250929", - "anthropic", - "session-shared", - 2.75, - None, - None, - ), - ], - &GroupBy::Session, - ); + fn create_opencode_sqlite_db(db_path: &std::path::Path) -> rusqlite::Connection { + let conn = rusqlite::Connection::open(db_path).unwrap(); + conn.execute_batch( + "CREATE TABLE message ( + id TEXT PRIMARY KEY, + session_id TEXT NOT NULL, + data TEXT NOT NULL + );", + ) + .unwrap(); + conn + } - assert_eq!(entries.len(), 1); - assert_eq!(entries[0].session_id.as_deref(), Some("session-shared")); - assert_eq!(entries[0].model, "claude-sonnet-4-5"); - assert!((entries[0].cost - 4.0).abs() < f64::EPSILON); - assert_eq!(entries[0].message_count, 2); - assert!(entries[0].workspace_key.is_none()); - assert!(entries[0].workspace_label.is_none()); - // Session grouping does not merge_clients into a comma list. - assert!(entries[0].merged_clients.is_none()); + fn create_opencode_v2_sqlite_db(db_path: &std::path::Path) -> rusqlite::Connection { + let conn = rusqlite::Connection::open(db_path).unwrap(); + conn.execute_batch( + "CREATE TABLE message ( + id TEXT PRIMARY KEY, + session_id TEXT NOT NULL, + data TEXT NOT NULL + ); + CREATE TABLE session ( + id TEXT PRIMARY KEY, + directory TEXT NOT NULL + ); + CREATE TABLE session_message ( + id TEXT PRIMARY KEY, + session_id TEXT NOT NULL, + type TEXT NOT NULL, + data TEXT NOT NULL + );", + ) + .unwrap(); + conn } - #[test] - fn test_session_grouping_separates_different_sessions() { - let entries = aggregate_model_usage_entries( - vec![ - make_workspace_message("codex", "gpt-5", "openai", "session-a", 1.0, None, None), - make_workspace_message("codex", "gpt-5", "openai", "session-b", 2.0, None, None), - ], - &GroupBy::Session, - ); + fn create_hermes_sqlite_db(db_path: &std::path::Path) -> rusqlite::Connection { + let conn = rusqlite::Connection::open(db_path).unwrap(); + conn.execute_batch( + "CREATE TABLE sessions ( + id TEXT PRIMARY KEY, + source TEXT NOT NULL, + model TEXT, + started_at REAL NOT NULL, + message_count INTEGER DEFAULT 0, + input_tokens INTEGER DEFAULT 0, + output_tokens INTEGER DEFAULT 0, + cache_read_tokens INTEGER DEFAULT 0, + cache_write_tokens INTEGER DEFAULT 0, + reasoning_tokens INTEGER DEFAULT 0, + billing_provider TEXT, + estimated_cost_usd REAL, + actual_cost_usd REAL + );", + ) + .unwrap(); + conn + } - assert_eq!(entries.len(), 2); - let session_ids: HashSet<_> = entries - .iter() - .map(|e| e.session_id.as_deref().unwrap()) - .collect(); - assert_eq!(session_ids, HashSet::from(["session-a", "session-b"])); + fn create_zed_sqlite_db(db_path: &std::path::Path) -> rusqlite::Connection { + let conn = rusqlite::Connection::open(db_path).unwrap(); + conn.execute_batch( + "CREATE TABLE threads ( + id TEXT PRIMARY KEY, + summary TEXT NOT NULL, + updated_at TEXT NOT NULL, + data_type TEXT NOT NULL, + data BLOB NOT NULL + );", + ) + .unwrap(); + conn } - #[test] - fn test_client_session_grouping_keeps_clients_separate() { - // Same session_id seen by two different clients (unusual in practice - // but possible if parsers collide on an id space). ClientSession - // must yield two rows; Session would yield one (covered above). - let entries = aggregate_model_usage_entries( - vec![ - make_workspace_message( - "claude", - "claude-sonnet-4-5-20250929", - "anthropic", - "session-shared", - 1.0, - None, - None, - ), - make_workspace_message( - "amp", - "claude-sonnet-4-5-20250929", - "anthropic", - "session-shared", - 3.0, - None, - None, - ), - ], - &GroupBy::ClientSession, + fn insert_zed_thread(conn: &rusqlite::Connection, id: &str, model: &str) { + let payload = format!( + r#"{{ + "version": "0.3.0", + "title": "Test thread", + "updated_at": "2026-05-01T12:30:00Z", + "request_token_usage": {{ + "turn-1": {{ + "input_tokens": 42, + "output_tokens": 7, + "cache_creation_input_tokens": 3, + "cache_read_input_tokens": 5 + }} + }}, + "model": {{ + "provider": "zed.dev", + "model": "{model}" + }}, + "imported": false + }}"# ); - - assert_eq!(entries.len(), 2); - for entry in &entries { - assert_eq!(entry.session_id.as_deref(), Some("session-shared")); - assert!(entry.merged_clients.is_none()); - } - let by_client: HashSet<_> = entries.iter().map(|e| e.client.as_str()).collect(); - assert_eq!(by_client, HashSet::from(["claude", "amp"])); + conn.execute( + "INSERT INTO threads (id, summary, updated_at, data_type, data) VALUES (?1, ?2, ?3, ?4, ?5)", + rusqlite::params![id, "Test thread", "2026-05-01T12:30:00Z", "json", payload.as_bytes()], + ) + .unwrap(); } - #[test] - fn test_non_session_grouping_does_not_populate_session_id() { - // Defensive: only Session/ClientSession variants should set the - // session_id field on ModelUsage — every other group_by must leave - // it None so the camelCase JSON output omits it via - // `skip_serializing_if = "Option::is_none"`. - for group_by in &[ - GroupBy::Model, - GroupBy::ClientModel, - GroupBy::ClientProviderModel, - GroupBy::WorkspaceModel, - ] { - let entries = aggregate_model_usage_entries( - vec![make_workspace_message( - "codex", - "gpt-5", + fn insert_hermes_session( + conn: &rusqlite::Connection, + id: &str, + model: &str, + message_count: i64, + input_tokens: i64, + output_tokens: i64, + actual_cost_usd: f64, + ) { + conn.execute( + "INSERT INTO sessions ( + id, source, model, started_at, message_count, + input_tokens, output_tokens, cache_read_tokens, cache_write_tokens, reasoning_tokens, + billing_provider, estimated_cost_usd, actual_cost_usd + ) VALUES (?1, 'cli', ?2, 1775001102.0, ?3, ?4, ?5, 0, 0, 0, 'anthropic', NULL, ?6)", + rusqlite::params![ + id, + model, + message_count, + input_tokens, + output_tokens, + actual_cost_usd + ], + ) + .unwrap(); + } + + #[test] + fn test_normalize_model_for_grouping() { + assert_eq!( + normalize_model_for_grouping("claude-opus-4-5-20251101"), + "claude-opus-4-5" + ); + assert_eq!( + normalize_model_for_grouping("claude-sonnet-4-5-20250929"), + "claude-sonnet-4-5" + ); + assert_eq!( + normalize_model_for_grouping("claude-sonnet-4-20250514"), + "claude-sonnet-4" + ); + + assert_eq!( + normalize_model_for_grouping("claude-opus-4.5"), + "claude-opus-4-5" + ); + assert_eq!( + normalize_model_for_grouping("claude-sonnet-4.5"), + "claude-sonnet-4-5" + ); + assert_eq!( + normalize_model_for_grouping("claude-opus-4.6"), + "claude-opus-4-6" + ); + assert_eq!( + normalize_model_for_grouping("anthropic/claude-4-6-sonnet"), + "claude-sonnet-4-6" + ); + assert_eq!( + normalize_model_for_grouping("anthropic/claude-4-5-haiku"), + "claude-haiku-4-5" + ); + assert_eq!( + normalize_model_for_grouping("anthropic/claude-4-6-opus"), + "claude-opus-4-6" + ); + + assert_eq!(normalize_model_for_grouping("gpt-5.2"), "gpt-5.2"); + assert_eq!(normalize_model_for_grouping("gpt-5.4(xhigh)"), "gpt-5.4"); + assert_eq!(normalize_model_for_grouping("gpt-5.4(high)"), "gpt-5.4"); + assert_eq!(normalize_model_for_grouping("gpt-5.4(minimal)"), "gpt-5.4"); + assert_eq!(normalize_model_for_grouping("gpt-5.4(auto)"), "gpt-5.4"); + assert_eq!(normalize_model_for_grouping("gpt-5.4(none)"), "gpt-5.4"); + assert_eq!( + normalize_model_for_grouping("gpt-5.4(weirdgarbage)"), + "gpt-5.4(weirdgarbage)" + ); + assert_eq!( + normalize_model_for_grouping("claude-sonnet-4.5(high)"), + "claude-sonnet-4-5" + ); + assert_eq!( + normalize_model_for_grouping("gemini-3-pro(auto)"), + "gemini-3-pro" + ); + assert_eq!( + normalize_model_for_grouping("gemini-2.5-pro"), + "gemini-2.5-pro" + ); + + assert_eq!( + normalize_model_for_grouping("claude-opus-4-5-high"), + "claude-opus-4-5-high" + ); + assert_eq!( + normalize_model_for_grouping("claude-opus-4-5-thinking-high"), + "claude-opus-4-5-thinking-high" + ); + assert_eq!( + normalize_model_for_grouping("claude-sonnet-4-5-high"), + "claude-sonnet-4-5-high" + ); + + assert_eq!( + normalize_model_for_grouping("claude-4-sonnet"), + "claude-4-sonnet" + ); + assert_eq!( + normalize_model_for_grouping("claude-4-opus-thinking"), + "claude-4-opus-thinking" + ); + + assert_eq!(normalize_model_for_grouping("big-pickle"), "big-pickle"); + assert_eq!(normalize_model_for_grouping("grok-code"), "grok-code"); + + assert_eq!( + normalize_model_for_grouping("claude-opus-4.5-20251101"), + "claude-opus-4-5" + ); + } + + fn test_alias_map(pairs: &[(&str, &str)]) -> ModelAliasMap { + ModelAliasMap { + entries: pairs + .iter() + .map(|(k, v)| ((*k).to_string(), (*v).to_string())) + .collect::>(), + } + } + + #[test] + fn model_aliases_fold_grouping_only_not_canonical_or_pricing() { + let _guard = crate::model_alias::lock_global_alias_tests(); + clear_model_aliases(); + + // Alias a channel-specific spelling onto the display/group key. + set_model_aliases(&test_alias_map(&[("claude-opus-4-8-cc", "claude-opus-4-8")])); + + // Grouping sees the folded label. + assert_eq!( + normalize_model_for_grouping("claude-opus-4-8-cc"), + "claude-opus-4-8" + ); + // Raw identity path stays alias-free (syntactic only). + assert_eq!( + canonical_model_id("claude-opus-4-8-cc"), + "claude-opus-4-8-cc" + ); + + // Pricing resolves the *raw* message model id, never the grouping label. + // Custom pricing is registered only under the raw channel spelling; the + // display/group key has no rate. After the alias fold, cost still uses + // the raw path and remains non-zero. + let mut custom = HashMap::new(); + custom.insert( + "claude-opus-4-8-cc".to_string(), + pricing::ModelPricing { + input_cost_per_token: Some(0.001), + output_cost_per_token: Some(0.002), + ..Default::default() + }, + ); + let service = pricing::PricingService::new_with_custom( + pricing::custom::CustomPricing::from_models(custom), + HashMap::new(), + HashMap::new(), + ); + let tokens = TokenBreakdown { + input: 1000, + output: 500, + cache_read: 0, + cache_write: 0, + reasoning: 0, + }; + let raw_cost = service.calculate_cost_with_provider("claude-opus-4-8-cc", None, &tokens); + let group_label_cost = + service.calculate_cost_with_provider("claude-opus-4-8", None, &tokens); + assert!(raw_cost > 0.0, "raw id must remain the pricing key"); + assert_eq!( + group_label_cost, 0.0, + "grouping label must not be used as the pricing key" + ); + + let mut msg = UnifiedMessage::new( + "claude", + "claude-opus-4-8-cc", + "anthropic", + "s1", + 1_733_011_200_000, + tokens.clone(), + 0.0, + ); + apply_pricing_if_available(&mut msg, Some(&service)); + assert!( + (msg.cost - raw_cost).abs() < 1e-12, + "apply_pricing must keep using the raw message model_id (got {}, expected {})", + msg.cost, + raw_cost + ); + // Message identity is never rewritten by aliases. + assert_eq!(msg.model_id, "claude-opus-4-8-cc"); + + // Model-report grouping merges channel variants under the alias label, + // while pre-computed costs simply sum (aliases never reprice). + let mut other = UnifiedMessage::new( + "claude", + "claude-opus-4-8", + "anthropic", + "s2", + 1_733_011_200_001, + tokens, + 1.5, + ); + other.mark_estimated_cost(); + msg.cost = raw_cost; + msg.mark_estimated_cost(); + let entries = aggregate_model_usage_entries(vec![msg, other], &GroupBy::Model); + assert_eq!( + entries.len(), + 1, + "alias must fold both variants into one bucket" + ); + assert_eq!(entries[0].model, "claude-opus-4-8"); + assert!( + (entries[0].cost - (raw_cost + 1.5)).abs() < 1e-9, + "merged cost is the sum of already-costed buckets, not a reprice" + ); + + // Graph/export client contribution keeps the alias-free raw identity. + let graph = fold_messages_streaming(&[UnifiedMessage::new( + "claude", + "claude-opus-4-8-cc", + "anthropic", + "s1", + 1_733_011_200_000, + TokenBreakdown { + input: 10, + output: 5, + cache_read: 0, + cache_write: 0, + reasoning: 0, + }, + 0.25, + )]); + assert_eq!(graph.len(), 1); + assert_eq!(graph[0].clients.len(), 1); + assert_eq!(graph[0].clients[0].model_id, "claude-opus-4-8-cc"); + + clear_model_aliases(); + } + + #[test] + fn model_alias_reload_invalidates_usage_data_consumers() { + let _guard = crate::model_alias::lock_global_alias_tests(); + clear_model_aliases(); + + static FIRES: AtomicUsize = AtomicUsize::new(0); + register_usage_data_invalidation_hook(|| { + FIRES.fetch_add(1, Ordering::SeqCst); + }); + let baseline_fires = FIRES.load(Ordering::SeqCst); + let gen0 = model_alias_generation(); + + set_model_aliases(&test_alias_map(&[("claude-opus-4-8-cc", "claude-opus-4-8")])); + assert_eq!( + normalize_model_for_grouping("claude-opus-4-8-cc"), + "claude-opus-4-8" + ); + let gen1 = model_alias_generation(); + assert!(gen1 > gen0); + + // Reload replaces the map; consumers see the new fold on the next report. + set_model_aliases(&test_alias_map(&[("gpt-5.5-cc", "gpt-5.5")])); + assert_eq!( + normalize_model_for_grouping("claude-opus-4-8-cc"), + "claude-opus-4-8-cc", + "previous alias must not survive a reload" + ); + assert_eq!(normalize_model_for_grouping("gpt-5.5-cc"), "gpt-5.5"); + let gen2 = model_alias_generation(); + assert!(gen2 > gen1); + + clear_model_aliases(); + assert_eq!(normalize_model_for_grouping("gpt-5.5-cc"), "gpt-5.5-cc"); + assert!(model_alias_generation() > gen2); + + let after_fires = FIRES.load(Ordering::SeqCst); + assert!( + after_fires >= baseline_fires + 3, + "set/reload/clear must each fire the usage-data invalidation hook \ + (baseline={baseline_fires}, after={after_fires})" + ); + } + + #[test] + fn aggregate_model_usage_entries_uses_fold_start_alias_snapshot() { + // Codex P2: multi-message report folds must keep one alias config for + // the whole fold. `aggregate_model_usage_entries` snapshots at start; + // prove the snapshot pattern it uses is stable under mid-fold reload, + // then that the real aggregator still merges under the fold-start map. + let _guard = crate::model_alias::lock_global_alias_tests(); + clear_model_aliases(); + set_model_aliases(&test_alias_map(&[("alias-a", "canonical-b")])); + + // Same capture point as aggregate_model_usage_entries. + let aliases = snapshot_grouping_aliases(); + assert_eq!( + aliases.fold(normalize_syntactic("alias-a")), + "canonical-b", + "fold start sees A→B" + ); + + // Mid-fold mutation of the process-wide map. + set_model_aliases(&test_alias_map(&[("alias-a", "canonical-other")])); + assert_eq!( + normalize_model_for_grouping("alias-a"), + "canonical-other", + "live per-call path sees the reloaded map" + ); + // Every message in the fold keeps the snapshotted label. + for _ in 0..3 { + assert_eq!( + aliases.fold(normalize_syntactic("alias-a")), + "canonical-b", + "snapshot must ignore mid-fold reload for the rest of the fold" + ); + } + + // Re-install fold-start aliases and run the real aggregator: both + // channel variants merge under the snapshotted canonical label. + set_model_aliases(&test_alias_map(&[("alias-a", "canonical-b")])); + let tokens = TokenBreakdown { + input: 10, + output: 5, + cache_read: 0, + cache_write: 0, + reasoning: 0, + }; + let msg_a = UnifiedMessage::new( + "claude", + "alias-a", + "anthropic", + "s1", + 1_733_011_200_000, + tokens.clone(), + 1.0, + ); + let msg_b = UnifiedMessage::new( + "claude", + "canonical-b", + "anthropic", + "s2", + 1_733_011_200_001, + tokens, + 2.0, + ); + let entries = aggregate_model_usage_entries(vec![msg_a, msg_b], &GroupBy::Model); + assert_eq!(entries.len(), 1, "fold-start alias must merge both variants"); + assert_eq!(entries[0].model, "canonical-b"); + assert!((entries[0].cost - 3.0).abs() < 1e-12); + + clear_model_aliases(); + } + + #[test] + fn test_group_by_from_str_valid_values() { + assert_eq!(GroupBy::from_str("model").unwrap(), GroupBy::Model); + assert_eq!( + GroupBy::from_str("client,model").unwrap(), + GroupBy::ClientModel + ); + assert_eq!( + GroupBy::from_str("client-model").unwrap(), + GroupBy::ClientModel + ); + assert_eq!( + GroupBy::from_str("client,provider,model").unwrap(), + GroupBy::ClientProviderModel + ); + assert_eq!( + GroupBy::from_str("client-provider-model").unwrap(), + GroupBy::ClientProviderModel + ); + assert_eq!( + GroupBy::from_str("workspace,model").unwrap(), + GroupBy::WorkspaceModel + ); + assert_eq!( + GroupBy::from_str("workspace-model").unwrap(), + GroupBy::WorkspaceModel + ); + assert_eq!(GroupBy::from_str("session").unwrap(), GroupBy::Session); + assert_eq!( + GroupBy::from_str("session,model").unwrap(), + GroupBy::Session + ); + assert_eq!( + GroupBy::from_str("session-model").unwrap(), + GroupBy::Session + ); + assert_eq!( + GroupBy::from_str("client,session").unwrap(), + GroupBy::ClientSession + ); + assert_eq!( + GroupBy::from_str("client,session,model").unwrap(), + GroupBy::ClientSession + ); + assert_eq!( + GroupBy::from_str("client-session-model").unwrap(), + GroupBy::ClientSession + ); + assert!(GroupBy::from_str("unknown").is_err()); + } + + #[test] + fn test_group_by_default_is_client_model() { + assert_eq!(GroupBy::default(), GroupBy::ClientModel); + } + + #[test] + fn test_group_by_display_round_trips_with_from_str() { + let variants = [ + GroupBy::Model, + GroupBy::ClientModel, + GroupBy::ClientProviderModel, + GroupBy::WorkspaceModel, + GroupBy::Session, + GroupBy::ClientSession, + ]; + + for variant in variants { + let rendered = variant.to_string(); + let parsed = GroupBy::from_str(&rendered).unwrap(); + assert_eq!(parsed, variant); + } + } + + #[test] + fn test_group_by_from_str_whitespace_handling() { + assert_eq!( + GroupBy::from_str("client, model").unwrap(), + GroupBy::ClientModel + ); + assert_eq!(GroupBy::from_str(" model ").unwrap(), GroupBy::Model); + assert_eq!( + GroupBy::from_str("client , provider , model").unwrap(), + GroupBy::ClientProviderModel + ); + assert_eq!( + GroupBy::from_str("workspace, model").unwrap(), + GroupBy::WorkspaceModel + ); + } + + #[test] + fn test_model_usage_performance_uses_only_timed_positive_token_messages() { + let mut timed = make_workspace_message( + "opencode", + "gpt-5.4", + "openai", + "session-1", + 0.0, + None, + None, + ); + timed.tokens = TokenBreakdown { + input: 100, + output: 50, + cache_read: 25, + cache_write: 0, + reasoning: 25, + }; + timed.duration_ms = Some(400); + + let mut untimed = make_workspace_message( + "opencode", + "gpt-5.4", + "openai", + "session-2", + 0.0, + None, + None, + ); + untimed.tokens = TokenBreakdown { + input: 300, + output: 0, + cache_read: 0, + cache_write: 0, + reasoning: 0, + }; + + let entries = aggregate_model_usage_entries(vec![timed, untimed], &GroupBy::ClientModel); + + assert_eq!(entries.len(), 1); + let performance = &entries[0].performance; + assert_eq!(performance.total_duration_ms, 400); + assert_eq!(performance.timed_tokens, 200); + assert_eq!(performance.sample_count, 1); + assert_eq!(performance.ms_per_1k_tokens, Some(2000.0)); + assert!((performance.token_coverage - 0.4).abs() < f64::EPSILON); + } + + #[test] + fn test_model_usage_performance_is_null_without_duration_samples() { + let entries = aggregate_model_usage_entries( + vec![make_workspace_message( + "claude", + "claude-sonnet-4-5", + "anthropic", + "session-1", + 0.0, + None, + None, + )], + &GroupBy::ClientModel, + ); + + assert_eq!(entries.len(), 1); + assert_eq!(entries[0].performance.ms_per_1k_tokens, None); + assert_eq!(entries[0].performance.total_duration_ms, 0); + assert_eq!(entries[0].performance.timed_tokens, 0); + assert_eq!(entries[0].performance.token_coverage, 0.0); + } + + #[test] + fn test_workspace_model_grouping_merges_same_workspace_and_model() { + let entries = aggregate_model_usage_entries( + vec![ + make_workspace_message( + "claude", + "claude-sonnet-4-5-20250929", + "anthropic", + "session-1", + 1.25, + Some("/repo-a"), + Some("repo-a"), + ), + make_workspace_message( + "qwen", + "claude-sonnet-4-5-20250929", + "anthropic", + "session-2", + 2.75, + Some("/repo-a"), + Some("repo-a"), + ), + ], + &GroupBy::WorkspaceModel, + ); + + assert_eq!(entries.len(), 1); + assert_eq!(entries[0].model, "claude-sonnet-4-5"); + assert_eq!(entries[0].workspace_key.as_deref(), Some("/repo-a")); + assert_eq!(entries[0].workspace_label.as_deref(), Some("repo-a")); + assert_eq!(entries[0].cost, 4.0); + assert_eq!(entries[0].message_count, 2); + assert_eq!(entries[0].merged_clients.as_deref(), Some("claude, qwen")); + } + + #[test] + fn test_model_grouping_merges_anthropic_prefixed_claude_variant_with_canonical_model() { + let entries = aggregate_model_usage_entries( + vec![ + make_workspace_message( + "claude", + "anthropic/claude-4-6-sonnet", + "anthropic", + "session-1", + 1.25, + Some("/repo-a"), + Some("repo-a"), + ), + make_workspace_message( + "claude", + "claude-sonnet-4-6", + "anthropic", + "session-2", + 2.75, + Some("/repo-b"), + Some("repo-b"), + ), + ], + &GroupBy::ClientModel, + ); + + assert_eq!(entries.len(), 1); + assert_eq!(entries[0].model, "claude-sonnet-4-6"); + assert_eq!(entries[0].input, 20); + assert_eq!(entries[0].output, 10); + assert_eq!(entries[0].cost, 4.0); + assert_eq!(entries[0].message_count, 2); + } + + #[test] + fn test_workspace_model_grouping_separates_different_workspaces() { + let entries = aggregate_model_usage_entries( + vec![ + make_workspace_message( + "claude", + "claude-sonnet-4-5-20250929", + "anthropic", + "session-1", + 1.0, + Some("/repo-a"), + Some("repo-a"), + ), + make_workspace_message( + "claude", + "claude-sonnet-4-5-20250929", + "anthropic", + "session-2", + 2.0, + Some("/repo-b"), + Some("repo-b"), + ), + ], + &GroupBy::WorkspaceModel, + ); + + assert_eq!(entries.len(), 2); + let labels: HashSet<_> = entries + .iter() + .map(|entry| entry.workspace_label.as_deref().unwrap()) + .collect(); + assert_eq!(labels, HashSet::from(["repo-a", "repo-b"])); + } + + #[test] + fn test_workspace_model_grouping_uses_unknown_bucket_without_workspace_metadata() { + let entries = aggregate_model_usage_entries( + vec![ + make_workspace_message( + "claude", + "claude-sonnet-4-5-20250929", + "anthropic", + "session-1", + 1.0, + None, + None, + ), + make_workspace_message( + "claude", + "claude-sonnet-4-5-20250929", + "anthropic", + "session-2", + "2.0".parse().unwrap(), + None, + None, + ), + ], + &GroupBy::WorkspaceModel, + ); + + assert_eq!(entries.len(), 1); + assert_eq!(entries[0].workspace_key, None); + assert_eq!( + entries[0].workspace_label.as_deref(), + Some(UNKNOWN_WORKSPACE_LABEL) + ); + assert_eq!(entries[0].message_count, 2); + assert_eq!(entries[0].cost, 3.0); + } + + #[test] + fn test_parsed_round_trip_preserves_workspace_metadata() { + let mut unified = UnifiedMessage::new( + "qwen", + "qwen3.5-plus", + "qwen", + "session-1", + 1_742_390_400_000, + TokenBreakdown { + input: 10, + output: 5, + cache_read: 2, + cache_write: 0, + reasoning: 1, + }, + 1.25, + ); + unified.set_workspace( + Some("//server/share/demo-workspace".to_string()), + Some("demo-workspace".to_string()), + ); + unified.duration_ms = Some(2500); + + let parsed = unified_to_parsed(&unified); + let round_tripped = parsed_to_unified(&parsed, 2.5); + + assert_eq!( + round_tripped.workspace_key.as_deref(), + Some("//server/share/demo-workspace") + ); + assert_eq!( + round_tripped.workspace_label.as_deref(), + Some("demo-workspace") + ); + assert_eq!(round_tripped.cost, 2.5); + assert_eq!(round_tripped.duration_ms, Some(2500)); + } + + #[test] + fn test_workspace_model_grouping_keeps_real_unknown_workspace_separate() { + let entries = aggregate_model_usage_entries( + vec![ + make_workspace_message( + "claude", + "claude-sonnet-4-5-20250929", + "anthropic", + "session-1", + 1.0, + Some("unknown-workspace"), + Some("unknown-workspace"), + ), + make_workspace_message( + "claude", + "claude-sonnet-4-5-20250929", + "anthropic", + "session-2", + 2.0, + None, + None, + ), + ], + &GroupBy::WorkspaceModel, + ); + + assert_eq!(entries.len(), 2); + assert!(entries.iter().any(|entry| { + entry.workspace_key.as_deref() == Some("unknown-workspace") + && entry.workspace_label.as_deref() == Some("unknown-workspace") + && (entry.cost - 1.0).abs() < f64::EPSILON + })); + assert!(entries.iter().any(|entry| { + entry.workspace_key.is_none() + && entry.workspace_label.as_deref() == Some(UNKNOWN_WORKSPACE_LABEL) + && (entry.cost - 2.0).abs() < f64::EPSILON + })); + } + + #[test] + fn test_session_grouping_merges_same_session_and_model() { + // Two messages with the same session_id + same model — should collapse + // into one row regardless of the client that produced them, because + // GroupBy::Session keys on (session_id, model) only. + let entries = aggregate_model_usage_entries( + vec![ + make_workspace_message( + "claude", + "claude-sonnet-4-5-20250929", + "anthropic", + "session-shared", + 1.25, + None, + None, + ), + make_workspace_message( + "amp", + "claude-sonnet-4-5-20250929", + "anthropic", + "session-shared", + 2.75, + None, + None, + ), + ], + &GroupBy::Session, + ); + + assert_eq!(entries.len(), 1); + assert_eq!(entries[0].session_id.as_deref(), Some("session-shared")); + assert_eq!(entries[0].model, "claude-sonnet-4-5"); + assert!((entries[0].cost - 4.0).abs() < f64::EPSILON); + assert_eq!(entries[0].message_count, 2); + assert!(entries[0].workspace_key.is_none()); + assert!(entries[0].workspace_label.is_none()); + // Session grouping does not merge_clients into a comma list. + assert!(entries[0].merged_clients.is_none()); + } + + #[test] + fn test_session_grouping_separates_different_sessions() { + let entries = aggregate_model_usage_entries( + vec![ + make_workspace_message("codex", "gpt-5", "openai", "session-a", 1.0, None, None), + make_workspace_message("codex", "gpt-5", "openai", "session-b", 2.0, None, None), + ], + &GroupBy::Session, + ); + + assert_eq!(entries.len(), 2); + let session_ids: HashSet<_> = entries + .iter() + .map(|e| e.session_id.as_deref().unwrap()) + .collect(); + assert_eq!(session_ids, HashSet::from(["session-a", "session-b"])); + } + + #[test] + fn test_client_session_grouping_keeps_clients_separate() { + // Same session_id seen by two different clients (unusual in practice + // but possible if parsers collide on an id space). ClientSession + // must yield two rows; Session would yield one (covered above). + let entries = aggregate_model_usage_entries( + vec![ + make_workspace_message( + "claude", + "claude-sonnet-4-5-20250929", + "anthropic", + "session-shared", + 1.0, + None, + None, + ), + make_workspace_message( + "amp", + "claude-sonnet-4-5-20250929", + "anthropic", + "session-shared", + 3.0, + None, + None, + ), + ], + &GroupBy::ClientSession, + ); + + assert_eq!(entries.len(), 2); + for entry in &entries { + assert_eq!(entry.session_id.as_deref(), Some("session-shared")); + assert!(entry.merged_clients.is_none()); + } + let by_client: HashSet<_> = entries.iter().map(|e| e.client.as_str()).collect(); + assert_eq!(by_client, HashSet::from(["claude", "amp"])); + } + + #[test] + fn test_non_session_grouping_does_not_populate_session_id() { + // Defensive: only Session/ClientSession variants should set the + // session_id field on ModelUsage — every other group_by must leave + // it None so the camelCase JSON output omits it via + // `skip_serializing_if = "Option::is_none"`. + for group_by in &[ + GroupBy::Model, + GroupBy::ClientModel, + GroupBy::ClientProviderModel, + GroupBy::WorkspaceModel, + ] { + let entries = aggregate_model_usage_entries( + vec![make_workspace_message( + "codex", + "gpt-5", "openai", "session-x", 1.0, @@ -5170,1715 +7500,3252 @@ mod tests { )], group_by, ); - assert_eq!(entries.len(), 1); - assert!( - entries[0].session_id.is_none(), - "session_id leaked into {:?} grouping", - group_by + assert_eq!(entries.len(), 1); + assert!( + entries[0].session_id.is_none(), + "session_id leaked into {:?} grouping", + group_by + ); + } + } + + #[test] + fn test_retain_for_requested_clients_keeps_original_client_matches() { + let requested: HashSet<&str> = HashSet::from(["opencode"]); + assert!(retain_for_requested_clients( + "opencode", + "gpt-4o", + "anthropic", + &requested + )); + assert!(!retain_for_requested_clients( + "claude", + "gpt-4o", + "anthropic", + &requested + )); + } + + #[test] + fn test_retain_for_requested_clients_accepts_synthetic_gateway_traffic() { + let requested: HashSet<&str> = HashSet::from(["synthetic"]); + assert!(retain_for_requested_clients( + "opencode", + "hf:deepseek-ai/DeepSeek-V3-0324", + "unknown", + &requested + )); + assert!(retain_for_requested_clients( + "synthetic", + "deepseek-v3-0324", + "synthetic", + &requested + )); + assert!(!retain_for_requested_clients( + "opencode", + "gpt-4o", + "anthropic", + &requested + )); + } + + #[test] + fn test_retain_for_requested_clients_preserves_kilo_split() { + let kilocode_only: HashSet<&str> = HashSet::from(["kilocode"]); + assert!(retain_for_requested_clients( + "kilocode", + "gpt-5", + "openai", + &kilocode_only + )); + assert!(!retain_for_requested_clients( + "kilo", + "gpt-5", + "openai", + &kilocode_only + )); + + let kilo_only: HashSet<&str> = HashSet::from(["kilo"]); + assert!(retain_for_requested_clients( + "kilo", "gpt-5", "openai", &kilo_only + )); + assert!(!retain_for_requested_clients( + "kilocode", "gpt-5", "openai", &kilo_only + )); + } + + #[test] + #[serial_test::serial] + fn test_cursor_parse_path_reprices_zero_cost_composer_1_5_rows() { + let cache_home = tempfile::TempDir::new().unwrap(); + let _env = EnvGuard::set(&[ + ("HOME", cache_home.path().as_os_str()), + ("TOKSCALE_CONFIG_DIR", cache_home.path().as_os_str()), + ]); + let temp_dir = tempfile::TempDir::new().unwrap(); + let cursor_cache_dir = temp_dir.path().join(".config/tokscale/cursor-cache"); + std::fs::create_dir_all(&cursor_cache_dir).unwrap(); + + let csv = r#"Date,Kind,Model,Max Mode,Input (w/ Cache Write),Input (w/o Cache Write),Cache Read,Output Tokens,Total Tokens,Cost +"2026-03-04T12:00:00.000Z","Included","Composer 1.5","No","1200","1000","5000","2000","8000","0""#; + std::fs::write(cursor_cache_dir.join("usage.csv"), csv).unwrap(); + + let pricing = pricing::PricingService::new(HashMap::new(), HashMap::new()); + let messages = parse_all_messages_with_pricing( + temp_dir.path().to_str().unwrap(), + &["cursor".to_string()], + Some(&pricing), + ); + + assert_eq!(messages.len(), 1); + assert_eq!(messages[0].client, "cursor"); + assert_eq!(messages[0].model_id, "Composer 1.5"); + assert!(messages[0].cost > 0.0); + } + + fn write_kimi_repeated_status_fixture(source_home: &std::path::Path) { + let session_dir = source_home.join(".kimi/sessions/group-1/session-1"); + std::fs::create_dir_all(&session_dir).unwrap(); + std::fs::write( + session_dir.join("wire.jsonl"), + r#"{"type": "metadata", "protocol_version": "1.3"} +{"timestamp": 1770983410.0, "message": {"type": "StatusUpdate", "payload": {"token_usage": {"input_other": 10, "output": 1, "input_cache_read": 0, "input_cache_creation": 0}, "message_id": "msg-progressive"}}} +{"timestamp": 1770983420.0, "message": {"type": "StatusUpdate", "payload": {"token_usage": {"input_other": 20, "output": 2, "input_cache_read": 0, "input_cache_creation": 0}, "message_id": "msg-progressive"}}} +{"timestamp": 1770983430.0, "message": {"type": "StatusUpdate", "payload": {"token_usage": {"input_other": 5, "output": 1, "input_cache_read": 0, "input_cache_creation": 0}, "message_id": "msg-distinct"}}} +{"timestamp": 1770983440.0, "message": {"type": "StatusUpdate", "payload": {"token_usage": {"input_other": 7, "output": 1, "input_cache_read": 0, "input_cache_creation": 0}}}} +{"timestamp": 1770983450.0, "message": {"type": "StatusUpdate", "payload": {"token_usage": {"input_other": 8, "output": 1, "input_cache_read": 0, "input_cache_creation": 0}}}}"#, + ) + .unwrap(); + } + + #[test] + #[serial_test::serial] + fn test_parse_all_messages_with_pricing_kimi_deduplicates_repeated_status_updates() { + let cache_home = tempfile::TempDir::new().unwrap(); + let source_home = tempfile::TempDir::new().unwrap(); + let _env = EnvGuard::set(&[ + ("HOME", cache_home.path().as_os_str()), + ("TOKSCALE_CONFIG_DIR", cache_home.path().as_os_str()), + ]); + + { + write_kimi_repeated_status_fixture(source_home.path()); + + let messages = parse_all_messages_with_pricing( + source_home.path().to_str().unwrap(), + &["kimi".to_string()], + None, + ); + + assert_eq!(messages.len(), 4); + assert_eq!(messages.iter().map(|m| m.tokens.input).sum::(), 40); + assert_eq!(messages.iter().map(|m| m.tokens.output).sum::(), 5); + } + } + + #[test] + #[serial_test::serial] + fn test_parse_local_clients_kimi_deduplicates_repeated_status_updates() { + let cache_home = tempfile::TempDir::new().unwrap(); + let source_home = tempfile::TempDir::new().unwrap(); + let _env = EnvGuard::set(&[ + ("HOME", cache_home.path().as_os_str()), + ("TOKSCALE_CONFIG_DIR", cache_home.path().as_os_str()), + ]); + + { + write_kimi_repeated_status_fixture(source_home.path()); + + let parsed = parse_local_clients(LocalParseOptions { + home_dir: Some(source_home.path().to_str().unwrap().to_string()), + use_env_roots: false, + clients: Some(vec!["kimi".to_string()]), + since: None, + until: None, + year: None, + scanner_settings: scanner::ScannerSettings::default(), + modified_after: None, + }) + .unwrap(); + + assert_eq!(parsed.counts.get(ClientId::Kimi), 4); + assert_eq!(parsed.messages.len(), 4); + assert_eq!(parsed.messages.iter().map(|m| m.input).sum::(), 40); + assert_eq!(parsed.messages.iter().map(|m| m.output).sum::(), 5); + } + } + + // Regression: the streaming driver must NOT share one dedup set across + // different clients. kimi and codebuff both emit raw upstream message ids + // as dedup_key with no client namespace, so a shared set would let one + // client's key suppress an identical key from the other. Here both a kimi + // message and a codebuff message carry dedup_key "COLLIDE"; both must + // survive. With a single shared `seen_keys` (the pre-fix behaviour) the + // second lane's message is silently dropped and this fails. + #[test] + #[serial_test::serial] + fn test_streaming_driver_does_not_dedup_across_clients() { + let cache_home = tempfile::TempDir::new().unwrap(); + let source_home = tempfile::TempDir::new().unwrap(); + let _env = EnvGuard::set(&[ + ("HOME", cache_home.path().as_os_str()), + ("TOKSCALE_CONFIG_DIR", cache_home.path().as_os_str()), + ]); + + { + // kimi: one StatusUpdate carrying message_id "COLLIDE". + let kimi_dir = source_home.path().join(".kimi/sessions/g/s"); + std::fs::create_dir_all(&kimi_dir).unwrap(); + std::fs::write( + kimi_dir.join("wire.jsonl"), + r#"{"type": "metadata", "protocol_version": "1.3"} +{"timestamp": 1770983410.0, "message": {"type": "StatusUpdate", "payload": {"token_usage": {"input_other": 100, "output": 50, "input_cache_read": 0, "input_cache_creation": 0}, "message_id": "COLLIDE"}}}"#, + ) + .unwrap(); + + // codebuff: one assistant message whose upstream id is "COLLIDE". + let cb_dir = source_home.path().join(".config/manicode/projects/proj"); + std::fs::create_dir_all(&cb_dir).unwrap(); + std::fs::write( + cb_dir.join("chat-messages.json"), + r#"[{"role":"assistant","id":"COLLIDE","metadata":{"model":"claude-sonnet-4","usage":{"inputTokens":200,"outputTokens":80}},"credits":0.02}]"#, + ) + .unwrap(); + + let mut seen: Vec = Vec::new(); + scan_messages_streaming( + source_home.path().to_str().unwrap(), + &["kimi".to_string(), "codebuff".to_string()], + None, + false, + &scanner::ScannerSettings::default(), + &|_m: &UnifiedMessage| true, + &mut |m: &UnifiedMessage| seen.push(m.client.clone()), + ); + + assert!( + seen.iter().any(|c| c == "kimi"), + "kimi message with shared dedup_key must survive: {seen:?}" + ); + assert!( + seen.iter().any(|c| c == "codebuff"), + "codebuff message with shared dedup_key must survive: {seen:?}" + ); + } + } + + // M2 (codex fork-replay): the parser-level fork dedup (#649/#681) must also + // collapse replayed parent token_count rows through OUR streaming report + // path (scan_messages_streaming), not just the materialized + // parse_all_messages_with_pricing path the upstream tests exercise. Without + // the fork-parent-scoped dedup key, each fork's replayed parent rows survive + // per child and inflate codex totals. + #[test] + #[serial_test::serial] + fn test_streaming_codex_collapses_parent_replay_across_forks() { + let cache_home = tempfile::TempDir::new().unwrap(); + let source_home = tempfile::TempDir::new().unwrap(); + let _env = EnvGuard::set(&[ + ("HOME", cache_home.path().as_os_str()), + ("TOKSCALE_CONFIG_DIR", cache_home.path().as_os_str()), + ]); + + { + write_codex_parent_replay_fixture(source_home.path()); + + let mut input_sum = 0i64; + let mut output_sum = 0i64; + let mut count = 0usize; + scan_messages_streaming( + source_home.path().to_str().unwrap(), + &["codex".to_string()], + None, + false, + &scanner::ScannerSettings::default(), + &|_m: &UnifiedMessage| true, + &mut |m: &UnifiedMessage| { + input_sum += m.tokens.input; + output_sum += m.tokens.output; + count += 1; + }, + ); + + // Same collapse as the materialized + // test_parse_all_messages_with_pricing_codex_deduplicates_parent_replay_across_forks: + // the parent's two turns plus the single own-turn shared (by identical + // cumulative total) across the two forks. Without #649/#681 the + // replayed parent rows would survive per fork and inflate this. + assert_eq!(count, 3, "replayed parent rows must collapse to 3 messages"); + assert_eq!(input_sum, 140); + assert_eq!(output_sum, 14); + } + } + + // Issue #6: the agents report must dedup the simple_lane! clients + // (copilot/codebuff/kimi/…) like the model/graph/hourly reports. Here + // codebuff emits the SAME upstream message id "DUP" in two different + // project files. The OLD materialized path (parse_local_unified_messages) + // never gated codebuff, so it counts both; the streaming-backed + // get_agents_report keeps one — matching get_model_report. Repointing + // get_agents_report at the old path makes the parity assertion FAIL (RED). + #[test] + #[serial_test::serial] + fn test_agents_report_dedups_like_model_report_issue6() { + let cache_home = tempfile::TempDir::new().unwrap(); + let source_home = tempfile::TempDir::new().unwrap(); + let _env = EnvGuard::set(&[ + ("HOME", cache_home.path().as_os_str()), + ("TOKSCALE_CONFIG_DIR", cache_home.path().as_os_str()), + ]); + // Hermetic: cache-only pricing + temp HOME → no network, pricing None. + let _pricing = EnvGuard::set(&[("TOKSCALE_PRICING_CACHE_ONLY", std::ffi::OsStr::new("1"))]); + + { + let write_codebuff = |proj: &str| { + let dir = source_home + .path() + .join(format!(".config/manicode/projects/{proj}")); + std::fs::create_dir_all(&dir).unwrap(); + std::fs::write( + dir.join("chat-messages.json"), + r#"[{"role":"assistant","id":"DUP","metadata":{"model":"claude-sonnet-4","usage":{"inputTokens":200,"outputTokens":80}},"credits":0.02}]"#, + ) + .unwrap(); + }; + write_codebuff("projA"); + write_codebuff("projB"); + + let home = source_home.path().to_str().unwrap().to_string(); + let clients = Some(vec!["codebuff".to_string()]); + let rt = tokio::runtime::Builder::new_current_thread() + .enable_all() + .build() + .unwrap(); + + let agents = rt + .block_on(get_agents_report(ReportOptions { + home_dir: Some(home.clone()), + use_env_roots: false, + clients: clients.clone(), + ..Default::default() + })) + .unwrap(); + let model = rt + .block_on(get_model_report(ReportOptions { + home_dir: Some(home.clone()), + use_env_roots: false, + clients: clients.clone(), + ..Default::default() + })) + .unwrap(); + let old = rt + .block_on(parse_local_unified_messages(LocalParseOptions { + home_dir: Some(home.clone()), + use_env_roots: false, + clients: clients.clone(), + since: None, + until: None, + year: None, + scanner_settings: scanner::ScannerSettings::default(), + modified_after: None, + })) + .unwrap(); + let old_total: i32 = old.iter().map(|m| m.message_count.max(0)).sum(); + + assert_eq!(old_total, 2, "old materialized path must NOT dedup codebuff"); + assert_eq!(model.total_messages, 1, "model report dedups codebuff"); + assert_eq!(agents.total_messages, 1, "agents report must dedup codebuff"); + assert_eq!( + agents.total_messages, model.total_messages, + "issue #6: agents must agree with the model report" + ); + assert_ne!( + old_total, model.total_messages, + "the old path diverged from the model report (the #6 bug)" + ); + assert!( + (agents.total_cost - model.total_cost).abs() < 1e-9, + "agents/model cost parity (agents={}, model={})", + agents.total_cost, + model.total_cost + ); + } + } + + // Preservation: with no duplicate dedup_keys (and only parse_local==true + // clients), the streaming-backed agents report produces the SAME numbers the + // old materialized path did. codebuff + kimi, distinct ids, no agent + // attribution → a single "Main" bucket. + #[test] + #[serial_test::serial] + fn test_agents_report_preserves_numbers_without_duplicates() { + let cache_home = tempfile::TempDir::new().unwrap(); + let source_home = tempfile::TempDir::new().unwrap(); + let _env = EnvGuard::set(&[ + ("HOME", cache_home.path().as_os_str()), + ("TOKSCALE_CONFIG_DIR", cache_home.path().as_os_str()), + ]); + let _pricing = EnvGuard::set(&[("TOKSCALE_PRICING_CACHE_ONLY", std::ffi::OsStr::new("1"))]); + + { + let cb_dir = source_home.path().join(".config/manicode/projects/proj"); + std::fs::create_dir_all(&cb_dir).unwrap(); + std::fs::write( + cb_dir.join("chat-messages.json"), + r#"[{"role":"assistant","id":"A","metadata":{"model":"claude-sonnet-4","usage":{"inputTokens":200,"outputTokens":80}},"credits":0.02}]"#, + ) + .unwrap(); + let kimi_dir = source_home.path().join(".kimi/sessions/g/s"); + std::fs::create_dir_all(&kimi_dir).unwrap(); + std::fs::write( + kimi_dir.join("wire.jsonl"), + "{\"type\": \"metadata\", \"protocol_version\": \"1.3\"}\n{\"timestamp\": 1770983410.0, \"message\": {\"type\": \"StatusUpdate\", \"payload\": {\"token_usage\": {\"input_other\": 100, \"output\": 50, \"input_cache_read\": 0, \"input_cache_creation\": 0}, \"message_id\": \"K\"}}}", + ) + .unwrap(); + + let home = source_home.path().to_str().unwrap().to_string(); + let clients = Some(vec!["codebuff".to_string(), "kimi".to_string()]); + let rt = tokio::runtime::Builder::new_current_thread() + .enable_all() + .build() + .unwrap(); + + let agents = rt + .block_on(get_agents_report(ReportOptions { + home_dir: Some(home.clone()), + use_env_roots: false, + clients: clients.clone(), + ..Default::default() + })) + .unwrap(); + + assert_eq!( + agents.entries.len(), + 1, + "no agent attribution → a single Main bucket" + ); + let main = &agents.entries[0]; + assert_eq!(main.agent, "Main"); + assert_eq!(main.messages, 2); + // BTreeSet → sorted, both clients fold into Main. + assert_eq!(main.clients, vec!["codebuff".to_string(), "kimi".to_string()]); + + // Byte-for-byte equivalence with the old materialized path for the + // non-duplicate case (both parse identically; only dedup differs). + let old = rt + .block_on(parse_local_unified_messages(LocalParseOptions { + home_dir: Some(home.clone()), + use_env_roots: false, + clients, + since: None, + until: None, + year: None, + scanner_settings: scanner::ScannerSettings::default(), + modified_after: None, + })) + .unwrap(); + let old_input: i64 = old.iter().map(|m| m.tokens.input).sum(); + let old_output: i64 = old.iter().map(|m| m.tokens.output).sum(); + let old_cache_read: i64 = old.iter().map(|m| m.tokens.cache_read).sum(); + let old_cache_write: i64 = old.iter().map(|m| m.tokens.cache_write).sum(); + let old_reasoning: i64 = old.iter().map(|m| m.tokens.reasoning).sum(); + let old_messages: i32 = old.iter().map(|m| m.message_count.max(0)).sum(); + let old_cost: f64 = old.iter().map(|m| m.cost).sum(); + + assert_eq!(main.input, old_input); + assert_eq!(main.output, old_output); + assert_eq!(main.cache_read, old_cache_read); + assert_eq!(main.cache_write, old_cache_write); + assert_eq!(main.reasoning, old_reasoning); + assert_eq!(main.messages, old_messages); + assert!((agents.total_cost - old_cost).abs() < 1e-9); + // Sanity: codebuff contributes its known tokens. + assert!(main.input >= 200 && main.output >= 80); + } + } + + // Issue #36: the client selection must be applied at the STREAMING SCAN, + // not by a downstream membership filter over the pre-aggregated buckets. + // codebuff (200/80) and kimi (100/50) fold into ONE shared "Main" agent + // bucket; the FFI/DashboardModel now thread the selection into + // ReportOptions.clients, so filtering to codebuff yields ONLY its 200/80 — + // NOT the mixed 300/130. The old approach (unfiltered report + a Swift + // membership filter over whole buckets) kept the entire shared bucket and + // would read 300/130 here, so this test is RED against it and GREEN now. + #[test] + #[serial_test::serial] + fn test_agents_report_client_filter_scopes_shared_bucket_issue36() { + let cache_home = tempfile::TempDir::new().unwrap(); + let source_home = tempfile::TempDir::new().unwrap(); + let _env = EnvGuard::set(&[ + ("HOME", cache_home.path().as_os_str()), + ("TOKSCALE_CONFIG_DIR", cache_home.path().as_os_str()), + ]); + let _pricing = EnvGuard::set(&[("TOKSCALE_PRICING_CACHE_ONLY", std::ffi::OsStr::new("1"))]); + + { + let cb_dir = source_home.path().join(".config/manicode/projects/proj"); + std::fs::create_dir_all(&cb_dir).unwrap(); + std::fs::write( + cb_dir.join("chat-messages.json"), + r#"[{"role":"assistant","id":"A","metadata":{"model":"claude-sonnet-4","usage":{"inputTokens":200,"outputTokens":80}},"credits":0.02}]"#, + ) + .unwrap(); + let kimi_dir = source_home.path().join(".kimi/sessions/g/s"); + std::fs::create_dir_all(&kimi_dir).unwrap(); + std::fs::write( + kimi_dir.join("wire.jsonl"), + "{\"type\": \"metadata\", \"protocol_version\": \"1.3\"}\n{\"timestamp\": 1770983410.0, \"message\": {\"type\": \"StatusUpdate\", \"payload\": {\"token_usage\": {\"input_other\": 100, \"output\": 50, \"input_cache_read\": 0, \"input_cache_creation\": 0}, \"message_id\": \"K\"}}}", + ) + .unwrap(); + + let home = source_home.path().to_str().unwrap().to_string(); + let rt = tokio::runtime::Builder::new_current_thread() + .enable_all() + .build() + .unwrap(); + let run = |clients: Option>| { + rt.block_on(get_agents_report(ReportOptions { + home_dir: Some(home.clone()), + use_env_roots: false, + clients, + ..Default::default() + })) + .unwrap() + }; + + // All clients: one shared "Main" bucket carrying the mixed total. + let all = run(None); + assert_eq!(all.entries.len(), 1, "codebuff + kimi share one Main bucket"); + assert_eq!(all.entries[0].agent, "Main"); + assert_eq!(all.entries[0].input, 300, "mixed bucket = 200 + 100"); + assert_eq!(all.entries[0].output, 130, "mixed bucket = 80 + 50"); + assert_eq!(all.total_messages, 2); + + // Filtered to codebuff: the SAME shared bucket, scoped at the scan + // to codebuff's contribution alone — proves the FFI-level filter, + // not a whole-bucket membership keep (which would still read 300). + let filtered = run(Some(vec!["codebuff".to_string()])); + assert_eq!(filtered.entries.len(), 1); + assert_eq!(filtered.entries[0].agent, "Main"); + assert_eq!( + filtered.entries[0].input, 200, + "filtered = codebuff only, not the mixed 300" + ); + assert_eq!(filtered.entries[0].output, 80); + assert_eq!(filtered.entries[0].clients, vec!["codebuff".to_string()]); + assert_eq!(filtered.total_messages, 1, "kimi's message is gone"); + } + } + + // Issue #36 (round 3): a cc-mirror variant id (`cc-mirror/kimi-code`) is + // produced during CLAUDE-lane parsing, not by a scanner lane of its own — + // `ClientId::from_str("cc-mirror/kimi-code")` is None. The two-level split + // must (a) map the variant to its producing `claude` lane so the scan finds + // it, and (b) keep ONLY the exact requested ids at fold time so requesting + // the variant returns just the variant, and requesting `claude` returns + // plain claude WITHOUT the variant (which the graph/daily/models surface as + // its own client id). RED against the old lane-only filter: requesting the + // variant returned empty, and requesting claude swept the variant in. + #[test] + #[serial_test::serial] + fn test_agents_report_cc_mirror_variant_slice_issue36() { + let cache_home = tempfile::TempDir::new().unwrap(); + let source_home = tempfile::TempDir::new().unwrap(); + let _env = EnvGuard::set(&[ + ("HOME", cache_home.path().as_os_str()), + ("TOKSCALE_CONFIG_DIR", cache_home.path().as_os_str()), + ]); + let _pricing = EnvGuard::set(&[("TOKSCALE_PRICING_CACHE_ONLY", std::ffi::OsStr::new("1"))]); + + { + // Plain claude session (client "claude"): 100 in / 50 out. + let claude_dir = source_home.path().join(".claude/projects/myproject"); + std::fs::create_dir_all(&claude_dir).unwrap(); + std::fs::write( + claude_dir.join("conversation.jsonl"), + r#"{"type":"assistant","timestamp":"2024-12-01T10:00:00.000Z","requestId":"req_plain","message":{"id":"msg_plain","model":"claude-3-5-sonnet","usage":{"input_tokens":100,"output_tokens":50}}}"#, + ) + .unwrap(); + // cc-mirror variant (client "cc-mirror/kimi-code"): 300 in / 70 out. + let variant_dir = source_home.path().join(".cc-mirror/kimi-code"); + let config_dir = variant_dir.join("config"); + let project_dir = config_dir.join("projects/proj"); + std::fs::create_dir_all(&project_dir).unwrap(); + std::fs::write( + variant_dir.join("variant.json"), + serde_json::json!({ + "name": "kimi-code", + "provider": "kimi", + "configDir": config_dir, + }) + .to_string(), + ) + .unwrap(); + std::fs::write( + project_dir.join("session.jsonl"), + r#"{"type":"assistant","timestamp":"2024-12-01T11:00:00.000Z","requestId":"req_variant","message":{"id":"msg_variant","model":"claude-3-5-sonnet","usage":{"input_tokens":300,"output_tokens":70}}}"#, + ) + .unwrap(); + + let home = source_home.path().to_str().unwrap().to_string(); + let rt = tokio::runtime::Builder::new_current_thread() + .enable_all() + .build() + .unwrap(); + let run = |clients: Option>| { + let report = rt + .block_on(get_agents_report(ReportOptions { + home_dir: Some(home.clone()), + use_env_roots: false, + clients, + ..Default::default() + })) + .unwrap(); + let input: i64 = report.entries.iter().map(|e| e.input).sum(); + (report.total_messages, input) + }; + + // All clients: both messages fold into "Main". + assert_eq!(run(None), (2, 400), "all: plain(100) + variant(300)"); + + // Variant slice: the claude lane is scanned (so the variant is + // found), then narrowed to exactly the variant id. + assert_eq!( + run(Some(vec!["cc-mirror/kimi-code".to_string()])), + (1, 300), + "variant slice = just the variant (300), not empty" + ); + + // Claude slice: plain claude ONLY — the distinct variant is excluded. + assert_eq!( + run(Some(vec!["claude".to_string()])), + (1, 100), + "claude slice = plain claude (100), not the mixed 400" ); } } + // Agent bucketing + fold arithmetic in isolation (no fixtures): normalized + // names, the "Main" fallback, plain `+=` token sums, and message_count.max(0). #[test] - fn test_retain_for_requested_clients_keeps_original_client_matches() { - let requested: HashSet<&str> = HashSet::from(["opencode"]); - assert!(retain_for_requested_clients( - "opencode", - "gpt-4o", - "anthropic", - &requested - )); - assert!(!retain_for_requested_clients( - "claude", - "gpt-4o", - "anthropic", - &requested - )); + fn test_agent_bucket_key_and_accumulator() { + let msg = |agent: Option<&str>| { + let mut m = UnifiedMessage::new_with_agent( + "codebuff", + "m", + "p", + "s", + 0, + TokenBreakdown { + input: 10, + output: 5, + cache_read: 2, + cache_write: 1, + reasoning: 3, + }, + 0.5, + agent.map(|a| a.to_string()), + ); + m.message_count = 2; + m + }; + + assert_eq!(agent_bucket_key(&msg(None)), "Main"); + assert_eq!(agent_bucket_key(&msg(Some(" "))), "Main"); + assert_eq!(agent_bucket_key(&msg(Some("OmO"))), "Sisyphus"); + + let mut acc = AgentAccumulator::default(); + acc.add(&msg(None)); + let mut negative = msg(None); + negative.message_count = -3; // .max(0) clamp → contributes 0 messages + acc.add(&negative); + + assert_eq!(acc.input, 20); + assert_eq!(acc.output, 10); + assert_eq!(acc.cache_read, 4); + assert_eq!(acc.cache_write, 2); + assert_eq!(acc.reasoning, 6); + assert!((acc.cost - 1.0).abs() < 1e-9); + assert_eq!(acc.messages, 2, "message_count.max(0): 2 + 0"); + assert!(acc.clients.contains("codebuff")); } #[test] - fn test_retain_for_requested_clients_accepts_synthetic_gateway_traffic() { - let requested: HashSet<&str> = HashSet::from(["synthetic"]); - assert!(retain_for_requested_clients( - "opencode", - "hf:deepseek-ai/DeepSeek-V3-0324", - "unknown", - &requested - )); - assert!(retain_for_requested_clients( - "synthetic", - "deepseek-v3-0324", - "synthetic", - &requested - )); - assert!(!retain_for_requested_clients( - "opencode", - "gpt-4o", - "anthropic", - &requested - )); + fn agent_accumulator_saturates_overflowing_token_folds() { + // Vendor-local sibling sweep alongside #823: AgentAccumulator::add is + // its own per-field CROSS-MESSAGE fold (agents streaming report), not + // one of the 6 sites #823 covers. An antigravity-cli row can carry an + // i64::MAX bucket after the untrusted-varint clamp, so two such rows + // folded into one agent bucket with plain `+=` overflow (debug panic / + // release wrap) before any saturating grand total runs. + let make = || { + UnifiedMessage::new( + "antigravity-cli", + "gemini-3-pro", + "antigravity", + "session-overflow", + 1_733_011_200_000, + TokenBreakdown { + input: i64::MAX, + output: 0, + cache_read: i64::MAX, + cache_write: 0, + reasoning: 0, + }, + 0.0, + ) + }; + + let mut acc = AgentAccumulator::default(); + acc.add(&make()); + acc.add(&make()); + + assert_eq!(acc.input, i64::MAX); + assert_eq!(acc.cache_read, i64::MAX); } #[test] - fn test_retain_for_requested_clients_preserves_kilo_split() { - let kilocode_only: HashSet<&str> = HashSet::from(["kilocode"]); - assert!(retain_for_requested_clients( - "kilocode", - "gpt-5", - "openai", - &kilocode_only - )); - assert!(!retain_for_requested_clients( - "kilo", - "gpt-5", - "openai", - &kilocode_only - )); + fn test_agent_bucket_key_copilot_uses_copilot_normalizer() { + // #724/#751: copilot messages carry a raw OTEL agent id. Our agents + // report must prettify it with the copilot-specific normalizer (the + // prettification upstream does in its CLI), while other clients keep the + // generic normalization. + let msg = |client: &str, agent: &str| { + UnifiedMessage::new_with_agent( + client, + "m", + "p", + "s", + 0, + TokenBreakdown::default(), + 0.0, + Some(agent.to_string()), + ) + }; - let kilo_only: HashSet<&str> = HashSet::from(["kilo"]); - assert!(retain_for_requested_clients( - "kilo", "gpt-5", "openai", &kilo_only - )); - assert!(!retain_for_requested_clients( - "kilocode", "gpt-5", "openai", &kilo_only - )); + // Copilot: raw OTEL ids resolve to their pretty display form. + assert_eq!( + agent_bucket_key(&msg("copilot", "github.copilot.default")), + "GitHub Copilot" + ); + assert_eq!( + agent_bucket_key(&msg("copilot", "Plugin:code-review-team:api-reviewer")), + "Code Review Team: API Reviewer" + ); + + // A non-copilot client with the same raw id must NOT get the + // copilot-specific prettification (proves the branch is client-scoped). + assert_ne!( + agent_bucket_key(&msg("codebuff", "github.copilot.default")), + "GitHub Copilot" + ); } #[test] #[serial_test::serial] - fn test_cursor_parse_path_reprices_zero_cost_composer_1_5_rows() { + fn test_source_cache_refreshes_stale_date_on_cache_hit() { let cache_home = tempfile::TempDir::new().unwrap(); - let _env = EnvGuard::set(&[ - ("HOME", cache_home.path().as_os_str()), - ("TOKSCALE_CONFIG_DIR", cache_home.path().as_os_str()), - ]); - let temp_dir = tempfile::TempDir::new().unwrap(); - let cursor_cache_dir = temp_dir.path().join(".config/tokscale/cursor-cache"); - std::fs::create_dir_all(&cursor_cache_dir).unwrap(); + let source_home = tempfile::TempDir::new().unwrap(); + let _env = opencode_test_env(cache_home.path(), source_home.path()); - let csv = r#"Date,Kind,Model,Max Mode,Input (w/ Cache Write),Input (w/o Cache Write),Cache Read,Output Tokens,Total Tokens,Cost -"2026-03-04T12:00:00.000Z","Included","Composer 1.5","No","1200","1000","5000","2000","8000","0""#; - std::fs::write(cursor_cache_dir.join("usage.csv"), csv).unwrap(); + { + let message_dir = scanner_fixture_path( + source_home.path(), + ".local/share/opencode/storage/message/project-1", + ); + std::fs::create_dir_all(&message_dir).unwrap(); + let path = message_dir.join("msg_001.json"); + std::fs::write( + &path, + r#"{"id":"msg-1","sessionID":"session-1","role":"assistant","modelID":"accounts/fireworks/models/deepseek-v3-0324","providerID":"fireworks","cost":0,"tokens":{"input":10,"output":5,"reasoning":0,"cache":{"read":0,"write":0}},"time":{"created":1733011200000}}"#, + ) + .unwrap(); - let pricing = pricing::PricingService::new(HashMap::new(), HashMap::new()); - let messages = parse_all_messages_with_pricing( - temp_dir.path().to_str().unwrap(), - &["cursor".to_string()], - Some(&pricing), - ); + let fingerprint = message_cache::SourceFingerprint::from_path(&path).unwrap(); + let mut stale_message = UnifiedMessage::new( + "opencode", + "accounts/fireworks/models/deepseek-v3-0324", + "fireworks", + "session-1", + 1_733_011_200_000, + TokenBreakdown { + input: 10, + output: 5, + cache_read: 0, + cache_write: 0, + reasoning: 0, + }, + 0.0, + ); + stale_message.date = "1900-01-01".to_string(); - assert_eq!(messages.len(), 1); - assert_eq!(messages[0].client, "cursor"); - assert_eq!(messages[0].model_id, "Composer 1.5"); - assert!(messages[0].cost > 0.0); - } + let mut cache = message_cache::SourceMessageCache::default(); + cache.insert(message_cache::CachedSourceEntry::new( + message_cache::CacheIdentity::for_client(ClientId::OpenCode), + &path, + fingerprint, + vec![stale_message], + Vec::new(), + None, + )); + cache.save_if_dirty(); + + let messages = parse_all_messages_with_pricing( + source_home.path().to_str().unwrap(), + &["opencode".to_string()], + None, + ); - fn write_kimi_repeated_status_fixture(source_home: &std::path::Path) { - let session_dir = source_home.join(".kimi/sessions/group-1/session-1"); - std::fs::create_dir_all(&session_dir).unwrap(); - std::fs::write( - session_dir.join("wire.jsonl"), - r#"{"type": "metadata", "protocol_version": "1.3"} -{"timestamp": 1770983410.0, "message": {"type": "StatusUpdate", "payload": {"token_usage": {"input_other": 10, "output": 1, "input_cache_read": 0, "input_cache_creation": 0}, "message_id": "msg-progressive"}}} -{"timestamp": 1770983420.0, "message": {"type": "StatusUpdate", "payload": {"token_usage": {"input_other": 20, "output": 2, "input_cache_read": 0, "input_cache_creation": 0}, "message_id": "msg-progressive"}}} -{"timestamp": 1770983430.0, "message": {"type": "StatusUpdate", "payload": {"token_usage": {"input_other": 5, "output": 1, "input_cache_read": 0, "input_cache_creation": 0}, "message_id": "msg-distinct"}}} -{"timestamp": 1770983440.0, "message": {"type": "StatusUpdate", "payload": {"token_usage": {"input_other": 7, "output": 1, "input_cache_read": 0, "input_cache_creation": 0}}}} -{"timestamp": 1770983450.0, "message": {"type": "StatusUpdate", "payload": {"token_usage": {"input_other": 8, "output": 1, "input_cache_read": 0, "input_cache_creation": 0}}}}"#, - ) - .unwrap(); + assert_eq!(messages.len(), 1); + assert_ne!(messages[0].date, "1900-01-01"); + assert_eq!( + messages[0].date, + UnifiedMessage::new( + "opencode", + "accounts/fireworks/models/deepseek-v3-0324", + "fireworks", + "session-1", + 1_733_011_200_000, + TokenBreakdown { + input: 10, + output: 5, + cache_read: 0, + cache_write: 0, + reasoning: 0, + }, + 0.0, + ) + .date + ); + } } + #[cfg(unix)] #[test] #[serial_test::serial] - fn test_parse_all_messages_with_pricing_kimi_deduplicates_repeated_status_updates() { + fn test_empty_parse_results_are_not_cached_for_optional_file_sources() { + use std::os::unix::fs::PermissionsExt; + let cache_home = tempfile::TempDir::new().unwrap(); let source_home = tempfile::TempDir::new().unwrap(); - let _env = EnvGuard::set(&[ - ("HOME", cache_home.path().as_os_str()), - ("TOKSCALE_CONFIG_DIR", cache_home.path().as_os_str()), - ]); + let _env = opencode_test_env(cache_home.path(), source_home.path()); { - write_kimi_repeated_status_fixture(source_home.path()); + let message_dir = scanner_fixture_path( + source_home.path(), + ".local/share/opencode/storage/message/project-1", + ); + std::fs::create_dir_all(&message_dir).unwrap(); + let path = message_dir.join("msg_001.json"); + std::fs::write( + &path, + r#"{"id":"msg-1","sessionID":"session-1","role":"assistant","modelID":"accounts/fireworks/models/deepseek-v3-0324","providerID":"fireworks","cost":0,"tokens":{"input":10,"output":5,"reasoning":0,"cache":{"read":0,"write":0}},"time":{"created":1733011200000}}"#, + ) + .unwrap(); - let messages = parse_all_messages_with_pricing( + let mut permissions = std::fs::metadata(&path).unwrap().permissions(); + permissions.set_mode(0o000); + std::fs::set_permissions(&path, permissions).unwrap(); + + let first_messages = parse_all_messages_with_pricing( source_home.path().to_str().unwrap(), - &["kimi".to_string()], + &["opencode".to_string()], None, ); + assert!(first_messages.is_empty()); - assert_eq!(messages.len(), 4); - assert_eq!(messages.iter().map(|m| m.tokens.input).sum::(), 40); - assert_eq!(messages.iter().map(|m| m.tokens.output).sum::(), 5); + let cache = message_cache::SourceMessageCache::load(); + assert!(cache + .get( + message_cache::CacheIdentity::for_client(ClientId::OpenCode), + &path + ) + .is_none()); + + let mut readable_permissions = std::fs::metadata(&path).unwrap().permissions(); + readable_permissions.set_mode(0o644); + std::fs::set_permissions(&path, readable_permissions).unwrap(); + + let second_messages = parse_all_messages_with_pricing( + source_home.path().to_str().unwrap(), + &["opencode".to_string()], + None, + ); + assert_eq!(second_messages.len(), 1); } } #[test] #[serial_test::serial] - fn test_parse_local_clients_kimi_deduplicates_repeated_status_updates() { + fn test_empty_cache_hits_are_reparsed_for_optional_file_sources() { let cache_home = tempfile::TempDir::new().unwrap(); let source_home = tempfile::TempDir::new().unwrap(); - let _env = EnvGuard::set(&[ - ("HOME", cache_home.path().as_os_str()), - ("TOKSCALE_CONFIG_DIR", cache_home.path().as_os_str()), - ]); + let _env = opencode_test_env(cache_home.path(), source_home.path()); { - write_kimi_repeated_status_fixture(source_home.path()); - - let parsed = parse_local_clients(LocalParseOptions { - home_dir: Some(source_home.path().to_str().unwrap().to_string()), - use_env_roots: false, - clients: Some(vec!["kimi".to_string()]), - since: None, - until: None, - year: None, - scanner_settings: scanner::ScannerSettings::default(), - modified_after: None, - }) + let message_dir = scanner_fixture_path( + source_home.path(), + ".local/share/opencode/storage/message/project-1", + ); + std::fs::create_dir_all(&message_dir).unwrap(); + let source_path = message_dir.join("msg_001.json"); + std::fs::write( + &source_path, + r#"{"id":"msg-1","sessionID":"session-1","role":"assistant","modelID":"accounts/fireworks/models/deepseek-v3-0324","providerID":"fireworks","cost":0,"tokens":{"input":10,"output":5,"reasoning":0,"cache":{"read":0,"write":0}},"time":{"created":1733011200000}}"#, + ) .unwrap(); + let cache_path = scanner::scan_all_clients_with_env_strategy( + source_home.path().to_str().unwrap(), + &["opencode".to_string()], + true, + ) + .get(ClientId::OpenCode) + .first() + .cloned() + .expect("scanner must find the OpenCode fixture"); - assert_eq!(parsed.counts.get(ClientId::Kimi), 4); - assert_eq!(parsed.messages.len(), 4); - assert_eq!(parsed.messages.iter().map(|m| m.input).sum::(), 40); - assert_eq!(parsed.messages.iter().map(|m| m.output).sum::(), 5); + let fingerprint = message_cache::SourceFingerprint::from_path(&cache_path).unwrap(); + let mut cache = message_cache::SourceMessageCache::default(); + cache.insert(message_cache::CachedSourceEntry::new( + message_cache::CacheIdentity::for_client(ClientId::OpenCode), + &cache_path, + fingerprint, + Vec::new(), + Vec::new(), + None, + )); + cache.save_if_dirty(); + + let messages = parse_all_messages_with_pricing( + source_home.path().to_str().unwrap(), + &["opencode".to_string()], + None, + ); + assert_eq!(messages.len(), 1); + + let loaded = message_cache::SourceMessageCache::load(); + let repaired_entry = loaded + .get( + message_cache::CacheIdentity::for_client(ClientId::OpenCode), + &cache_path, + ) + .unwrap(); + assert_eq!(repaired_entry.messages.len(), 1); } } - // Regression: the streaming driver must NOT share one dedup set across - // different clients. kimi and codebuff both emit raw upstream message ids - // as dedup_key with no client namespace, so a shared set would let one - // client's key suppress an identical key from the other. Here both a kimi - // message and a codebuff message carry dedup_key "COLLIDE"; both must - // survive. With a single shared `seen_keys` (the pre-fix behaviour) the - // second lane's message is silently dropped and this fails. #[test] #[serial_test::serial] - fn test_streaming_driver_does_not_dedup_across_clients() { + fn test_opencode_aliases_preserve_cross_store_identity_across_lanes() { let cache_home = tempfile::TempDir::new().unwrap(); let source_home = tempfile::TempDir::new().unwrap(); - let _env = EnvGuard::set(&[ - ("HOME", cache_home.path().as_os_str()), - ("TOKSCALE_CONFIG_DIR", cache_home.path().as_os_str()), - ]); + let _env = opencode_test_env(cache_home.path(), source_home.path()); - { - // kimi: one StatusUpdate carrying message_id "COLLIDE". - let kimi_dir = source_home.path().join(".kimi/sessions/g/s"); - std::fs::create_dir_all(&kimi_dir).unwrap(); - std::fs::write( - kimi_dir.join("wire.jsonl"), - r#"{"type": "metadata", "protocol_version": "1.3"} -{"timestamp": 1770983410.0, "message": {"type": "StatusUpdate", "payload": {"token_usage": {"input_other": 100, "output": 50, "input_cache_read": 0, "input_cache_creation": 0}, "message_id": "COLLIDE"}}}"#, + let db_dir = source_home.path().join(".local/share/opencode"); + std::fs::create_dir_all(&db_dir).unwrap(); + let db_path = db_dir.join("opencode-next.db"); + let conn = create_opencode_v2_sqlite_db(&db_path); + let v1_payload = build_opencode_sqlite_payload( + 1_700_000_000_000.0, + 1_700_000_000_500.0, + 100, + 10, + 1, + 5, + 2, + 0.0, + ); + conn.execute( + "INSERT INTO message (id, session_id, data) VALUES (?1, ?2, ?3)", + rusqlite::params!["legacy-message", "session-overlap", &v1_payload], + ) + .unwrap(); + let v2_payload = build_opencode_sqlite_payload( + 1_700_000_000_000.0, + 1_700_000_000_500.0, + 100, + 10, + 1, + 5, + 2, + 0.02, + ) + .replacen('{', r#"{"id":"v2-embedded","#, 1); + conn.execute( + "INSERT INTO session_message (id, session_id, type, data) VALUES (?1, ?2, ?3, ?4)", + rusqlite::params!["v2-message", "session-overlap", "assistant", &v2_payload], + ) + .unwrap(); + let deferred_payload = build_opencode_sqlite_payload( + 1_700_000_010_000.0, + 1_700_000_010_500.0, + 20, + 5, + 0, + 0, + 0, + 0.0, + ); + let deferred_v2_payload = + deferred_payload.replacen('{', r#"{"id":"shared-order","#, 1); + conn.execute( + "INSERT INTO message (id, session_id, data) VALUES (?1, ?2, ?3)", + rusqlite::params!["legacy-order", "session-order", &deferred_payload], + ) + .unwrap(); + conn.execute( + "INSERT INTO session_message (id, session_id, type, data) VALUES (?1, ?2, ?3, ?4)", + rusqlite::params![ + "v2-order", + "session-order", + "assistant", + &deferred_v2_payload + ], + ) + .unwrap(); + drop(conn); + + let sibling_db = db_dir.join("opencode-stable.db"); + let sibling_conn = create_opencode_v2_sqlite_db(&sibling_db); + sibling_conn + .execute( + "INSERT INTO session_message (id, session_id, type, data) VALUES (?1, ?2, ?3, ?4)", + rusqlite::params![ + "v2-sibling", + "session-overlap", + "assistant", + &v2_payload + ], ) .unwrap(); - - // codebuff: one assistant message whose upstream id is "COLLIDE". - let cb_dir = source_home.path().join(".config/manicode/projects/proj"); - std::fs::create_dir_all(&cb_dir).unwrap(); - std::fs::write( - cb_dir.join("chat-messages.json"), - r#"[{"role":"assistant","id":"COLLIDE","metadata":{"model":"claude-sonnet-4","usage":{"inputTokens":200,"outputTokens":80}},"credits":0.02}]"#, + let provider_payload = build_opencode_sqlite_payload( + 1_700_000_010_000.0, + 1_700_000_010_500.0, + 10, + 5, + 0, + 0, + 0, + 0.02, + ) + .replacen('{', r#"{"id":"shared-order","#, 1); + sibling_conn + .execute( + "INSERT INTO session_message (id, session_id, type, data) VALUES (?1, ?2, ?3, ?4)", + rusqlite::params![ + "v2-order-provider", + "session-order", + "assistant", + &provider_payload + ], ) .unwrap(); + drop(sibling_conn); + + let json_dir = db_dir.join("storage/message/project-overlap"); + std::fs::create_dir_all(&json_dir).unwrap(); + let migrated_json_payload = build_opencode_sqlite_payload( + 1_700_000_000_000.0, + 1_700_000_000_500.0, + 100, + 10, + 1, + 5, + 2, + 0.03, + ); + std::fs::write( + json_dir.join("legacy-message.json"), + migrated_json_payload, + ) + .unwrap(); + let json_payload = build_opencode_sqlite_payload( + 1_700_000_010_000.0, + 1_700_000_010_500.0, + 30, + 5, + 0, + 0, + 0, + 0.03, + ); + std::fs::write(json_dir.join("legacy-order.json"), json_payload).unwrap(); + + let clients = vec!["opencode".to_string()]; + let materialized = parse_all_messages_with_pricing_with_env_strategy( + source_home.path().to_str().unwrap(), + &clients, + None, + false, + &scanner::ScannerSettings::default(), + ); + assert_eq!(materialized.len(), 3); + let mut materialized_inputs: Vec<_> = materialized + .iter() + .map(|message| message.tokens.input) + .collect(); + materialized_inputs.sort_unstable(); + assert_eq!(materialized_inputs, vec![10, 20, 100]); + let migrated = materialized + .iter() + .find(|message| message.tokens.input == 100) + .unwrap(); + assert_eq!(migrated.dedup_key.as_deref(), Some("v2-embedded")); + assert_eq!(migrated.dedup_aliases, vec!["legacy-message"]); + assert_eq!(migrated.cost, 0.02); + assert_eq!(migrated.cost_source, CostSource::ProviderReported); + let warm = parse_all_messages_with_pricing_with_env_strategy( + source_home.path().to_str().unwrap(), + &clients, + None, + false, + &scanner::ScannerSettings::default(), + ); + assert_eq!(warm, materialized, "cache hits must retain OpenCode aliases"); - let mut seen: Vec = Vec::new(); - scan_messages_streaming( - source_home.path().to_str().unwrap(), - &["kimi".to_string(), "codebuff".to_string()], - None, - false, - &scanner::ScannerSettings::default(), - &|_m: &UnifiedMessage| true, - &mut |m: &UnifiedMessage| seen.push(m.client.clone()), - ); + let mut streamed = Vec::new(); + scan_messages_streaming( + source_home.path().to_str().unwrap(), + &clients, + None, + false, + &scanner::ScannerSettings::default(), + &|_| true, + &mut |message| streamed.push(message.clone()), + ); + assert_eq!(streamed.len(), 3); + let mut streamed_inputs: Vec<_> = streamed + .iter() + .map(|message| message.tokens.input) + .collect(); + streamed_inputs.sort_unstable(); + assert_eq!(streamed_inputs, vec![10, 20, 100]); + let migrated = streamed + .iter() + .find(|message| message.tokens.input == 100) + .unwrap(); + assert_eq!(migrated.dedup_key.as_deref(), Some("v2-embedded")); + assert_eq!(migrated.dedup_aliases, vec!["legacy-message"]); + assert_eq!(migrated.cost, 0.02); + assert_eq!(migrated.cost_source, CostSource::ProviderReported); - assert!( - seen.iter().any(|c| c == "kimi"), - "kimi message with shared dedup_key must survive: {seen:?}" - ); - assert!( - seen.iter().any(|c| c == "codebuff"), - "codebuff message with shared dedup_key must survive: {seen:?}" - ); - } + let counted = parse_local_clients(LocalParseOptions { + home_dir: Some(source_home.path().to_string_lossy().into_owned()), + use_env_roots: false, + clients: Some(clients), + ..Default::default() + }) + .unwrap(); + assert_eq!(counted.counts.get(ClientId::OpenCode), 3); + assert_eq!(counted.messages.len(), 3); + assert_eq!( + counted + .messages + .iter() + .map(|message| message.input) + .sum::(), + 130 + ); } - // M2 (codex fork-replay): the parser-level fork dedup (#649/#681) must also - // collapse replayed parent token_count rows through OUR streaming report - // path (scan_messages_streaming), not just the materialized - // parse_all_messages_with_pricing path the upstream tests exercise. Without - // the fork-parent-scoped dedup key, each fork's replayed parent rows survive - // per child and inflate codex totals. #[test] #[serial_test::serial] - fn test_streaming_codex_collapses_parent_replay_across_forks() { + fn test_schema_29_hybrid_opencode_cache_rebuilds_with_v2_rows_across_lanes() { + #[derive(serde::Serialize)] + struct Schema29Store { + schema_version: u32, + entries: Vec, + } + let cache_home = tempfile::TempDir::new().unwrap(); let source_home = tempfile::TempDir::new().unwrap(); - let _env = EnvGuard::set(&[ - ("HOME", cache_home.path().as_os_str()), - ("TOKSCALE_CONFIG_DIR", cache_home.path().as_os_str()), - ]); + let _env = opencode_test_env(cache_home.path(), source_home.path()); + let _pricing_env = + EnvGuard::set(&[("TOKSCALE_PRICING_CACHE_ONLY", std::ffi::OsStr::new("1"))]); - { - write_codex_parent_replay_fixture(source_home.path()); + let db_dir = source_home.path().join(".local/share/opencode"); + std::fs::create_dir_all(&db_dir).unwrap(); + let db_path = db_dir.join("opencode-next.db"); + let conn = create_opencode_v2_sqlite_db(&db_path); + conn.execute( + "INSERT INTO session (id, directory) VALUES (?1, ?2), (?3, ?4)", + rusqlite::params!["session-v1", "/workspace/v1", "session-v2", "/workspace/v2"], + ) + .unwrap(); + let v1_data = r#"{ + "id": "message-v1", + "role": "assistant", + "modelID": "claude-sonnet-4", + "providerID": "anthropic", + "mode": "build", + "cost": 0.0, + "tokens": { "input": 100, "output": 10, "reasoning": 1, "cache": { "read": 5, "write": 2 } }, + "time": { "created": 1700000000000.0, "completed": 1700000000500.0 } + }"#; + let v2_data = r#"{ + "id": "message-v1", + "model": { "id": "claude-sonnet-4", "providerID": "anthropic" }, + "agent": "build", + "cost": 0.02, + "tokens": { "input": 200, "output": 20, "reasoning": 2, "cache": { "read": 10, "write": 4 } }, + "time": { "created": 1700000001000.0, "completed": 1700000001750.0 } + }"#; + conn.execute( + "INSERT INTO message (id, session_id, data) VALUES (?1, ?2, ?3)", + rusqlite::params!["row-v1", "session-v1", v1_data], + ) + .unwrap(); + conn.execute( + "INSERT INTO session_message (id, session_id, type, data) VALUES (?1, ?2, ?3, ?4)", + rusqlite::params!["row-v2", "session-v2", "assistant", v2_data], + ) + .unwrap(); + drop(conn); - let mut input_sum = 0i64; - let mut output_sum = 0i64; - let mut count = 0usize; - scan_messages_streaming( + let json_dir = db_dir.join("storage/message/project-v1"); + std::fs::create_dir_all(&json_dir).unwrap(); + std::fs::write( + json_dir.join("message-v1.json"), + r#"{ + "id": "message-v1", + "sessionID": "session-v1", + "role": "assistant", + "modelID": "claude-sonnet-4", + "providerID": "anthropic", + "mode": "build", + "cost": 0.01, + "tokens": { "input": 100, "output": 10, "reasoning": 1, "cache": { "read": 5, "write": 2 } }, + "time": { "created": 1700000000000.0, "completed": 1700000000500.0 } + }"#, + ) + .unwrap(); + + let fingerprint = message_cache::SourceFingerprint::from_sqlite_path(&db_path).unwrap(); + let mut stale_message = UnifiedMessage::new_with_dedup( + "opencode", + "claude-sonnet-4", + "anthropic", + "session-v1", + 1_700_000_000_000, + TokenBreakdown { + input: 100, + output: 10, + cache_read: 5, + cache_write: 2, + reasoning: 1, + }, + 0.0, + Some("message-v1".to_string()), + ); + stale_message.duration_ms = Some(500); + let stale_store = Schema29Store { + schema_version: 29, + entries: vec![message_cache::CachedSourceEntry::new( + message_cache::CacheIdentity::for_client(ClientId::OpenCode), + &db_path, + fingerprint.clone(), + vec![stale_message], + Vec::new(), + None, + )], + }; + let cache_file = crate::paths::get_cache_dir().join("source-message-cache.bin"); + std::fs::create_dir_all(cache_file.parent().unwrap()).unwrap(); + let writer = std::io::BufWriter::new(std::fs::File::create(&cache_file).unwrap()); + bincode::options() + .serialize_into(writer, &stale_store) + .unwrap(); + + assert_eq!( + message_cache::SourceFingerprint::from_sqlite_path(&db_path).unwrap(), + fingerprint, + "the schema-29 v1-only entry must match the unchanged hybrid database" + ); + assert!( + message_cache::SourceMessageCache::load().entries.is_empty(), + "schema-29 cache entries must be rejected before parsing v2 rows" + ); + + let clients = vec!["opencode".to_string()]; + let scan_result = scan_test_sources( + source_home.path(), + &clients, + &scanner::ScannerSettings::default(), + ); + let cache_db_path = scanner_spelling(&scan_result.opencode_dbs, &db_path); + let parse_materialized = || { + parse_all_messages_with_pricing_with_env_strategy( source_home.path().to_str().unwrap(), - &["codex".to_string()], + &clients, None, false, &scanner::ScannerSettings::default(), - &|_m: &UnifiedMessage| true, - &mut |m: &UnifiedMessage| { - input_sum += m.tokens.input; - output_sum += m.tokens.output; - count += 1; - }, - ); - - // Same collapse as the materialized - // test_parse_all_messages_with_pricing_codex_deduplicates_parent_replay_across_forks: - // the parent's two turns plus the single own-turn shared (by identical - // cumulative total) across the two forks. Without #649/#681 the - // replayed parent rows would survive per fork and inflate this. - assert_eq!(count, 3, "replayed parent rows must collapse to 3 messages"); - assert_eq!(input_sum, 140); - assert_eq!(output_sum, 14); - } - } - - // Issue #6: the agents report must dedup the simple_lane! clients - // (copilot/codebuff/kimi/…) like the model/graph/hourly reports. Here - // codebuff emits the SAME upstream message id "DUP" in two different - // project files. The OLD materialized path (parse_local_unified_messages) - // never gated codebuff, so it counts both; the streaming-backed - // get_agents_report keeps one — matching get_model_report. Repointing - // get_agents_report at the old path makes the parity assertion FAIL (RED). - #[test] - #[serial_test::serial] - fn test_agents_report_dedups_like_model_report_issue6() { - let cache_home = tempfile::TempDir::new().unwrap(); - let source_home = tempfile::TempDir::new().unwrap(); - let _env = EnvGuard::set(&[ - ("HOME", cache_home.path().as_os_str()), - ("TOKSCALE_CONFIG_DIR", cache_home.path().as_os_str()), - ]); - // Hermetic: cache-only pricing + temp HOME → no network, pricing None. - let _pricing = EnvGuard::set(&[("TOKSCALE_PRICING_CACHE_ONLY", std::ffi::OsStr::new("1"))]); - - { - let write_codebuff = |proj: &str| { - let dir = source_home - .path() - .join(format!(".config/manicode/projects/{proj}")); - std::fs::create_dir_all(&dir).unwrap(); - std::fs::write( - dir.join("chat-messages.json"), - r#"[{"role":"assistant","id":"DUP","metadata":{"model":"claude-sonnet-4","usage":{"inputTokens":200,"outputTokens":80}},"credits":0.02}]"#, + ) + }; + let mut cold = parse_materialized(); + let mut warm = parse_materialized(); + cold.sort_by(|left, right| { + (&left.dedup_key, left.tokens.input).cmp(&(&right.dedup_key, right.tokens.input)) + }); + warm.sort_by(|left, right| { + (&left.dedup_key, left.tokens.input).cmp(&(&right.dedup_key, right.tokens.input)) + }); + assert_eq!(cold.len(), 2); + assert_eq!( + cold.iter() + .map(|message| message.dedup_key.as_deref()) + .collect::>(), + vec![Some("message-v1"), Some("message-v1")], + "same embedded ids with incompatible timestamps and tokens must remain distinct" + ); + assert_eq!( + cold.iter() + .map(|message| (&message.dedup_key, &message.tokens)) + .collect::>(), + warm.iter() + .map(|message| (&message.dedup_key, &message.tokens)) + .collect::>(), + "the rebuilt current-schema entry must preserve v1+v2 output on a warm hit" + ); + let rebuilt_cache = message_cache::SourceMessageCache::load(); + assert_eq!( + rebuilt_cache + .get( + message_cache::CacheIdentity::for_client(ClientId::OpenCode), + &cache_db_path ) - .unwrap(); - }; - write_codebuff("projA"); - write_codebuff("projB"); - - let home = source_home.path().to_str().unwrap().to_string(); - let clients = Some(vec!["codebuff".to_string()]); - let rt = tokio::runtime::Builder::new_current_thread() - .enable_all() - .build() - .unwrap(); + .unwrap() + .messages + .len(), + 2 + ); - let agents = rt - .block_on(get_agents_report(ReportOptions { - home_dir: Some(home.clone()), - use_env_roots: false, - clients: clients.clone(), - ..Default::default() - })) - .unwrap(); - let model = rt - .block_on(get_model_report(ReportOptions { - home_dir: Some(home.clone()), - use_env_roots: false, - clients: clients.clone(), - ..Default::default() - })) - .unwrap(); - let old = rt - .block_on(parse_local_unified_messages(LocalParseOptions { - home_dir: Some(home.clone()), - use_env_roots: false, - clients: clients.clone(), - since: None, - until: None, - year: None, - scanner_settings: scanner::ScannerSettings::default(), - modified_after: None, - })) - .unwrap(); - let old_total: i32 = old.iter().map(|m| m.message_count.max(0)).sum(); + let mut streamed = Vec::new(); + scan_messages_streaming( + source_home.path().to_str().unwrap(), + &clients, + None, + false, + &scanner::ScannerSettings::default(), + &|_| true, + &mut |message| streamed.push(message.clone()), + ); + streamed.sort_by(|left, right| { + (&left.dedup_key, left.tokens.input).cmp(&(&right.dedup_key, right.tokens.input)) + }); + assert_eq!( + streamed + .iter() + .map(|message| (&message.dedup_key, &message.tokens)) + .collect::>(), + cold.iter() + .map(|message| (&message.dedup_key, &message.tokens)) + .collect::>() + ); - assert_eq!(old_total, 2, "old materialized path must NOT dedup codebuff"); - assert_eq!(model.total_messages, 1, "model report dedups codebuff"); - assert_eq!(agents.total_messages, 1, "agents report must dedup codebuff"); - assert_eq!( - agents.total_messages, model.total_messages, - "issue #6: agents must agree with the model report" - ); - assert_ne!( - old_total, model.total_messages, - "the old path diverged from the model report (the #6 bug)" - ); - assert!( - (agents.total_cost - model.total_cost).abs() < 1e-9, - "agents/model cost parity (agents={}, model={})", - agents.total_cost, - model.total_cost - ); - } + let counted = parse_local_clients(LocalParseOptions { + home_dir: Some(source_home.path().to_string_lossy().into_owned()), + use_env_roots: false, + clients: Some(clients.clone()), + ..Default::default() + }) + .unwrap(); + assert_eq!(counted.counts.get(ClientId::OpenCode), 2); + assert_eq!(counted.messages.len(), 2); + assert_eq!( + counted + .messages + .iter() + .map(|message| message.input) + .sum::(), + 300 + ); + + let runtime = tokio::runtime::Builder::new_current_thread() + .enable_all() + .build() + .unwrap(); + let options = ReportOptions { + home_dir: Some(source_home.path().to_string_lossy().into_owned()), + use_env_roots: false, + clients: Some(clients), + ..Default::default() + }; + let model = runtime.block_on(get_model_report(options.clone())).unwrap(); + let monthly = runtime + .block_on(get_monthly_report(options.clone())) + .unwrap(); + let hourly = runtime + .block_on(get_hourly_report(options.clone())) + .unwrap(); + let agents = runtime.block_on(get_agents_report(options)).unwrap(); + assert_eq!(model.total_messages, 2); + assert_eq!( + ( + model.total_input, + model.total_output, + model.total_cache_read, + model.total_cache_write, + model + .entries + .iter() + .map(|entry| entry.reasoning) + .sum::(), + ), + (300, 30, 15, 6, 3) + ); + assert_eq!( + monthly.entries.iter().fold((0, 0, 0, 0), |totals, entry| ( + totals.0 + entry.input, + totals.1 + entry.output, + totals.2 + entry.cache_read, + totals.3 + entry.cache_write, + ),), + (300, 30, 15, 6) + ); + assert_eq!( + hourly + .entries + .iter() + .fold((0, 0, 0, 0, 0), |totals, entry| ( + totals.0 + entry.input, + totals.1 + entry.output, + totals.2 + entry.cache_read, + totals.3 + entry.cache_write, + totals.4 + entry.reasoning, + ),), + (300, 30, 15, 6, 3) + ); + assert_eq!( + agents + .entries + .iter() + .fold((0, 0, 0, 0, 0), |totals, entry| ( + totals.0 + entry.input, + totals.1 + entry.output, + totals.2 + entry.cache_read, + totals.3 + entry.cache_write, + totals.4 + entry.reasoning, + ),), + (300, 30, 15, 6, 3) + ); + assert_eq!( + monthly + .entries + .iter() + .map(|entry| entry.message_count) + .sum::(), + model.total_messages + ); + assert_eq!( + hourly + .entries + .iter() + .map(|entry| entry.message_count) + .sum::(), + model.total_messages + ); + assert_eq!(agents.total_messages, model.total_messages); } - // Preservation: with no duplicate dedup_keys (and only parse_local==true - // clients), the streaming-backed agents report produces the SAME numbers the - // old materialized path did. codebuff + kimi, distinct ids, no agent - // attribution → a single "Main" bucket. #[test] #[serial_test::serial] - fn test_agents_report_preserves_numbers_without_duplicates() { + fn m16_schema_30_jcode_cache_rebuilds_start_anchor_across_lanes() { + #[derive(serde::Serialize)] + struct Schema30Store { + schema_version: u32, + entries: Vec, + } + let cache_home = tempfile::TempDir::new().unwrap(); let source_home = tempfile::TempDir::new().unwrap(); let _env = EnvGuard::set(&[ ("HOME", cache_home.path().as_os_str()), ("TOKSCALE_CONFIG_DIR", cache_home.path().as_os_str()), ]); - let _pricing = EnvGuard::set(&[("TOKSCALE_PRICING_CACHE_ONLY", std::ffi::OsStr::new("1"))]); + let _pricing_env = + EnvGuard::set(&[("TOKSCALE_PRICING_CACHE_ONLY", std::ffi::OsStr::new("1"))]); - { - let cb_dir = source_home.path().join(".config/manicode/projects/proj"); - std::fs::create_dir_all(&cb_dir).unwrap(); - std::fs::write( - cb_dir.join("chat-messages.json"), - r#"[{"role":"assistant","id":"A","metadata":{"model":"claude-sonnet-4","usage":{"inputTokens":200,"outputTokens":80}},"credits":0.02}]"#, - ) + let sessions_dir = source_home.path().join(".jcode/sessions"); + std::fs::create_dir_all(&sessions_dir).unwrap(); + let source_path = sessions_dir.join("session_m16.json"); + std::fs::write( + &source_path, + r#"{"id":"session_m16","provider_key":"cliproxyapi","model":"claude-sonnet-4","working_dir":"/workspace/m16","messages":[{"id":"u1","role":"user","timestamp":"2026-06-16T12:00:00Z"},{"id":"a1","role":"assistant","timestamp":"2026-06-16T12:00:01Z","token_usage":{"input_tokens":1200,"output_tokens":300},"tool_duration_ms":1000}]}"#, + ) + .unwrap(); + + let start = sessions::utils::parse_timestamp_str("2026-06-16T12:00:00Z").unwrap(); + let end = sessions::utils::parse_timestamp_str("2026-06-16T12:00:01Z").unwrap(); + let fingerprint = message_cache::SourceFingerprint::from_jcode_path(&source_path).unwrap(); + let mut stale_message = UnifiedMessage::new_with_dedup( + "jcode", + "claude-sonnet-4", + "cliproxyapi", + "session_m16", + end, + TokenBreakdown { + input: 1200, + output: 300, + ..Default::default() + }, + 0.0, + Some("stale-schema-30".to_string()), + ); + stale_message.duration_ms = Some(end - start); + let stale_store = Schema30Store { + schema_version: 30, + entries: vec![message_cache::CachedSourceEntry::new( + message_cache::CacheIdentity::for_client(ClientId::Jcode), + &source_path, + fingerprint.clone(), + vec![stale_message], + Vec::new(), + None, + )], + }; + let cache_file = crate::paths::get_cache_dir().join("source-message-cache.bin"); + std::fs::create_dir_all(cache_file.parent().unwrap()).unwrap(); + let writer = std::io::BufWriter::new(std::fs::File::create(&cache_file).unwrap()); + bincode::options() + .serialize_into(writer, &stale_store) .unwrap(); - let kimi_dir = source_home.path().join(".kimi/sessions/g/s"); - std::fs::create_dir_all(&kimi_dir).unwrap(); - std::fs::write( - kimi_dir.join("wire.jsonl"), - "{\"type\": \"metadata\", \"protocol_version\": \"1.3\"}\n{\"timestamp\": 1770983410.0, \"message\": {\"type\": \"StatusUpdate\", \"payload\": {\"token_usage\": {\"input_other\": 100, \"output\": 50, \"input_cache_read\": 0, \"input_cache_creation\": 0}, \"message_id\": \"K\"}}}", + + assert_eq!( + message_cache::SourceFingerprint::from_jcode_path(&source_path).unwrap(), + fingerprint, + "the source fingerprint must stay unchanged across the schema-only rebuild" + ); + assert!( + message_cache::SourceMessageCache::load().entries.is_empty(), + "schema-30 entries must be rejected before corrected Jcode output is loaded" + ); + + let clients = vec!["jcode".to_string()]; + let scan_result = scan_test_sources( + source_home.path(), + &clients, + &scanner::ScannerSettings::default(), + ); + let cache_source_path = scanner_spelling(scan_result.get(ClientId::Jcode), &source_path); + let parse_materialized = || { + parse_all_messages_with_pricing_with_env_strategy( + source_home.path().to_str().unwrap(), + &clients, + None, + false, + &scanner::ScannerSettings::default(), ) + }; + let cold = parse_materialized(); + let warm = parse_materialized(); + assert_eq!(cold.len(), 1); + assert_eq!( + ( + cold[0].timestamp, + cold[0].duration_ms, + cold[0].tokens.input, + cold[0].tokens.output, + ), + (start, Some(end - start), 1200, 300) + ); + assert_eq!( + warm.iter() + .map(|message| ( + message.timestamp, + message.duration_ms, + message.tokens.clone() + )) + .collect::>(), + cold.iter() + .map(|message| ( + message.timestamp, + message.duration_ms, + message.tokens.clone() + )) + .collect::>() + ); + let rebuilt_cache = message_cache::SourceMessageCache::load(); + assert_eq!( + rebuilt_cache + .get( + message_cache::CacheIdentity::for_client(ClientId::Jcode), + &cache_source_path + ) + .unwrap() + .messages + .len(), + 1 + ); + + let mut streamed = Vec::new(); + scan_messages_streaming( + source_home.path().to_str().unwrap(), + &clients, + None, + false, + &scanner::ScannerSettings::default(), + &|_| true, + &mut |message| streamed.push(message.clone()), + ); + assert_eq!(streamed.len(), 1); + assert_eq!( + ( + streamed[0].timestamp, + streamed[0].duration_ms, + streamed[0].tokens.input, + streamed[0].tokens.output, + ), + (start, Some(end - start), 1200, 300) + ); + + let counted = parse_local_clients(LocalParseOptions { + home_dir: Some(source_home.path().to_string_lossy().into_owned()), + use_env_roots: false, + clients: Some(clients.clone()), + ..Default::default() + }) + .unwrap(); + assert_eq!(counted.counts.get(ClientId::Jcode), 1); + assert_eq!(counted.messages.len(), 1); + assert_eq!( + (counted.messages[0].input, counted.messages[0].output), + (1200, 300) + ); + + let runtime = tokio::runtime::Builder::new_current_thread() + .enable_all() + .build() + .unwrap(); + let options = ReportOptions { + home_dir: Some(source_home.path().to_string_lossy().into_owned()), + use_env_roots: false, + clients: Some(clients), + ..Default::default() + }; + let model = runtime.block_on(get_model_report(options.clone())).unwrap(); + let monthly = runtime + .block_on(get_monthly_report(options.clone())) + .unwrap(); + let hourly = runtime + .block_on(get_hourly_report(options.clone())) .unwrap(); + let agents = runtime.block_on(get_agents_report(options)).unwrap(); + assert_eq!( + (model.total_messages, model.total_input, model.total_output), + (1, 1200, 300) + ); + assert_eq!( + monthly.entries.iter().fold((0, 0, 0), |totals, entry| ( + totals.0 + entry.message_count, + totals.1 + entry.input, + totals.2 + entry.output, + )), + (1, 1200, 300) + ); + assert_eq!( + hourly.entries.iter().fold((0, 0, 0), |totals, entry| ( + totals.0 + entry.message_count, + totals.1 + entry.input, + totals.2 + entry.output, + )), + (1, 1200, 300) + ); + assert_eq!(agents.total_messages, 1); + assert_eq!( + agents.entries.iter().fold((0, 0), |totals, entry| ( + totals.0 + entry.input, + totals.1 + entry.output, + )), + (1200, 300) + ); + } - let home = source_home.path().to_str().unwrap().to_string(); - let clients = Some(vec!["codebuff".to_string(), "kimi".to_string()]); - let rt = tokio::runtime::Builder::new_current_thread() - .enable_all() - .build() - .unwrap(); + #[test] + #[serial_test::serial] + fn test_sqlite_source_cache_invalidates_on_wal_change() { + let cache_home = tempfile::TempDir::new().unwrap(); + let source_home = tempfile::TempDir::new().unwrap(); + let _env = opencode_test_env(cache_home.path(), source_home.path()); - let agents = rt - .block_on(get_agents_report(ReportOptions { - home_dir: Some(home.clone()), - use_env_roots: false, - clients: clients.clone(), - ..Default::default() - })) + { + let db_dir = source_home.path().join(".local/share/opencode"); + std::fs::create_dir_all(&db_dir).unwrap(); + let db_path = db_dir.join("opencode.db"); + + let conn = rusqlite::Connection::open(&db_path).unwrap(); + let journal_mode: String = conn + .query_row("PRAGMA journal_mode=WAL;", [], |row| row.get(0)) .unwrap(); + assert_eq!(journal_mode.to_lowercase(), "wal"); + conn.execute_batch( + "PRAGMA wal_autocheckpoint=0; + CREATE TABLE message ( + id TEXT PRIMARY KEY, + session_id TEXT NOT NULL, + data TEXT NOT NULL + );", + ) + .unwrap(); - assert_eq!( - agents.entries.len(), - 1, - "no agent attribution → a single Main bucket" - ); - let main = &agents.entries[0]; - assert_eq!(main.agent, "Main"); - assert_eq!(main.messages, 2); - // BTreeSet → sorted, both clients fold into Main. - assert_eq!(main.clients, vec!["codebuff".to_string(), "kimi".to_string()]); + let row_one = r#"{ + "role": "assistant", + "modelID": "claude-sonnet-4", + "providerID": "anthropic", + "tokens": { "input": 100, "output": 50, "reasoning": 0, "cache": { "read": 0, "write": 0 } }, + "time": { "created": 1700000000000.0 } + }"#; + let row_two = r#"{ + "role": "assistant", + "modelID": "claude-sonnet-4", + "providerID": "anthropic", + "tokens": { "input": 120, "output": 60, "reasoning": 0, "cache": { "read": 0, "write": 0 } }, + "time": { "created": 1700000001000.0 } + }"#; - // Byte-for-byte equivalence with the old materialized path for the - // non-duplicate case (both parse identically; only dedup differs). - let old = rt - .block_on(parse_local_unified_messages(LocalParseOptions { - home_dir: Some(home.clone()), - use_env_roots: false, - clients, - since: None, - until: None, - year: None, - scanner_settings: scanner::ScannerSettings::default(), - modified_after: None, - })) - .unwrap(); - let old_input: i64 = old.iter().map(|m| m.tokens.input).sum(); - let old_output: i64 = old.iter().map(|m| m.tokens.output).sum(); - let old_cache_read: i64 = old.iter().map(|m| m.tokens.cache_read).sum(); - let old_cache_write: i64 = old.iter().map(|m| m.tokens.cache_write).sum(); - let old_reasoning: i64 = old.iter().map(|m| m.tokens.reasoning).sum(); - let old_messages: i32 = old.iter().map(|m| m.message_count.max(0)).sum(); - let old_cost: f64 = old.iter().map(|m| m.cost).sum(); + conn.execute( + "INSERT INTO message (id, session_id, data) VALUES (?1, ?2, ?3)", + rusqlite::params!["msg-1", "session-1", row_one], + ) + .unwrap(); - assert_eq!(main.input, old_input); - assert_eq!(main.output, old_output); - assert_eq!(main.cache_read, old_cache_read); - assert_eq!(main.cache_write, old_cache_write); - assert_eq!(main.reasoning, old_reasoning); - assert_eq!(main.messages, old_messages); - assert!((agents.total_cost - old_cost).abs() < 1e-9); - // Sanity: codebuff contributes its known tokens. - assert!(main.input >= 200 && main.output >= 80); + let first_messages = parse_all_messages_with_pricing( + source_home.path().to_str().unwrap(), + &["opencode".to_string()], + None, + ); + assert_eq!(first_messages.len(), 1); + + conn.execute( + "INSERT INTO message (id, session_id, data) VALUES (?1, ?2, ?3)", + rusqlite::params!["msg-2", "session-1", row_two], + ) + .unwrap(); + assert!(db_path.with_extension("db-wal").exists()); + + let refreshed_messages = parse_all_messages_with_pricing( + source_home.path().to_str().unwrap(), + &["opencode".to_string()], + None, + ); + assert_eq!(refreshed_messages.len(), 2); } } - // Issue #36: the client selection must be applied at the STREAMING SCAN, - // not by a downstream membership filter over the pre-aggregated buckets. - // codebuff (200/80) and kimi (100/50) fold into ONE shared "Main" agent - // bucket; the FFI/DashboardModel now thread the selection into - // ReportOptions.clients, so filtering to codebuff yields ONLY its 200/80 — - // NOT the mixed 300/130. The old approach (unfiltered report + a Swift - // membership filter over whole buckets) kept the entire shared bucket and - // would read 300/130 here, so this test is RED against it and GREEN now. #[test] #[serial_test::serial] - fn test_agents_report_client_filter_scopes_shared_bucket_issue36() { + fn test_parse_all_messages_dedups_across_channel_suffixed_opencode_dbs() { + // Regression guard: a session that appears in both `opencode.db` and + // `opencode-.db` (e.g. the user switches channels mid-session) + // must only be counted once. let cache_home = tempfile::TempDir::new().unwrap(); let source_home = tempfile::TempDir::new().unwrap(); - let _env = EnvGuard::set(&[ - ("HOME", cache_home.path().as_os_str()), - ("TOKSCALE_CONFIG_DIR", cache_home.path().as_os_str()), - ]); - let _pricing = EnvGuard::set(&[("TOKSCALE_PRICING_CACHE_ONLY", std::ffi::OsStr::new("1"))]); + let _env = opencode_test_env(cache_home.path(), source_home.path()); { - let cb_dir = source_home.path().join(".config/manicode/projects/proj"); - std::fs::create_dir_all(&cb_dir).unwrap(); - std::fs::write( - cb_dir.join("chat-messages.json"), - r#"[{"role":"assistant","id":"A","metadata":{"model":"claude-sonnet-4","usage":{"inputTokens":200,"outputTokens":80}},"credits":0.02}]"#, + let db_dir = source_home.path().join(".local/share/opencode"); + std::fs::create_dir_all(&db_dir).unwrap(); + + let schema = "PRAGMA journal_mode=WAL; + PRAGMA wal_autocheckpoint=0; + CREATE TABLE message ( + id TEXT PRIMARY KEY, + session_id TEXT NOT NULL, + data TEXT NOT NULL + );"; + let row = |input: u64, ts: u64| { + format!( + r#"{{ + "role": "assistant", + "modelID": "claude-sonnet-4", + "providerID": "anthropic", + "tokens": {{ "input": {input}, "output": 10, "reasoning": 0, "cache": {{ "read": 0, "write": 0 }} }}, + "time": {{ "created": {ts}.0 }} + }}"# + ) + }; + + let default_db = db_dir.join("opencode.db"); + let conn = rusqlite::Connection::open(&default_db).unwrap(); + conn.execute_batch(schema).unwrap(); + conn.execute( + "INSERT INTO message (id, session_id, data) VALUES (?1, ?2, ?3)", + rusqlite::params![ + "shared-msg", + "session-shared", + row(100, 1_700_000_000_000u64) + ], ) .unwrap(); - let kimi_dir = source_home.path().join(".kimi/sessions/g/s"); - std::fs::create_dir_all(&kimi_dir).unwrap(); - std::fs::write( - kimi_dir.join("wire.jsonl"), - "{\"type\": \"metadata\", \"protocol_version\": \"1.3\"}\n{\"timestamp\": 1770983410.0, \"message\": {\"type\": \"StatusUpdate\", \"payload\": {\"token_usage\": {\"input_other\": 100, \"output\": 50, \"input_cache_read\": 0, \"input_cache_creation\": 0}, \"message_id\": \"K\"}}}", + conn.execute( + "INSERT INTO message (id, session_id, data) VALUES (?1, ?2, ?3)", + rusqlite::params![ + "latest-only", + "session-latest", + row(200, 1_700_000_001_000u64) + ], ) .unwrap(); + drop(conn); - let home = source_home.path().to_str().unwrap().to_string(); - let rt = tokio::runtime::Builder::new_current_thread() - .enable_all() - .build() - .unwrap(); - let run = |clients: Option>| { - rt.block_on(get_agents_report(ReportOptions { - home_dir: Some(home.clone()), - use_env_roots: false, - clients, - ..Default::default() - })) - .unwrap() - }; + let stable_db = db_dir.join("opencode-stable.db"); + let conn = rusqlite::Connection::open(&stable_db).unwrap(); + conn.execute_batch(schema).unwrap(); + conn.execute( + "INSERT INTO message (id, session_id, data) VALUES (?1, ?2, ?3)", + rusqlite::params![ + "shared-msg", + "session-shared", + row(100, 1_700_000_000_000u64) + ], + ) + .unwrap(); + conn.execute( + "INSERT INTO message (id, session_id, data) VALUES (?1, ?2, ?3)", + rusqlite::params![ + "stable-only", + "session-stable", + row(300, 1_700_000_002_000u64) + ], + ) + .unwrap(); + drop(conn); - // All clients: one shared "Main" bucket carrying the mixed total. - let all = run(None); - assert_eq!(all.entries.len(), 1, "codebuff + kimi share one Main bucket"); - assert_eq!(all.entries[0].agent, "Main"); - assert_eq!(all.entries[0].input, 300, "mixed bucket = 200 + 100"); - assert_eq!(all.entries[0].output, 130, "mixed bucket = 80 + 50"); - assert_eq!(all.total_messages, 2); + let messages = parse_all_messages_with_pricing( + source_home.path().to_str().unwrap(), + &["opencode".to_string()], + None, + ); + assert_eq!( + messages.len(), + 3, + "expected 3 unique messages (shared + latest-only + stable-only), got {}", + messages.len() + ); + let mut ids: Vec = messages + .iter() + .filter_map(|m| m.dedup_key.clone()) + .collect(); + ids.sort(); + assert_eq!(ids, vec!["latest-only", "shared-msg", "stable-only"]); - // Filtered to codebuff: the SAME shared bucket, scoped at the scan - // to codebuff's contribution alone — proves the FFI-level filter, - // not a whole-bucket membership keep (which would still read 300). - let filtered = run(Some(vec!["codebuff".to_string()])); - assert_eq!(filtered.entries.len(), 1); - assert_eq!(filtered.entries[0].agent, "Main"); + let messages_warm = parse_all_messages_with_pricing( + source_home.path().to_str().unwrap(), + &["opencode".to_string()], + None, + ); assert_eq!( - filtered.entries[0].input, 200, - "filtered = codebuff only, not the mixed 300" + messages_warm.len(), + 3, + "warm cache must also dedup shared message across channel dbs" ); - assert_eq!(filtered.entries[0].output, 80); - assert_eq!(filtered.entries[0].clients, vec!["codebuff".to_string()]); - assert_eq!(filtered.total_messages, 1, "kimi's message is gone"); } } - // Issue #36 (round 3): a cc-mirror variant id (`cc-mirror/kimi-code`) is - // produced during CLAUDE-lane parsing, not by a scanner lane of its own — - // `ClientId::from_str("cc-mirror/kimi-code")` is None. The two-level split - // must (a) map the variant to its producing `claude` lane so the scan finds - // it, and (b) keep ONLY the exact requested ids at fold time so requesting - // the variant returns just the variant, and requesting `claude` returns - // plain claude WITHOUT the variant (which the graph/daily/models surface as - // its own client id). RED against the old lane-only filter: requesting the - // variant returned empty, and requesting claude swept the variant in. #[test] #[serial_test::serial] - fn test_agents_report_cc_mirror_variant_slice_issue36() { + fn test_parse_all_messages_with_pricing_opencode_sqlite_deduplicates_forked_history() { let cache_home = tempfile::TempDir::new().unwrap(); let source_home = tempfile::TempDir::new().unwrap(); - let _env = EnvGuard::set(&[ - ("HOME", cache_home.path().as_os_str()), - ("TOKSCALE_CONFIG_DIR", cache_home.path().as_os_str()), - ]); - let _pricing = EnvGuard::set(&[("TOKSCALE_PRICING_CACHE_ONLY", std::ffi::OsStr::new("1"))]); + let _env = opencode_test_env(cache_home.path(), source_home.path()); { - // Plain claude session (client "claude"): 100 in / 50 out. - let claude_dir = source_home.path().join(".claude/projects/myproject"); - std::fs::create_dir_all(&claude_dir).unwrap(); - std::fs::write( - claude_dir.join("conversation.jsonl"), - r#"{"type":"assistant","timestamp":"2024-12-01T10:00:00.000Z","requestId":"req_plain","message":{"id":"msg_plain","model":"claude-3-5-sonnet","usage":{"input_tokens":100,"output_tokens":50}}}"#, - ) - .unwrap(); - // cc-mirror variant (client "cc-mirror/kimi-code"): 300 in / 70 out. - let variant_dir = source_home.path().join(".cc-mirror/kimi-code"); - let config_dir = variant_dir.join("config"); - let project_dir = config_dir.join("projects/proj"); - std::fs::create_dir_all(&project_dir).unwrap(); - std::fs::write( - variant_dir.join("variant.json"), - serde_json::json!({ - "name": "kimi-code", - "provider": "kimi", - "configDir": config_dir, - }) - .to_string(), - ) - .unwrap(); - std::fs::write( - project_dir.join("session.jsonl"), - r#"{"type":"assistant","timestamp":"2024-12-01T11:00:00.000Z","requestId":"req_variant","message":{"id":"msg_variant","model":"claude-3-5-sonnet","usage":{"input_tokens":300,"output_tokens":70}}}"#, - ) - .unwrap(); + let db_dir = source_home.path().join(".local/share/opencode"); + std::fs::create_dir_all(&db_dir).unwrap(); + let db_path = db_dir.join("opencode.db"); + let conn = create_opencode_sqlite_db(&db_path); - let home = source_home.path().to_str().unwrap().to_string(); - let rt = tokio::runtime::Builder::new_current_thread() - .enable_all() - .build() - .unwrap(); - let run = |clients: Option>| { - let report = rt - .block_on(get_agents_report(ReportOptions { - home_dir: Some(home.clone()), - use_env_roots: false, - clients, - ..Default::default() - })) - .unwrap(); - let input: i64 = report.entries.iter().map(|e| e.input).sum(); - (report.total_messages, input) - }; + let msg_a = build_opencode_sqlite_payload( + 1_700_000_000_000.0, + 1_700_000_000_500.0, + 100, + 50, + 0, + 10, + 5, + 0.01, + ); + let msg_b = build_opencode_sqlite_payload( + 1_700_000_001_000.0, + 1_700_000_001_500.0, + 200, + 80, + 10, + 20, + 0, + 0.02, + ); + let msg_c = build_opencode_sqlite_payload( + 1_700_000_002_000.0, + 1_700_000_002_500.0, + 300, + 120, + 15, + 0, + 0, + 0.03, + ); - // All clients: both messages fold into "Main". - assert_eq!(run(None), (2, 400), "all: plain(100) + variant(300)"); + for (id, session_id, payload) in [ + ("root_a", "root", msg_a.as_str()), + ("root_b", "root", msg_b.as_str()), + ("fork_a_copy", "fork", msg_a.as_str()), + ("fork_b_copy", "fork", msg_b.as_str()), + ("fork_c_new", "fork", msg_c.as_str()), + ] { + conn.execute( + "INSERT INTO message (id, session_id, data) VALUES (?1, ?2, ?3)", + rusqlite::params![id, session_id, payload], + ) + .unwrap(); + } + drop(conn); - // Variant slice: the claude lane is scanned (so the variant is - // found), then narrowed to exactly the variant id. - assert_eq!( - run(Some(vec!["cc-mirror/kimi-code".to_string()])), - (1, 300), - "variant slice = just the variant (300), not empty" + let messages = parse_all_messages_with_pricing( + source_home.path().to_str().unwrap(), + &["opencode".to_string()], + None, ); - // Claude slice: plain claude ONLY — the distinct variant is excluded. - assert_eq!( - run(Some(vec!["claude".to_string()])), - (1, 100), - "claude slice = plain claude (100), not the mixed 400" - ); + assert_eq!(messages.len(), 3); + assert_eq!(messages.iter().map(|m| m.tokens.input).sum::(), 600); + assert_eq!(messages.iter().map(|m| m.tokens.output).sum::(), 250); + assert_eq!(messages.iter().map(|m| m.cost).sum::(), 0.06); } } - // Agent bucketing + fold arithmetic in isolation (no fixtures): normalized - // names, the "Main" fallback, plain `+=` token sums, and message_count.max(0). #[test] - fn test_agent_bucket_key_and_accumulator() { - let msg = |agent: Option<&str>| { - let mut m = UnifiedMessage::new_with_agent( - "codebuff", - "m", - "p", - "s", + #[serial_test::serial] + fn test_parse_local_clients_opencode_sqlite_counts_deduplicated_forked_history() { + let cache_home = tempfile::TempDir::new().unwrap(); + let source_home = tempfile::TempDir::new().unwrap(); + let _env = opencode_test_env(cache_home.path(), source_home.path()); + + { + let db_dir = source_home.path().join(".local/share/opencode"); + std::fs::create_dir_all(&db_dir).unwrap(); + let db_path = db_dir.join("opencode.db"); + let conn = create_opencode_sqlite_db(&db_path); + + let msg_a = build_opencode_sqlite_payload( + 1_700_000_000_000.0, + 1_700_000_000_500.0, + 100, + 50, 0, - TokenBreakdown { - input: 10, - output: 5, - cache_read: 2, - cache_write: 1, - reasoning: 3, - }, - 0.5, - agent.map(|a| a.to_string()), + 10, + 5, + 0.01, + ); + let msg_b = build_opencode_sqlite_payload( + 1_700_000_001_000.0, + 1_700_000_001_500.0, + 200, + 80, + 10, + 20, + 0, + 0.02, + ); + let msg_c = build_opencode_sqlite_payload( + 1_700_000_002_000.0, + 1_700_000_002_500.0, + 300, + 120, + 15, + 0, + 0, + 0.03, ); - m.message_count = 2; - m - }; - assert_eq!(agent_bucket_key(&msg(None)), "Main"); - assert_eq!(agent_bucket_key(&msg(Some(" "))), "Main"); - assert_eq!(agent_bucket_key(&msg(Some("OmO"))), "Sisyphus"); + for (id, session_id, payload) in [ + ("root_a", "root", msg_a.as_str()), + ("root_b", "root", msg_b.as_str()), + ("fork_a_copy", "fork", msg_a.as_str()), + ("fork_b_copy", "fork", msg_b.as_str()), + ("fork_c_new", "fork", msg_c.as_str()), + ] { + conn.execute( + "INSERT INTO message (id, session_id, data) VALUES (?1, ?2, ?3)", + rusqlite::params![id, session_id, payload], + ) + .unwrap(); + } + drop(conn); - let mut acc = AgentAccumulator::default(); - acc.add(&msg(None)); - let mut negative = msg(None); - negative.message_count = -3; // .max(0) clamp → contributes 0 messages - acc.add(&negative); + let parsed = parse_local_clients(LocalParseOptions { + home_dir: Some(source_home.path().to_str().unwrap().to_string()), + use_env_roots: false, + clients: Some(vec!["opencode".to_string()]), + since: None, + until: None, + year: None, + scanner_settings: scanner::ScannerSettings::default(), + modified_after: None, + }) + .unwrap(); - assert_eq!(acc.input, 20); - assert_eq!(acc.output, 10); - assert_eq!(acc.cache_read, 4); - assert_eq!(acc.cache_write, 2); - assert_eq!(acc.reasoning, 6); - assert!((acc.cost - 1.0).abs() < 1e-9); - assert_eq!(acc.messages, 2, "message_count.max(0): 2 + 0"); - assert!(acc.clients.contains("codebuff")); + assert_eq!(parsed.counts.get(ClientId::OpenCode), 3); + assert_eq!(parsed.messages.len(), 3); + assert_eq!(parsed.messages.iter().map(|m| m.input).sum::(), 600); + assert_eq!(parsed.messages.iter().map(|m| m.output).sum::(), 250); + } } - #[test] - fn agent_accumulator_saturates_overflowing_token_folds() { - // Vendor-local sibling sweep alongside #823: AgentAccumulator::add is - // its own per-field CROSS-MESSAGE fold (agents streaming report), not - // one of the 6 sites #823 covers. An antigravity-cli row can carry an - // i64::MAX bucket after the untrusted-varint clamp, so two such rows - // folded into one agent bucket with plain `+=` overflow (debug panic / - // release wrap) before any saturating grand total runs. - let make = || { - UnifiedMessage::new( - "antigravity-cli", - "gemini-3-pro", - "antigravity", - "session-overflow", - 1_733_011_200_000, - TokenBreakdown { - input: i64::MAX, - output: 0, - cache_read: i64::MAX, - cache_write: 0, - reasoning: 0, - }, - 0.0, - ) - }; + /// Regression fixture for Codex sessions that are live-only, archive-only, + /// or briefly present in both roots while the CLI moves a transcript. + fn write_codex_sessions_and_archived_sessions_fixture(source_home: &std::path::Path) { + let sessions_dir = source_home.join(".codex/sessions"); + let archived_dir = source_home.join(".codex/archived_sessions"); + std::fs::create_dir_all(&sessions_dir).unwrap(); + std::fs::create_dir_all(&archived_dir).unwrap(); - let mut acc = AgentAccumulator::default(); - acc.add(&make()); - acc.add(&make()); + std::fs::write( + sessions_dir.join("live-only.jsonl"), + concat!( + r#"{"timestamp":"2026-06-25T10:00:00Z","type":"session_meta","payload":{"id":"33333333-3333-7333-8333-333333333333","source":"interactive","model_provider":"openai","cwd":"/repo"}}"#, + "\n", + r#"{"timestamp":"2026-06-25T10:00:01Z","type":"turn_context","payload":{"model":"gpt-5.5"}}"#, + "\n", + r#"{"timestamp":"2026-06-25T10:00:02Z","type":"event_msg","payload":{"type":"token_count","info":{"total_token_usage":{"input_tokens":50,"output_tokens":5,"total_tokens":55},"last_token_usage":{"input_tokens":50,"output_tokens":5,"total_tokens":55}}}}"#, + "\n" + ), + ) + .unwrap(); - assert_eq!(acc.input, i64::MAX); - assert_eq!(acc.cache_read, i64::MAX); + std::fs::write( + archived_dir.join("archived-only.jsonl"), + concat!( + r#"{"timestamp":"2026-06-20T09:00:00Z","type":"session_meta","payload":{"id":"44444444-4444-7444-8444-444444444444","source":"interactive","model_provider":"openai","cwd":"/repo"}}"#, + "\n", + r#"{"timestamp":"2026-06-20T09:00:01Z","type":"turn_context","payload":{"model":"gpt-5.5"}}"#, + "\n", + r#"{"timestamp":"2026-06-20T09:00:02Z","type":"event_msg","payload":{"type":"token_count","info":{"total_token_usage":{"input_tokens":70,"output_tokens":7,"total_tokens":77},"last_token_usage":{"input_tokens":70,"output_tokens":7,"total_tokens":77}}}}"#, + "\n" + ), + ) + .unwrap(); + + let shared_content = concat!( + r#"{"timestamp":"2026-06-22T08:00:00Z","type":"session_meta","payload":{"id":"55555555-5555-7555-8555-555555555555","source":"interactive","model_provider":"openai","cwd":"/repo"}}"#, + "\n", + r#"{"timestamp":"2026-06-22T08:00:01Z","type":"turn_context","payload":{"model":"gpt-5.5"}}"#, + "\n", + r#"{"timestamp":"2026-06-22T08:00:02Z","type":"event_msg","payload":{"type":"token_count","info":{"total_token_usage":{"input_tokens":30,"output_tokens":3,"total_tokens":33},"last_token_usage":{"input_tokens":30,"output_tokens":3,"total_tokens":33}}}}"#, + "\n" + ); + std::fs::write( + sessions_dir.join("shared-in-sessions.jsonl"), + shared_content, + ) + .unwrap(); + std::fs::write( + archived_dir.join("shared-in-archived.jsonl"), + shared_content, + ) + .unwrap(); + } + + fn scan_test_sources( + source_home: &Path, + clients: &[String], + scanner_settings: &scanner::ScannerSettings, + ) -> scanner::ScanResult { + scanner::scan_all_clients_with_scanner_settings( + source_home.to_str().unwrap(), + clients, + false, + scanner_settings, + ) + } + + // Cache keys preserve the scanner-returned path spelling. On Windows an + // explicit-home resolver can use `/` while TempDir::join produces `\`, so + // cache fixtures must seed and query with the production scanner spelling. + fn scanner_spelling(paths: &[PathBuf], expected: &Path) -> PathBuf { + let expected = std::fs::canonicalize(expected).unwrap(); + paths + .iter() + .find(|path| std::fs::canonicalize(path).is_ok_and(|path| path == expected)) + .cloned() + .expect("scanner must return the expected physical source") + } + + fn with_isolated_tokscale_cache( + cache_home: &std::path::Path, + action: impl FnOnce() -> T, + ) -> T { + let _env = EnvGuard::set(&[ + ("HOME", cache_home.as_os_str()), + ("TOKSCALE_CONFIG_DIR", cache_home.as_os_str()), + ("TOKSCALE_PRICING_CACHE_ONLY", std::ffi::OsStr::new("1")), + ]); + action() + } + + fn shard_cache_has_entries(cache_home: &std::path::Path, namespace: &str) -> bool { + std::fs::read_dir( + cache_home + .join("cache/source-message-cache-v2") + .join(namespace), + ) + .is_ok_and(|entries| { + entries + .filter_map(Result::ok) + .any(|entry| entry.path().extension().is_some_and(|ext| ext == "bin")) + }) } - #[test] - fn test_agent_bucket_key_copilot_uses_copilot_normalizer() { - // #724/#751: copilot messages carry a raw OTEL agent id. Our agents - // report must prettify it with the copilot-specific normalizer (the - // prettification upstream does in its CLI), while other clients keep the - // generic normalization. - let msg = |client: &str, agent: &str| { - UnifiedMessage::new_with_agent( - client, - "m", - "p", - "s", - 0, - TokenBreakdown::default(), - 0.0, - Some(agent.to_string()), - ) + #[test] + #[serial_test::serial] + fn test_claude_shipping_report_cold_start_persists_shard() { + let source_home = tempfile::TempDir::new().unwrap(); + let cache_home = tempfile::TempDir::new().unwrap(); + let transcripts_dir = source_home.path().join(".claude/transcripts"); + std::fs::create_dir_all(&transcripts_dir).unwrap(); + let source = transcripts_dir.join("ses_123456789012345678901234567.jsonl"); + std::fs::write( + &source, + concat!( + r#"{"type":"user","timestamp":"2026-04-01T10:00:00.000Z","message":{"content":"prompt"}}"#, + "\n", + r#"{"type":"assistant","timestamp":"2026-04-01T10:00:01.000Z","requestId":"req-shipping","message":{"id":"msg-shipping","model":"claude-sonnet-4","usage":{"input_tokens":123,"output_tokens":45}}}"#, + "\n", + ), + ) + .unwrap(); + + let runtime = tokio::runtime::Builder::new_current_thread() + .enable_all() + .build() + .unwrap(); + let report_options = || ReportOptions { + home_dir: Some(source_home.path().to_str().unwrap().to_string()), + use_env_roots: false, + clients: Some(vec!["claude".to_string()]), + ..Default::default() }; - - // Copilot: raw OTEL ids resolve to their pretty display form. - assert_eq!( - agent_bucket_key(&msg("copilot", "github.copilot.default")), - "GitHub Copilot" - ); - assert_eq!( - agent_bucket_key(&msg("copilot", "Plugin:code-review-team:api-reviewer")), - "Code Review Team: API Reviewer" + let cold = with_isolated_tokscale_cache(cache_home.path(), || { + runtime + .block_on(get_model_report(report_options())) + .unwrap() + }); + assert_eq!(cold.total_messages, 1); + assert!(shard_cache_has_entries(cache_home.path(), "claude")); + let clients = ["claude".to_string()]; + let scan_result = scan_test_sources( + source_home.path(), + &clients, + &scanner::ScannerSettings::default(), ); + let cache_source = scanner_spelling(scan_result.get(ClientId::Claude), &source); + with_isolated_tokscale_cache(cache_home.path(), || { + let cache = message_cache::SourceMessageCache::load(); + assert!(cache + .get( + message_cache::CacheIdentity::for_client(ClientId::Claude), + &cache_source, + ) + .is_some_and(|entry| entry.messages.len() == 1)); + }); - // A non-copilot client with the same raw id must NOT get the - // copilot-specific prettification (proves the branch is client-scoped). - assert_ne!( - agent_bucket_key(&msg("codebuff", "github.copilot.default")), - "GitHub Copilot" - ); + let warm = with_isolated_tokscale_cache(cache_home.path(), || { + runtime + .block_on(get_model_report(report_options())) + .unwrap() + }); + assert_eq!(warm.total_messages, cold.total_messages); + assert_eq!(warm.total_input, cold.total_input); + assert_eq!(warm.total_output, cold.total_output); } #[test] #[serial_test::serial] - fn test_source_cache_refreshes_stale_date_on_cache_hit() { - let cache_home = tempfile::TempDir::new().unwrap(); + fn test_claude_shipping_warm_cache_reuses_persisted_parent_paths() { let source_home = tempfile::TempDir::new().unwrap(); - let _env = opencode_test_env(cache_home.path(), source_home.path()); + let cache_home = tempfile::TempDir::new().unwrap(); + let project_dir = source_home.path().join(".claude/projects/project-one"); + std::fs::create_dir_all(&project_dir).unwrap(); + let source = project_dir.join("agent-child.jsonl"); + std::fs::write( + &source, + concat!( + r#"{"type":"assistant","isSidechain":true,"sessionId":"flat-parent","agentId":"child","timestamp":"2026-01-01T00:00:00Z","requestId":"req-source","message":{"id":"msg-source","model":"claude-sonnet-4","usage":{"input_tokens":123,"output_tokens":45}}}"#, + "\n" + ), + ) + .unwrap(); + let clients = ["claude".to_string()]; + let scan_result = scan_test_sources( + source_home.path(), + &clients, + &scanner::ScannerSettings::default(), + ); + let cache_source = scanner_spelling(scan_result.get(ClientId::Claude), &source); - { - let message_dir = scanner_fixture_path( - source_home.path(), - ".local/share/opencode/storage/message/project-1", - ); - std::fs::create_dir_all(&message_dir).unwrap(); - let path = message_dir.join("msg_001.json"); - std::fs::write( - &path, - r#"{"id":"msg-1","sessionID":"session-1","role":"assistant","modelID":"accounts/fireworks/models/deepseek-v3-0324","providerID":"fireworks","cost":0,"tokens":{"input":10,"output":5,"reasoning":0,"cache":{"read":0,"write":0}},"time":{"created":1733011200000}}"#, - ) + let mut fingerprint = message_cache::SourceFingerprint::from_claude_code_path_with_home( + &cache_source, + Some(source_home.path()), + ) + .unwrap(); + let persisted_parent = fingerprint + .related_files + .iter_mut() + .find(|related| related.suffix.starts_with("parent-session-")) .unwrap(); + persisted_parent.path = + message_cache::CachedPath::from_path(&project_dir.join("cached-only-parent.jsonl")); + assert!(!persisted_parent.exists); - let fingerprint = message_cache::SourceFingerprint::from_path(&path).unwrap(); - let mut stale_message = UnifiedMessage::new( - "opencode", - "accounts/fireworks/models/deepseek-v3-0324", - "fireworks", - "session-1", - 1_733_011_200_000, - TokenBreakdown { - input: 10, - output: 5, - cache_read: 0, - cache_write: 0, - reasoning: 0, - }, - 0.0, - ); - stale_message.date = "1900-01-01".to_string(); - + let sentinel = UnifiedMessage::new_with_dedup( + "claude", + "cached-sentinel", + "anthropic", + "cached-session", + 1_767_225_600_000, + TokenBreakdown { + input: 999, + output: 1, + ..Default::default() + }, + 0.0, + Some("cached-sentinel".to_string()), + ); + with_isolated_tokscale_cache(cache_home.path(), || { let mut cache = message_cache::SourceMessageCache::default(); cache.insert(message_cache::CachedSourceEntry::new( - &path, + message_cache::CacheIdentity::for_client(ClientId::Claude), + &cache_source, fingerprint, - vec![stale_message], + vec![sentinel], Vec::new(), None, )); cache.save_if_dirty(); + }); - let messages = parse_all_messages_with_pricing( + let mut streamed = Vec::new(); + with_isolated_tokscale_cache(cache_home.path(), || { + scan_messages_streaming( source_home.path().to_str().unwrap(), - &["opencode".to_string()], + &clients, None, + false, + &scanner::ScannerSettings::default(), + &|_| true, + &mut |message| streamed.push(message.clone()), ); + }); - assert_eq!(messages.len(), 1); - assert_ne!(messages[0].date, "1900-01-01"); + assert_eq!(streamed.len(), 1); + assert_eq!(streamed[0].model_id, "cached-sentinel"); + assert_eq!(streamed[0].tokens.input, 999); + } + + #[test] + #[serial_test::serial] + fn test_zed_shipping_report_cold_start_persists_shard() { + let source_home = tempfile::TempDir::new().unwrap(); + let cache_home = tempfile::TempDir::new().unwrap(); + let threads_dir = source_home.path().join("custom-zed/threads"); + std::fs::create_dir_all(&threads_dir).unwrap(); + let source = threads_dir.join("threads.db"); + let connection = create_zed_sqlite_db(&source); + insert_zed_thread(&connection, "zed-shipping", "claude-sonnet-4-5"); + drop(connection); + + let runtime = tokio::runtime::Builder::new_current_thread() + .enable_all() + .build() + .unwrap(); + let report_options = || { + let mut extra_scan_paths = std::collections::BTreeMap::new(); + extra_scan_paths.insert("zed".to_string(), vec![threads_dir.clone()]); + ReportOptions { + home_dir: Some(source_home.path().to_str().unwrap().to_string()), + use_env_roots: false, + clients: Some(vec!["zed".to_string()]), + scanner_settings: scanner::ScannerSettings { + extra_scan_paths, + ..Default::default() + }, + ..Default::default() + } + }; + let cold = with_isolated_tokscale_cache(cache_home.path(), || { + runtime + .block_on(get_model_report(report_options())) + .unwrap() + }); + assert_eq!(cold.total_messages, 1); + assert!(shard_cache_has_entries(cache_home.path(), "zed")); + with_isolated_tokscale_cache(cache_home.path(), || { + let cache = message_cache::SourceMessageCache::load(); + assert!(cache + .get( + message_cache::CacheIdentity::for_client(ClientId::Zed), + &source, + ) + .is_some_and(|entry| entry.messages.len() == 1)); + }); + + let warm = with_isolated_tokscale_cache(cache_home.path(), || { + runtime + .block_on(get_model_report(report_options())) + .unwrap() + }); + assert_eq!(warm.total_messages, cold.total_messages); + assert_eq!(warm.total_input, cold.total_input); + assert_eq!(warm.total_output, cold.total_output); + } + + fn write_codex_duration_prefix_fixture( + source_home: &std::path::Path, + ) -> (std::path::PathBuf, String) { + let lines = include_str!("../tests/fixtures/codex_duration_timing.jsonl") + .lines() + .collect::>(); + let sessions_dir = source_home.join(".codex/sessions"); + std::fs::create_dir_all(&sessions_dir).unwrap(); + let source = sessions_dir.join("codex-duration-timing.jsonl"); + std::fs::write(&source, format!("{}\n", lines[..5].join("\n"))).unwrap(); + (source, format!("{}\n", lines[5..].join("\n"))) + } + + #[test] + #[serial_test::serial] + fn test_codex_non_overlapping_durations_survive_incremental_and_streaming_caches() { + let materialized_source_home = tempfile::TempDir::new().unwrap(); + let streaming_source_home = tempfile::TempDir::new().unwrap(); + let materialized_cache = tempfile::TempDir::new().unwrap(); + let streaming_cache = tempfile::TempDir::new().unwrap(); + let (source, suffix) = + write_codex_duration_prefix_fixture(materialized_source_home.path()); + let (streaming_source, streaming_suffix) = + write_codex_duration_prefix_fixture(streaming_source_home.path()); + + let home = materialized_source_home + .path() + .to_str() + .unwrap() + .to_string(); + let streaming_home = streaming_source_home + .path() + .to_str() + .unwrap() + .to_string(); + let clients = vec!["codex".to_string()]; + let streaming_scan = scan_test_sources( + streaming_source_home.path(), + &clients, + &scanner::ScannerSettings::default(), + ); + let cache_streaming_source = + scanner_spelling(streaming_scan.get(ClientId::Codex), &streaming_source); + let parse_materialized = || { + parse_all_messages_with_pricing_with_env_strategy( + &home, + &clients, + None, + false, + &scanner::ScannerSettings::default(), + ) + }; + + let prefix = with_isolated_tokscale_cache(materialized_cache.path(), parse_materialized); + assert_eq!(prefix.len(), 1); + assert_eq!(prefix[0].duration_ms, Some(1_000)); + + let mut source_file = std::fs::OpenOptions::new() + .append(true) + .open(&source) + .unwrap(); + source_file.write_all(suffix.as_bytes()).unwrap(); + source_file.flush().unwrap(); + drop(source_file); + + let materialized_incremental = + with_isolated_tokscale_cache(materialized_cache.path(), parse_materialized); + let materialized_warm = + with_isolated_tokscale_cache(materialized_cache.path(), parse_materialized); + for (phase, messages) in [ + ("incremental", &materialized_incremental), + ("warm", &materialized_warm), + ] { + assert_eq!(messages.len(), 3, "{phase} materialized messages"); assert_eq!( - messages[0].date, - UnifiedMessage::new( - "opencode", - "accounts/fireworks/models/deepseek-v3-0324", - "fireworks", - "session-1", - 1_733_011_200_000, - TokenBreakdown { - input: 10, - output: 5, - cache_read: 0, - cache_write: 0, - reasoning: 0, - }, - 0.0, + messages + .iter() + .map(|message| message.duration_ms) + .collect::>(), + vec![Some(1_000), Some(4_000), Some(2_000)], + "{phase} materialized durations" + ); + } + + let runtime = tokio::runtime::Builder::new_current_thread() + .enable_all() + .build() + .unwrap(); + let report_options = || ReportOptions { + home_dir: Some(streaming_home.clone()), + use_env_roots: false, + clients: Some(clients.clone()), + ..Default::default() + }; + let streaming_cold = with_isolated_tokscale_cache(streaming_cache.path(), || { + runtime + .block_on(get_model_report(report_options())) + .unwrap() + }); + assert_eq!(streaming_cold.total_messages, 1); + assert!( + shard_cache_has_entries(streaming_cache.path(), "codex"), + "the cold shipping report must persist its own Codex shard" + ); + with_isolated_tokscale_cache(streaming_cache.path(), || { + let cache = message_cache::SourceMessageCache::load(); + let incremental = cache + .get( + message_cache::CacheIdentity::for_client(ClientId::Codex), + &cache_streaming_source, ) - .date + .and_then(|entry| entry.codex_incremental.as_ref()) + .expect("cold shipping report must persist Codex incremental state"); + assert_eq!( + incremental.consumed_offset, + std::fs::metadata(&streaming_source).unwrap().len() ); - } - } - - #[cfg(unix)] - #[test] - #[serial_test::serial] - fn test_empty_parse_results_are_not_cached_for_optional_file_sources() { - use std::os::unix::fs::PermissionsExt; - - let cache_home = tempfile::TempDir::new().unwrap(); - let source_home = tempfile::TempDir::new().unwrap(); - let _env = opencode_test_env(cache_home.path(), source_home.path()); + }); - { - let message_dir = scanner_fixture_path( - source_home.path(), - ".local/share/opencode/storage/message/project-1", - ); - std::fs::create_dir_all(&message_dir).unwrap(); - let path = message_dir.join("msg_001.json"); - std::fs::write( - &path, - r#"{"id":"msg-1","sessionID":"session-1","role":"assistant","modelID":"accounts/fireworks/models/deepseek-v3-0324","providerID":"fireworks","cost":0,"tokens":{"input":10,"output":5,"reasoning":0,"cache":{"read":0,"write":0}},"time":{"created":1733011200000}}"#, - ) + let mut streaming_source_file = std::fs::OpenOptions::new() + .append(true) + .open(&streaming_source) .unwrap(); + streaming_source_file + .write_all(streaming_suffix.as_bytes()) + .unwrap(); + streaming_source_file.flush().unwrap(); + drop(streaming_source_file); - let mut permissions = std::fs::metadata(&path).unwrap().permissions(); - permissions.set_mode(0o000); - std::fs::set_permissions(&path, permissions).unwrap(); - - let first_messages = parse_all_messages_with_pricing( - source_home.path().to_str().unwrap(), - &["opencode".to_string()], - None, - ); - assert!(first_messages.is_empty()); - + let streaming_incremental = + with_isolated_tokscale_cache(streaming_cache.path(), || { + runtime + .block_on(get_model_report(report_options())) + .unwrap() + }); + with_isolated_tokscale_cache(streaming_cache.path(), || { let cache = message_cache::SourceMessageCache::load(); - assert!(cache.get(&path).is_none()); - - let mut readable_permissions = std::fs::metadata(&path).unwrap().permissions(); - readable_permissions.set_mode(0o644); - std::fs::set_permissions(&path, readable_permissions).unwrap(); + let incremental = cache + .get( + message_cache::CacheIdentity::for_client(ClientId::Codex), + &cache_streaming_source, + ) + .and_then(|entry| entry.codex_incremental.as_ref()) + .expect("incremental shipping report must refresh Codex state"); + assert_eq!( + incremental.consumed_offset, + std::fs::metadata(&streaming_source).unwrap().len() + ); + }); + let streaming_warm = with_isolated_tokscale_cache(streaming_cache.path(), || { + runtime + .block_on(get_model_report(report_options())) + .unwrap() + }); - let second_messages = parse_all_messages_with_pricing( - source_home.path().to_str().unwrap(), - &["opencode".to_string()], - None, + for (phase, report) in [ + ("incremental", &streaming_incremental), + ("warm", &streaming_warm), + ] { + assert_eq!(report.total_messages, 3, "{phase} streaming messages"); + assert_eq!(report.entries.len(), 1, "{phase} model groups"); + let performance = &report.entries[0].performance; + assert_eq!( + performance.total_duration_ms, 7_000, + "{phase} total duration" + ); + assert_eq!(performance.timed_tokens, 170, "{phase} timed tokens"); + assert_eq!(performance.sample_count, 3, "{phase} samples"); + assert_eq!(performance.token_coverage, 1.0, "{phase} coverage"); + let expected_ms_per_1k = 7_000.0 * 1_000.0 / 170.0; + assert!( + (performance.ms_per_1k_tokens.unwrap() - expected_ms_per_1k).abs() < f64::EPSILON, + "{phase} milliseconds per 1K tokens" ); - assert_eq!(second_messages.len(), 1); } } #[test] #[serial_test::serial] - fn test_empty_cache_hits_are_reparsed_for_optional_file_sources() { - let cache_home = tempfile::TempDir::new().unwrap(); + fn test_codex_archive_roots_are_exact_once_across_all_consumers() { let source_home = tempfile::TempDir::new().unwrap(); - let _env = opencode_test_env(cache_home.path(), source_home.path()); + let materialized_cache = tempfile::TempDir::new().unwrap(); + let streaming_cache = tempfile::TempDir::new().unwrap(); + let count_cache = tempfile::TempDir::new().unwrap(); - { - let message_dir = scanner_fixture_path( - source_home.path(), - ".local/share/opencode/storage/message/project-1", - ); - std::fs::create_dir_all(&message_dir).unwrap(); - let source_path = message_dir.join("msg_001.json"); - std::fs::write( - &source_path, - r#"{"id":"msg-1","sessionID":"session-1","role":"assistant","modelID":"accounts/fireworks/models/deepseek-v3-0324","providerID":"fireworks","cost":0,"tokens":{"input":10,"output":5,"reasoning":0,"cache":{"read":0,"write":0}},"time":{"created":1733011200000}}"#, - ) - .unwrap(); - let cache_path = scanner::scan_all_clients_with_env_strategy( - source_home.path().to_str().unwrap(), - &["opencode".to_string()], - true, - ) - .get(ClientId::OpenCode) - .first() - .cloned() - .expect("scanner must find the OpenCode fixture"); + write_codex_sessions_and_archived_sessions_fixture(source_home.path()); - let fingerprint = message_cache::SourceFingerprint::from_path(&cache_path).unwrap(); - let mut cache = message_cache::SourceMessageCache::default(); - cache.insert(message_cache::CachedSourceEntry::new( - &cache_path, - fingerprint, - Vec::new(), - Vec::new(), + let home = source_home.path().to_str().unwrap().to_string(); + let clients = vec!["codex".to_string()]; + let materialized = with_isolated_tokscale_cache(materialized_cache.path(), || { + parse_all_messages_with_pricing_with_env_strategy( + &home, + &clients, None, - )); - cache.save_if_dirty(); + false, + &scanner::ScannerSettings::default(), + ) + }); - let messages = parse_all_messages_with_pricing( - source_home.path().to_str().unwrap(), - &["opencode".to_string()], - None, - ); - assert_eq!(messages.len(), 1); + assert_eq!(materialized.len(), 3); + let session_ids: HashSet<_> = materialized + .iter() + .map(|message| message.session_id.as_str()) + .collect(); + assert!(session_ids.contains("live-only")); + assert!(session_ids.contains("archived-only")); + assert_eq!( + materialized + .iter() + .map(|message| message.tokens.input) + .sum::(), + 150, + ); + assert_eq!( + materialized + .iter() + .map(|message| message.tokens.output) + .sum::(), + 15, + ); - let loaded = message_cache::SourceMessageCache::load(); - let repaired_entry = loaded.get(&cache_path).unwrap(); - assert_eq!(repaired_entry.messages.len(), 1); - } - } + let runtime = tokio::runtime::Builder::new_current_thread() + .enable_all() + .build() + .unwrap(); + let report_options = || ReportOptions { + home_dir: Some(home.clone()), + use_env_roots: false, + clients: Some(clients.clone()), + ..Default::default() + }; + let streaming_cold = with_isolated_tokscale_cache(streaming_cache.path(), || { + runtime + .block_on(get_model_report(report_options())) + .unwrap() + }); - #[test] - #[serial_test::serial] - fn test_sqlite_source_cache_invalidates_on_wal_change() { - let cache_home = tempfile::TempDir::new().unwrap(); - let source_home = tempfile::TempDir::new().unwrap(); - let _env = opencode_test_env(cache_home.path(), source_home.path()); + assert!( + shard_cache_has_entries(streaming_cache.path(), "codex"), + "the cold shipping report must persist the cache used by the warm pass", + ); + let streaming_warm = with_isolated_tokscale_cache(streaming_cache.path(), || { + runtime + .block_on(get_model_report(report_options())) + .unwrap() + }); + for (phase, streaming) in [("cold", &streaming_cold), ("warm", &streaming_warm)] { + assert_eq!(streaming.total_messages, 3, "{phase} streaming messages"); + assert_eq!(streaming.total_input, 150, "{phase} streaming input"); + assert_eq!(streaming.total_output, 15, "{phase} streaming output"); + } - { - let db_dir = source_home.path().join(".local/share/opencode"); - std::fs::create_dir_all(&db_dir).unwrap(); - let db_path = db_dir.join("opencode.db"); + let counted = with_isolated_tokscale_cache(count_cache.path(), || { + parse_local_clients(LocalParseOptions { + home_dir: Some(home), + use_env_roots: false, + clients: Some(clients), + since: None, + until: None, + year: None, + scanner_settings: scanner::ScannerSettings::default(), + modified_after: None, + }) + .unwrap() + }); + assert_eq!(counted.counts.get(ClientId::Codex), 3); + assert_eq!(counted.messages.len(), 3); + assert_eq!( + counted + .messages + .iter() + .map(|message| message.input) + .sum::(), + 150, + ); + assert_eq!( + counted + .messages + .iter() + .map(|message| message.output) + .sum::(), + 15, + ); + } - let conn = rusqlite::Connection::open(&db_path).unwrap(); - let journal_mode: String = conn - .query_row("PRAGMA journal_mode=WAL;", [], |row| row.get(0)) - .unwrap(); - assert_eq!(journal_mode.to_lowercase(), "wal"); - conn.execute_batch( - "PRAGMA wal_autocheckpoint=0; - CREATE TABLE message ( - id TEXT PRIMARY KEY, - session_id TEXT NOT NULL, - data TEXT NOT NULL - );", - ) - .unwrap(); + fn write_codex_forked_history_fixture(source_home: &std::path::Path) { + let codex_dir = source_home.join(".codex/sessions"); + std::fs::create_dir_all(&codex_dir).unwrap(); + std::fs::write( + codex_dir.join("parent.jsonl"), + concat!( + r#"{"timestamp":"2026-04-30T10:00:00Z","type":"session_meta","payload":{"id":"parent-session","source":"interactive","model_provider":"openai","cwd":"/Users/alice/root"}}"#, + "\n", + r#"{"timestamp":"2026-04-30T10:00:01Z","type":"turn_context","payload":{"model":"gpt-5.2"}}"#, + "\n", + r#"{"timestamp":"2026-04-30T10:00:02Z","type":"event_msg","payload":{"type":"token_count","info":{"total_token_usage":{"input_tokens":50,"cached_input_tokens":10,"output_tokens":15,"total_tokens":65},"last_token_usage":{"input_tokens":50,"cached_input_tokens":10,"output_tokens":15,"total_tokens":65}}}}"#, + "\n", + r#"{"timestamp":"2026-04-30T10:00:03Z","type":"event_msg","payload":{"type":"token_count","info":{"total_token_usage":{"input_tokens":100,"cached_input_tokens":20,"output_tokens":30,"total_tokens":130},"last_token_usage":{"input_tokens":50,"cached_input_tokens":10,"output_tokens":15,"total_tokens":65}}}}"#, + "\n" + ), + ) + .unwrap(); + std::fs::write( + codex_dir.join("fork.jsonl"), + concat!( + r#"{"timestamp":"2026-04-30T10:01:00Z","type":"session_meta","payload":{"id":"fork-session","source":{"subagent":{"thread_spawn":{"parent_thread_id":"parent-session","depth":1}}},"model_provider":"openai","cwd":"/Users/alice/root-worktree"}}"#, + "\n", + r#"{"timestamp":"2026-04-30T10:01:01Z","type":"event_msg","payload":{"type":"token_count","info":{"total_token_usage":{"input_tokens":100,"cached_input_tokens":20,"output_tokens":30,"total_tokens":130},"last_token_usage":{"input_tokens":100,"cached_input_tokens":20,"output_tokens":30,"total_tokens":130}}}}"#, + "\n", + r#"{"timestamp":"2026-04-30T10:01:02Z","type":"turn_context","payload":{"model":"gpt-5.2"}}"#, + "\n", + r#"{"timestamp":"2026-04-30T10:01:03Z","type":"event_msg","payload":{"type":"token_count","info":{"total_token_usage":{"input_tokens":50,"cached_input_tokens":10,"output_tokens":15,"total_tokens":65},"last_token_usage":{"input_tokens":50,"cached_input_tokens":10,"output_tokens":15,"total_tokens":65}}}}"#, + "\n", + r#"{"timestamp":"2026-04-30T10:01:04Z","type":"event_msg","payload":{"type":"token_count","info":{"total_token_usage":{"input_tokens":100,"cached_input_tokens":20,"output_tokens":30,"total_tokens":130},"last_token_usage":{"input_tokens":50,"cached_input_tokens":10,"output_tokens":15,"total_tokens":65}}}}"#, + "\n", + r#"{"timestamp":"2026-04-30T10:01:05Z","type":"event_msg","payload":{"type":"token_count","info":{"total_token_usage":{"input_tokens":110,"cached_input_tokens":22,"output_tokens":33,"total_tokens":143},"last_token_usage":{"input_tokens":10,"cached_input_tokens":2,"output_tokens":3,"total_tokens":13}}}}"#, + "\n" + ), + ) + .unwrap(); + } - let row_one = r#"{ - "role": "assistant", - "modelID": "claude-sonnet-4", - "providerID": "anthropic", - "tokens": { "input": 100, "output": 50, "reasoning": 0, "cache": { "read": 0, "write": 0 } }, - "time": { "created": 1700000000000.0 } - }"#; - let row_two = r#"{ - "role": "assistant", - "modelID": "claude-sonnet-4", - "providerID": "anthropic", - "tokens": { "input": 120, "output": 60, "reasoning": 0, "cache": { "read": 0, "write": 0 } }, - "time": { "created": 1700000001000.0 } - }"#; + fn write_codex_parent_replay_fixture(source_home: &std::path::Path) { + let codex_dir = source_home.join(".codex/sessions"); + std::fs::create_dir_all(&codex_dir).unwrap(); + std::fs::write( + codex_dir.join("parent.jsonl"), + concat!( + r#"{"timestamp":"2026-05-24T20:00:00Z","type":"session_meta","payload":{"id":"019e5b00-0000-7000-8000-000000000001","source":"vscode","model_provider":"openai","cwd":"/repo"}}"#, + "\n", + r#"{"timestamp":"2026-05-24T20:00:01Z","type":"turn_context","payload":{"turn_id":"019e5b00-0001-7000-8000-000000000001","model":"gpt-5.5","cwd":"/repo"}}"#, + "\n", + r#"{"timestamp":"2026-05-24T20:00:02Z","type":"event_msg","payload":{"type":"token_count","info":{"total_token_usage":{"input_tokens":100,"output_tokens":10,"total_tokens":110},"last_token_usage":{"input_tokens":100,"output_tokens":10,"total_tokens":110}}}}"#, + "\n", + r#"{"timestamp":"2026-05-24T20:00:03Z","type":"event_msg","payload":{"type":"token_count","info":{"total_token_usage":{"input_tokens":130,"output_tokens":13,"total_tokens":143},"last_token_usage":{"input_tokens":30,"output_tokens":3,"total_tokens":33}}}}"#, + "\n" + ), + ) + .unwrap(); - conn.execute( - "INSERT INTO message (id, session_id, data) VALUES (?1, ?2, ?3)", - rusqlite::params!["msg-1", "session-1", row_one], + for (filename, child_id, child_turn_id, timestamp) in [ + ( + "child-a.jsonl", + "019e5c03-1e99-7000-8000-000000000001", + "019e5c03-6425-7000-8000-000000000001", + "2026-05-24T21:00:00Z", + ), + ( + "child-b.jsonl", + "019e5c04-1e99-7000-8000-000000000001", + "019e5c04-6425-7000-8000-000000000001", + "2026-05-24T22:00:00Z", + ), + ] { + std::fs::write( + codex_dir.join(filename), + format!( + concat!( + r#"{{"timestamp":"{timestamp}","type":"session_meta","payload":{{"id":"{child_id}","forked_from_id":"019e5b00-0000-7000-8000-000000000001","source":{{"subagent":{{"thread_spawn":{{"parent_thread_id":"019e5b00-0000-7000-8000-000000000001","depth":1}}}}}},"model_provider":"openai","agent_nickname":"worker","cwd":"/repo"}}}}"#, + "\n", + r#"{{"timestamp":"{timestamp}","type":"session_meta","payload":{{"id":"019e5b00-0000-7000-8000-000000000001","source":"vscode","model_provider":"openai","cwd":"/repo"}}}}"#, + "\n", + r#"{{"timestamp":"{timestamp}","type":"turn_context","payload":{{"turn_id":"019e5b00-0001-7000-8000-000000000001","model":"gpt-5.5","cwd":"/repo"}}}}"#, + "\n", + r#"{{"timestamp":"{timestamp}","type":"event_msg","payload":{{"type":"token_count","info":{{"total_token_usage":{{"input_tokens":100,"output_tokens":10,"total_tokens":110}},"last_token_usage":{{"input_tokens":100,"output_tokens":10,"total_tokens":110}}}}}}}}"#, + "\n", + r#"{{"timestamp":"{timestamp}","type":"event_msg","payload":{{"type":"token_count","info":{{"total_token_usage":{{"input_tokens":130,"output_tokens":13,"total_tokens":143}},"last_token_usage":{{"input_tokens":30,"output_tokens":3,"total_tokens":33}}}}}}}}"#, + "\n", + r#"{{"timestamp":"{timestamp}","type":"event_msg","payload":{{"type":"task_started","turn_id":"{child_turn_id}"}}}}"#, + "\n", + r#"{{"timestamp":"{timestamp}","type":"turn_context","payload":{{"turn_id":"{child_turn_id}","model":"gpt-5.5","cwd":"/repo"}}}}"#, + "\n", + r#"{{"timestamp":"{timestamp}","type":"event_msg","payload":{{"type":"token_count","info":{{"total_token_usage":{{"input_tokens":140,"output_tokens":14,"total_tokens":154}},"last_token_usage":{{"input_tokens":10,"output_tokens":1,"total_tokens":11}}}}}}}}"#, + "\n", + ), + timestamp = timestamp, + child_id = child_id, + child_turn_id = child_turn_id, + ), ) .unwrap(); + } + } - let first_messages = parse_all_messages_with_pricing( - source_home.path().to_str().unwrap(), - &["opencode".to_string()], - None, - ); - assert_eq!(first_messages.len(), 1); + fn write_codex_user_fork_replay_fixture(source_home: &std::path::Path) { + let sessions_dir = source_home.join(".codex/sessions/2026/01/02"); + let archived_dir = source_home.join(".codex/archived_sessions"); + std::fs::create_dir_all(&sessions_dir).unwrap(); + std::fs::create_dir_all(&archived_dir).unwrap(); - conn.execute( - "INSERT INTO message (id, session_id, data) VALUES (?1, ?2, ?3)", - rusqlite::params!["msg-2", "session-1", row_two], - ) - .unwrap(); - assert!(db_path.with_extension("db-wal").exists()); + std::fs::write( + archived_dir.join("rollout-2026-01-02T03-04-05-11111111-1111-7111-8111-111111111111.jsonl"), + concat!( + r#"{"timestamp":"2026-01-02T03:04:05Z","type":"session_meta","payload":{"id":"11111111-1111-7111-8111-111111111111","source":"vscode","thread_source":"user","model_provider":"openai","cwd":"/repo"}}"#, + "\n", + r#"{"timestamp":"2026-01-02T03:04:06Z","type":"turn_context","payload":{"turn_id":"11111111-3333-7333-8333-333333333333","model":"gpt-5.5","cwd":"/repo"}}"#, + "\n", + r#"{"timestamp":"2026-01-02T03:04:07Z","type":"event_msg","payload":{"type":"token_count","info":{"total_token_usage":{"input_tokens":1000,"cached_input_tokens":400,"output_tokens":100,"total_tokens":1100},"last_token_usage":{"input_tokens":1000,"cached_input_tokens":400,"output_tokens":100,"total_tokens":1100}}}}"#, + "\n", + r#"{"timestamp":"2026-01-02T03:04:08Z","type":"event_msg","payload":{"type":"token_count","info":{"total_token_usage":{"input_tokens":1200,"cached_input_tokens":450,"output_tokens":120,"total_tokens":1320},"last_token_usage":{"input_tokens":200,"cached_input_tokens":50,"output_tokens":20,"total_tokens":220}}}}"#, + "\n" + ), + ) + .unwrap(); - let refreshed_messages = parse_all_messages_with_pricing( - source_home.path().to_str().unwrap(), - &["opencode".to_string()], - None, - ); - assert_eq!(refreshed_messages.len(), 2); - } + std::fs::write( + sessions_dir.join("rollout-2026-01-02T03-10-00-22222222-2222-7222-8222-222222222222.jsonl"), + concat!( + r#"{"timestamp":"2026-01-02T03:10:00Z","type":"session_meta","payload":{"id":"22222222-2222-7222-8222-222222222222","forked_from_id":"11111111-1111-7111-8111-111111111111","source":"vscode","thread_source":"user","model_provider":"openai","cwd":"/repo"}}"#, + "\n", + r#"{"timestamp":"2026-01-02T03:10:00Z","type":"session_meta","payload":{"id":"11111111-1111-7111-8111-111111111111","source":"vscode","thread_source":"user","model_provider":"openai","cwd":"/repo"}}"#, + "\n", + r#"{"timestamp":"2026-01-02T03:10:00Z","type":"turn_context","payload":{"turn_id":"11111111-3333-7333-8333-333333333333","model":"gpt-5.5","cwd":"/repo"}}"#, + "\n", + r#"{"timestamp":"2026-01-02T03:10:00Z","type":"event_msg","payload":{"type":"token_count","info":{"total_token_usage":{"input_tokens":1000,"cached_input_tokens":400,"output_tokens":100,"total_tokens":1100},"last_token_usage":{"input_tokens":1000,"cached_input_tokens":400,"output_tokens":100,"total_tokens":1100}}}}"#, + "\n", + r#"{"timestamp":"2026-01-02T03:10:00Z","type":"event_msg","payload":{"type":"token_count","info":{"total_token_usage":{"input_tokens":1200,"cached_input_tokens":450,"output_tokens":120,"total_tokens":1320},"last_token_usage":{"input_tokens":200,"cached_input_tokens":50,"output_tokens":20,"total_tokens":220}}}}"#, + "\n", + r#"{"timestamp":"2026-01-02T03:10:30Z","type":"turn_context","payload":{"turn_id":"22222222-4444-7444-8444-444444444444","model":"gpt-5.5","cwd":"/repo"}}"#, + "\n", + r#"{"timestamp":"2026-01-02T03:10:30Z","type":"session_meta","payload":{"id":"22222222-2222-7222-8222-222222222222","forked_from_id":"11111111-1111-7111-8111-111111111111","source":"vscode","thread_source":"user","model_provider":"openai","cwd":"/repo"}}"#, + "\n", + r#"{"timestamp":"2026-01-02T03:10:53Z","type":"event_msg","payload":{"type":"token_count","info":{"total_token_usage":{"input_tokens":1500,"cached_input_tokens":500,"output_tokens":150,"total_tokens":1650},"last_token_usage":{"input_tokens":300,"cached_input_tokens":50,"output_tokens":30,"total_tokens":330}}}}"#, + "\n" + ), + ) + .unwrap(); } #[test] #[serial_test::serial] - fn test_parse_all_messages_dedups_across_channel_suffixed_opencode_dbs() { - // Regression guard: a session that appears in both `opencode.db` and - // `opencode-.db` (e.g. the user switches channels mid-session) - // must only be counted once. + fn test_parse_all_messages_with_pricing_codex_deduplicates_forked_history() { let cache_home = tempfile::TempDir::new().unwrap(); let source_home = tempfile::TempDir::new().unwrap(); - let _env = opencode_test_env(cache_home.path(), source_home.path()); + let _env = EnvGuard::set(&[ + ("HOME", cache_home.path().as_os_str()), + ("TOKSCALE_CONFIG_DIR", cache_home.path().as_os_str()), + ]); { - let db_dir = source_home.path().join(".local/share/opencode"); - std::fs::create_dir_all(&db_dir).unwrap(); - - let schema = "PRAGMA journal_mode=WAL; - PRAGMA wal_autocheckpoint=0; - CREATE TABLE message ( - id TEXT PRIMARY KEY, - session_id TEXT NOT NULL, - data TEXT NOT NULL - );"; - let row = |input: u64, ts: u64| { - format!( - r#"{{ - "role": "assistant", - "modelID": "claude-sonnet-4", - "providerID": "anthropic", - "tokens": {{ "input": {input}, "output": 10, "reasoning": 0, "cache": {{ "read": 0, "write": 0 }} }}, - "time": {{ "created": {ts}.0 }} - }}"# - ) - }; - - let default_db = db_dir.join("opencode.db"); - let conn = rusqlite::Connection::open(&default_db).unwrap(); - conn.execute_batch(schema).unwrap(); - conn.execute( - "INSERT INTO message (id, session_id, data) VALUES (?1, ?2, ?3)", - rusqlite::params![ - "shared-msg", - "session-shared", - row(100, 1_700_000_000_000u64) - ], - ) - .unwrap(); - conn.execute( - "INSERT INTO message (id, session_id, data) VALUES (?1, ?2, ?3)", - rusqlite::params![ - "latest-only", - "session-latest", - row(200, 1_700_000_001_000u64) - ], - ) - .unwrap(); - drop(conn); - - let stable_db = db_dir.join("opencode-stable.db"); - let conn = rusqlite::Connection::open(&stable_db).unwrap(); - conn.execute_batch(schema).unwrap(); - conn.execute( - "INSERT INTO message (id, session_id, data) VALUES (?1, ?2, ?3)", - rusqlite::params![ - "shared-msg", - "session-shared", - row(100, 1_700_000_000_000u64) - ], - ) - .unwrap(); - conn.execute( - "INSERT INTO message (id, session_id, data) VALUES (?1, ?2, ?3)", - rusqlite::params![ - "stable-only", - "session-stable", - row(300, 1_700_000_002_000u64) - ], - ) - .unwrap(); - drop(conn); + write_codex_forked_history_fixture(source_home.path()); let messages = parse_all_messages_with_pricing( source_home.path().to_str().unwrap(), - &["opencode".to_string()], + &["codex".to_string()], None, ); - assert_eq!( - messages.len(), - 3, - "expected 3 unique messages (shared + latest-only + stable-only), got {}", - messages.len() - ); - let mut ids: Vec = messages - .iter() - .filter_map(|m| m.dedup_key.clone()) - .collect(); - ids.sort(); - assert_eq!(ids, vec!["latest-only", "shared-msg", "stable-only"]); - let messages_warm = parse_all_messages_with_pricing( - source_home.path().to_str().unwrap(), - &["opencode".to_string()], - None, + assert_eq!(messages.len(), 3); + assert_eq!( + messages + .iter() + .map(|message| message.tokens.input) + .sum::(), + 88 ); assert_eq!( - messages_warm.len(), - 3, - "warm cache must also dedup shared message across channel dbs" + messages + .iter() + .map(|message| message.tokens.cache_read) + .sum::(), + 22 + ); + assert_eq!( + messages + .iter() + .map(|message| message.tokens.output) + .sum::(), + 33 ); } } #[test] #[serial_test::serial] - fn test_parse_all_messages_with_pricing_opencode_sqlite_deduplicates_forked_history() { + fn test_parse_all_messages_with_pricing_codex_keeps_user_fork_own_turn() { let cache_home = tempfile::TempDir::new().unwrap(); let source_home = tempfile::TempDir::new().unwrap(); - let _env = opencode_test_env(cache_home.path(), source_home.path()); + let _env = EnvGuard::set(&[ + ("HOME", cache_home.path().as_os_str()), + ("TOKSCALE_CONFIG_DIR", cache_home.path().as_os_str()), + ]); { - let db_dir = source_home.path().join(".local/share/opencode"); - std::fs::create_dir_all(&db_dir).unwrap(); - let db_path = db_dir.join("opencode.db"); - let conn = create_opencode_sqlite_db(&db_path); - - let msg_a = build_opencode_sqlite_payload( - 1_700_000_000_000.0, - 1_700_000_000_500.0, - 100, - 50, - 0, - 10, - 5, - 0.01, - ); - let msg_b = build_opencode_sqlite_payload( - 1_700_000_001_000.0, - 1_700_000_001_500.0, - 200, - 80, - 10, - 20, - 0, - 0.02, - ); - let msg_c = build_opencode_sqlite_payload( - 1_700_000_002_000.0, - 1_700_000_002_500.0, - 300, - 120, - 15, - 0, - 0, - 0.03, - ); - - for (id, session_id, payload) in [ - ("root_a", "root", msg_a.as_str()), - ("root_b", "root", msg_b.as_str()), - ("fork_a_copy", "fork", msg_a.as_str()), - ("fork_b_copy", "fork", msg_b.as_str()), - ("fork_c_new", "fork", msg_c.as_str()), - ] { - conn.execute( - "INSERT INTO message (id, session_id, data) VALUES (?1, ?2, ?3)", - rusqlite::params![id, session_id, payload], - ) - .unwrap(); - } - drop(conn); + write_codex_user_fork_replay_fixture(source_home.path()); let messages = parse_all_messages_with_pricing( source_home.path().to_str().unwrap(), - &["opencode".to_string()], + &["codex".to_string()], None, ); - assert_eq!(messages.len(), 3); - assert_eq!(messages.iter().map(|m| m.tokens.input).sum::(), 600); - assert_eq!(messages.iter().map(|m| m.tokens.output).sum::(), 250); - assert_eq!(messages.iter().map(|m| m.cost).sum::(), 0.06); + let session_ids: HashSet<_> = messages + .iter() + .map(|message| message.session_id.as_str()) + .collect(); + assert!(session_ids.contains( + "rollout-2026-01-02T03-10-00-22222222-2222-7222-8222-222222222222" + )); + assert_eq!(messages.iter().map(|m| m.tokens.input).sum::(), 1000); + assert_eq!(messages.iter().map(|m| m.tokens.cache_read).sum::(), 500); + assert_eq!(messages.iter().map(|m| m.tokens.output).sum::(), 150); } } #[test] #[serial_test::serial] - fn test_parse_local_clients_opencode_sqlite_counts_deduplicated_forked_history() { + fn test_parse_all_messages_with_pricing_codex_deduplicates_parent_replay_across_forks() { let cache_home = tempfile::TempDir::new().unwrap(); let source_home = tempfile::TempDir::new().unwrap(); - let _env = opencode_test_env(cache_home.path(), source_home.path()); + let _env = EnvGuard::set(&[ + ("HOME", cache_home.path().as_os_str()), + ("TOKSCALE_CONFIG_DIR", cache_home.path().as_os_str()), + ]); { - let db_dir = source_home.path().join(".local/share/opencode"); - std::fs::create_dir_all(&db_dir).unwrap(); - let db_path = db_dir.join("opencode.db"); - let conn = create_opencode_sqlite_db(&db_path); + write_codex_parent_replay_fixture(source_home.path()); - let msg_a = build_opencode_sqlite_payload( - 1_700_000_000_000.0, - 1_700_000_000_500.0, - 100, - 50, - 0, - 10, - 5, - 0.01, - ); - let msg_b = build_opencode_sqlite_payload( - 1_700_000_001_000.0, - 1_700_000_001_500.0, - 200, - 80, - 10, - 20, - 0, - 0.02, - ); - let msg_c = build_opencode_sqlite_payload( - 1_700_000_002_000.0, - 1_700_000_002_500.0, - 300, - 120, - 15, - 0, - 0, - 0.03, + let messages = parse_all_messages_with_pricing( + source_home.path().to_str().unwrap(), + &["codex".to_string()], + None, ); - for (id, session_id, payload) in [ - ("root_a", "root", msg_a.as_str()), - ("root_b", "root", msg_b.as_str()), - ("fork_a_copy", "fork", msg_a.as_str()), - ("fork_b_copy", "fork", msg_b.as_str()), - ("fork_c_new", "fork", msg_c.as_str()), - ] { - conn.execute( - "INSERT INTO message (id, session_id, data) VALUES (?1, ?2, ?3)", - rusqlite::params![id, session_id, payload], - ) - .unwrap(); - } - drop(conn); - - let parsed = parse_local_clients(LocalParseOptions { - home_dir: Some(source_home.path().to_str().unwrap().to_string()), - use_env_roots: false, - clients: Some(vec!["opencode".to_string()]), - since: None, - until: None, - year: None, - scanner_settings: scanner::ScannerSettings::default(), - modified_after: None, - }) - .unwrap(); - - assert_eq!(parsed.counts.get(ClientId::OpenCode), 3); - assert_eq!(parsed.messages.len(), 3); - assert_eq!(parsed.messages.iter().map(|m| m.input).sum::(), 600); - assert_eq!(parsed.messages.iter().map(|m| m.output).sum::(), 250); + // Parent contributes its two turns. The two forks each replay the + // parent history (skipped) and then emit one own turn that lands on + // the identical cumulative total (140/14). Sibling forks sharing a + // cumulative total is the signature of a replayed row, so the + // fork-parent-scoped dedup key collapses them into one. Real fork + // fan-out replays the same upstream totals into 10-100+ siblings; + // two distinct turns reaching a byte-identical cumulative vector by + // chance does not happen in practice because the cumulative encodes + // each fork's divergent context size. + assert_eq!(messages.len(), 3); + assert_eq!(messages.iter().map(|m| m.tokens.input).sum::(), 140); + assert_eq!(messages.iter().map(|m| m.tokens.output).sum::(), 14); } } - /// Regression fixture for Codex sessions that are live-only, archive-only, - /// or briefly present in both roots while the CLI moves a transcript. - fn write_codex_sessions_and_archived_sessions_fixture(source_home: &std::path::Path) { - let sessions_dir = source_home.join(".codex/sessions"); - let archived_dir = source_home.join(".codex/archived_sessions"); - std::fs::create_dir_all(&sessions_dir).unwrap(); - std::fs::create_dir_all(&archived_dir).unwrap(); - + fn write_codex_twin_token_count_fixture(source_home: &std::path::Path) { + // Single session with two turns whose `last_token_usage` deltas are + // byte-identical but emitted at different timestamps. The fork-dedup + // key includes the cumulative total, so both turns must survive even + // when a user happens to send two turns producing the same per-turn + // delta. + let codex_dir = source_home.join(".codex/sessions"); + std::fs::create_dir_all(&codex_dir).unwrap(); std::fs::write( - sessions_dir.join("live-only.jsonl"), + codex_dir.join("twin-deltas.jsonl"), concat!( - r#"{"timestamp":"2026-06-25T10:00:00Z","type":"session_meta","payload":{"id":"33333333-3333-7333-8333-333333333333","source":"interactive","model_provider":"openai","cwd":"/repo"}}"#, - "\n", - r#"{"timestamp":"2026-06-25T10:00:01Z","type":"turn_context","payload":{"model":"gpt-5.5"}}"#, + r#"{"timestamp":"2026-04-30T11:00:00Z","type":"session_meta","payload":{"id":"twin-session","source":"interactive","model_provider":"openai","cwd":"/Users/alice/root"}}"#, "\n", - r#"{"timestamp":"2026-06-25T10:00:02Z","type":"event_msg","payload":{"type":"token_count","info":{"total_token_usage":{"input_tokens":50,"output_tokens":5,"total_tokens":55},"last_token_usage":{"input_tokens":50,"output_tokens":5,"total_tokens":55}}}}"#, - "\n" - ), - ) - .unwrap(); - - std::fs::write( - archived_dir.join("archived-only.jsonl"), - concat!( - r#"{"timestamp":"2026-06-20T09:00:00Z","type":"session_meta","payload":{"id":"44444444-4444-7444-8444-444444444444","source":"interactive","model_provider":"openai","cwd":"/repo"}}"#, + r#"{"timestamp":"2026-04-30T11:00:01Z","type":"turn_context","payload":{"model":"gpt-5.2"}}"#, "\n", - r#"{"timestamp":"2026-06-20T09:00:01Z","type":"turn_context","payload":{"model":"gpt-5.5"}}"#, + r#"{"timestamp":"2026-04-30T11:00:02Z","type":"event_msg","payload":{"type":"token_count","info":{"total_token_usage":{"input_tokens":10,"cached_input_tokens":2,"output_tokens":3},"last_token_usage":{"input_tokens":10,"cached_input_tokens":2,"output_tokens":3}}}}"#, "\n", - r#"{"timestamp":"2026-06-20T09:00:02Z","type":"event_msg","payload":{"type":"token_count","info":{"total_token_usage":{"input_tokens":70,"output_tokens":7,"total_tokens":77},"last_token_usage":{"input_tokens":70,"output_tokens":7,"total_tokens":77}}}}"#, + r#"{"timestamp":"2026-04-30T11:00:03Z","type":"event_msg","payload":{"type":"token_count","info":{"total_token_usage":{"input_tokens":20,"cached_input_tokens":4,"output_tokens":6},"last_token_usage":{"input_tokens":10,"cached_input_tokens":2,"output_tokens":3}}}}"#, "\n" ), ) .unwrap(); + } - let shared_content = concat!( - r#"{"timestamp":"2026-06-22T08:00:00Z","type":"session_meta","payload":{"id":"55555555-5555-7555-8555-555555555555","source":"interactive","model_provider":"openai","cwd":"/repo"}}"#, - "\n", - r#"{"timestamp":"2026-06-22T08:00:01Z","type":"turn_context","payload":{"model":"gpt-5.5"}}"#, - "\n", - r#"{"timestamp":"2026-06-22T08:00:02Z","type":"event_msg","payload":{"type":"token_count","info":{"total_token_usage":{"input_tokens":30,"output_tokens":3,"total_tokens":33},"last_token_usage":{"input_tokens":30,"output_tokens":3,"total_tokens":33}}}}"#, - "\n" - ); - std::fs::write( - sessions_dir.join("shared-in-sessions.jsonl"), - shared_content, - ) - .unwrap(); - std::fs::write( - archived_dir.join("shared-in-archived.jsonl"), - shared_content, - ) - .unwrap(); + #[test] + #[serial_test::serial] + fn test_parse_all_messages_with_pricing_codex_keeps_twin_token_counts_at_distinct_timestamps() { + let cache_home = tempfile::TempDir::new().unwrap(); + let source_home = tempfile::TempDir::new().unwrap(); + let _env = EnvGuard::set(&[ + ("HOME", cache_home.path().as_os_str()), + ("TOKSCALE_CONFIG_DIR", cache_home.path().as_os_str()), + ]); + + { + write_codex_twin_token_count_fixture(source_home.path()); + + let messages = parse_all_messages_with_pricing( + source_home.path().to_str().unwrap(), + &["codex".to_string()], + None, + ); + + assert_eq!( + messages.len(), + 2, + "two turns with identical token deltas at distinct timestamps must both survive dedup", + ); + assert_eq!( + messages + .iter() + .map(|message| message.tokens.input) + .sum::(), + 16, + "input tokens normalize cache_read out of input: 2 turns × (10 - 2) = 16", + ); + assert_eq!( + messages + .iter() + .map(|message| message.tokens.output) + .sum::(), + 6, + ); + assert_eq!( + messages + .iter() + .map(|message| message.tokens.cache_read) + .sum::(), + 4, + ); + } } - fn with_isolated_tokscale_cache( - cache_home: &std::path::Path, - action: impl FnOnce() -> T, - ) -> T { + #[test] + #[serial_test::serial] + fn test_parse_local_clients_codex_counts_deduplicated_forked_history() { + let cache_home = tempfile::TempDir::new().unwrap(); + let source_home = tempfile::TempDir::new().unwrap(); let _env = EnvGuard::set(&[ - ("HOME", cache_home.as_os_str()), - ("TOKSCALE_CONFIG_DIR", cache_home.as_os_str()), - ("TOKSCALE_PRICING_CACHE_ONLY", std::ffi::OsStr::new("1")), + ("HOME", cache_home.path().as_os_str()), + ("TOKSCALE_CONFIG_DIR", cache_home.path().as_os_str()), ]); - action() - } - fn write_codex_duration_prefix_fixture( - source_home: &std::path::Path, - ) -> (std::path::PathBuf, String) { - let lines = include_str!("../tests/fixtures/codex_duration_timing.jsonl") - .lines() - .collect::>(); - let sessions_dir = source_home.join(".codex/sessions"); - std::fs::create_dir_all(&sessions_dir).unwrap(); - let source = sessions_dir.join("codex-duration-timing.jsonl"); - std::fs::write(&source, format!("{}\n", lines[..5].join("\n"))).unwrap(); - (source, format!("{}\n", lines[5..].join("\n"))) + { + write_codex_forked_history_fixture(source_home.path()); + + let parsed = parse_local_clients(LocalParseOptions { + home_dir: Some(source_home.path().to_str().unwrap().to_string()), + use_env_roots: false, + clients: Some(vec!["codex".to_string()]), + since: None, + until: None, + year: None, + scanner_settings: scanner::ScannerSettings::default(), + modified_after: None, + }) + .unwrap(); + + assert_eq!(parsed.counts.get(ClientId::Codex), 3); + assert_eq!(parsed.messages.len(), 3); + assert_eq!( + parsed + .messages + .iter() + .map(|message| message.input) + .sum::(), + 88 + ); + assert_eq!( + parsed + .messages + .iter() + .map(|message| message.cache_read) + .sum::(), + 22 + ); + assert_eq!( + parsed + .messages + .iter() + .map(|message| message.output) + .sum::(), + 33 + ); + } } #[test] #[serial_test::serial] - fn test_codex_non_overlapping_durations_survive_incremental_and_streaming_caches() { + fn test_codex_cache_reparses_from_zero_when_incremental_prefix_is_stale() { + let cache_home = tempfile::TempDir::new().unwrap(); + let fresh_cache_home = tempfile::TempDir::new().unwrap(); let source_home = tempfile::TempDir::new().unwrap(); - let materialized_cache = tempfile::TempDir::new().unwrap(); - let streaming_cache = tempfile::TempDir::new().unwrap(); - let (source, suffix) = write_codex_duration_prefix_fixture(source_home.path()); + let _env = EnvGuard::set(&[ + ("HOME", cache_home.path().as_os_str()), + ("TOKSCALE_CONFIG_DIR", cache_home.path().as_os_str()), + ]); - let home = source_home.path().to_str().unwrap().to_string(); - let clients = vec!["codex".to_string()]; - let parse_materialized = || { - parse_all_messages_with_pricing_with_env_strategy( - &home, - &clients, - None, - false, - &scanner::ScannerSettings::default(), + { + let codex_dir = scanner_fixture_path(source_home.path(), ".codex/sessions"); + std::fs::create_dir_all(&codex_dir).unwrap(); + let path = codex_dir.join("session.jsonl"); + std::fs::write( + &path, + concat!( + r#"{"type":"turn_context","payload":{"model":"gpt-5.4"}}"#, + "\n", + r#"{"type":"event_msg","payload":{"type":"token_count","info":{"total_token_usage":{"input_tokens":10,"cached_input_tokens":2,"output_tokens":3},"last_token_usage":{"input_tokens":10,"cached_input_tokens":2,"output_tokens":3}}}}"#, + "\n" + ), ) - }; + .unwrap(); - let prefix = with_isolated_tokscale_cache(materialized_cache.path(), parse_materialized); - assert_eq!(prefix.len(), 1); - assert_eq!(prefix[0].duration_ms, Some(1_000)); + let initial_messages = parse_all_messages_with_pricing( + source_home.path().to_str().unwrap(), + &["codex".to_string()], + None, + ); + assert_eq!(initial_messages.len(), 1); + assert_eq!(initial_messages[0].model_id, "gpt-5.4"); + assert!(message_cache::SourceMessageCache::load() + .get( + message_cache::CacheIdentity::for_client(ClientId::Codex), + &path + ) + .and_then(|entry| entry.codex_incremental.as_ref()) + .is_some()); - let mut source_file = std::fs::OpenOptions::new() - .append(true) - .open(&source) + std::thread::sleep(std::time::Duration::from_millis(5)); + std::fs::write( + &path, + concat!( + r#"{"type":"turn_context","payload":{"model":"gpt-5.5"}}"#, + "\n", + r#"{"type":"event_msg","payload":{"type":"token_count","info":{"total_token_usage":{"input_tokens":10,"cached_input_tokens":2,"output_tokens":3},"last_token_usage":{"input_tokens":10,"cached_input_tokens":2,"output_tokens":3}}}}"#, + "\n", + r#"{"type":"event_msg","payload":{"type":"token_count","info":{"total_token_usage":{"input_tokens":15,"cached_input_tokens":3,"output_tokens":5},"last_token_usage":{"input_tokens":5,"cached_input_tokens":1,"output_tokens":2}}}}"#, + "\n" + ), + ) .unwrap(); - source_file.write_all(suffix.as_bytes()).unwrap(); - source_file.flush().unwrap(); - drop(source_file); - let materialized_incremental = - with_isolated_tokscale_cache(materialized_cache.path(), parse_materialized); - let materialized_warm = - with_isolated_tokscale_cache(materialized_cache.path(), parse_materialized); - for (phase, messages) in [ - ("incremental", &materialized_incremental), - ("warm", &materialized_warm), - ] { - assert_eq!(messages.len(), 3, "{phase} materialized messages"); - assert_eq!( - messages - .iter() - .map(|message| message.duration_ms) - .collect::>(), - vec![Some(1_000), Some(4_000), Some(2_000)], - "{phase} materialized durations" + let warm_messages = parse_all_messages_with_pricing( + source_home.path().to_str().unwrap(), + &["codex".to_string()], + None, + ); + let _fresh_env = EnvGuard::set(&[ + ("HOME", fresh_cache_home.path().as_os_str()), + ("TOKSCALE_CONFIG_DIR", fresh_cache_home.path().as_os_str()), + ]); + let fresh_messages = parse_all_messages_with_pricing( + source_home.path().to_str().unwrap(), + &["codex".to_string()], + None, ); + + assert_eq!(warm_messages, fresh_messages); + assert_eq!(warm_messages.len(), 2); + assert!(warm_messages + .iter() + .all(|message| message.model_id == "gpt-5.5")); } + } - let runtime = tokio::runtime::Builder::new_current_thread() - .enable_all() - .build() + #[test] + #[serial_test::serial] + fn test_source_cache_keeps_untimestamped_rows_in_sync_after_append() { + let cache_home = tempfile::TempDir::new().unwrap(); + let fresh_cache_home = tempfile::TempDir::new().unwrap(); + let source_home = tempfile::TempDir::new().unwrap(); + let _env = EnvGuard::set(&[ + ("HOME", cache_home.path().as_os_str()), + ("TOKSCALE_CONFIG_DIR", cache_home.path().as_os_str()), + ]); + + { + let codex_dir = scanner_fixture_path(source_home.path(), ".codex/sessions"); + std::fs::create_dir_all(&codex_dir).unwrap(); + let path = codex_dir.join("session.jsonl"); + std::fs::write( + &path, + concat!( + r#"{"type":"turn_context","payload":{"model":"gpt-5.4"}}"#, + "\n", + r#"{"type":"event_msg","payload":{"type":"token_count","info":{"total_token_usage":{"input_tokens":10,"cached_input_tokens":2,"output_tokens":3},"last_token_usage":{"input_tokens":10,"cached_input_tokens":2,"output_tokens":3}}}}"#, + "\n" + ), + ) .unwrap(); - let report_options = || ReportOptions { - home_dir: Some(home.clone()), - use_env_roots: false, - clients: Some(clients.clone()), - ..Default::default() - }; - let streaming_cold = with_isolated_tokscale_cache(streaming_cache.path(), || { - runtime - .block_on(get_model_report(report_options())) - .unwrap() - }); - // The cold Codex streaming lane does not persist SourceMessageCache. - // Seed this isolated cache through the public materialized path so the - // second report exercises the shipping cache-hit branch. - let cache_seed = with_isolated_tokscale_cache(streaming_cache.path(), parse_materialized); - assert_eq!(cache_seed.len(), 3); - assert!( - streaming_cache - .path() - .join("cache/source-message-cache.bin") - .is_file(), - "materialized seed must persist the cache used by the warm report" - ); - let streaming_warm = with_isolated_tokscale_cache(streaming_cache.path(), || { - runtime - .block_on(get_model_report(report_options())) - .unwrap() - }); + let first_messages = parse_all_messages_with_pricing( + source_home.path().to_str().unwrap(), + &["codex".to_string()], + None, + ); + assert_eq!(first_messages.len(), 1); + + std::thread::sleep(std::time::Duration::from_millis(5)); + let mut file = std::fs::OpenOptions::new() + .append(true) + .open(&path) + .unwrap(); + file.write_all( + concat!( + r#"{"type":"event_msg","payload":{"type":"token_count","info":{"total_token_usage":{"input_tokens":15,"cached_input_tokens":3,"output_tokens":5},"last_token_usage":{"input_tokens":5,"cached_input_tokens":1,"output_tokens":2}}}}"#, + "\n" + ) + .as_bytes(), + ) + .unwrap(); + file.flush().unwrap(); - for (phase, report) in [("cold", &streaming_cold), ("warm", &streaming_warm)] { - assert_eq!(report.total_messages, 3, "{phase} streaming messages"); - assert_eq!(report.entries.len(), 1, "{phase} model groups"); - let performance = &report.entries[0].performance; - assert_eq!( - performance.total_duration_ms, 7_000, - "{phase} total duration" + let warm_messages = parse_all_messages_with_pricing( + source_home.path().to_str().unwrap(), + &["codex".to_string()], + None, ); - assert_eq!(performance.timed_tokens, 170, "{phase} timed tokens"); - assert_eq!(performance.sample_count, 3, "{phase} samples"); - assert_eq!(performance.token_coverage, 1.0, "{phase} coverage"); - let expected_ms_per_1k = 7_000.0 * 1_000.0 / 170.0; - assert!( - (performance.ms_per_1k_tokens.unwrap() - expected_ms_per_1k).abs() < f64::EPSILON, - "{phase} milliseconds per 1K tokens" + let _fresh_env = EnvGuard::set(&[ + ("HOME", fresh_cache_home.path().as_os_str()), + ("TOKSCALE_CONFIG_DIR", fresh_cache_home.path().as_os_str()), + ]); + let fresh_messages = parse_all_messages_with_pricing( + source_home.path().to_str().unwrap(), + &["codex".to_string()], + None, ); + + assert_eq!(warm_messages, fresh_messages); } } #[test] #[serial_test::serial] - fn test_codex_archive_roots_are_exact_once_across_all_consumers() { + fn test_source_cache_matches_cold_parse_after_malformed_json_append() { + let cache_home = tempfile::TempDir::new().unwrap(); + let fresh_cache_home = tempfile::TempDir::new().unwrap(); let source_home = tempfile::TempDir::new().unwrap(); - let materialized_cache = tempfile::TempDir::new().unwrap(); - let streaming_cache = tempfile::TempDir::new().unwrap(); - let count_cache = tempfile::TempDir::new().unwrap(); - - write_codex_sessions_and_archived_sessions_fixture(source_home.path()); + let _env = EnvGuard::set(&[ + ("HOME", cache_home.path().as_os_str()), + ("TOKSCALE_CONFIG_DIR", cache_home.path().as_os_str()), + ]); - let home = source_home.path().to_str().unwrap().to_string(); - let clients = vec!["codex".to_string()]; - let materialized = with_isolated_tokscale_cache(materialized_cache.path(), || { - parse_all_messages_with_pricing_with_env_strategy( - &home, - &clients, - None, - false, - &scanner::ScannerSettings::default(), + { + let codex_dir = scanner_fixture_path(source_home.path(), ".codex/sessions"); + std::fs::create_dir_all(&codex_dir).unwrap(); + let path = codex_dir.join("session.jsonl"); + std::fs::write( + &path, + concat!( + r#"{"type":"turn_context","payload":{"model":"gpt-5.4"}}"#, + "\n", + r#"{"type":"event_msg","payload":{"type":"token_count","info":{"total_token_usage":{"input_tokens":10,"cached_input_tokens":2,"output_tokens":3},"last_token_usage":{"input_tokens":10,"cached_input_tokens":2,"output_tokens":3}}}}"#, + "\n", + r#"{"type":"event_msg","payload":{"type":"token_count","info":{"total_token_usage":{"input_tokens":999""#, + "\n" + ), ) - }); + .unwrap(); - assert_eq!(materialized.len(), 3); - let session_ids: HashSet<_> = materialized - .iter() - .map(|message| message.session_id.as_str()) - .collect(); - assert!(session_ids.contains("live-only")); - assert!(session_ids.contains("archived-only")); - assert_eq!( - materialized - .iter() - .map(|message| message.tokens.input) - .sum::(), - 150, - ); - assert_eq!( - materialized - .iter() - .map(|message| message.tokens.output) - .sum::(), - 15, - ); + let initial_messages = parse_all_messages_with_pricing( + source_home.path().to_str().unwrap(), + &["codex".to_string()], + None, + ); + assert_eq!(initial_messages.len(), 1); - let runtime = tokio::runtime::Builder::new_current_thread() - .enable_all() - .build() + std::thread::sleep(std::time::Duration::from_millis(5)); + let mut file = std::fs::OpenOptions::new() + .append(true) + .open(&path) + .unwrap(); + file.write_all( + concat!( + r#"{"type":"event_msg","payload":{"type":"token_count","info":{"total_token_usage":{"input_tokens":15,"cached_input_tokens":3,"output_tokens":5},"last_token_usage":{"input_tokens":5,"cached_input_tokens":1,"output_tokens":2}}}}"#, + "\n" + ) + .as_bytes(), + ) .unwrap(); - let report_options = || ReportOptions { - home_dir: Some(home.clone()), - use_env_roots: false, - clients: Some(clients.clone()), - ..Default::default() - }; - let streaming_cold = with_isolated_tokscale_cache(streaming_cache.path(), || { - runtime - .block_on(get_model_report(report_options())) - .unwrap() - }); + file.flush().unwrap(); - // Codex's cold streaming lane does not populate SourceMessageCache, so - // seed the same isolated cache through the public materialized path - // before exercising the streaming cache-hit branch. - let cache_seed = with_isolated_tokscale_cache(streaming_cache.path(), || { - parse_all_messages_with_pricing_with_env_strategy( - &home, - &clients, + let warm_messages = parse_all_messages_with_pricing( + source_home.path().to_str().unwrap(), + &["codex".to_string()], None, - false, - &scanner::ScannerSettings::default(), - ) - }); - assert_eq!(cache_seed.len(), 3); - assert!( - streaming_cache - .path() - .join("cache/source-message-cache.bin") - .is_file(), - "materialized seed must persist the cache used by the warm pass", - ); - let streaming_warm = with_isolated_tokscale_cache(streaming_cache.path(), || { - runtime - .block_on(get_model_report(report_options())) - .unwrap() - }); - for (phase, streaming) in [("cold", &streaming_cold), ("warm", &streaming_warm)] { - assert_eq!(streaming.total_messages, 3, "{phase} streaming messages"); - assert_eq!(streaming.total_input, 150, "{phase} streaming input"); - assert_eq!(streaming.total_output, 15, "{phase} streaming output"); - } + ); + assert!(message_cache::SourceMessageCache::load() + .get( + message_cache::CacheIdentity::for_client(ClientId::Codex), + &path + ) + .is_none()); - let counted = with_isolated_tokscale_cache(count_cache.path(), || { - parse_local_clients(LocalParseOptions { - home_dir: Some(home), - use_env_roots: false, - clients: Some(clients), - since: None, - until: None, - year: None, - scanner_settings: scanner::ScannerSettings::default(), - modified_after: None, - }) - .unwrap() - }); - assert_eq!(counted.counts.get(ClientId::Codex), 3); - assert_eq!(counted.messages.len(), 3); - assert_eq!( - counted - .messages - .iter() - .map(|message| message.input) - .sum::(), - 150, - ); - assert_eq!( - counted - .messages - .iter() - .map(|message| message.output) - .sum::(), - 15, - ); - } + let _fresh_env = EnvGuard::set(&[ + ("HOME", fresh_cache_home.path().as_os_str()), + ("TOKSCALE_CONFIG_DIR", fresh_cache_home.path().as_os_str()), + ]); + let fresh_messages = parse_all_messages_with_pricing( + source_home.path().to_str().unwrap(), + &["codex".to_string()], + None, + ); - fn write_codex_forked_history_fixture(source_home: &std::path::Path) { - let codex_dir = source_home.join(".codex/sessions"); - std::fs::create_dir_all(&codex_dir).unwrap(); - std::fs::write( - codex_dir.join("parent.jsonl"), - concat!( - r#"{"timestamp":"2026-04-30T10:00:00Z","type":"session_meta","payload":{"id":"parent-session","source":"interactive","model_provider":"openai","cwd":"/Users/alice/root"}}"#, - "\n", - r#"{"timestamp":"2026-04-30T10:00:01Z","type":"turn_context","payload":{"model":"gpt-5.2"}}"#, - "\n", - r#"{"timestamp":"2026-04-30T10:00:02Z","type":"event_msg","payload":{"type":"token_count","info":{"total_token_usage":{"input_tokens":50,"cached_input_tokens":10,"output_tokens":15,"total_tokens":65},"last_token_usage":{"input_tokens":50,"cached_input_tokens":10,"output_tokens":15,"total_tokens":65}}}}"#, - "\n", - r#"{"timestamp":"2026-04-30T10:00:03Z","type":"event_msg","payload":{"type":"token_count","info":{"total_token_usage":{"input_tokens":100,"cached_input_tokens":20,"output_tokens":30,"total_tokens":130},"last_token_usage":{"input_tokens":50,"cached_input_tokens":10,"output_tokens":15,"total_tokens":65}}}}"#, - "\n" - ), - ) - .unwrap(); - std::fs::write( - codex_dir.join("fork.jsonl"), - concat!( - r#"{"timestamp":"2026-04-30T10:01:00Z","type":"session_meta","payload":{"id":"fork-session","source":{"subagent":{"thread_spawn":{"parent_thread_id":"parent-session","depth":1}}},"model_provider":"openai","cwd":"/Users/alice/root-worktree"}}"#, - "\n", - r#"{"timestamp":"2026-04-30T10:01:01Z","type":"event_msg","payload":{"type":"token_count","info":{"total_token_usage":{"input_tokens":100,"cached_input_tokens":20,"output_tokens":30,"total_tokens":130},"last_token_usage":{"input_tokens":100,"cached_input_tokens":20,"output_tokens":30,"total_tokens":130}}}}"#, - "\n", - r#"{"timestamp":"2026-04-30T10:01:02Z","type":"turn_context","payload":{"model":"gpt-5.2"}}"#, - "\n", - r#"{"timestamp":"2026-04-30T10:01:03Z","type":"event_msg","payload":{"type":"token_count","info":{"total_token_usage":{"input_tokens":50,"cached_input_tokens":10,"output_tokens":15,"total_tokens":65},"last_token_usage":{"input_tokens":50,"cached_input_tokens":10,"output_tokens":15,"total_tokens":65}}}}"#, - "\n", - r#"{"timestamp":"2026-04-30T10:01:04Z","type":"event_msg","payload":{"type":"token_count","info":{"total_token_usage":{"input_tokens":100,"cached_input_tokens":20,"output_tokens":30,"total_tokens":130},"last_token_usage":{"input_tokens":50,"cached_input_tokens":10,"output_tokens":15,"total_tokens":65}}}}"#, - "\n", - r#"{"timestamp":"2026-04-30T10:01:05Z","type":"event_msg","payload":{"type":"token_count","info":{"total_token_usage":{"input_tokens":110,"cached_input_tokens":22,"output_tokens":33,"total_tokens":143},"last_token_usage":{"input_tokens":10,"cached_input_tokens":2,"output_tokens":3,"total_tokens":13}}}}"#, - "\n" - ), - ) - .unwrap(); + assert_eq!(warm_messages, fresh_messages); + } } - fn write_codex_parent_replay_fixture(source_home: &std::path::Path) { - let codex_dir = source_home.join(".codex/sessions"); - std::fs::create_dir_all(&codex_dir).unwrap(); - std::fs::write( - codex_dir.join("parent.jsonl"), - concat!( - r#"{"timestamp":"2026-05-24T20:00:00Z","type":"session_meta","payload":{"id":"019e5b00-0000-7000-8000-000000000001","source":"vscode","model_provider":"openai","cwd":"/repo"}}"#, - "\n", - r#"{"timestamp":"2026-05-24T20:00:01Z","type":"turn_context","payload":{"turn_id":"019e5b00-0001-7000-8000-000000000001","model":"gpt-5.5","cwd":"/repo"}}"#, - "\n", - r#"{"timestamp":"2026-05-24T20:00:02Z","type":"event_msg","payload":{"type":"token_count","info":{"total_token_usage":{"input_tokens":100,"output_tokens":10,"total_tokens":110},"last_token_usage":{"input_tokens":100,"output_tokens":10,"total_tokens":110}}}}"#, - "\n", - r#"{"timestamp":"2026-05-24T20:00:03Z","type":"event_msg","payload":{"type":"token_count","info":{"total_token_usage":{"input_tokens":130,"output_tokens":13,"total_tokens":143},"last_token_usage":{"input_tokens":30,"output_tokens":3,"total_tokens":33}}}}"#, - "\n" - ), - ) - .unwrap(); + #[test] + #[serial_test::serial] + fn test_exact_hit_codex_cache_repairs_fallback_timestamps_without_incremental_state() { + let cache_home = tempfile::TempDir::new().unwrap(); + let source_home = tempfile::TempDir::new().unwrap(); + let _env = EnvGuard::set(&[ + ("HOME", cache_home.path().as_os_str()), + ("TOKSCALE_CONFIG_DIR", cache_home.path().as_os_str()), + ]); - for (filename, child_id, child_turn_id, timestamp) in [ - ( - "child-a.jsonl", - "019e5c03-1e99-7000-8000-000000000001", - "019e5c03-6425-7000-8000-000000000001", - "2026-05-24T21:00:00Z", - ), - ( - "child-b.jsonl", - "019e5c04-1e99-7000-8000-000000000001", - "019e5c04-6425-7000-8000-000000000001", - "2026-05-24T22:00:00Z", - ), - ] { + { + let session_dir = scanner_fixture_path(source_home.path(), ".codex/sessions"); + std::fs::create_dir_all(&session_dir).unwrap(); + let path = session_dir.join("session.jsonl"); std::fs::write( - codex_dir.join(filename), - format!( - concat!( - r#"{{"timestamp":"{timestamp}","type":"session_meta","payload":{{"id":"{child_id}","forked_from_id":"019e5b00-0000-7000-8000-000000000001","source":{{"subagent":{{"thread_spawn":{{"parent_thread_id":"019e5b00-0000-7000-8000-000000000001","depth":1}}}}}},"model_provider":"openai","agent_nickname":"worker","cwd":"/repo"}}}}"#, - "\n", - r#"{{"timestamp":"{timestamp}","type":"session_meta","payload":{{"id":"019e5b00-0000-7000-8000-000000000001","source":"vscode","model_provider":"openai","cwd":"/repo"}}}}"#, - "\n", - r#"{{"timestamp":"{timestamp}","type":"turn_context","payload":{{"turn_id":"019e5b00-0001-7000-8000-000000000001","model":"gpt-5.5","cwd":"/repo"}}}}"#, - "\n", - r#"{{"timestamp":"{timestamp}","type":"event_msg","payload":{{"type":"token_count","info":{{"total_token_usage":{{"input_tokens":100,"output_tokens":10,"total_tokens":110}},"last_token_usage":{{"input_tokens":100,"output_tokens":10,"total_tokens":110}}}}}}}}"#, - "\n", - r#"{{"timestamp":"{timestamp}","type":"event_msg","payload":{{"type":"token_count","info":{{"total_token_usage":{{"input_tokens":130,"output_tokens":13,"total_tokens":143}},"last_token_usage":{{"input_tokens":30,"output_tokens":3,"total_tokens":33}}}}}}}}"#, - "\n", - r#"{{"timestamp":"{timestamp}","type":"event_msg","payload":{{"type":"task_started","turn_id":"{child_turn_id}"}}}}"#, - "\n", - r#"{{"timestamp":"{timestamp}","type":"turn_context","payload":{{"turn_id":"{child_turn_id}","model":"gpt-5.5","cwd":"/repo"}}}}"#, - "\n", - r#"{{"timestamp":"{timestamp}","type":"event_msg","payload":{{"type":"token_count","info":{{"total_token_usage":{{"input_tokens":140,"output_tokens":14,"total_tokens":154}},"last_token_usage":{{"input_tokens":10,"output_tokens":1,"total_tokens":11}}}}}}}}"#, - "\n", - ), - timestamp = timestamp, - child_id = child_id, - child_turn_id = child_turn_id, + &path, + concat!( + r#"{"type":"turn_context","payload":{"model":"gpt-5.4"}}"#, + "\n", + r#"{"type":"event_msg","payload":{"type":"token_count","info":{"total_token_usage":{"input_tokens":10,"cached_input_tokens":2,"output_tokens":3},"last_token_usage":{"input_tokens":10,"cached_input_tokens":2,"output_tokens":3}}}}"#, + "\n" ), ) .unwrap(); - } - } - fn write_codex_user_fork_replay_fixture(source_home: &std::path::Path) { - let sessions_dir = source_home.join(".codex/sessions/2026/01/02"); - let archived_dir = source_home.join(".codex/archived_sessions"); - std::fs::create_dir_all(&sessions_dir).unwrap(); - std::fs::create_dir_all(&archived_dir).unwrap(); + let expected = crate::sessions::codex::parse_codex_file(&path); + assert_eq!(expected.len(), 1); - std::fs::write( - archived_dir.join("rollout-2026-01-02T03-04-05-11111111-1111-7111-8111-111111111111.jsonl"), - concat!( - r#"{"timestamp":"2026-01-02T03:04:05Z","type":"session_meta","payload":{"id":"11111111-1111-7111-8111-111111111111","source":"vscode","thread_source":"user","model_provider":"openai","cwd":"/repo"}}"#, - "\n", - r#"{"timestamp":"2026-01-02T03:04:06Z","type":"turn_context","payload":{"turn_id":"11111111-3333-7333-8333-333333333333","model":"gpt-5.5","cwd":"/repo"}}"#, - "\n", - r#"{"timestamp":"2026-01-02T03:04:07Z","type":"event_msg","payload":{"type":"token_count","info":{"total_token_usage":{"input_tokens":1000,"cached_input_tokens":400,"output_tokens":100,"total_tokens":1100},"last_token_usage":{"input_tokens":1000,"cached_input_tokens":400,"output_tokens":100,"total_tokens":1100}}}}"#, - "\n", - r#"{"timestamp":"2026-01-02T03:04:08Z","type":"event_msg","payload":{"type":"token_count","info":{"total_token_usage":{"input_tokens":1200,"cached_input_tokens":450,"output_tokens":120,"total_tokens":1320},"last_token_usage":{"input_tokens":200,"cached_input_tokens":50,"output_tokens":20,"total_tokens":220}}}}"#, - "\n" - ), - ) - .unwrap(); + let fingerprint = message_cache::SourceFingerprint::from_path(&path).unwrap(); + let mut stale_message = expected[0].clone(); + stale_message.timestamp = 0; + stale_message.date = "1900-01-01".to_string(); - std::fs::write( - sessions_dir.join("rollout-2026-01-02T03-10-00-22222222-2222-7222-8222-222222222222.jsonl"), - concat!( - r#"{"timestamp":"2026-01-02T03:10:00Z","type":"session_meta","payload":{"id":"22222222-2222-7222-8222-222222222222","forked_from_id":"11111111-1111-7111-8111-111111111111","source":"vscode","thread_source":"user","model_provider":"openai","cwd":"/repo"}}"#, - "\n", - r#"{"timestamp":"2026-01-02T03:10:00Z","type":"session_meta","payload":{"id":"11111111-1111-7111-8111-111111111111","source":"vscode","thread_source":"user","model_provider":"openai","cwd":"/repo"}}"#, - "\n", - r#"{"timestamp":"2026-01-02T03:10:00Z","type":"turn_context","payload":{"turn_id":"11111111-3333-7333-8333-333333333333","model":"gpt-5.5","cwd":"/repo"}}"#, - "\n", - r#"{"timestamp":"2026-01-02T03:10:00Z","type":"event_msg","payload":{"type":"token_count","info":{"total_token_usage":{"input_tokens":1000,"cached_input_tokens":400,"output_tokens":100,"total_tokens":1100},"last_token_usage":{"input_tokens":1000,"cached_input_tokens":400,"output_tokens":100,"total_tokens":1100}}}}"#, - "\n", - r#"{"timestamp":"2026-01-02T03:10:00Z","type":"event_msg","payload":{"type":"token_count","info":{"total_token_usage":{"input_tokens":1200,"cached_input_tokens":450,"output_tokens":120,"total_tokens":1320},"last_token_usage":{"input_tokens":200,"cached_input_tokens":50,"output_tokens":20,"total_tokens":220}}}}"#, - "\n", - r#"{"timestamp":"2026-01-02T03:10:30Z","type":"turn_context","payload":{"turn_id":"22222222-4444-7444-8444-444444444444","model":"gpt-5.5","cwd":"/repo"}}"#, - "\n", - r#"{"timestamp":"2026-01-02T03:10:30Z","type":"session_meta","payload":{"id":"22222222-2222-7222-8222-222222222222","forked_from_id":"11111111-1111-7111-8111-111111111111","source":"vscode","thread_source":"user","model_provider":"openai","cwd":"/repo"}}"#, - "\n", - r#"{"timestamp":"2026-01-02T03:10:53Z","type":"event_msg","payload":{"type":"token_count","info":{"total_token_usage":{"input_tokens":1500,"cached_input_tokens":500,"output_tokens":150,"total_tokens":1650},"last_token_usage":{"input_tokens":300,"cached_input_tokens":50,"output_tokens":30,"total_tokens":330}}}}"#, - "\n" - ), - ) - .unwrap(); + let mut cache = message_cache::SourceMessageCache::default(); + cache.insert(message_cache::CachedSourceEntry::new( + message_cache::CacheIdentity::for_client(ClientId::Codex), + &path, + fingerprint, + vec![stale_message], + vec![0], + None, + )); + cache.save_if_dirty(); + + let messages = parse_all_messages_with_pricing( + source_home.path().to_str().unwrap(), + &["codex".to_string()], + None, + ); + + assert_eq!(messages, expected); + } } #[test] #[serial_test::serial] - fn test_parse_all_messages_with_pricing_codex_deduplicates_forked_history() { + fn test_codex_cache_repairs_fallback_timestamps_after_source_mtime_change() { let cache_home = tempfile::TempDir::new().unwrap(); + let fresh_cache_home = tempfile::TempDir::new().unwrap(); let source_home = tempfile::TempDir::new().unwrap(); let _env = EnvGuard::set(&[ ("HOME", cache_home.path().as_os_str()), @@ -6886,42 +10753,51 @@ mod tests { ]); { - write_codex_forked_history_fixture(source_home.path()); + let session_dir = scanner_fixture_path(source_home.path(), ".codex/sessions"); + std::fs::create_dir_all(&session_dir).unwrap(); + let path = session_dir.join("session.jsonl"); + let contents = concat!( + r#"{"type":"turn_context","payload":{"model":"gpt-5.4"}}"#, + "\n", + r#"{"type":"event_msg","payload":{"type":"token_count","info":{"total_token_usage":{"input_tokens":10,"cached_input_tokens":2,"output_tokens":3},"last_token_usage":{"input_tokens":10,"cached_input_tokens":2,"output_tokens":3}}}}"#, + "\n" + ); + std::fs::write(&path, contents).unwrap(); - let messages = parse_all_messages_with_pricing( + let initial_messages = parse_all_messages_with_pricing( source_home.path().to_str().unwrap(), &["codex".to_string()], None, ); + assert_eq!(initial_messages.len(), 1); - assert_eq!(messages.len(), 3); - assert_eq!( - messages - .iter() - .map(|message| message.tokens.input) - .sum::(), - 88 - ); - assert_eq!( - messages - .iter() - .map(|message| message.tokens.cache_read) - .sum::(), - 22 + std::thread::sleep(std::time::Duration::from_millis(20)); + std::fs::write(&path, contents).unwrap(); + + let warm_messages = parse_all_messages_with_pricing( + source_home.path().to_str().unwrap(), + &["codex".to_string()], + None, ); - assert_eq!( - messages - .iter() - .map(|message| message.tokens.output) - .sum::(), - 33 + + let _fresh_env = EnvGuard::set(&[ + ("HOME", fresh_cache_home.path().as_os_str()), + ("TOKSCALE_CONFIG_DIR", fresh_cache_home.path().as_os_str()), + ]); + let fresh_messages = parse_all_messages_with_pricing( + source_home.path().to_str().unwrap(), + &["codex".to_string()], + None, ); + + assert_eq!(warm_messages, fresh_messages); + assert_ne!(warm_messages[0].timestamp, initial_messages[0].timestamp); } } #[test] #[serial_test::serial] - fn test_parse_all_messages_with_pricing_codex_keeps_user_fork_own_turn() { + fn test_full_log_parse_preserves_valid_messages_before_invalid_line_error() { let cache_home = tempfile::TempDir::new().unwrap(); let source_home = tempfile::TempDir::new().unwrap(); let _env = EnvGuard::set(&[ @@ -6930,31 +10806,47 @@ mod tests { ]); { - write_codex_user_fork_replay_fixture(source_home.path()); + let session_dir = scanner_fixture_path(source_home.path(), ".codex/sessions"); + std::fs::create_dir_all(&session_dir).unwrap(); + let path = session_dir.join("session.jsonl"); + + let mut file = std::fs::File::create(&path).unwrap(); + file.write_all( + concat!( + r#"{"type":"turn_context","payload":{"model":"gpt-5.4"}}"#, + "\n", + r#"{"type":"event_msg","payload":{"type":"token_count","info":{"total_token_usage":{"input_tokens":10,"cached_input_tokens":2,"output_tokens":3},"last_token_usage":{"input_tokens":10,"cached_input_tokens":2,"output_tokens":3}}}}"#, + "\n" + ) + .as_bytes(), + ) + .unwrap(); + file.write_all(&[0xff, b'\n']).unwrap(); + file.flush().unwrap(); let messages = parse_all_messages_with_pricing( source_home.path().to_str().unwrap(), &["codex".to_string()], None, ); + assert_eq!(messages.len(), 1); + assert_eq!(messages[0].model_id, "gpt-5.4"); - let session_ids: HashSet<_> = messages - .iter() - .map(|message| message.session_id.as_str()) - .collect(); - assert!(session_ids.contains( - "rollout-2026-01-02T03-10-00-22222222-2222-7222-8222-222222222222" - )); - assert_eq!(messages.iter().map(|m| m.tokens.input).sum::(), 1000); - assert_eq!(messages.iter().map(|m| m.tokens.cache_read).sum::(), 500); - assert_eq!(messages.iter().map(|m| m.tokens.output).sum::(), 150); + let cache = message_cache::SourceMessageCache::load(); + assert!(cache + .get( + message_cache::CacheIdentity::for_client(ClientId::Codex), + &path + ) + .is_none()); } } #[test] #[serial_test::serial] - fn test_parse_all_messages_with_pricing_codex_deduplicates_parent_replay_across_forks() { + fn test_codex_cache_does_not_persist_unknown_before_later_turn_context() { let cache_home = tempfile::TempDir::new().unwrap(); + let fresh_cache_home = tempfile::TempDir::new().unwrap(); let source_home = tempfile::TempDir::new().unwrap(); let _env = EnvGuard::set(&[ ("HOME", cache_home.path().as_os_str()), @@ -6962,57 +10854,84 @@ mod tests { ]); { - write_codex_parent_replay_fixture(source_home.path()); + let session_dir = scanner_fixture_path(source_home.path(), ".codex/sessions"); + std::fs::create_dir_all(&session_dir).unwrap(); + let path = session_dir.join("session.jsonl"); + std::fs::write( + &path, + concat!( + r#"{"type":"session_meta","payload":{"source":"interactive","model_provider":"openai"}}"#, + "\n", + r#"{"timestamp":"2026-04-27T10:00:00Z","type":"event_msg","payload":{"type":"token_count","info":{"total_token_usage":{"input_tokens":10,"cached_input_tokens":2,"output_tokens":3},"last_token_usage":{"input_tokens":10,"cached_input_tokens":2,"output_tokens":3}}}}"#, + "\n" + ), + ) + .unwrap(); - let messages = parse_all_messages_with_pricing( + let initial_messages = parse_all_messages_with_pricing( + source_home.path().to_str().unwrap(), + &["codex".to_string()], + None, + ); + assert_eq!(initial_messages.len(), 1); + assert_eq!(initial_messages[0].model_id, "unknown"); + assert!(message_cache::SourceMessageCache::load() + .get( + message_cache::CacheIdentity::for_client(ClientId::Codex), + &path + ) + .is_none()); + + std::thread::sleep(std::time::Duration::from_millis(5)); + let mut file = std::fs::OpenOptions::new() + .append(true) + .open(&path) + .unwrap(); + file.write_all( + concat!( + r#"{"timestamp":"2026-04-27T10:00:04Z","type":"turn_context","payload":{"model":"gpt-5.5"}}"#, + "\n" + ) + .as_bytes(), + ) + .unwrap(); + file.flush().unwrap(); + + let resumed_messages = parse_all_messages_with_pricing( + source_home.path().to_str().unwrap(), + &["codex".to_string()], + None, + ); + + let _fresh_env = EnvGuard::set(&[ + ("HOME", fresh_cache_home.path().as_os_str()), + ("TOKSCALE_CONFIG_DIR", fresh_cache_home.path().as_os_str()), + ]); + let fresh_messages = parse_all_messages_with_pricing( source_home.path().to_str().unwrap(), &["codex".to_string()], None, ); - // Parent contributes its two turns. The two forks each replay the - // parent history (skipped) and then emit one own turn that lands on - // the identical cumulative total (140/14). Sibling forks sharing a - // cumulative total is the signature of a replayed row, so the - // fork-parent-scoped dedup key collapses them into one. Real fork - // fan-out replays the same upstream totals into 10-100+ siblings; - // two distinct turns reaching a byte-identical cumulative vector by - // chance does not happen in practice because the cumulative encodes - // each fork's divergent context size. - assert_eq!(messages.len(), 3); - assert_eq!(messages.iter().map(|m| m.tokens.input).sum::(), 140); - assert_eq!(messages.iter().map(|m| m.tokens.output).sum::(), 14); - } - } + assert_eq!(resumed_messages, fresh_messages); + assert_eq!(resumed_messages.len(), 1); + assert_eq!(resumed_messages[0].model_id, "gpt-5.5"); - fn write_codex_twin_token_count_fixture(source_home: &std::path::Path) { - // Single session with two turns whose `last_token_usage` deltas are - // byte-identical but emitted at different timestamps. The fork-dedup - // key includes the cumulative total, so both turns must survive even - // when a user happens to send two turns producing the same per-turn - // delta. - let codex_dir = source_home.join(".codex/sessions"); - std::fs::create_dir_all(&codex_dir).unwrap(); - std::fs::write( - codex_dir.join("twin-deltas.jsonl"), - concat!( - r#"{"timestamp":"2026-04-30T11:00:00Z","type":"session_meta","payload":{"id":"twin-session","source":"interactive","model_provider":"openai","cwd":"/Users/alice/root"}}"#, - "\n", - r#"{"timestamp":"2026-04-30T11:00:01Z","type":"turn_context","payload":{"model":"gpt-5.2"}}"#, - "\n", - r#"{"timestamp":"2026-04-30T11:00:02Z","type":"event_msg","payload":{"type":"token_count","info":{"total_token_usage":{"input_tokens":10,"cached_input_tokens":2,"output_tokens":3},"last_token_usage":{"input_tokens":10,"cached_input_tokens":2,"output_tokens":3}}}}"#, - "\n", - r#"{"timestamp":"2026-04-30T11:00:03Z","type":"event_msg","payload":{"type":"token_count","info":{"total_token_usage":{"input_tokens":20,"cached_input_tokens":4,"output_tokens":6},"last_token_usage":{"input_tokens":10,"cached_input_tokens":2,"output_tokens":3}}}}"#, - "\n" - ), - ) - .unwrap(); + drop(_fresh_env); + assert!(message_cache::SourceMessageCache::load() + .get( + message_cache::CacheIdentity::for_client(ClientId::Codex), + &path + ) + .is_some()); + } } #[test] #[serial_test::serial] - fn test_parse_all_messages_with_pricing_codex_keeps_twin_token_counts_at_distinct_timestamps() { + fn test_codex_cache_skips_non_newline_terminated_resume_prefix() { let cache_home = tempfile::TempDir::new().unwrap(); + let fresh_cache_home = tempfile::TempDir::new().unwrap(); let source_home = tempfile::TempDir::new().unwrap(); let _env = EnvGuard::set(&[ ("HOME", cache_home.path().as_os_str()), @@ -7020,669 +10939,817 @@ mod tests { ]); { - write_codex_twin_token_count_fixture(source_home.path()); + let session_dir = scanner_fixture_path(source_home.path(), ".codex/sessions"); + std::fs::create_dir_all(&session_dir).unwrap(); + let path = session_dir.join("session.jsonl"); + std::fs::write( + &path, + concat!( + r#"{"type":"turn_context","payload":{"model":"gpt-5.4"}}"#, + "\n", + r#"{"type":"event_msg","payload":{"type":"token_count","info":{"total_token_usage":{"input_tokens":10,"cached_input_tokens":2,"output_tokens":3},"last_token_usage":{"input_tokens":10,"cached_input_tokens":2,"output_tokens":3}}}}"# + ), + ) + .unwrap(); - let messages = parse_all_messages_with_pricing( + let initial_messages = parse_all_messages_with_pricing( source_home.path().to_str().unwrap(), &["codex".to_string()], None, ); + assert_eq!(initial_messages.len(), 1); + assert!(message_cache::SourceMessageCache::load() + .get( + message_cache::CacheIdentity::for_client(ClientId::Codex), + &path + ) + .is_none()); - assert_eq!( - messages.len(), - 2, - "two turns with identical token deltas at distinct timestamps must both survive dedup", - ); - assert_eq!( - messages - .iter() - .map(|message| message.tokens.input) - .sum::(), - 16, - "input tokens normalize cache_read out of input: 2 turns × (10 - 2) = 16", - ); - assert_eq!( - messages - .iter() - .map(|message| message.tokens.output) - .sum::(), - 6, + std::thread::sleep(std::time::Duration::from_millis(5)); + let mut file = std::fs::OpenOptions::new() + .append(true) + .open(&path) + .unwrap(); + file.write_all( + concat!( + "\n", + r#"{"type":"event_msg","payload":{"type":"token_count","info":{"total_token_usage":{"input_tokens":15,"cached_input_tokens":3,"output_tokens":5},"last_token_usage":{"input_tokens":5,"cached_input_tokens":1,"output_tokens":2}}}}"#, + "\n" + ) + .as_bytes(), + ) + .unwrap(); + file.flush().unwrap(); + + let warm_messages = parse_all_messages_with_pricing( + source_home.path().to_str().unwrap(), + &["codex".to_string()], + None, ); - assert_eq!( - messages - .iter() - .map(|message| message.tokens.cache_read) - .sum::(), - 4, + + let _fresh_env = EnvGuard::set(&[ + ("HOME", fresh_cache_home.path().as_os_str()), + ("TOKSCALE_CONFIG_DIR", fresh_cache_home.path().as_os_str()), + ]); + let fresh_messages = parse_all_messages_with_pricing( + source_home.path().to_str().unwrap(), + &["codex".to_string()], + None, ); + + assert_eq!(warm_messages, fresh_messages); + assert_eq!(warm_messages.len(), 2); } } #[test] #[serial_test::serial] - fn test_parse_local_clients_codex_counts_deduplicated_forked_history() { - let cache_home = tempfile::TempDir::new().unwrap(); + fn test_source_cache_does_not_reuse_priced_cost_without_pricing_service() { + let temp_home = tempfile::TempDir::new().unwrap(); let source_home = tempfile::TempDir::new().unwrap(); let _env = EnvGuard::set(&[ - ("HOME", cache_home.path().as_os_str()), - ("TOKSCALE_CONFIG_DIR", cache_home.path().as_os_str()), + ("HOME", temp_home.path().as_os_str()), + ("TOKSCALE_CONFIG_DIR", temp_home.path().as_os_str()), ]); - { - write_codex_forked_history_fixture(source_home.path()); + let cursor_cache_dir = source_home.path().join(".config/tokscale/cursor-cache"); + std::fs::create_dir_all(&cursor_cache_dir).unwrap(); - let parsed = parse_local_clients(LocalParseOptions { - home_dir: Some(source_home.path().to_str().unwrap().to_string()), - use_env_roots: false, - clients: Some(vec!["codex".to_string()]), - since: None, - until: None, - year: None, - scanner_settings: scanner::ScannerSettings::default(), - modified_after: None, - }) - .unwrap(); + let csv = r#"Date,Kind,Model,Max Mode,Input (w/ Cache Write),Input (w/o Cache Write),Cache Read,Output Tokens,Total Tokens,Cost +"2026-03-04T12:00:00.000Z","Included","Composer 1.5","No","1200","1000","5000","2000","8000","0""#; + std::fs::write(cursor_cache_dir.join("usage.csv"), csv).unwrap(); - assert_eq!(parsed.counts.get(ClientId::Codex), 3); - assert_eq!(parsed.messages.len(), 3); - assert_eq!( - parsed - .messages - .iter() - .map(|message| message.input) - .sum::(), - 88 + let mut litellm = HashMap::new(); + litellm.insert( + "Composer 1.5".into(), + pricing::ModelPricing { + input_cost_per_token: Some(0.001), + output_cost_per_token: Some(0.002), + cache_read_input_token_cost: Some(0.0005), + ..Default::default() + }, ); - assert_eq!( - parsed - .messages - .iter() - .map(|message| message.cache_read) - .sum::(), - 22 + let pricing = pricing::PricingService::new(litellm, HashMap::new()); + + let repriced_messages = parse_all_messages_with_pricing( + source_home.path().to_str().unwrap(), + &["cursor".to_string()], + Some(&pricing), ); - assert_eq!( - parsed - .messages - .iter() - .map(|message| message.output) - .sum::(), - 33 + assert_eq!(repriced_messages.len(), 1); + assert!(repriced_messages[0].cost > 0.0); + + let cached_messages = parse_all_messages_with_pricing( + source_home.path().to_str().unwrap(), + &["cursor".to_string()], + None, ); + + assert_eq!(cached_messages.len(), 1); + assert_eq!(cached_messages[0].cost, 0.0); } } #[test] - #[serial_test::serial] - fn test_codex_cache_reparses_from_zero_when_incremental_prefix_is_stale() { - let cache_home = tempfile::TempDir::new().unwrap(); - let fresh_cache_home = tempfile::TempDir::new().unwrap(); - let source_home = tempfile::TempDir::new().unwrap(); - let _env = EnvGuard::set(&[ - ("HOME", cache_home.path().as_os_str()), - ("TOKSCALE_CONFIG_DIR", cache_home.path().as_os_str()), - ]); + fn test_apply_pricing_if_available_keeps_existing_cost_without_pricing() { + let mut msg = UnifiedMessage::new_with_agent( + "roocode", + "gpt-4o", + "provider", + "session-1", + 1_733_011_200_000, + TokenBreakdown { + input: 10, + output: 5, + cache_read: 0, + cache_write: 0, + reasoning: 0, + }, + 0.42, + Some("planner".to_string()), + ); + + apply_pricing_if_available(&mut msg, None); + + assert_eq!(msg.cost, 0.42); + } + + #[test] + fn test_apply_pricing_if_available_overrides_cost_when_pricing_exists() { + let mut litellm = HashMap::new(); + litellm.insert( + "gpt-4o".into(), + pricing::ModelPricing { + input_cost_per_token: Some(0.001), + output_cost_per_token: Some(0.002), + ..Default::default() + }, + ); + let pricing = pricing::PricingService::new(litellm, HashMap::new()); + + let mut msg = UnifiedMessage::new( + "codex", + "gpt-4o", + "provider", + "session-1", + 1_733_011_200_000, + TokenBreakdown { + input: 10, + output: 5, + cache_read: 0, + cache_write: 0, + reasoning: 0, + }, + 0.0, + ); - { - let codex_dir = scanner_fixture_path(source_home.path(), ".codex/sessions"); - std::fs::create_dir_all(&codex_dir).unwrap(); - let path = codex_dir.join("session.jsonl"); - std::fs::write( - &path, - concat!( - r#"{"type":"turn_context","payload":{"model":"gpt-5.4"}}"#, - "\n", - r#"{"type":"event_msg","payload":{"type":"token_count","info":{"total_token_usage":{"input_tokens":10,"cached_input_tokens":2,"output_tokens":3},"last_token_usage":{"input_tokens":10,"cached_input_tokens":2,"output_tokens":3}}}}"#, - "\n" - ), - ) - .unwrap(); + apply_pricing_if_available(&mut msg, Some(&pricing)); - let initial_messages = parse_all_messages_with_pricing( - source_home.path().to_str().unwrap(), - &["codex".to_string()], - None, - ); - assert_eq!(initial_messages.len(), 1); - assert_eq!(initial_messages[0].model_id, "gpt-5.4"); - assert!(message_cache::SourceMessageCache::load() - .get(&path) - .and_then(|entry| entry.codex_incremental.as_ref()) - .is_some()); + assert_eq!(msg.cost, 0.02); + assert_eq!(msg.cost_source, CostSource::Estimated); + } - std::thread::sleep(std::time::Duration::from_millis(5)); - std::fs::write( - &path, - concat!( - r#"{"type":"turn_context","payload":{"model":"gpt-5.5"}}"#, - "\n", - r#"{"type":"event_msg","payload":{"type":"token_count","info":{"total_token_usage":{"input_tokens":10,"cached_input_tokens":2,"output_tokens":3},"last_token_usage":{"input_tokens":10,"cached_input_tokens":2,"output_tokens":3}}}}"#, - "\n", - r#"{"type":"event_msg","payload":{"type":"token_count","info":{"total_token_usage":{"input_tokens":15,"cached_input_tokens":3,"output_tokens":5},"last_token_usage":{"input_tokens":5,"cached_input_tokens":1,"output_tokens":2}}}}"#, - "\n" - ), - ) - .unwrap(); + #[test] + fn test_apply_pricing_if_available_preserves_provider_reported_cost() { + let mut litellm = HashMap::new(); + litellm.insert( + "gpt-4o".into(), + pricing::ModelPricing { + input_cost_per_token: Some(0.001), + output_cost_per_token: Some(0.002), + ..Default::default() + }, + ); + let pricing = pricing::PricingService::new(litellm, HashMap::new()); + let mut msg = UnifiedMessage::new( + "opencode", + "gpt-4o", + "openai", + "session-1", + 1_733_011_200_000, + TokenBreakdown { + input: 10, + output: 5, + cache_read: 0, + cache_write: 0, + reasoning: 0, + }, + 0.42, + ); + msg.mark_provider_reported_cost(); - let warm_messages = parse_all_messages_with_pricing( - source_home.path().to_str().unwrap(), - &["codex".to_string()], - None, - ); - let _fresh_env = EnvGuard::set(&[ - ("HOME", fresh_cache_home.path().as_os_str()), - ("TOKSCALE_CONFIG_DIR", fresh_cache_home.path().as_os_str()), - ]); - let fresh_messages = parse_all_messages_with_pricing( - source_home.path().to_str().unwrap(), - &["codex".to_string()], - None, - ); + apply_pricing_if_available(&mut msg, Some(&pricing)); - assert_eq!(warm_messages, fresh_messages); - assert_eq!(warm_messages.len(), 2); - assert!(warm_messages - .iter() - .all(|message| message.model_id == "gpt-5.5")); - } + assert_eq!(msg.cost, 0.42); + assert_eq!(msg.cost_source, CostSource::ProviderReported); } #[test] #[serial_test::serial] - fn test_source_cache_keeps_untimestamped_rows_in_sync_after_append() { + fn test_cost_provenance_matches_materialized_and_streaming_lanes() { let cache_home = tempfile::TempDir::new().unwrap(); - let fresh_cache_home = tempfile::TempDir::new().unwrap(); let source_home = tempfile::TempDir::new().unwrap(); - let _env = EnvGuard::set(&[ - ("HOME", cache_home.path().as_os_str()), - ("TOKSCALE_CONFIG_DIR", cache_home.path().as_os_str()), - ]); + let _env = opencode_test_env(cache_home.path(), source_home.path()); - { - let codex_dir = scanner_fixture_path(source_home.path(), ".codex/sessions"); - std::fs::create_dir_all(&codex_dir).unwrap(); - let path = codex_dir.join("session.jsonl"); - std::fs::write( - &path, - concat!( - r#"{"type":"turn_context","payload":{"model":"gpt-5.4"}}"#, - "\n", - r#"{"type":"event_msg","payload":{"type":"token_count","info":{"total_token_usage":{"input_tokens":10,"cached_input_tokens":2,"output_tokens":3},"last_token_usage":{"input_tokens":10,"cached_input_tokens":2,"output_tokens":3}}}}"#, - "\n" + let opencode_data_dir = source_home.path().join(".local/share/opencode"); + std::fs::create_dir_all(&opencode_data_dir).unwrap(); + let opencode_db = + rusqlite::Connection::open(opencode_data_dir.join("opencode.db")).unwrap(); + opencode_db.execute_batch("CREATE TABLE message (id TEXT PRIMARY KEY, session_id TEXT NOT NULL, data TEXT NOT NULL);").unwrap(); + let sqlite_rows = [ + ( + "sqlite-a", + r#"{"id":"json-authoritative","role":"assistant","modelID":"gpt-4o","providerID":"openai","cost":0.0,"tokens":{"input":20,"output":5,"reasoning":0,"cache":{"read":0,"write":0}},"time":{"created":1733011200000}}"#, + ), + ( + "sqlite-b", + r#"{"id":"sqlite-authoritative","role":"assistant","modelID":"gpt-4o","providerID":"openai","cost":0.06,"tokens":{"input":20,"output":5,"reasoning":0,"cache":{"read":0,"write":0}},"time":{"created":1733011201000}}"#, + ), + ( + "sqlite-c", + r#"{"id":"both-authoritative","role":"assistant","modelID":"gpt-4o","providerID":"openai","cost":0.07,"tokens":{"input":20,"output":5,"reasoning":0,"cache":{"read":0,"write":0}},"time":{"created":1733011202000}}"#, + ), + ( + "sqlite-d", + r#"{"id":"both-estimated","role":"assistant","modelID":"gpt-4o","providerID":"openai","cost":0.0,"tokens":{"input":40,"output":5,"reasoning":0,"cache":{"read":0,"write":0}},"time":{"created":1733011203000}}"#, + ), + ]; + for (row_id, data) in sqlite_rows { + opencode_db + .execute( + "INSERT INTO message (id, session_id, data) VALUES (?1, ?2, ?3)", + rusqlite::params![row_id, "oc-session", data], + ) + .unwrap(); + } + drop(opencode_db); + + let opencode_dir = opencode_data_dir.join("storage/message/project-1"); + std::fs::create_dir_all(&opencode_dir).unwrap(); + let json_rows = [ + ( + "a-json-authoritative.json", + r#"{"id":"json-authoritative","sessionID":"oc-session","role":"assistant","modelID":"gpt-4o","providerID":"openai","cost":0.05,"tokens":{"input":10,"output":5,"reasoning":0,"cache":{"read":0,"write":0}},"time":{"created":1733011200000}}"#, + ), + ( + "b-json-estimated.json", + r#"{"id":"sqlite-authoritative","sessionID":"oc-session","role":"assistant","modelID":"gpt-4o","providerID":"openai","cost":0.0,"tokens":{"input":10,"output":5,"reasoning":0,"cache":{"read":0,"write":0}},"time":{"created":1733011201000}}"#, + ), + ( + "c-json-authoritative.json", + r#"{"id":"both-authoritative","sessionID":"oc-session","role":"assistant","modelID":"gpt-4o","providerID":"openai","cost":0.08,"tokens":{"input":10,"output":5,"reasoning":0,"cache":{"read":0,"write":0}},"time":{"created":1733011202000}}"#, + ), + ( + "d-json-estimated.json", + r#"{"id":"both-estimated","sessionID":"oc-session","role":"assistant","modelID":"gpt-4o","providerID":"openai","cost":0.0,"tokens":{"input":10,"output":5,"reasoning":0,"cache":{"read":0,"write":0}},"time":{"created":1733011203000}}"#, + ), + ]; + for (name, data) in json_rows { + std::fs::write(opencode_dir.join(name), data).unwrap(); + } + + let mut litellm = HashMap::new(); + litellm.insert( + "openai/gpt-4o".to_string(), + pricing::ModelPricing { + input_cost_per_token: Some(0.01), + output_cost_per_token: Some(0.02), + ..Default::default() + }, + ); + let pricing = pricing::PricingService::new(litellm, HashMap::new()); + let clients = vec!["opencode".to_string()]; + let materialized = parse_all_messages_with_pricing_with_env_strategy( + source_home.path().to_str().unwrap(), + &clients, + Some(&pricing), + false, + &scanner::ScannerSettings::default(), + ); + let mut streamed = Vec::new(); + scan_messages_streaming( + source_home.path().to_str().unwrap(), + &clients, + Some(&pricing), + false, + &scanner::ScannerSettings::default(), + &|_| true, + &mut |message| streamed.push(message.clone()), + ); + + let summarize = |messages: Vec| { + let mut rows: Vec<(String, f64, CostSource)> = messages + .into_iter() + .map(|message| { + ( + message.dedup_key.unwrap_or_default(), + message.cost, + message.cost_source, + ) + }) + .collect(); + rows.sort_by(|left, right| left.0.cmp(&right.0)); + rows + }; + let materialized = summarize(materialized); + let streamed = summarize(streamed); + assert_eq!(materialized, streamed); + assert_eq!( + materialized, + vec![ + ( + "both-authoritative".to_string(), + 0.07, + CostSource::ProviderReported ), - ) - .unwrap(); + ("both-estimated".to_string(), 0.5, CostSource::Estimated), + ( + "json-authoritative".to_string(), + 0.05, + CostSource::ProviderReported + ), + ( + "sqlite-authoritative".to_string(), + 0.06, + CostSource::ProviderReported + ), + ] + ); + } - let first_messages = parse_all_messages_with_pricing( - source_home.path().to_str().unwrap(), - &["codex".to_string()], - None, - ); - assert_eq!(first_messages.len(), 1); + #[test] + fn test_apply_pricing_if_available_applies_zed_hosted_markup() { + let mut litellm = HashMap::new(); + litellm.insert( + "claude-sonnet-4-5".into(), + pricing::ModelPricing { + input_cost_per_token: Some(0.001), + output_cost_per_token: Some(0.002), + ..Default::default() + }, + ); + let pricing = pricing::PricingService::new(litellm, HashMap::new()); - std::thread::sleep(std::time::Duration::from_millis(5)); - let mut file = std::fs::OpenOptions::new() - .append(true) - .open(&path) - .unwrap(); - file.write_all( - concat!( - r#"{"type":"event_msg","payload":{"type":"token_count","info":{"total_token_usage":{"input_tokens":15,"cached_input_tokens":3,"output_tokens":5},"last_token_usage":{"input_tokens":5,"cached_input_tokens":1,"output_tokens":2}}}}"#, - "\n" - ) - .as_bytes(), - ) - .unwrap(); - file.flush().unwrap(); + let mut msg = UnifiedMessage::new( + "zed", + "claude-sonnet-4-5", + crate::sessions::zed::ZED_HOSTED_PROVIDER, + "session-1", + 1_733_011_200_000, + TokenBreakdown { + input: 10, + output: 5, + cache_read: 0, + cache_write: 0, + reasoning: 0, + }, + 0.0, + ); - let warm_messages = parse_all_messages_with_pricing( - source_home.path().to_str().unwrap(), - &["codex".to_string()], - None, - ); - let _fresh_env = EnvGuard::set(&[ - ("HOME", fresh_cache_home.path().as_os_str()), - ("TOKSCALE_CONFIG_DIR", fresh_cache_home.path().as_os_str()), - ]); - let fresh_messages = parse_all_messages_with_pricing( - source_home.path().to_str().unwrap(), - &["codex".to_string()], - None, - ); + apply_pricing_if_available(&mut msg, Some(&pricing)); - assert_eq!(warm_messages, fresh_messages); - } + assert!((msg.cost - 0.022).abs() < 1e-12); } #[test] - #[serial_test::serial] - fn test_source_cache_matches_cold_parse_after_malformed_json_append() { - let cache_home = tempfile::TempDir::new().unwrap(); - let fresh_cache_home = tempfile::TempDir::new().unwrap(); - let source_home = tempfile::TempDir::new().unwrap(); - let _env = EnvGuard::set(&[ - ("HOME", cache_home.path().as_os_str()), - ("TOKSCALE_CONFIG_DIR", cache_home.path().as_os_str()), - ]); + fn test_apply_pricing_if_available_skips_zed_markup_for_non_zed_client() { + // Non-zed client with provider_id "zed.dev" must not receive the +10% + // markup. The multiplier is gated on (client == "zed" AND provider). + let mut litellm = HashMap::new(); + litellm.insert( + "claude-sonnet-4-5".into(), + pricing::ModelPricing { + input_cost_per_token: Some(0.001), + output_cost_per_token: Some(0.002), + ..Default::default() + }, + ); + let pricing = pricing::PricingService::new(litellm, HashMap::new()); - { - let codex_dir = scanner_fixture_path(source_home.path(), ".codex/sessions"); - std::fs::create_dir_all(&codex_dir).unwrap(); - let path = codex_dir.join("session.jsonl"); - std::fs::write( - &path, - concat!( - r#"{"type":"turn_context","payload":{"model":"gpt-5.4"}}"#, - "\n", - r#"{"type":"event_msg","payload":{"type":"token_count","info":{"total_token_usage":{"input_tokens":10,"cached_input_tokens":2,"output_tokens":3},"last_token_usage":{"input_tokens":10,"cached_input_tokens":2,"output_tokens":3}}}}"#, - "\n", - r#"{"type":"event_msg","payload":{"type":"token_count","info":{"total_token_usage":{"input_tokens":999""#, - "\n" - ), - ) - .unwrap(); + let mut msg = UnifiedMessage::new( + "claudecode", + "claude-sonnet-4-5", + crate::sessions::zed::ZED_HOSTED_PROVIDER, + "session-1", + 1_733_011_200_000, + TokenBreakdown { + input: 10, + output: 5, + cache_read: 0, + cache_write: 0, + reasoning: 0, + }, + 0.0, + ); - let initial_messages = parse_all_messages_with_pricing( - source_home.path().to_str().unwrap(), - &["codex".to_string()], - None, - ); - assert_eq!(initial_messages.len(), 1); + apply_pricing_if_available(&mut msg, Some(&pricing)); - std::thread::sleep(std::time::Duration::from_millis(5)); - let mut file = std::fs::OpenOptions::new() - .append(true) - .open(&path) - .unwrap(); - file.write_all( - concat!( - r#"{"type":"event_msg","payload":{"type":"token_count","info":{"total_token_usage":{"input_tokens":15,"cached_input_tokens":3,"output_tokens":5},"last_token_usage":{"input_tokens":5,"cached_input_tokens":1,"output_tokens":2}}}}"#, - "\n" - ) - .as_bytes(), - ) - .unwrap(); - file.flush().unwrap(); + // 10 * 0.001 + 5 * 0.002 = 0.020, no markup. + assert!((msg.cost - 0.020).abs() < 1e-12); + } - let warm_messages = parse_all_messages_with_pricing( - source_home.path().to_str().unwrap(), - &["codex".to_string()], - None, - ); - assert!(message_cache::SourceMessageCache::load() - .get(&path) - .is_none()); + #[test] + fn test_apply_pricing_if_available_skips_zed_markup_for_byok_provider() { + // A Zed message whose provider_id is the upstream provider directly + // (BYOK / non-hosted path) must not be marked up — the user is paying + // the upstream API directly, not through Zed. + let mut litellm = HashMap::new(); + litellm.insert( + "claude-sonnet-4-5".into(), + pricing::ModelPricing { + input_cost_per_token: Some(0.001), + output_cost_per_token: Some(0.002), + ..Default::default() + }, + ); + let pricing = pricing::PricingService::new(litellm, HashMap::new()); - let _fresh_env = EnvGuard::set(&[ - ("HOME", fresh_cache_home.path().as_os_str()), - ("TOKSCALE_CONFIG_DIR", fresh_cache_home.path().as_os_str()), - ]); - let fresh_messages = parse_all_messages_with_pricing( - source_home.path().to_str().unwrap(), - &["codex".to_string()], - None, - ); + let mut msg = UnifiedMessage::new( + "zed", + "claude-sonnet-4-5", + "anthropic", + "session-1", + 1_733_011_200_000, + TokenBreakdown { + input: 10, + output: 5, + cache_read: 0, + cache_write: 0, + reasoning: 0, + }, + 0.0, + ); - assert_eq!(warm_messages, fresh_messages); - } + apply_pricing_if_available(&mut msg, Some(&pricing)); + + assert!((msg.cost - 0.020).abs() < 1e-12); } #[test] - #[serial_test::serial] - fn test_exact_hit_codex_cache_repairs_fallback_timestamps_without_incremental_state() { - let cache_home = tempfile::TempDir::new().unwrap(); - let source_home = tempfile::TempDir::new().unwrap(); - let _env = EnvGuard::set(&[ - ("HOME", cache_home.path().as_os_str()), - ("TOKSCALE_CONFIG_DIR", cache_home.path().as_os_str()), - ]); + fn test_apply_pricing_if_available_uses_reasoning_for_gemini() { + let mut litellm = HashMap::new(); + litellm.insert( + "gemini-2.5-pro".into(), + pricing::ModelPricing { + input_cost_per_token: Some(0.001), + output_cost_per_token: Some(0.002), + ..Default::default() + }, + ); + let pricing = pricing::PricingService::new(litellm, HashMap::new()); - { - let session_dir = scanner_fixture_path(source_home.path(), ".codex/sessions"); - std::fs::create_dir_all(&session_dir).unwrap(); - let path = session_dir.join("session.jsonl"); - std::fs::write( - &path, - concat!( - r#"{"type":"turn_context","payload":{"model":"gpt-5.4"}}"#, - "\n", - r#"{"type":"event_msg","payload":{"type":"token_count","info":{"total_token_usage":{"input_tokens":10,"cached_input_tokens":2,"output_tokens":3},"last_token_usage":{"input_tokens":10,"cached_input_tokens":2,"output_tokens":3}}}}"#, - "\n" - ), - ) - .unwrap(); + let mut msg = UnifiedMessage::new( + "gemini", + "gemini-2.5-pro", + "google", + "session-1", + 1_733_011_200_000, + TokenBreakdown { + input: 10, + output: 5, + cache_read: 0, + cache_write: 0, + reasoning: 7, + }, + 0.0, + ); - let expected = crate::sessions::codex::parse_codex_file(&path); - assert_eq!(expected.len(), 1); + apply_pricing_if_available(&mut msg, Some(&pricing)); - let fingerprint = message_cache::SourceFingerprint::from_path(&path).unwrap(); - let mut stale_message = expected[0].clone(); - stale_message.timestamp = 0; - stale_message.date = "1900-01-01".to_string(); + assert_eq!(msg.cost, 0.034); + } - let mut cache = message_cache::SourceMessageCache::default(); - cache.insert(message_cache::CachedSourceEntry::new( - &path, - fingerprint, - vec![stale_message], - vec![0], - None, - )); - cache.save_if_dirty(); + #[test] + fn test_apply_pricing_if_available_uses_cache_read_pricing_for_gemini() { + let mut litellm = HashMap::new(); + litellm.insert( + "gemini-2.5-pro".into(), + pricing::ModelPricing { + input_cost_per_token: Some(0.001), + output_cost_per_token: Some(0.002), + cache_read_input_token_cost: Some(0.0001), + ..Default::default() + }, + ); + let pricing = pricing::PricingService::new(litellm, HashMap::new()); - let messages = parse_all_messages_with_pricing( - source_home.path().to_str().unwrap(), - &["codex".to_string()], - None, - ); + let mut msg = UnifiedMessage::new( + "gemini", + "gemini-2.5-pro", + "google", + "session-1", + 1_733_011_200_000, + TokenBreakdown { + input: 10, + output: 5, + cache_read: 7, + cache_write: 0, + reasoning: 3, + }, + 0.0, + ); - assert_eq!(messages, expected); - } + apply_pricing_if_available(&mut msg, Some(&pricing)); + + assert_eq!(msg.cost, 0.0267); } #[test] - #[serial_test::serial] - fn test_codex_cache_repairs_fallback_timestamps_after_source_mtime_change() { - let cache_home = tempfile::TempDir::new().unwrap(); - let fresh_cache_home = tempfile::TempDir::new().unwrap(); - let source_home = tempfile::TempDir::new().unwrap(); - let _env = EnvGuard::set(&[ - ("HOME", cache_home.path().as_os_str()), - ("TOKSCALE_CONFIG_DIR", cache_home.path().as_os_str()), - ]); - - { - let session_dir = scanner_fixture_path(source_home.path(), ".codex/sessions"); - std::fs::create_dir_all(&session_dir).unwrap(); - let path = session_dir.join("session.jsonl"); - let contents = concat!( - r#"{"type":"turn_context","payload":{"model":"gpt-5.4"}}"#, - "\n", - r#"{"type":"event_msg","payload":{"type":"token_count","info":{"total_token_usage":{"input_tokens":10,"cached_input_tokens":2,"output_tokens":3},"last_token_usage":{"input_tokens":10,"cached_input_tokens":2,"output_tokens":3}}}}"#, - "\n" - ); - std::fs::write(&path, contents).unwrap(); - - let initial_messages = parse_all_messages_with_pricing( - source_home.path().to_str().unwrap(), - &["codex".to_string()], - None, - ); - assert_eq!(initial_messages.len(), 1); - - std::thread::sleep(std::time::Duration::from_millis(20)); - std::fs::write(&path, contents).unwrap(); + fn test_apply_pricing_if_available_uses_market_rate_for_free_variant() { + let mut openrouter = HashMap::new(); + openrouter.insert( + "z-ai/glm-4.7".into(), + pricing::ModelPricing { + input_cost_per_token: Some(0.001), + output_cost_per_token: Some(0.002), + ..Default::default() + }, + ); + let pricing = pricing::PricingService::new(HashMap::new(), openrouter); - let warm_messages = parse_all_messages_with_pricing( - source_home.path().to_str().unwrap(), - &["codex".to_string()], - None, - ); + let mut msg = UnifiedMessage::new( + "opencode", + "glm-4.7-free", + "modal", + "session-1", + 1_733_011_200_000, + TokenBreakdown { + input: 10, + output: 5, + cache_read: 0, + cache_write: 0, + reasoning: 0, + }, + 0.0, + ); - let _fresh_env = EnvGuard::set(&[ - ("HOME", fresh_cache_home.path().as_os_str()), - ("TOKSCALE_CONFIG_DIR", fresh_cache_home.path().as_os_str()), - ]); - let fresh_messages = parse_all_messages_with_pricing( - source_home.path().to_str().unwrap(), - &["codex".to_string()], - None, - ); + apply_pricing_if_available(&mut msg, Some(&pricing)); - assert_eq!(warm_messages, fresh_messages); - assert_ne!(warm_messages[0].timestamp, initial_messages[0].timestamp); - } + assert_eq!(msg.cost, 0.02); } #[test] - #[serial_test::serial] - fn test_full_log_parse_preserves_valid_messages_before_invalid_line_error() { - let cache_home = tempfile::TempDir::new().unwrap(); - let source_home = tempfile::TempDir::new().unwrap(); - let _env = EnvGuard::set(&[ - ("HOME", cache_home.path().as_os_str()), - ("TOKSCALE_CONFIG_DIR", cache_home.path().as_os_str()), - ]); - - { - let session_dir = scanner_fixture_path(source_home.path(), ".codex/sessions"); - std::fs::create_dir_all(&session_dir).unwrap(); - let path = session_dir.join("session.jsonl"); + fn test_apply_pricing_if_available_prefers_provider_aware_match() { + let mut litellm = HashMap::new(); + litellm.insert( + "xai/grok-code-fast-1-0825".into(), + pricing::ModelPricing { + input_cost_per_token: Some(0.001), + output_cost_per_token: Some(0.002), + ..Default::default() + }, + ); + litellm.insert( + "azure_ai/grok-code-fast-1".into(), + pricing::ModelPricing { + input_cost_per_token: Some(0.01), + output_cost_per_token: Some(0.02), + ..Default::default() + }, + ); + let pricing = pricing::PricingService::new(litellm, HashMap::new()); - let mut file = std::fs::File::create(&path).unwrap(); - file.write_all( - concat!( - r#"{"type":"turn_context","payload":{"model":"gpt-5.4"}}"#, - "\n", - r#"{"type":"event_msg","payload":{"type":"token_count","info":{"total_token_usage":{"input_tokens":10,"cached_input_tokens":2,"output_tokens":3},"last_token_usage":{"input_tokens":10,"cached_input_tokens":2,"output_tokens":3}}}}"#, - "\n" - ) - .as_bytes(), - ) - .unwrap(); - file.write_all(&[0xff, b'\n']).unwrap(); - file.flush().unwrap(); + let mut msg = UnifiedMessage::new( + "opencode", + "grok-code", + "azure", + "session-1", + 1_733_011_200_000, + TokenBreakdown { + input: 10, + output: 5, + cache_read: 0, + cache_write: 0, + reasoning: 0, + }, + 0.0, + ); - let messages = parse_all_messages_with_pricing( - source_home.path().to_str().unwrap(), - &["codex".to_string()], - None, - ); - assert_eq!(messages.len(), 1); - assert_eq!(messages[0].model_id, "gpt-5.4"); + apply_pricing_if_available(&mut msg, Some(&pricing)); - let cache = message_cache::SourceMessageCache::load(); - assert!(cache.get(&path).is_none()); - } + assert_eq!(msg.cost, 0.2); } #[test] - #[serial_test::serial] - fn test_codex_cache_does_not_persist_unknown_before_later_turn_context() { - let cache_home = tempfile::TempDir::new().unwrap(); - let fresh_cache_home = tempfile::TempDir::new().unwrap(); - let source_home = tempfile::TempDir::new().unwrap(); - let _env = EnvGuard::set(&[ - ("HOME", cache_home.path().as_os_str()), - ("TOKSCALE_CONFIG_DIR", cache_home.path().as_os_str()), - ]); - - { - let session_dir = scanner_fixture_path(source_home.path(), ".codex/sessions"); - std::fs::create_dir_all(&session_dir).unwrap(); - let path = session_dir.join("session.jsonl"); - std::fs::write( - &path, - concat!( - r#"{"type":"session_meta","payload":{"source":"interactive","model_provider":"openai"}}"#, - "\n", - r#"{"timestamp":"2026-04-27T10:00:00Z","type":"event_msg","payload":{"type":"token_count","info":{"total_token_usage":{"input_tokens":10,"cached_input_tokens":2,"output_tokens":3},"last_token_usage":{"input_tokens":10,"cached_input_tokens":2,"output_tokens":3}}}}"#, - "\n" - ), - ) - .unwrap(); + fn test_apply_pricing_if_available_uses_nested_reseller_exact_match() { + let mut litellm = HashMap::new(); + litellm.insert( + "gpt-4".into(), + pricing::ModelPricing { + input_cost_per_token: Some(0.001), + output_cost_per_token: Some(0.002), + ..Default::default() + }, + ); + litellm.insert( + "azure/openai/gpt-4".into(), + pricing::ModelPricing { + input_cost_per_token: Some(0.01), + output_cost_per_token: Some(0.02), + ..Default::default() + }, + ); + let pricing = pricing::PricingService::new(litellm, HashMap::new()); - let initial_messages = parse_all_messages_with_pricing( - source_home.path().to_str().unwrap(), - &["codex".to_string()], - None, - ); - assert_eq!(initial_messages.len(), 1); - assert_eq!(initial_messages[0].model_id, "unknown"); - assert!(message_cache::SourceMessageCache::load() - .get(&path) - .is_none()); + let mut msg = UnifiedMessage::new( + "opencode", + "gpt-4", + "azure", + "session-1", + 1_733_011_200_000, + TokenBreakdown { + input: 10, + output: 5, + cache_read: 0, + cache_write: 0, + reasoning: 0, + }, + 0.0, + ); - std::thread::sleep(std::time::Duration::from_millis(5)); - let mut file = std::fs::OpenOptions::new() - .append(true) - .open(&path) - .unwrap(); - file.write_all( - concat!( - r#"{"timestamp":"2026-04-27T10:00:04Z","type":"turn_context","payload":{"model":"gpt-5.5"}}"#, - "\n" - ) - .as_bytes(), - ) - .unwrap(); - file.flush().unwrap(); + apply_pricing_if_available(&mut msg, Some(&pricing)); - let resumed_messages = parse_all_messages_with_pricing( - source_home.path().to_str().unwrap(), - &["codex".to_string()], - None, - ); + assert_eq!(msg.cost, 0.2); + } - let _fresh_env = EnvGuard::set(&[ - ("HOME", fresh_cache_home.path().as_os_str()), - ("TOKSCALE_CONFIG_DIR", fresh_cache_home.path().as_os_str()), - ]); - let fresh_messages = parse_all_messages_with_pricing( - source_home.path().to_str().unwrap(), - &["codex".to_string()], - None, - ); + #[test] + fn test_apply_pricing_if_available_keeps_scoped_fireworks_cost_without_exact_pricing() { + let mut litellm = HashMap::new(); + litellm.insert( + "fireworks_ai/accounts/fireworks/models/deepseek-r1-0528-distill-qwen3-8b".into(), + pricing::ModelPricing { + input_cost_per_token: Some(0.0000002), + output_cost_per_token: Some(0.0000002), + ..Default::default() + }, + ); - assert_eq!(resumed_messages, fresh_messages); - assert_eq!(resumed_messages.len(), 1); - assert_eq!(resumed_messages[0].model_id, "gpt-5.5"); + let mut openrouter = HashMap::new(); + openrouter.insert( + "deepseek/deepseek-v4-pro".into(), + pricing::ModelPricing { + input_cost_per_token: Some(0.000001), + output_cost_per_token: Some(0.000002), + ..Default::default() + }, + ); - drop(_fresh_env); - assert!(message_cache::SourceMessageCache::load() - .get(&path) - .is_some()); - } + let pricing = pricing::PricingService::new(litellm, openrouter); + let mut msg = UnifiedMessage::new( + "opencode", + "accounts/fireworks/models/deepseek-v4-pro", + "fireworks", + "session-1", + 1_733_011_200_000, + TokenBreakdown { + input: 10, + output: 5, + cache_read: 0, + cache_write: 0, + reasoning: 0, + }, + 0.123, + ); + + apply_pricing_if_available(&mut msg, Some(&pricing)); + + assert_eq!(msg.cost, 0.123); } #[test] - #[serial_test::serial] - fn test_codex_cache_skips_non_newline_terminated_resume_prefix() { - let cache_home = tempfile::TempDir::new().unwrap(); - let fresh_cache_home = tempfile::TempDir::new().unwrap(); - let source_home = tempfile::TempDir::new().unwrap(); - let _env = EnvGuard::set(&[ - ("HOME", cache_home.path().as_os_str()), - ("TOKSCALE_CONFIG_DIR", cache_home.path().as_os_str()), - ]); - - { - let session_dir = scanner_fixture_path(source_home.path(), ".codex/sessions"); - std::fs::create_dir_all(&session_dir).unwrap(); - let path = session_dir.join("session.jsonl"); - std::fs::write( - &path, - concat!( - r#"{"type":"turn_context","payload":{"model":"gpt-5.4"}}"#, - "\n", - r#"{"type":"event_msg","payload":{"type":"token_count","info":{"total_token_usage":{"input_tokens":10,"cached_input_tokens":2,"output_tokens":3},"last_token_usage":{"input_tokens":10,"cached_input_tokens":2,"output_tokens":3}}}}"# - ), - ) - .unwrap(); + fn test_apply_pricing_if_available_prefers_provider_specific_exact_match_over_plain_exact() { + let mut litellm = HashMap::new(); + litellm.insert( + "gemini-2.5-pro".into(), + pricing::ModelPricing { + input_cost_per_token: Some(0.001), + output_cost_per_token: Some(0.002), + cache_creation_input_token_cost: None, + ..Default::default() + }, + ); - let initial_messages = parse_all_messages_with_pricing( - source_home.path().to_str().unwrap(), - &["codex".to_string()], - None, - ); - assert_eq!(initial_messages.len(), 1); - assert!(message_cache::SourceMessageCache::load() - .get(&path) - .is_none()); + let mut openrouter = HashMap::new(); + openrouter.insert( + "google/gemini-2.5-pro".into(), + pricing::ModelPricing { + input_cost_per_token: Some(0.001), + output_cost_per_token: Some(0.002), + cache_creation_input_token_cost: Some(0.01), + ..Default::default() + }, + ); - std::thread::sleep(std::time::Duration::from_millis(5)); - let mut file = std::fs::OpenOptions::new() - .append(true) - .open(&path) - .unwrap(); - file.write_all( - concat!( - "\n", - r#"{"type":"event_msg","payload":{"type":"token_count","info":{"total_token_usage":{"input_tokens":15,"cached_input_tokens":3,"output_tokens":5},"last_token_usage":{"input_tokens":5,"cached_input_tokens":1,"output_tokens":2}}}}"#, - "\n" - ) - .as_bytes(), - ) - .unwrap(); - file.flush().unwrap(); + let pricing = pricing::PricingService::new(litellm, openrouter); - let warm_messages = parse_all_messages_with_pricing( - source_home.path().to_str().unwrap(), - &["codex".to_string()], - None, - ); + let mut msg = UnifiedMessage::new( + "opencode", + "gemini-2.5-pro", + "google", + "session-1", + 1_733_011_200_000, + TokenBreakdown { + input: 10, + output: 5, + cache_read: 0, + cache_write: 3, + reasoning: 0, + }, + 0.0, + ); - let _fresh_env = EnvGuard::set(&[ - ("HOME", fresh_cache_home.path().as_os_str()), - ("TOKSCALE_CONFIG_DIR", fresh_cache_home.path().as_os_str()), - ]); - let fresh_messages = parse_all_messages_with_pricing( - source_home.path().to_str().unwrap(), - &["codex".to_string()], - None, - ); + apply_pricing_if_available(&mut msg, Some(&pricing)); - assert_eq!(warm_messages, fresh_messages); - assert_eq!(warm_messages.len(), 2); - } + assert_eq!(msg.cost, 0.05); } #[test] - #[serial_test::serial] - fn test_source_cache_does_not_reuse_priced_cost_without_pricing_service() { - let temp_home = tempfile::TempDir::new().unwrap(); - let source_home = tempfile::TempDir::new().unwrap(); - let _env = EnvGuard::set(&[ - ("HOME", temp_home.path().as_os_str()), - ("TOKSCALE_CONFIG_DIR", temp_home.path().as_os_str()), - ]); - { - let cursor_cache_dir = source_home.path().join(".config/tokscale/cursor-cache"); - std::fs::create_dir_all(&cursor_cache_dir).unwrap(); + fn test_apply_pricing_if_available_normalizes_openai_codex_provider() { + let mut litellm = HashMap::new(); + litellm.insert( + "openai/gpt-5.2-preview".into(), + pricing::ModelPricing { + input_cost_per_token: Some(0.01), + output_cost_per_token: Some(0.02), + ..Default::default() + }, + ); + litellm.insert( + "google/gpt-5.2-preview-max".into(), + pricing::ModelPricing { + input_cost_per_token: Some(0.1), + output_cost_per_token: Some(0.2), + ..Default::default() + }, + ); + let pricing = pricing::PricingService::new(litellm, HashMap::new()); - let csv = r#"Date,Kind,Model,Max Mode,Input (w/ Cache Write),Input (w/o Cache Write),Cache Read,Output Tokens,Total Tokens,Cost -"2026-03-04T12:00:00.000Z","Included","Composer 1.5","No","1200","1000","5000","2000","8000","0""#; - std::fs::write(cursor_cache_dir.join("usage.csv"), csv).unwrap(); + let mut msg = UnifiedMessage::new( + "openclaw", + "gpt-5.2", + "openai-codex", + "session-1", + 1_733_011_200_000, + TokenBreakdown { + input: 10, + output: 5, + cache_read: 0, + cache_write: 0, + reasoning: 0, + }, + 0.0, + ); - let mut litellm = HashMap::new(); - litellm.insert( - "Composer 1.5".into(), - pricing::ModelPricing { - input_cost_per_token: Some(0.001), - output_cost_per_token: Some(0.002), - cache_read_input_token_cost: Some(0.0005), - ..Default::default() - }, - ); - let pricing = pricing::PricingService::new(litellm, HashMap::new()); + apply_pricing_if_available(&mut msg, Some(&pricing)); - let repriced_messages = parse_all_messages_with_pricing( - source_home.path().to_str().unwrap(), - &["cursor".to_string()], - Some(&pricing), - ); - assert_eq!(repriced_messages.len(), 1); - assert!(repriced_messages[0].cost > 0.0); + assert_eq!(msg.cost, 0.2); + } - let cached_messages = parse_all_messages_with_pricing( - source_home.path().to_str().unwrap(), - &["cursor".to_string()], - None, - ); + #[test] + fn test_apply_pricing_if_available_prices_claude_code_gpt_5_3_codex() { + let pricing = pricing::PricingService::new(HashMap::new(), HashMap::new()); + + let mut msg = UnifiedMessage::new( + "claude", + "gpt-5.3-codex", + "openai", + "session-1", + 1_776_000_000_000, + TokenBreakdown { + input: 1_000_000, + output: 100_000, + cache_read: 50_000, + cache_write: 0, + reasoning: 0, + }, + 0.0, + ); - assert_eq!(cached_messages.len(), 1); - assert_eq!(cached_messages[0].cost, 0.0); - } + apply_pricing_if_available(&mut msg, Some(&pricing)); + + let expected = 1.75 + 1.4 + 0.00875; + assert!((msg.cost - expected).abs() < 1e-12); } #[test] - fn test_apply_pricing_if_available_keeps_existing_cost_without_pricing() { - let mut msg = UnifiedMessage::new_with_agent( - "roocode", - "gpt-4o", - "provider", + fn test_apply_pricing_if_available_prices_claude_code_minimax_model() { + let mut litellm = HashMap::new(); + litellm.insert( + "minimax/minimax-m2.1".into(), + pricing::ModelPricing { + input_cost_per_token: Some(0.01), + output_cost_per_token: Some(0.02), + ..Default::default() + }, + ); + let pricing = pricing::PricingService::new(litellm, HashMap::new()); + + let mut msg = UnifiedMessage::new( + "claude", + "MiniMax-M2.1", + "minimax", "session-1", - 1_733_011_200_000, + 1_776_000_000_000, TokenBreakdown { input: 10, output: 5, @@ -7690,37 +11757,36 @@ mod tests { cache_write: 0, reasoning: 0, }, - 0.42, - Some("planner".to_string()), + 0.0, ); - apply_pricing_if_available(&mut msg, None); + apply_pricing_if_available(&mut msg, Some(&pricing)); - assert_eq!(msg.cost, 0.42); + assert_eq!(msg.cost, 0.2); } #[test] - fn test_apply_pricing_if_available_overrides_cost_when_pricing_exists() { - let mut litellm = HashMap::new(); - litellm.insert( - "gpt-4o".into(), + fn test_apply_pricing_if_available_prices_kimi_k2p6_alias() { + let mut openrouter = HashMap::new(); + openrouter.insert( + "moonshotai/kimi-k2.6".into(), pricing::ModelPricing { - input_cost_per_token: Some(0.001), - output_cost_per_token: Some(0.002), + input_cost_per_token: Some(9.5e-7), + output_cost_per_token: Some(0.000004), ..Default::default() }, ); - let pricing = pricing::PricingService::new(litellm, HashMap::new()); + let pricing = pricing::PricingService::new(HashMap::new(), openrouter); let mut msg = UnifiedMessage::new( - "codex", - "gpt-4o", - "provider", + "kimi", + "k2p6", + "kimi-for-coding", "session-1", - 1_733_011_200_000, + 1_776_000_000_000, TokenBreakdown { - input: 10, - output: 5, + input: 1_000_000, + output: 250_000, cache_read: 0, cache_write: 0, reasoning: 0, @@ -7730,25 +11796,29 @@ mod tests { apply_pricing_if_available(&mut msg, Some(&pricing)); - assert_eq!(msg.cost, 0.02); - assert_eq!(msg.cost_source, CostSource::Estimated); + let expected = 1_000_000.0 * 9.5e-7 + 250_000.0 * 0.000004; + assert!((msg.cost - expected).abs() < 1e-12); + assert!(msg.cost > 0.0); } #[test] - fn test_apply_pricing_if_available_preserves_provider_reported_cost() { - let mut litellm = HashMap::new(); - litellm.insert( - "gpt-4o".into(), + fn test_select_local_parse_pricing_prefers_fresh_service_for_new_models() { + let mut fresh_litellm = HashMap::new(); + fresh_litellm.insert( + "gpt-5.4".into(), pricing::ModelPricing { - input_cost_per_token: Some(0.001), - output_cost_per_token: Some(0.002), + input_cost_per_token: Some(0.000002), + output_cost_per_token: Some(0.00001), ..Default::default() }, ); - let pricing = pricing::PricingService::new(litellm, HashMap::new()); + let fresh = Arc::new(pricing::PricingService::new(fresh_litellm, HashMap::new())); + let stale = pricing::PricingService::new(HashMap::new(), HashMap::new()); + let selected = select_local_parse_pricing(Ok(Arc::clone(&fresh)), || Some(stale)).unwrap(); + let mut msg = UnifiedMessage::new( "opencode", - "gpt-4o", + "gpt-5.4", "openai", "session-1", 1_733_011_200_000, @@ -7759,1051 +11829,1258 @@ mod tests { cache_write: 0, reasoning: 0, }, - 0.42, + 0.0, ); - msg.mark_provider_reported_cost(); - apply_pricing_if_available(&mut msg, Some(&pricing)); + apply_pricing_if_available(&mut msg, Some(selected.as_ref())); - assert_eq!(msg.cost, 0.42); - assert_eq!(msg.cost_source, CostSource::ProviderReported); + assert!(msg.cost > 0.0); } #[test] - #[serial_test::serial] - fn test_cost_provenance_matches_materialized_and_streaming_lanes() { - let cache_home = tempfile::TempDir::new().unwrap(); - let source_home = tempfile::TempDir::new().unwrap(); - let _env = opencode_test_env(cache_home.path(), source_home.path()); + fn test_select_local_parse_pricing_falls_back_to_stale_cache_on_fetch_error() { + let mut stale_litellm = HashMap::new(); + stale_litellm.insert( + "gpt-5.2".into(), + pricing::ModelPricing { + input_cost_per_token: Some(0.00000175), + output_cost_per_token: Some(0.000014), + ..Default::default() + }, + ); + let stale = pricing::PricingService::new(stale_litellm, HashMap::new()); - let opencode_data_dir = source_home.path().join(".local/share/opencode"); - std::fs::create_dir_all(&opencode_data_dir).unwrap(); - let opencode_db = - rusqlite::Connection::open(opencode_data_dir.join("opencode.db")).unwrap(); - opencode_db.execute_batch("CREATE TABLE message (id TEXT PRIMARY KEY, session_id TEXT NOT NULL, data TEXT NOT NULL);").unwrap(); - let sqlite_rows = [ - ( - "sqlite-a", - r#"{"id":"json-authoritative","role":"assistant","modelID":"gpt-4o","providerID":"openai","cost":0.0,"tokens":{"input":20,"output":5,"reasoning":0,"cache":{"read":0,"write":0}},"time":{"created":1733011200000}}"#, - ), - ( - "sqlite-b", - r#"{"id":"sqlite-authoritative","role":"assistant","modelID":"gpt-4o","providerID":"openai","cost":0.06,"tokens":{"input":20,"output":5,"reasoning":0,"cache":{"read":0,"write":0}},"time":{"created":1733011201000}}"#, + let selected = + select_local_parse_pricing(Err("network failed".to_string()), || Some(stale)).unwrap(); + + assert!(selected.lookup_with_source("gpt-5.2", None).is_some()); + } + + #[test] + fn test_select_local_parse_pricing_does_not_evaluate_stale_fallback_on_fresh_success() { + let fresh = Arc::new(pricing::PricingService::new(HashMap::new(), HashMap::new())); + let mut stale_called = false; + + let selected = select_local_parse_pricing(Ok(Arc::clone(&fresh)), || { + stale_called = true; + None + }) + .unwrap(); + + assert!(Arc::ptr_eq(&selected, &fresh)); + assert!(!stale_called); + } + + #[test] + fn test_dedupe_latest_trae_messages_keeps_latest_timestamp_for_session() { + let messages = vec![ + make_trae_message( + "session-stable", + 1_700_000_002_000, + Some("trae:session-stable:1_700_000_002"), + 0.2, ), - ( - "sqlite-c", - r#"{"id":"both-authoritative","role":"assistant","modelID":"gpt-4o","providerID":"openai","cost":0.07,"tokens":{"input":20,"output":5,"reasoning":0,"cache":{"read":0,"write":0}},"time":{"created":1733011202000}}"#, + make_trae_message( + "session-stable", + 1_700_000_003_000, + Some("trae:session-stable:1_700_000_003"), + 0.3, ), - ( - "sqlite-d", - r#"{"id":"both-estimated","role":"assistant","modelID":"gpt-4o","providerID":"openai","cost":0.0,"tokens":{"input":40,"output":5,"reasoning":0,"cache":{"read":0,"write":0}},"time":{"created":1733011203000}}"#, + make_trae_message( + "session-other", + 1_700_000_001_000, + Some("trae:session-other:1_700_000_001"), + 0.1, ), ]; - for (row_id, data) in sqlite_rows { - opencode_db - .execute( - "INSERT INTO message (id, session_id, data) VALUES (?1, ?2, ?3)", - rusqlite::params![row_id, "oc-session", data], - ) - .unwrap(); - } - drop(opencode_db); - let opencode_dir = opencode_data_dir.join("storage/message/project-1"); - std::fs::create_dir_all(&opencode_dir).unwrap(); - let json_rows = [ - ( - "a-json-authoritative.json", - r#"{"id":"json-authoritative","sessionID":"oc-session","role":"assistant","modelID":"gpt-4o","providerID":"openai","cost":0.05,"tokens":{"input":10,"output":5,"reasoning":0,"cache":{"read":0,"write":0}},"time":{"created":1733011200000}}"#, - ), - ( - "b-json-estimated.json", - r#"{"id":"sqlite-authoritative","sessionID":"oc-session","role":"assistant","modelID":"gpt-4o","providerID":"openai","cost":0.0,"tokens":{"input":10,"output":5,"reasoning":0,"cache":{"read":0,"write":0}},"time":{"created":1733011201000}}"#, + let deduped = dedupe_latest_trae_messages(messages); + + assert_eq!(deduped.len(), 2); + let stable = deduped + .iter() + .find(|msg| msg.session_id == "session-stable") + .expect("session-stable should remain after dedupe"); + assert_eq!(stable.timestamp, 1_700_000_003_000); + assert_eq!(stable.cost, 0.3); + assert_eq!( + stable.dedup_key.as_deref(), + Some("trae:session-stable:1_700_000_003") + ); + } + + #[test] + fn test_dedupe_latest_trae_messages_tiebreaks_by_dedup_key() { + let messages = vec![ + make_trae_message( + "session-stable", + 1_700_000_010_000, + Some("dedupe-key-a"), + 0.2, ), - ( - "c-json-authoritative.json", - r#"{"id":"both-authoritative","sessionID":"oc-session","role":"assistant","modelID":"gpt-4o","providerID":"openai","cost":0.08,"tokens":{"input":10,"output":5,"reasoning":0,"cache":{"read":0,"write":0}},"time":{"created":1733011202000}}"#, + make_trae_message( + "session-stable", + 1_700_000_010_000, + Some("dedupe-key-z"), + 0.4, ), - ( - "d-json-estimated.json", - r#"{"id":"both-estimated","sessionID":"oc-session","role":"assistant","modelID":"gpt-4o","providerID":"openai","cost":0.0,"tokens":{"input":10,"output":5,"reasoning":0,"cache":{"read":0,"write":0}},"time":{"created":1733011203000}}"#, + make_trae_message( + "session-stable", + 1_700_000_009_000, + Some("dedupe-key-m"), + 0.1, ), ]; - for (name, data) in json_rows { - std::fs::write(opencode_dir.join(name), data).unwrap(); - } - let mut litellm = HashMap::new(); - litellm.insert( - "openai/gpt-4o".to_string(), - pricing::ModelPricing { - input_cost_per_token: Some(0.01), - output_cost_per_token: Some(0.02), - ..Default::default() - }, - ); - let pricing = pricing::PricingService::new(litellm, HashMap::new()); - let clients = vec!["opencode".to_string()]; - let materialized = parse_all_messages_with_pricing_with_env_strategy( - source_home.path().to_str().unwrap(), - &clients, + let deduped = dedupe_latest_trae_messages(messages); + + assert_eq!(deduped.len(), 1); + assert_eq!(deduped[0].timestamp, 1_700_000_010_000); + assert_eq!(deduped[0].dedup_key.as_deref(), Some("dedupe-key-z")); + assert_eq!(deduped[0].cost, 0.4); + } + + #[test] + #[serial_test::serial] + fn test_parse_all_messages_with_pricing_keeps_gateway_message_under_synthetic_filter() { + let cache_home = tempfile::TempDir::new().unwrap(); + let temp_dir = tempfile::TempDir::new().unwrap(); + let _env = opencode_test_env(cache_home.path(), temp_dir.path()); + let message_dir = temp_dir + .path() + .join(".local/share/opencode/storage/message/project-1"); + std::fs::create_dir_all(&message_dir).unwrap(); + std::fs::write( + message_dir.join("msg_001.json"), + r#"{"id":"msg-1","sessionID":"session-1","role":"assistant","modelID":"hf:deepseek-ai/DeepSeek-V3-0324","providerID":"unknown","cost":0,"tokens":{"input":10,"output":5,"reasoning":0,"cache":{"read":0,"write":0}},"time":{"created":1733011200000}}"#, + ) + .unwrap(); + + let pricing = pricing::PricingService::new(HashMap::new(), HashMap::new()); + let messages = parse_all_messages_with_pricing( + temp_dir.path().to_str().unwrap(), + &["synthetic".to_string()], Some(&pricing), - false, - &scanner::ScannerSettings::default(), ); - let mut streamed = Vec::new(); - scan_messages_streaming( - source_home.path().to_str().unwrap(), - &clients, + + assert_eq!(messages.len(), 1); + assert_eq!(messages[0].client, "opencode"); + assert_eq!(messages[0].model_id, "deepseek-v3-0324"); + assert_eq!(messages[0].provider_id, "synthetic"); + } + + #[test] + fn test_parse_local_clients_preserves_gateway_message_client_counts() { + let temp_dir = tempfile::TempDir::new().unwrap(); + let message_dir = temp_dir + .path() + .join(".local/share/opencode/storage/message/project-1"); + std::fs::create_dir_all(&message_dir).unwrap(); + std::fs::write( + message_dir.join("msg_001.json"), + r#"{"id":"msg-1","sessionID":"session-1","role":"assistant","modelID":"accounts/fireworks/models/deepseek-v3-0324","providerID":"fireworks","cost":0,"tokens":{"input":10,"output":5,"reasoning":0,"cache":{"read":0,"write":0}},"time":{"created":1733011200000}}"#, + ) + .unwrap(); + + let parsed = parse_local_clients(LocalParseOptions { + home_dir: Some(temp_dir.path().to_str().unwrap().to_string()), + use_env_roots: false, + clients: Some(vec!["opencode".to_string(), "synthetic".to_string()]), + since: None, + until: None, + year: None, + scanner_settings: scanner::ScannerSettings::default(), + modified_after: None, + }) + .unwrap(); + + assert_eq!(parsed.counts.get(ClientId::OpenCode), 1); + assert_eq!(parsed.messages.len(), 1); + assert_eq!(parsed.messages[0].client, "opencode"); + assert_eq!(parsed.messages[0].model_id, "deepseek-v3-0324"); + // opencode canonicalizes the raw "fireworks" gateway id to "fireworks_ai" (#760). + assert_eq!(parsed.messages[0].provider_id, "fireworks_ai"); + } + + #[test] + #[serial_test::serial] + fn test_parse_all_messages_fireworks_provider_kept_under_synthetic_only_filter() { + let cache_home = tempfile::TempDir::new().unwrap(); + let temp_dir = tempfile::TempDir::new().unwrap(); + let _env = opencode_test_env(cache_home.path(), temp_dir.path()); + let message_dir = temp_dir + .path() + .join(".local/share/opencode/storage/message/project-1"); + std::fs::create_dir_all(&message_dir).unwrap(); + std::fs::write( + message_dir.join("msg_001.json"), + r#"{"id":"msg-1","sessionID":"session-1","role":"assistant","modelID":"accounts/fireworks/models/deepseek-v3-0324","providerID":"fireworks","cost":0.1,"tokens":{"input":10,"output":5,"reasoning":0,"cache":{"read":0,"write":0}},"time":{"created":1733011200000}}"#, + ) + .unwrap(); + + let pricing = pricing::PricingService::new(HashMap::new(), HashMap::new()); + let messages = parse_all_messages_with_pricing( + temp_dir.path().to_str().unwrap(), + &["synthetic".to_string()], Some(&pricing), - false, - &scanner::ScannerSettings::default(), - &|_| true, - &mut |message| streamed.push(message.clone()), ); - let summarize = |messages: Vec| { - let mut rows: Vec<(String, f64, CostSource)> = messages - .into_iter() - .map(|message| { - ( - message.dedup_key.unwrap_or_default(), - message.cost, - message.cost_source, - ) - }) - .collect(); - rows.sort_by(|left, right| left.0.cmp(&right.0)); - rows - }; - let materialized = summarize(materialized); - let streamed = summarize(streamed); - assert_eq!(materialized, streamed); assert_eq!( - materialized, - vec![ - ( - "both-authoritative".to_string(), - 0.07, - CostSource::ProviderReported - ), - ("both-estimated".to_string(), 0.5, CostSource::Estimated), - ( - "json-authoritative".to_string(), - 0.05, - CostSource::ProviderReported - ), - ( - "sqlite-authoritative".to_string(), - 0.06, - CostSource::ProviderReported - ), - ] + messages.len(), + 1, + "fireworks gateway message must not be dropped when filtering for synthetic" ); + assert_eq!(messages[0].client, "opencode"); + assert_eq!(messages[0].model_id, "deepseek-v3-0324"); + // opencode canonicalizes the raw "fireworks" gateway id to "fireworks_ai" (#760). + assert_eq!(messages[0].provider_id, "fireworks_ai"); } #[test] - fn test_apply_pricing_if_available_applies_zed_hosted_markup() { - let mut litellm = HashMap::new(); - litellm.insert( - "claude-sonnet-4-5".into(), - pricing::ModelPricing { - input_cost_per_token: Some(0.001), - output_cost_per_token: Some(0.002), - ..Default::default() - }, - ); - let pricing = pricing::PricingService::new(litellm, HashMap::new()); - - let mut msg = UnifiedMessage::new( - "zed", - "claude-sonnet-4-5", - crate::sessions::zed::ZED_HOSTED_PROVIDER, - "session-1", - 1_733_011_200_000, - TokenBreakdown { - input: 10, - output: 5, - cache_read: 0, - cache_write: 0, - reasoning: 0, - }, - 0.0, - ); + fn test_parse_local_clients_fireworks_provider_kept_under_synthetic_only_filter() { + let temp_dir = tempfile::TempDir::new().unwrap(); + let message_dir = temp_dir + .path() + .join(".local/share/opencode/storage/message/project-1"); + std::fs::create_dir_all(&message_dir).unwrap(); + std::fs::write( + message_dir.join("msg_001.json"), + r#"{"id":"msg-1","sessionID":"session-1","role":"assistant","modelID":"accounts/fireworks/models/deepseek-v3-0324","providerID":"fireworks","cost":0.1,"tokens":{"input":10,"output":5,"reasoning":0,"cache":{"read":0,"write":0}},"time":{"created":1733011200000}}"#, + ) + .unwrap(); - apply_pricing_if_available(&mut msg, Some(&pricing)); + let parsed = parse_local_clients(LocalParseOptions { + home_dir: Some(temp_dir.path().to_str().unwrap().to_string()), + use_env_roots: false, + clients: Some(vec!["synthetic".to_string()]), + since: None, + until: None, + year: None, + scanner_settings: scanner::ScannerSettings::default(), + modified_after: None, + }) + .unwrap(); - assert!((msg.cost - 0.022).abs() < 1e-12); + assert_eq!( + parsed.messages.len(), + 1, + "fireworks gateway message must not be dropped when filtering for synthetic only" + ); + assert_eq!(parsed.messages[0].client, "opencode"); + assert_eq!(parsed.messages[0].model_id, "deepseek-v3-0324"); + // opencode canonicalizes the raw "fireworks" gateway id to "fireworks_ai" (#760). + assert_eq!(parsed.messages[0].provider_id, "fireworks_ai"); } #[test] - fn test_apply_pricing_if_available_skips_zed_markup_for_non_zed_client() { - // Non-zed client with provider_id "zed.dev" must not receive the +10% - // markup. The multiplier is gated on (client == "zed" AND provider). - let mut litellm = HashMap::new(); - litellm.insert( - "claude-sonnet-4-5".into(), - pricing::ModelPricing { - input_cost_per_token: Some(0.001), - output_cost_per_token: Some(0.002), - ..Default::default() - }, - ); - let pricing = pricing::PricingService::new(litellm, HashMap::new()); - - let mut msg = UnifiedMessage::new( - "claudecode", - "claude-sonnet-4-5", - crate::sessions::zed::ZED_HOSTED_PROVIDER, - "session-1", - 1_733_011_200_000, - TokenBreakdown { - input: 10, - output: 5, - cache_read: 0, - cache_write: 0, - reasoning: 0, - }, - 0.0, - ); + fn test_parse_local_clients_honors_scanner_settings_opencode_db_paths() { + // Regression guard: `parse_local_clients` used to call + // `scan_all_clients_with_env_strategy`, which silently dropped + // `options.scanner_settings`. Users with + // `scanner.opencodeDbPaths` pointing at an OPENCODE_DB outside the + // XDG data dir would see no rows through the clients/wrapped + // command paths even though model/monthly/graph reports honored + // the same config. + let temp_dir = tempfile::TempDir::new().unwrap(); + // Deliberately do not create ~/.local/share/opencode so nothing + // is auto-discoverable; the only db the scanner can find must + // come from `scanner_settings`. + let outside_dir = temp_dir.path().join("elsewhere"); + std::fs::create_dir_all(&outside_dir).unwrap(); + let external_db = outside_dir.join("opencode.db"); - apply_pricing_if_available(&mut msg, Some(&pricing)); + let conn = rusqlite::Connection::open(&external_db).unwrap(); + conn.execute_batch( + "PRAGMA journal_mode=WAL; + CREATE TABLE message ( + id TEXT PRIMARY KEY, + session_id TEXT NOT NULL, + data TEXT NOT NULL + );", + ) + .unwrap(); + conn.execute( + "INSERT INTO message (id, session_id, data) VALUES (?1, ?2, ?3)", + rusqlite::params![ + "ext-msg-1", + "ext-session", + r#"{ + "role": "assistant", + "modelID": "claude-sonnet-4", + "providerID": "anthropic", + "tokens": { "input": 42, "output": 7, "reasoning": 0, "cache": { "read": 0, "write": 0 } }, + "time": { "created": 1700000000000.0 } + }"# + ], + ) + .unwrap(); + drop(conn); - // 10 * 0.001 + 5 * 0.002 = 0.020, no markup. - assert!((msg.cost - 0.020).abs() < 1e-12); - } + // Without scanner_settings: no rows (nothing auto-discoverable). + let parsed_default = parse_local_clients(LocalParseOptions { + home_dir: Some(temp_dir.path().to_str().unwrap().to_string()), + use_env_roots: false, + clients: Some(vec!["opencode".to_string()]), + since: None, + until: None, + year: None, + scanner_settings: scanner::ScannerSettings::default(), + modified_after: None, + }) + .unwrap(); + assert_eq!(parsed_default.counts.get(ClientId::OpenCode), 0); + assert!(parsed_default.messages.is_empty()); - #[test] - fn test_apply_pricing_if_available_skips_zed_markup_for_byok_provider() { - // A Zed message whose provider_id is the upstream provider directly - // (BYOK / non-hosted path) must not be marked up — the user is paying - // the upstream API directly, not through Zed. - let mut litellm = HashMap::new(); - litellm.insert( - "claude-sonnet-4-5".into(), - pricing::ModelPricing { - input_cost_per_token: Some(0.001), - output_cost_per_token: Some(0.002), + // With scanner_settings pointing at the external db: the user + // row must show up. + let parsed_with_settings = parse_local_clients(LocalParseOptions { + home_dir: Some(temp_dir.path().to_str().unwrap().to_string()), + use_env_roots: false, + clients: Some(vec!["opencode".to_string()]), + since: None, + until: None, + year: None, + scanner_settings: scanner::ScannerSettings { + opencode_db_paths: vec![external_db.clone()], ..Default::default() }, + modified_after: None, + }) + .unwrap(); + assert_eq!( + parsed_with_settings.counts.get(ClientId::OpenCode), + 1, + "scanner.opencodeDbPaths must reach the parse_local_clients path" ); - let pricing = pricing::PricingService::new(litellm, HashMap::new()); - - let mut msg = UnifiedMessage::new( - "zed", - "claude-sonnet-4-5", - "anthropic", - "session-1", - 1_733_011_200_000, - TokenBreakdown { - input: 10, - output: 5, - cache_read: 0, - cache_write: 0, - reasoning: 0, - }, - 0.0, - ); - - apply_pricing_if_available(&mut msg, Some(&pricing)); - - assert!((msg.cost - 0.020).abs() < 1e-12); + assert_eq!(parsed_with_settings.messages.len(), 1); + assert_eq!(parsed_with_settings.messages[0].client, "opencode"); + assert_eq!(parsed_with_settings.messages[0].model_id, "claude-sonnet-4"); } #[test] - fn test_apply_pricing_if_available_uses_reasoning_for_gemini() { - let mut litellm = HashMap::new(); - litellm.insert( - "gemini-2.5-pro".into(), - pricing::ModelPricing { - input_cost_per_token: Some(0.001), - output_cost_per_token: Some(0.002), - ..Default::default() - }, - ); - let pricing = pricing::PricingService::new(litellm, HashMap::new()); - - let mut msg = UnifiedMessage::new( - "gemini", - "gemini-2.5-pro", - "google", - "session-1", - 1_733_011_200_000, - TokenBreakdown { - input: 10, - output: 5, - cache_read: 0, - cache_write: 0, - reasoning: 7, - }, - 0.0, + fn test_parse_local_clients_honors_scanner_extra_scan_paths_for_hermes_profile_db() { + let temp_dir = tempfile::TempDir::new().unwrap(); + let profile_dir = temp_dir.path().join("external-hermes/director_planning"); + std::fs::create_dir_all(&profile_dir).unwrap(); + let profile_db = profile_dir.join("state.db"); + let conn = create_hermes_sqlite_db(&profile_db); + insert_hermes_session( + &conn, + "hermes-extra-session", + "claude-sonnet-4", + 2, + 100, + 25, + 0.07, ); + drop(conn); - apply_pricing_if_available(&mut msg, Some(&pricing)); - - assert_eq!(msg.cost, 0.034); - } + let parsed_default = parse_local_clients(LocalParseOptions { + home_dir: Some(temp_dir.path().to_str().unwrap().to_string()), + use_env_roots: false, + clients: Some(vec!["hermes".to_string()]), + since: None, + until: None, + year: None, + scanner_settings: scanner::ScannerSettings::default(), + modified_after: None, + }) + .unwrap(); + assert_eq!(parsed_default.counts.get(ClientId::Hermes), 0); + assert!(parsed_default.messages.is_empty()); - #[test] - fn test_apply_pricing_if_available_uses_cache_read_pricing_for_gemini() { - let mut litellm = HashMap::new(); - litellm.insert( - "gemini-2.5-pro".into(), - pricing::ModelPricing { - input_cost_per_token: Some(0.001), - output_cost_per_token: Some(0.002), - cache_read_input_token_cost: Some(0.0001), + let mut extra_scan_paths = std::collections::BTreeMap::new(); + extra_scan_paths.insert("hermes".to_string(), vec![profile_dir]); + let parsed_with_settings = parse_local_clients(LocalParseOptions { + home_dir: Some(temp_dir.path().to_str().unwrap().to_string()), + use_env_roots: false, + clients: Some(vec!["hermes".to_string()]), + since: None, + until: None, + year: None, + scanner_settings: scanner::ScannerSettings { + extra_scan_paths, ..Default::default() }, - ); - let pricing = pricing::PricingService::new(litellm, HashMap::new()); + modified_after: None, + }) + .unwrap(); - let mut msg = UnifiedMessage::new( - "gemini", - "gemini-2.5-pro", - "google", - "session-1", - 1_733_011_200_000, - TokenBreakdown { - input: 10, - output: 5, - cache_read: 7, - cache_write: 0, - reasoning: 3, - }, - 0.0, + assert_eq!(parsed_with_settings.counts.get(ClientId::Hermes), 2); + assert_eq!(parsed_with_settings.messages.len(), 1); + assert_eq!(parsed_with_settings.messages[0].client, "hermes"); + assert_eq!( + parsed_with_settings.messages[0].agent.as_deref(), + Some("Hermes Agent") ); - - apply_pricing_if_available(&mut msg, Some(&pricing)); - - assert_eq!(msg.cost, 0.0267); + assert_eq!( + parsed_with_settings.messages[0].session_id, + "hermes-extra-session" + ); + assert_eq!(parsed_with_settings.messages[0].model_id, "claude-sonnet-4"); + assert_eq!(parsed_with_settings.messages[0].input, 100); + assert_eq!(parsed_with_settings.messages[0].output, 25); } #[test] - fn test_apply_pricing_if_available_uses_market_rate_for_free_variant() { - let mut openrouter = HashMap::new(); - openrouter.insert( - "z-ai/glm-4.7".into(), - pricing::ModelPricing { - input_cost_per_token: Some(0.001), - output_cost_per_token: Some(0.002), - ..Default::default() - }, + #[serial_test::serial] + fn test_auto_discovered_hermes_profile_reaches_all_consumers() { + let source_home = tempfile::TempDir::new().unwrap(); + let materialized_cache = tempfile::TempDir::new().unwrap(); + let streaming_cache = tempfile::TempDir::new().unwrap(); + let count_cache = tempfile::TempDir::new().unwrap(); + let profile_dir = source_home.path().join(".hermes/profiles/research"); + std::fs::create_dir_all(&profile_dir).unwrap(); + let profile_db = profile_dir.join("state.db"); + let conn = create_hermes_sqlite_db(&profile_db); + insert_hermes_session( + &conn, + "hermes-auto-profile", + "claude-sonnet-4", + 2, + 100, + 25, + 0.07, ); - let pricing = pricing::PricingService::new(HashMap::new(), openrouter); + drop(conn); - let mut msg = UnifiedMessage::new( - "opencode", - "glm-4.7-free", - "modal", - "session-1", - 1_733_011_200_000, - TokenBreakdown { - input: 10, - output: 5, - cache_read: 0, - cache_write: 0, - reasoning: 0, - }, - 0.0, - ); + let home = source_home.path().to_str().unwrap().to_string(); + let clients = vec!["hermes".to_string()]; + let materialized = with_isolated_tokscale_cache(materialized_cache.path(), || { + parse_all_messages_with_pricing_with_env_strategy( + &home, + &clients, + None, + false, + &scanner::ScannerSettings::default(), + ) + }); + assert_eq!(materialized.len(), 1); + assert_eq!(materialized[0].session_id, "hermes-auto-profile"); + assert_eq!(materialized[0].tokens.input, 100); + assert_eq!(materialized[0].tokens.output, 25); - apply_pricing_if_available(&mut msg, Some(&pricing)); + let runtime = tokio::runtime::Builder::new_current_thread() + .enable_all() + .build() + .unwrap(); + let streaming = with_isolated_tokscale_cache(streaming_cache.path(), || { + runtime + .block_on(get_model_report(ReportOptions { + home_dir: Some(home.clone()), + use_env_roots: false, + clients: Some(clients.clone()), + ..Default::default() + })) + .unwrap() + }); + assert_eq!(streaming.total_messages, 2); + assert_eq!(streaming.total_input, 100); + assert_eq!(streaming.total_output, 25); - assert_eq!(msg.cost, 0.02); + let future_ms = std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .unwrap() + .as_millis() as u64 + + 3_600_000; + let counted = with_isolated_tokscale_cache(count_cache.path(), || { + parse_local_clients(LocalParseOptions { + home_dir: Some(home), + use_env_roots: false, + clients: Some(clients), + since: None, + until: None, + year: None, + scanner_settings: scanner::ScannerSettings::default(), + modified_after: Some(future_ms), + }) + .unwrap() + }); + assert_eq!(counted.counts.get(ClientId::Hermes), 2); + assert_eq!(counted.messages.len(), 1); + assert_eq!(counted.messages[0].session_id, "hermes-auto-profile"); + assert_eq!(counted.messages[0].input, 100); + assert_eq!(counted.messages[0].output, 25); } #[test] - fn test_apply_pricing_if_available_prefers_provider_aware_match() { - let mut litellm = HashMap::new(); - litellm.insert( - "xai/grok-code-fast-1-0825".into(), - pricing::ModelPricing { - input_cost_per_token: Some(0.001), - output_cost_per_token: Some(0.002), - ..Default::default() - }, - ); - litellm.insert( - "azure_ai/grok-code-fast-1".into(), - pricing::ModelPricing { - input_cost_per_token: Some(0.01), - output_cost_per_token: Some(0.02), - ..Default::default() - }, + fn test_modified_after_never_prunes_hermes_dbs_from_extra_scan_paths() { + // SQLite WAL writes may leave the main db file's mtime untouched, so + // `modified_after` must not prune Hermes/Zed dbs even when they come + // from user scan roots (the `files` lanes) rather than the default + // single-db path. A threshold in the future would prune any mtime. + let temp_dir = tempfile::TempDir::new().unwrap(); + let profile_dir = temp_dir.path().join("external-hermes/director_planning"); + std::fs::create_dir_all(&profile_dir).unwrap(); + let profile_db = profile_dir.join("state.db"); + let conn = create_hermes_sqlite_db(&profile_db); + insert_hermes_session( + &conn, + "hermes-wal-session", + "claude-sonnet-4", + 1, + 50, + 10, + 0.03, ); - let pricing = pricing::PricingService::new(litellm, HashMap::new()); + drop(conn); - let mut msg = UnifiedMessage::new( - "opencode", - "grok-code", - "azure", - "session-1", - 1_733_011_200_000, - TokenBreakdown { - input: 10, - output: 5, - cache_read: 0, - cache_write: 0, - reasoning: 0, + let future_ms = std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .unwrap() + .as_millis() as u64 + + 3_600_000; + let mut extra_scan_paths = std::collections::BTreeMap::new(); + extra_scan_paths.insert("hermes".to_string(), vec![profile_dir]); + let parsed = parse_local_clients(LocalParseOptions { + home_dir: Some(temp_dir.path().to_str().unwrap().to_string()), + use_env_roots: false, + clients: Some(vec!["hermes".to_string()]), + since: None, + until: None, + year: None, + scanner_settings: scanner::ScannerSettings { + extra_scan_paths, + ..Default::default() }, - 0.0, - ); - - apply_pricing_if_available(&mut msg, Some(&pricing)); + modified_after: Some(future_ms), + }) + .unwrap(); - assert_eq!(msg.cost, 0.2); + assert_eq!(parsed.counts.get(ClientId::Hermes), 1); + assert_eq!(parsed.messages.len(), 1); + assert_eq!(parsed.messages[0].session_id, "hermes-wal-session"); } #[test] - fn test_apply_pricing_if_available_uses_nested_reseller_exact_match() { - let mut litellm = HashMap::new(); - litellm.insert( - "gpt-4".into(), - pricing::ModelPricing { - input_cost_per_token: Some(0.001), - output_cost_per_token: Some(0.002), - ..Default::default() - }, - ); - litellm.insert( - "azure/openai/gpt-4".into(), - pricing::ModelPricing { - input_cost_per_token: Some(0.01), - output_cost_per_token: Some(0.02), - ..Default::default() - }, - ); - let pricing = pricing::PricingService::new(litellm, HashMap::new()); - - let mut msg = UnifiedMessage::new( - "opencode", - "gpt-4", - "azure", - "session-1", - 1_733_011_200_000, - TokenBreakdown { - input: 10, - output: 5, - cache_read: 0, - cache_write: 0, - reasoning: 0, - }, - 0.0, - ); + fn test_modified_after_never_prunes_antigravity_cli_dbs() { + // Antigravity CLI conversation `.db` files arrive via the generic + // `files` lane (a `*.db` glob), but they are SQLite — WAL writes may + // leave the main db mtime untouched, so they must be exempt from mtime + // pruning like Hermes/Zed. A plain-file client with the same old mtime + // is still pruned (control). + let temp_dir = tempfile::TempDir::new().unwrap(); + let cli_db = temp_dir.path().join("conv.db"); + std::fs::File::create(&cli_db).unwrap(); + let claude_log = temp_dir.path().join("session.jsonl"); + std::fs::File::create(&claude_log).unwrap(); - apply_pricing_if_available(&mut msg, Some(&pricing)); + let mut scan_result = scanner::ScanResult::default(); + scan_result + .get_mut(ClientId::AntigravityCli) + .push(cli_db.clone()); + scan_result + .get_mut(ClientId::Claude) + .push(claude_log.clone()); - assert_eq!(msg.cost, 0.2); - } + // A threshold in the future would prune any real on-disk mtime. + let future_ms = std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .unwrap() + .as_millis() as u64 + + 3_600_000; + crate::prune_scan_result_by_mtime(&mut scan_result, future_ms); - #[test] - fn test_apply_pricing_if_available_keeps_scoped_fireworks_cost_without_exact_pricing() { - let mut litellm = HashMap::new(); - litellm.insert( - "fireworks_ai/accounts/fireworks/models/deepseek-r1-0528-distill-qwen3-8b".into(), - pricing::ModelPricing { - input_cost_per_token: Some(0.0000002), - output_cost_per_token: Some(0.0000002), - ..Default::default() - }, + assert_eq!( + scan_result.get(ClientId::AntigravityCli), + std::slice::from_ref(&cli_db), + "Antigravity CLI .db (a WAL-mode SQLite source) must survive mtime pruning" ); - - let mut openrouter = HashMap::new(); - openrouter.insert( - "deepseek/deepseek-v4-pro".into(), - pricing::ModelPricing { - input_cost_per_token: Some(0.000001), - output_cost_per_token: Some(0.000002), - ..Default::default() - }, + assert!( + scan_result.get(ClientId::Claude).is_empty(), + "a plain-file client's stale log is still pruned" ); + } - let pricing = pricing::PricingService::new(litellm, openrouter); - let mut msg = UnifiedMessage::new( - "opencode", - "accounts/fireworks/models/deepseek-v4-pro", - "fireworks", - "session-1", - 1_733_011_200_000, - TokenBreakdown { - input: 10, - output: 5, - cache_read: 0, - cache_write: 0, - reasoning: 0, - }, - 0.123, - ); + /// Write a minimal Antigravity CLI conversation DB (one priced + /// `gen_metadata` row carrying `response_id`). The `trajectory_metadata_blob` + /// table is omitted on purpose — the parser tolerates its absence and falls + /// back to the file mtime for the timestamp. + fn write_antigravity_cli_db( + conversations_dir: &std::path::Path, + file_stem: &str, + response_id: &str, + ) { + fn encode_varint(mut value: u64) -> Vec { + let mut out = Vec::new(); + loop { + let mut byte = (value & 0x7f) as u8; + value >>= 7; + if value != 0 { + byte |= 0x80; + } + out.push(byte); + if value == 0 { + break; + } + } + out + } + fn enc_varint(field: u64, value: u64) -> Vec { + let mut out = encode_varint(field << 3); + out.extend(encode_varint(value)); + out + } + fn enc_len(field: u64, payload: &[u8]) -> Vec { + let mut out = encode_varint((field << 3) | 2); + out.extend(encode_varint(payload.len() as u64)); + out.extend_from_slice(payload); + out + } - apply_pricing_if_available(&mut msg, Some(&pricing)); + let mut usage = Vec::new(); + usage.extend(enc_varint(2, 500)); // new input + usage.extend(enc_varint(9, 300)); // output + usage.extend(enc_len(11, response_id.as_bytes())); // responseId + let mut chat_model = Vec::new(); + chat_model.extend(enc_len(4, &usage)); + chat_model.extend(enc_len(19, b"gemini-3-flash-a")); + let gen_blob = enc_len(1, &chat_model); - assert_eq!(msg.cost, 0.123); + std::fs::create_dir_all(conversations_dir).unwrap(); + let path = conversations_dir.join(format!("{file_stem}.db")); + let conn = rusqlite::Connection::open(&path).unwrap(); + conn.execute_batch("CREATE TABLE gen_metadata (idx integer, data blob, size integer);") + .unwrap(); + conn.execute( + "INSERT INTO gen_metadata (idx, data, size) VALUES (0, ?1, 0)", + rusqlite::params![gen_blob], + ) + .unwrap(); } + // Two independent Antigravity CLI conversation DBs that reuse the same + // responseId must both survive the streaming report path. responseIds are + // unique only within a conversation, so the cross-file dedup gate is + // namespaced by session; with a bare-responseId key (the pre-fix behaviour) + // the second conversation is silently dropped and this fails (count == 1). #[test] - fn test_apply_pricing_if_available_prefers_provider_specific_exact_match_over_plain_exact() { - let mut litellm = HashMap::new(); - litellm.insert( - "gemini-2.5-pro".into(), - pricing::ModelPricing { - input_cost_per_token: Some(0.001), - output_cost_per_token: Some(0.002), - cache_creation_input_token_cost: None, - ..Default::default() - }, - ); - - let mut openrouter = HashMap::new(); - openrouter.insert( - "google/gemini-2.5-pro".into(), - pricing::ModelPricing { - input_cost_per_token: Some(0.001), - output_cost_per_token: Some(0.002), - cache_creation_input_token_cost: Some(0.01), - ..Default::default() - }, - ); - - let pricing = pricing::PricingService::new(litellm, openrouter); + #[serial_test::serial] + fn test_streaming_antigravity_cli_keeps_colliding_response_ids_across_conversations() { + let cache_home = tempfile::TempDir::new().unwrap(); + let source_home = tempfile::TempDir::new().unwrap(); + let _env = EnvGuard::set(&[ + ("HOME", cache_home.path().as_os_str()), + ("TOKSCALE_CONFIG_DIR", cache_home.path().as_os_str()), + ]); - let mut msg = UnifiedMessage::new( - "opencode", - "gemini-2.5-pro", - "google", - "session-1", - 1_733_011_200_000, - TokenBreakdown { - input: 10, - output: 5, - cache_read: 0, - cache_write: 3, - reasoning: 0, - }, - 0.0, - ); + { + let conversations_dir = source_home + .path() + .join(".gemini/antigravity-cli/conversations"); + write_antigravity_cli_db(&conversations_dir, "conv-aaa", "SHARED"); + write_antigravity_cli_db(&conversations_dir, "conv-bbb", "SHARED"); - apply_pricing_if_available(&mut msg, Some(&pricing)); + let mut sessions: Vec = Vec::new(); + scan_messages_streaming( + source_home.path().to_str().unwrap(), + &["antigravity-cli".to_string()], + None, + false, + &scanner::ScannerSettings::default(), + &|_m: &UnifiedMessage| true, + &mut |m: &UnifiedMessage| sessions.push(m.session_id.clone()), + ); - assert_eq!(msg.cost, 0.05); + sessions.sort(); + assert_eq!( + sessions, + vec!["conv-aaa".to_string(), "conv-bbb".to_string()], + "both conversations reusing responseId \"SHARED\" must survive" + ); + } } + // jcode (`~/.jcode/sessions/session_*.json`) must be discovered by the + // generic scanner (EnvVar JCODE_HOME / .jcode root, `session_*.json` glob) + // and flow through the streaming lane with its authoritative per-message + // token_usage. #[test] - fn test_apply_pricing_if_available_normalizes_openai_codex_provider() { - let mut litellm = HashMap::new(); - litellm.insert( - "openai/gpt-5.2-preview".into(), - pricing::ModelPricing { - input_cost_per_token: Some(0.01), - output_cost_per_token: Some(0.02), - ..Default::default() - }, - ); - litellm.insert( - "google/gpt-5.2-preview-max".into(), - pricing::ModelPricing { - input_cost_per_token: Some(0.1), - output_cost_per_token: Some(0.2), - ..Default::default() - }, - ); - let pricing = pricing::PricingService::new(litellm, HashMap::new()); + #[serial_test::serial] + fn test_streaming_jcode_flows_through_lane() { + let cache_home = tempfile::TempDir::new().unwrap(); + let source_home = tempfile::TempDir::new().unwrap(); + let _env = EnvGuard::set(&[ + ("HOME", cache_home.path().as_os_str()), + ("TOKSCALE_CONFIG_DIR", cache_home.path().as_os_str()), + ]); - let mut msg = UnifiedMessage::new( - "openclaw", - "gpt-5.2", - "openai-codex", - "session-1", - 1_733_011_200_000, - TokenBreakdown { - input: 10, - output: 5, - cache_read: 0, - cache_write: 0, - reasoning: 0, - }, - 0.0, - ); + { + let sessions_dir = source_home.path().join(".jcode/sessions"); + std::fs::create_dir_all(&sessions_dir).unwrap(); + std::fs::write( + sessions_dir.join("session_test.json"), + r#"{"id":"session_test","provider_key":"cliproxyapi","model":"claude-sonnet-4","working_dir":"/x","messages":[{"id":"u1","role":"user","timestamp":"2026-06-16T12:00:00Z"},{"id":"a1","role":"assistant","timestamp":"2026-06-16T12:00:01Z","token_usage":{"input_tokens":1200,"output_tokens":300}}]}"#, + ) + .unwrap(); - apply_pricing_if_available(&mut msg, Some(&pricing)); + let mut input_sum = 0i64; + let mut count = 0usize; + scan_messages_streaming( + source_home.path().to_str().unwrap(), + &["jcode".to_string()], + None, + false, + &scanner::ScannerSettings::default(), + &|_m: &UnifiedMessage| true, + &mut |m: &UnifiedMessage| { + input_sum += m.tokens.input; + count += 1; + }, + ); - assert_eq!(msg.cost, 0.2); + assert_eq!(count, 1, "the jcode assistant message must flow through the streaming lane"); + assert_eq!(input_sum, 1200); + } } #[test] - fn test_apply_pricing_if_available_prices_claude_code_gpt_5_3_codex() { - let pricing = pricing::PricingService::new(HashMap::new(), HashMap::new()); - - let mut msg = UnifiedMessage::new( - "claude", - "gpt-5.3-codex", - "openai", - "session-1", - 1_776_000_000_000, - TokenBreakdown { - input: 1_000_000, - output: 100_000, - cache_read: 50_000, - cache_write: 0, - reasoning: 0, - }, - 0.0, - ); + #[serial_test::serial] + fn m21_sources_keep_materialized_streaming_count_and_report_parity() { + let source_home = tempfile::TempDir::new().unwrap(); + let materialized_cache = tempfile::TempDir::new().unwrap(); + let streaming_cache = tempfile::TempDir::new().unwrap(); + let report_cache = tempfile::TempDir::new().unwrap(); + let clients = vec![ + "kimi".to_string(), + "junie".to_string(), + "opencodereview".to_string(), + ]; - apply_pricing_if_available(&mut msg, Some(&pricing)); + let legacy_kimi = source_home + .path() + .join(".kimi/sessions/group/legacy-session/wire.jsonl"); + std::fs::create_dir_all(legacy_kimi.parent().unwrap()).unwrap(); + std::fs::write( + &legacy_kimi, + r#"{"timestamp":1770983410.0,"message":{"type":"StatusUpdate","payload":{"token_usage":{"input_other":7,"output":3},"message_id":"legacy-1"}}}"#, + ) + .unwrap(); - let expected = 1.75 + 1.4 + 0.00875; - assert!((msg.cost - expected).abs() < 1e-12); - } + let kimi_code = source_home + .path() + .join(".kimi-code/sessions/workspace/code-session/agents/main/wire.jsonl"); + std::fs::create_dir_all(kimi_code.parent().unwrap()).unwrap(); + let code_a = r#"{"type":"usage.record","model":"kimi-code/kimi-for-coding","usage":{"inputOther":100,"output":50,"inputCacheRead":10,"inputCacheCreation":0},"usageScope":"turn","time":1770983420000,"turnId":"turn-a"}"#; + let code_b = r#"{"type":"usage.record","model":"kimi-code/kimi-for-coding","usage":{"inputOther":100,"output":50,"inputCacheRead":10,"inputCacheCreation":0},"usageScope":"turn","time":1770983420001,"turnId":"turn-b"}"#; + std::fs::write(&kimi_code, format!("{code_a}\n{code_a}\n{code_b}\n")).unwrap(); + let kimi_code_replay = source_home + .path() + .join(".kimi-code/sessions/workspace/code-session/agents/reviewer/wire.jsonl"); + std::fs::create_dir_all(kimi_code_replay.parent().unwrap()).unwrap(); + std::fs::write(&kimi_code_replay, format!("{code_a}\n")).unwrap(); - #[test] - fn test_apply_pricing_if_available_prices_claude_code_minimax_model() { - let mut litellm = HashMap::new(); - litellm.insert( - "minimax/minimax-m2.1".into(), - pricing::ModelPricing { - input_cost_per_token: Some(0.01), - output_cost_per_token: Some(0.02), - ..Default::default() - }, - ); - let pricing = pricing::PricingService::new(litellm, HashMap::new()); + let junie = source_home + .path() + .join(".junie/sessions/session-260213-120000/events.jsonl"); + std::fs::create_dir_all(junie.parent().unwrap()).unwrap(); + let write_junie = |input: i64, cost: f64| { + let usage = serde_json::json!({ + "timestampMs": 1_770_983_430_000_i64, + "event": { + "agentEvent": { + "kind": "LlmResponseMetadataEvent", + "agent": { "name": "reviewer" }, + "modelUsage": [{ + "model": "gpt-5", + "provider": "openai", + "inputTokens": input, + "outputTokens": 50, + "time": 2_000, + "cost": cost + }] + } + } + }); + std::fs::write( + &junie, + format!("{}\n{usage}\n", r#"{"kind":"UserPromptEvent"}"#), + ) + .unwrap(); + }; + write_junie(100, 0.125); - let mut msg = UnifiedMessage::new( - "claude", - "MiniMax-M2.1", - "minimax", - "session-1", - 1_776_000_000_000, - TokenBreakdown { - input: 10, - output: 5, - cache_read: 0, - cache_write: 0, - reasoning: 0, - }, - 0.0, + let review = source_home + .path() + .join(".opencodereview/sessions/repo/review-session.jsonl"); + std::fs::create_dir_all(review.parent().unwrap()).unwrap(); + let review_contents = concat!( + r#"{"type":"session_start","cwd":"/work/repo"}"#, + "\n", + r#"{"type":"llm_response","timestamp":"2026-02-13T12:00:40Z","model":"gpt-4o","duration_ms":1500,"usage":{"prompt_tokens":20,"completion_tokens":5,"cache_read_tokens":1,"cache_write_tokens":2}}"#, + "\n" ); + std::fs::write(&review, review_contents).unwrap(); - apply_pricing_if_available(&mut msg, Some(&pricing)); - - assert_eq!(msg.cost, 0.2); - } + let home = source_home.path().to_string_lossy().into_owned(); + let run_materialized = || { + with_isolated_tokscale_cache(materialized_cache.path(), || { + let mut messages = parse_all_messages_with_pricing_with_env_strategy( + &home, + &clients, + None, + false, + &scanner::ScannerSettings::default(), + ); + messages.sort_by(|left, right| { + (&left.client, &left.session_id, &left.dedup_key).cmp(&( + &right.client, + &right.session_id, + &right.dedup_key, + )) + }); + messages + }) + }; + let run_streaming = || { + with_isolated_tokscale_cache(streaming_cache.path(), || { + let mut messages = Vec::new(); + scan_messages_streaming( + &home, + &clients, + None, + false, + &scanner::ScannerSettings::default(), + &|_| true, + &mut |message| messages.push(message.clone()), + ); + messages.sort_by(|left, right| { + (&left.client, &left.session_id, &left.dedup_key).cmp(&( + &right.client, + &right.session_id, + &right.dedup_key, + )) + }); + messages + }) + }; - #[test] - fn test_apply_pricing_if_available_prices_kimi_k2p6_alias() { - let mut openrouter = HashMap::new(); - openrouter.insert( - "moonshotai/kimi-k2.6".into(), - pricing::ModelPricing { - input_cost_per_token: Some(9.5e-7), - output_cost_per_token: Some(0.000004), - ..Default::default() - }, + let cold_materialized = run_materialized(); + let cold_streaming = run_streaming(); + assert_eq!(cold_materialized, cold_streaming); + assert_eq!(cold_materialized.len(), 5); + assert_eq!( + cold_materialized + .iter() + .filter(|message| message.client == "kimi") + .count(), + 3 ); - let pricing = pricing::PricingService::new(HashMap::new(), openrouter); + let junie_message = cold_materialized + .iter() + .find(|message| message.client == "junie") + .unwrap(); + assert_eq!(junie_message.cost_source, CostSource::ProviderReported); + assert!((junie_message.cost - 0.125).abs() < 1e-9); + assert!(junie_message.is_turn_start); + assert_eq!(junie_message.duration_ms, Some(2_000)); + let review_message = cold_materialized + .iter() + .find(|message| message.client == "opencodereview") + .unwrap(); + assert_eq!(review_message.duration_ms, Some(1_500)); + assert_eq!(review_message.workspace_label.as_deref(), Some("repo")); + assert_eq!(run_materialized(), cold_materialized); + assert_eq!(run_streaming(), cold_streaming); + + write_junie(200, 0.25); + let rewritten_materialized = run_materialized(); + let rewritten_streaming = run_streaming(); + assert_eq!(rewritten_materialized, rewritten_streaming); + let rewritten_junie = rewritten_materialized + .iter() + .find(|message| message.client == "junie") + .unwrap(); + assert_eq!(rewritten_junie.tokens.input, 200); + assert!((rewritten_junie.cost - 0.25).abs() < 1e-9); - let mut msg = UnifiedMessage::new( - "kimi", - "k2p6", - "kimi-for-coding", - "session-1", - 1_776_000_000_000, - TokenBreakdown { - input: 1_000_000, - output: 250_000, - cache_read: 0, - cache_write: 0, - reasoning: 0, - }, - 0.0, - ); + std::fs::remove_file(&review).unwrap(); + assert_eq!(run_materialized().len(), 4); + assert_eq!(run_streaming().len(), 4); + std::fs::write(&review, review_contents).unwrap(); - apply_pricing_if_available(&mut msg, Some(&pricing)); + let counted = parse_local_clients(LocalParseOptions { + home_dir: Some(home.clone()), + use_env_roots: false, + clients: Some(clients.clone()), + ..Default::default() + }) + .unwrap(); + assert_eq!(counted.counts.get(ClientId::Kimi), 3); + assert_eq!(counted.counts.get(ClientId::Junie), 1); + assert_eq!(counted.counts.get(ClientId::OpenCodeReview), 1); + assert_eq!(counted.messages.len(), 5); - let expected = 1_000_000.0 * 9.5e-7 + 250_000.0 * 0.000004; - assert!((msg.cost - expected).abs() < 1e-12); - assert!(msg.cost > 0.0); + with_isolated_tokscale_cache(report_cache.path(), || { + let runtime = tokio::runtime::Builder::new_current_thread() + .enable_all() + .build() + .unwrap(); + let options = ReportOptions { + home_dir: Some(home), + use_env_roots: false, + clients: Some(clients), + ..Default::default() + }; + let model = runtime.block_on(get_model_report(options.clone())).unwrap(); + let monthly = runtime + .block_on(get_monthly_report(options.clone())) + .unwrap(); + let hourly = runtime + .block_on(get_hourly_report(options.clone())) + .unwrap(); + let agents = runtime.block_on(get_agents_report(options)).unwrap(); + assert_eq!(model.total_messages, 5); + assert_eq!(model.total_input, 427); + assert_eq!(model.total_output, 158); + assert_eq!(model.total_cache_read, 21); + assert_eq!(model.total_cache_write, 2); + assert!((model.total_cost - 0.25).abs() < 1e-9); + assert_eq!( + monthly + .entries + .iter() + .map(|entry| entry.message_count) + .sum::(), + 5 + ); + assert_eq!( + hourly + .entries + .iter() + .map(|entry| entry.message_count) + .sum::(), + 5 + ); + assert_eq!(agents.total_messages, 5); + assert_eq!( + agents.entries.iter().map(|entry| entry.input).sum::(), + 427 + ); + }); } + // micode (`$XDG_DATA_HOME/micode/*.db`, WAL-mode SQLite) must be discovered + // via the generic `*.db` glob and flow through the streaming lane, keeping + // its authoritative per-message cost intact (MiMo models are unpriced, so + // apply_pricing leaves the embedded cost alone). #[test] - fn test_select_local_parse_pricing_prefers_fresh_service_for_new_models() { - let mut fresh_litellm = HashMap::new(); - fresh_litellm.insert( - "gpt-5.4".into(), - pricing::ModelPricing { - input_cost_per_token: Some(0.000002), - output_cost_per_token: Some(0.00001), - ..Default::default() - }, - ); - let fresh = Arc::new(pricing::PricingService::new(fresh_litellm, HashMap::new())); - let stale = pricing::PricingService::new(HashMap::new(), HashMap::new()); - let selected = select_local_parse_pricing(Ok(Arc::clone(&fresh)), || Some(stale)).unwrap(); + #[serial_test::serial] + fn test_streaming_micode_flows_with_authoritative_cost() { + let cache_home = tempfile::TempDir::new().unwrap(); + let source_home = tempfile::TempDir::new().unwrap(); + let _env = EnvGuard::set(&[ + ("HOME", cache_home.path().as_os_str()), + ("TOKSCALE_CONFIG_DIR", cache_home.path().as_os_str()), + ]); - let mut msg = UnifiedMessage::new( - "opencode", - "gpt-5.4", - "openai", - "session-1", - 1_733_011_200_000, - TokenBreakdown { - input: 10, - output: 5, - cache_read: 0, - cache_write: 0, - reasoning: 0, - }, - 0.0, - ); + { + let micode_dir = source_home.path().join(".local/share/mimocode"); + std::fs::create_dir_all(&micode_dir).unwrap(); + let db_path = micode_dir.join("test.db"); + { + let conn = rusqlite::Connection::open(&db_path).unwrap(); + conn.execute_batch( + "CREATE TABLE message (id TEXT PRIMARY KEY, session_id TEXT NOT NULL, data TEXT NOT NULL);", + ) + .unwrap(); + conn.execute( + "INSERT INTO message (id, session_id, data) VALUES (?1, ?2, ?3)", + rusqlite::params![ + "msg_001", + "ses_001", + r#"{"role":"assistant","modelID":"mimo-v2.5-pro","providerID":"mimo","cost":0.05,"tokens":{"input":1000,"output":500},"time":{"created":1700000000000.0,"completed":1700000001000.0}}"# + ], + ) + .unwrap(); + } - apply_pricing_if_available(&mut msg, Some(selected.as_ref())); + let mut cost_sum = 0.0f64; + let mut count = 0usize; + scan_messages_streaming( + source_home.path().to_str().unwrap(), + &["micode".to_string()], + None, + false, + &scanner::ScannerSettings::default(), + &|_m: &UnifiedMessage| true, + &mut |m: &UnifiedMessage| { + cost_sum += m.cost; + count += 1; + }, + ); - assert!(msg.cost > 0.0); + assert_eq!(count, 1, "the micode assistant message must flow through the streaming lane"); + assert!( + (cost_sum - 0.05).abs() < 1e-9, + "authoritative micode cost must survive pricing (got {cost_sum})" + ); + } } + // #742 Part 2: the micode lane is cost-guarded so MiMo Code's authoritative + // embedded cost is never overwritten by a recomputed tokens*rate when the + // model resolves to a price. reprice_lane_message(.., guard=true) reprices + // only when the embedded cost is absent (<= 0.0); guard=false is the old + // unconditional behavior that this fix replaces for micode. #[test] - fn test_select_local_parse_pricing_falls_back_to_stale_cache_on_fetch_error() { - let mut stale_litellm = HashMap::new(); - stale_litellm.insert( - "gpt-5.2".into(), + fn test_reprice_lane_message_guards_authoritative_micode_cost() { + // A pricing service that WOULD recompute a large cost for the MiMo model + // (1000*0.001 + 500*0.002 = 2.0, versus the embedded 0.05). + let mut litellm = HashMap::new(); + litellm.insert( + "mimo-v2.5-pro".into(), pricing::ModelPricing { - input_cost_per_token: Some(0.00000175), - output_cost_per_token: Some(0.000014), + input_cost_per_token: Some(0.001), + output_cost_per_token: Some(0.002), ..Default::default() }, ); - let stale = pricing::PricingService::new(stale_litellm, HashMap::new()); - - let selected = - select_local_parse_pricing(Err("network failed".to_string()), || Some(stale)).unwrap(); - - assert!(selected.lookup_with_source("gpt-5.2", None).is_some()); - } - - #[test] - fn test_select_local_parse_pricing_does_not_evaluate_stale_fallback_on_fresh_success() { - let fresh = Arc::new(pricing::PricingService::new(HashMap::new(), HashMap::new())); - let mut stale_called = false; - - let selected = select_local_parse_pricing(Ok(Arc::clone(&fresh)), || { - stale_called = true; - None - }) - .unwrap(); - - assert!(Arc::ptr_eq(&selected, &fresh)); - assert!(!stale_called); - } - - #[test] - fn test_dedupe_latest_trae_messages_keeps_latest_timestamp_for_session() { - let messages = vec![ - make_trae_message( - "session-stable", - 1_700_000_002_000, - Some("trae:session-stable:1_700_000_002"), - 0.2, - ), - make_trae_message( - "session-stable", - 1_700_000_003_000, - Some("trae:session-stable:1_700_000_003"), - 0.3, - ), - make_trae_message( - "session-other", - 1_700_000_001_000, - Some("trae:session-other:1_700_000_001"), - 0.1, - ), - ]; - - let deduped = dedupe_latest_trae_messages(messages); + let pricing = pricing::PricingService::new(litellm, HashMap::new()); + let recomputed = 1000.0 * 0.001 + 500.0 * 0.002; // 2.0 - assert_eq!(deduped.len(), 2); - let stable = deduped - .iter() - .find(|msg| msg.session_id == "session-stable") - .expect("session-stable should remain after dedupe"); - assert_eq!(stable.timestamp, 1_700_000_003_000); - assert_eq!(stable.cost, 0.3); - assert_eq!( - stable.dedup_key.as_deref(), - Some("trae:session-stable:1_700_000_003") + let make = |embedded_cost: f64| { + UnifiedMessage::new( + "micode", + "mimo-v2.5-pro", + "mimo", + "ses", + 1_700_000_000_000, + TokenBreakdown { + input: 1000, + output: 500, + cache_read: 0, + cache_write: 0, + reasoning: 0, + }, + embedded_cost, + ) + }; + + // guard=true + embedded cost present -> authoritative cost survives. + let mut guarded = make(0.05); + reprice_lane_message(&mut guarded, Some(&pricing), true); + assert!( + (guarded.cost - 0.05).abs() < 1e-9, + "cost-guarded reprice must keep the embedded 0.05, got {}", + guarded.cost ); - } - #[test] - fn test_dedupe_latest_trae_messages_tiebreaks_by_dedup_key() { - let messages = vec![ - make_trae_message( - "session-stable", - 1_700_000_010_000, - Some("dedupe-key-a"), - 0.2, - ), - make_trae_message( - "session-stable", - 1_700_000_010_000, - Some("dedupe-key-z"), - 0.4, - ), - make_trae_message( - "session-stable", - 1_700_000_009_000, - Some("dedupe-key-m"), - 0.1, - ), - ]; + // guard=false (old behavior) -> unconditionally overwritten by the recompute. + let mut unguarded = make(0.05); + reprice_lane_message(&mut unguarded, Some(&pricing), false); + assert!( + (unguarded.cost - recomputed).abs() < 1e-9, + "unguarded reprice overwrites the embedded cost with {recomputed}, got {}", + unguarded.cost + ); - let deduped = dedupe_latest_trae_messages(messages); + // guard=true + embedded cost absent (<= 0.0) -> still repriced (fallback). + let mut absent = make(0.0); + reprice_lane_message(&mut absent, Some(&pricing), true); + assert!( + (absent.cost - recomputed).abs() < 1e-9, + "a missing embedded cost must still be priced, got {}", + absent.cost + ); + } - assert_eq!(deduped.len(), 1); - assert_eq!(deduped[0].timestamp, 1_700_000_010_000); - assert_eq!(deduped[0].dedup_key.as_deref(), Some("dedupe-key-z")); - assert_eq!(deduped[0].cost, 0.4); + fn roo_task_root(home: &Path, client: ClientId) -> PathBuf { + let relative = match client { + ClientId::RooCode => ".config/Code/User/globalStorage/rooveterinaryinc.roo-cline/tasks", + ClientId::KiloCode => ".config/Code/User/globalStorage/kilocode.kilo-code/tasks", + ClientId::Cline => ".config/Code/User/globalStorage/saoudrizwan.claude-dev/tasks", + _ => panic!("not a Roo-family client: {client:?}"), + }; + home.join(relative) } - #[test] - #[serial_test::serial] - fn test_parse_all_messages_with_pricing_keeps_gateway_message_under_synthetic_filter() { - let cache_home = tempfile::TempDir::new().unwrap(); - let temp_dir = tempfile::TempDir::new().unwrap(); - let _env = opencode_test_env(cache_home.path(), temp_dir.path()); - let message_dir = temp_dir - .path() - .join(".local/share/opencode/storage/message/project-1"); - std::fs::create_dir_all(&message_dir).unwrap(); + fn write_roo_history_fixture(history: &Path, model: &str, agent: &str) { std::fs::write( - message_dir.join("msg_001.json"), - r#"{"id":"msg-1","sessionID":"session-1","role":"assistant","modelID":"hf:deepseek-ai/DeepSeek-V3-0324","providerID":"unknown","cost":0,"tokens":{"input":10,"output":5,"reasoning":0,"cache":{"read":0,"write":0}},"time":{"created":1733011200000}}"#, + history, + format!( + "{model}{agent}" + ), ) .unwrap(); - - let pricing = pricing::PricingService::new(HashMap::new(), HashMap::new()); - let messages = parse_all_messages_with_pricing( - temp_dir.path().to_str().unwrap(), - &["synthetic".to_string()], - Some(&pricing), - ); - - assert_eq!(messages.len(), 1); - assert_eq!(messages[0].client, "opencode"); - assert_eq!(messages[0].model_id, "deepseek-v3-0324"); - assert_eq!(messages[0].provider_id, "synthetic"); } - #[test] - fn test_parse_local_clients_preserves_gateway_message_client_counts() { - let temp_dir = tempfile::TempDir::new().unwrap(); - let message_dir = temp_dir - .path() - .join(".local/share/opencode/storage/message/project-1"); - std::fs::create_dir_all(&message_dir).unwrap(); + fn write_roo_task_fixture( + home: &Path, + client: ClientId, + task_id: &str, + model: &str, + agent: &str, + ) -> (PathBuf, PathBuf) { + let task_dir = roo_task_root(home, client).join(task_id); + std::fs::create_dir_all(&task_dir).unwrap(); + let ui_messages = task_dir.join("ui_messages.json"); std::fs::write( - message_dir.join("msg_001.json"), - r#"{"id":"msg-1","sessionID":"session-1","role":"assistant","modelID":"accounts/fireworks/models/deepseek-v3-0324","providerID":"fireworks","cost":0,"tokens":{"input":10,"output":5,"reasoning":0,"cache":{"read":0,"write":0}},"time":{"created":1733011200000}}"#, + &ui_messages, + r#"[{"type":"say","say":"api_req_started","ts":"2026-06-25T10:00:00Z","text":"{\"cost\":0.125,\"tokensIn\":100,\"tokensOut\":25,\"cacheReads\":10,\"cacheWrites\":5,\"apiProtocol\":\"anthropic\"}"}]"#, ) .unwrap(); - - let parsed = parse_local_clients(LocalParseOptions { - home_dir: Some(temp_dir.path().to_str().unwrap().to_string()), - use_env_roots: false, - clients: Some(vec!["opencode".to_string(), "synthetic".to_string()]), - since: None, - until: None, - year: None, - scanner_settings: scanner::ScannerSettings::default(), - modified_after: None, - }) - .unwrap(); - - assert_eq!(parsed.counts.get(ClientId::OpenCode), 1); - assert_eq!(parsed.messages.len(), 1); - assert_eq!(parsed.messages[0].client, "opencode"); - assert_eq!(parsed.messages[0].model_id, "deepseek-v3-0324"); - // opencode canonicalizes the raw "fireworks" gateway id to "fireworks_ai" (#760). - assert_eq!(parsed.messages[0].provider_id, "fireworks_ai"); + let history = sessions::roocode::history_path_for_ui_messages(&ui_messages); + write_roo_history_fixture(&history, model, agent); + (ui_messages, history) } #[test] #[serial_test::serial] - fn test_parse_all_messages_fireworks_provider_kept_under_synthetic_only_filter() { - let cache_home = tempfile::TempDir::new().unwrap(); - let temp_dir = tempfile::TempDir::new().unwrap(); - let _env = opencode_test_env(cache_home.path(), temp_dir.path()); - let message_dir = temp_dir - .path() - .join(".local/share/opencode/storage/message/project-1"); - std::fs::create_dir_all(&message_dir).unwrap(); - std::fs::write( - message_dir.join("msg_001.json"), - r#"{"id":"msg-1","sessionID":"session-1","role":"assistant","modelID":"accounts/fireworks/models/deepseek-v3-0324","providerID":"fireworks","cost":0.1,"tokens":{"input":10,"output":5,"reasoning":0,"cache":{"read":0,"write":0}},"time":{"created":1733011200000}}"#, - ) - .unwrap(); - - let pricing = pricing::PricingService::new(HashMap::new(), HashMap::new()); - let messages = parse_all_messages_with_pricing( - temp_dir.path().to_str().unwrap(), - &["synthetic".to_string()], - Some(&pricing), + fn test_roo_family_history_rewrite_refreshes_materialized_and_streaming_caches() { + let source_home = tempfile::TempDir::new().unwrap(); + let materialized_cache = tempfile::TempDir::new().unwrap(); + let streaming_cache = tempfile::TempDir::new().unwrap(); + let family = [ClientId::RooCode, ClientId::KiloCode, ClientId::Cline]; + let mut litellm = HashMap::new(); + litellm.insert( + "old-model".to_string(), + pricing::ModelPricing { + input_cost_per_token: Some(0.001), + output_cost_per_token: Some(0.002), + ..Default::default() + }, ); - - assert_eq!( - messages.len(), - 1, - "fireworks gateway message must not be dropped when filtering for synthetic" + litellm.insert( + "new-model-with-longer-id".to_string(), + pricing::ModelPricing { + input_cost_per_token: Some(0.01), + output_cost_per_token: Some(0.02), + ..Default::default() + }, ); - assert_eq!(messages[0].client, "opencode"); - assert_eq!(messages[0].model_id, "deepseek-v3-0324"); - // opencode canonicalizes the raw "fireworks" gateway id to "fireworks_ai" (#760). - assert_eq!(messages[0].provider_id, "fireworks_ai"); - } - - #[test] - fn test_parse_local_clients_fireworks_provider_kept_under_synthetic_only_filter() { - let temp_dir = tempfile::TempDir::new().unwrap(); - let message_dir = temp_dir - .path() - .join(".local/share/opencode/storage/message/project-1"); - std::fs::create_dir_all(&message_dir).unwrap(); - std::fs::write( - message_dir.join("msg_001.json"), - r#"{"id":"msg-1","sessionID":"session-1","role":"assistant","modelID":"accounts/fireworks/models/deepseek-v3-0324","providerID":"fireworks","cost":0.1,"tokens":{"input":10,"output":5,"reasoning":0,"cache":{"read":0,"write":0}},"time":{"created":1733011200000}}"#, - ) - .unwrap(); - - let parsed = parse_local_clients(LocalParseOptions { - home_dir: Some(temp_dir.path().to_str().unwrap().to_string()), - use_env_roots: false, - clients: Some(vec!["synthetic".to_string()]), - since: None, - until: None, - year: None, - scanner_settings: scanner::ScannerSettings::default(), - modified_after: None, - }) - .unwrap(); + let pricing = pricing::PricingService::new(litellm, HashMap::new()); + let mut histories = Vec::new(); + for client in family { + let task_id = format!("{}-task", client.as_str()); + let (_, history) = write_roo_task_fixture( + source_home.path(), + client, + &task_id, + "old-model", + "old-agent", + ); + histories.push(history); + } - assert_eq!( - parsed.messages.len(), - 1, - "fireworks gateway message must not be dropped when filtering for synthetic only" - ); - assert_eq!(parsed.messages[0].client, "opencode"); - assert_eq!(parsed.messages[0].model_id, "deepseek-v3-0324"); - // opencode canonicalizes the raw "fireworks" gateway id to "fireworks_ai" (#760). - assert_eq!(parsed.messages[0].provider_id, "fireworks_ai"); - } + let home = source_home.path().to_str().unwrap().to_string(); + let clients: Vec = family + .into_iter() + .map(|client| client.as_str().to_string()) + .collect(); + let run_materialized = || { + with_isolated_tokscale_cache(materialized_cache.path(), || { + parse_all_messages_with_pricing_with_env_strategy( + &home, + &clients, + Some(&pricing), + false, + &scanner::ScannerSettings::default(), + ) + }) + }; + let run_streaming = || { + with_isolated_tokscale_cache(streaming_cache.path(), || { + let mut messages = Vec::new(); + scan_messages_streaming( + &home, + &clients, + Some(&pricing), + false, + &scanner::ScannerSettings::default(), + &|_: &UnifiedMessage| true, + &mut |message: &UnifiedMessage| messages.push(message.clone()), + ); + messages + }) + }; - #[test] - fn test_parse_local_clients_honors_scanner_settings_opencode_db_paths() { - // Regression guard: `parse_local_clients` used to call - // `scan_all_clients_with_env_strategy`, which silently dropped - // `options.scanner_settings`. Users with - // `scanner.opencodeDbPaths` pointing at an OPENCODE_DB outside the - // XDG data dir would see no rows through the clients/wrapped - // command paths even though model/monthly/graph reports honored - // the same config. - let temp_dir = tempfile::TempDir::new().unwrap(); - // Deliberately do not create ~/.local/share/opencode so nothing - // is auto-discoverable; the only db the scanner can find must - // come from `scanner_settings`. - let outside_dir = temp_dir.path().join("elsewhere"); - std::fs::create_dir_all(&outside_dir).unwrap(); - let external_db = outside_dir.join("opencode.db"); + let materialized_before = run_materialized(); + let streaming_before = run_streaming(); + for (lane, messages) in [ + ("materialized", &materialized_before), + ("streaming", &streaming_before), + ] { + assert_eq!(messages.len(), 3, "{lane} seed message count"); + assert!(messages.iter().all(|message| { + message.model_id == "old-model" && message.agent.as_deref() == Some("old-agent") + })); + } - let conn = rusqlite::Connection::open(&external_db).unwrap(); - conn.execute_batch( - "PRAGMA journal_mode=WAL; - CREATE TABLE message ( - id TEXT PRIMARY KEY, - session_id TEXT NOT NULL, - data TEXT NOT NULL - );", - ) - .unwrap(); - conn.execute( - "INSERT INTO message (id, session_id, data) VALUES (?1, ?2, ?3)", - rusqlite::params![ - "ext-msg-1", - "ext-session", - r#"{ - "role": "assistant", - "modelID": "claude-sonnet-4", - "providerID": "anthropic", - "tokens": { "input": 42, "output": 7, "reasoning": 0, "cache": { "read": 0, "write": 0 } }, - "time": { "created": 1700000000000.0 } - }"# - ], - ) - .unwrap(); - drop(conn); + for history in &histories { + write_roo_history_fixture(history, "new-model-with-longer-id", "new-agent"); + } - // Without scanner_settings: no rows (nothing auto-discoverable). - let parsed_default = parse_local_clients(LocalParseOptions { - home_dir: Some(temp_dir.path().to_str().unwrap().to_string()), - use_env_roots: false, - clients: Some(vec!["opencode".to_string()]), - since: None, - until: None, - year: None, - scanner_settings: scanner::ScannerSettings::default(), - modified_after: None, - }) - .unwrap(); - assert_eq!(parsed_default.counts.get(ClientId::OpenCode), 0); - assert!(parsed_default.messages.is_empty()); + let materialized_after = run_materialized(); + let streaming_after = run_streaming(); + assert_eq!(materialized_after.len(), 3); + assert_eq!(streaming_after.len(), 3); + for client in family { + let client_name = client.as_str(); + let materialized = materialized_after + .iter() + .find(|message| message.client == client_name) + .unwrap(); + let streaming = streaming_after + .iter() + .find(|message| message.client == client_name) + .unwrap(); + let materialized_before = materialized_before + .iter() + .find(|message| message.client == client_name) + .unwrap(); - // With scanner_settings pointing at the external db: the user - // row must show up. - let parsed_with_settings = parse_local_clients(LocalParseOptions { - home_dir: Some(temp_dir.path().to_str().unwrap().to_string()), - use_env_roots: false, - clients: Some(vec!["opencode".to_string()]), - since: None, - until: None, - year: None, - scanner_settings: scanner::ScannerSettings { - opencode_db_paths: vec![external_db.clone()], - ..Default::default() - }, - modified_after: None, - }) - .unwrap(); - assert_eq!( - parsed_with_settings.counts.get(ClientId::OpenCode), - 1, - "scanner.opencodeDbPaths must reach the parse_local_clients path" - ); - assert_eq!(parsed_with_settings.messages.len(), 1); - assert_eq!(parsed_with_settings.messages[0].client, "opencode"); - assert_eq!(parsed_with_settings.messages[0].model_id, "claude-sonnet-4"); + assert_eq!(materialized.session_id, format!("{client_name}-task")); + assert_eq!(materialized.model_id, "new-model-with-longer-id"); + assert_eq!(materialized.agent.as_deref(), Some("new-agent")); + assert_eq!(materialized.tokens.input, 100); + assert_eq!(materialized.tokens.output, 25); + assert_eq!(materialized.tokens.cache_read, 10); + assert_eq!(materialized.tokens.cache_write, 5); + assert!( + materialized.cost > materialized_before.cost, + "{client_name} history model rewrite must refresh derived pricing" + ); + + assert_eq!(streaming.client, materialized.client); + assert_eq!(streaming.session_id, materialized.session_id); + assert_eq!(streaming.model_id, materialized.model_id); + assert_eq!(streaming.agent, materialized.agent); + assert_eq!(streaming.tokens.input, materialized.tokens.input); + assert_eq!(streaming.tokens.output, materialized.tokens.output); + assert_eq!(streaming.tokens.cache_read, materialized.tokens.cache_read); + assert_eq!( + streaming.tokens.cache_write, + materialized.tokens.cache_write + ); + assert!((streaming.cost - materialized.cost).abs() < 1e-9); + } } - #[test] - fn test_parse_local_clients_honors_scanner_extra_scan_paths_for_hermes_profile_db() { - let temp_dir = tempfile::TempDir::new().unwrap(); - let profile_dir = temp_dir.path().join("external-hermes/director_planning"); - std::fs::create_dir_all(&profile_dir).unwrap(); - let profile_db = profile_dir.join("state.db"); - let conn = create_hermes_sqlite_db(&profile_db); - insert_hermes_session( - &conn, - "hermes-extra-session", - "claude-sonnet-4", - 2, - 100, - 25, - 0.07, + fn latest_source_mtime_probes_roo_history(client: ClientId) { + let source_home = tempfile::TempDir::new().unwrap(); + let (ui_messages, history) = write_roo_task_fixture( + source_home.path(), + client, + "mtime-task", + "test-model", + "test-agent", ); - drop(conn); + let stale_time = + std::time::SystemTime::UNIX_EPOCH + std::time::Duration::from_secs(1_700_000_000); + let fresh_time = + std::time::SystemTime::UNIX_EPOCH + std::time::Duration::from_secs(1_700_086_400); + for (path, time) in [(&ui_messages, stale_time), (&history, fresh_time)] { + let file = std::fs::OpenOptions::new().write(true).open(path).unwrap(); + let Ok(()) = file.set_modified(time) else { + return; + }; + } - let parsed_default = parse_local_clients(LocalParseOptions { - home_dir: Some(temp_dir.path().to_str().unwrap().to_string()), + let token = latest_source_mtime_ms(&LocalParseOptions { + home_dir: Some(source_home.path().to_str().unwrap().to_string()), use_env_roots: false, - clients: Some(vec!["hermes".to_string()]), + clients: Some(vec![client.as_str().to_string()]), since: None, until: None, year: None, @@ -8811,2484 +13088,3161 @@ mod tests { modified_after: None, }) .unwrap(); - assert_eq!(parsed_default.counts.get(ClientId::Hermes), 0); - assert!(parsed_default.messages.is_empty()); - - let mut extra_scan_paths = std::collections::BTreeMap::new(); - extra_scan_paths.insert("hermes".to_string(), vec![profile_dir]); - let parsed_with_settings = parse_local_clients(LocalParseOptions { - home_dir: Some(temp_dir.path().to_str().unwrap().to_string()), - use_env_roots: false, - clients: Some(vec!["hermes".to_string()]), - since: None, - until: None, - year: None, - scanner_settings: scanner::ScannerSettings { - extra_scan_paths, - ..Default::default() - }, - modified_after: None, - }) - .unwrap(); - assert_eq!(parsed_with_settings.counts.get(ClientId::Hermes), 2); - assert_eq!(parsed_with_settings.messages.len(), 1); - assert_eq!(parsed_with_settings.messages[0].client, "hermes"); - assert_eq!( - parsed_with_settings.messages[0].agent.as_deref(), - Some("Hermes Agent") - ); assert_eq!( - parsed_with_settings.messages[0].session_id, - "hermes-extra-session" + token, + 1_700_086_400_000, + "{} change token must include the history sibling", + client.as_str() ); - assert_eq!(parsed_with_settings.messages[0].model_id, "claude-sonnet-4"); - assert_eq!(parsed_with_settings.messages[0].input, 100); - assert_eq!(parsed_with_settings.messages[0].output, 25); } #[test] - #[serial_test::serial] - fn test_auto_discovered_hermes_profile_reaches_all_consumers() { - let source_home = tempfile::TempDir::new().unwrap(); - let materialized_cache = tempfile::TempDir::new().unwrap(); - let streaming_cache = tempfile::TempDir::new().unwrap(); - let count_cache = tempfile::TempDir::new().unwrap(); - let profile_dir = source_home.path().join(".hermes/profiles/research"); - std::fs::create_dir_all(&profile_dir).unwrap(); - let profile_db = profile_dir.join("state.db"); - let conn = create_hermes_sqlite_db(&profile_db); - insert_hermes_session( - &conn, - "hermes-auto-profile", - "claude-sonnet-4", - 2, - 100, - 25, - 0.07, - ); - drop(conn); - - let home = source_home.path().to_str().unwrap().to_string(); - let clients = vec!["hermes".to_string()]; - let materialized = with_isolated_tokscale_cache(materialized_cache.path(), || { - parse_all_messages_with_pricing_with_env_strategy( - &home, - &clients, - None, - false, - &scanner::ScannerSettings::default(), - ) - }); - assert_eq!(materialized.len(), 1); - assert_eq!(materialized[0].session_id, "hermes-auto-profile"); - assert_eq!(materialized[0].tokens.input, 100); - assert_eq!(materialized[0].tokens.output, 25); + fn test_latest_source_mtime_ms_probes_roocode_history() { + latest_source_mtime_probes_roo_history(ClientId::RooCode); + } - let runtime = tokio::runtime::Builder::new_current_thread() - .enable_all() - .build() - .unwrap(); - let streaming = with_isolated_tokscale_cache(streaming_cache.path(), || { - runtime - .block_on(get_model_report(ReportOptions { - home_dir: Some(home.clone()), - use_env_roots: false, - clients: Some(clients.clone()), - ..Default::default() - })) - .unwrap() - }); - assert_eq!(streaming.total_messages, 2); - assert_eq!(streaming.total_input, 100); - assert_eq!(streaming.total_output, 25); + #[test] + fn test_latest_source_mtime_ms_probes_kilocode_history() { + latest_source_mtime_probes_roo_history(ClientId::KiloCode); + } - let future_ms = std::time::SystemTime::now() - .duration_since(std::time::UNIX_EPOCH) - .unwrap() - .as_millis() as u64 - + 3_600_000; - let counted = with_isolated_tokscale_cache(count_cache.path(), || { + #[test] + fn test_latest_source_mtime_ms_probes_cline_history() { + latest_source_mtime_probes_roo_history(ClientId::Cline); + } + + #[test] + #[serial_test::serial] + fn test_modified_after_keeps_roo_family_with_fresh_history() { + let source_home = tempfile::TempDir::new().unwrap(); + let cache_home = tempfile::TempDir::new().unwrap(); + let family = [ClientId::RooCode, ClientId::KiloCode, ClientId::Cline]; + let stale_time = + std::time::SystemTime::UNIX_EPOCH + std::time::Duration::from_secs(1_700_000_000); + let fresh_time = + std::time::SystemTime::UNIX_EPOCH + std::time::Duration::from_secs(1_700_086_400); + + for client in family { + let stale_id = format!("{}-stale", client.as_str()); + let active_id = format!("{}-active", client.as_str()); + let (stale_ui, stale_history) = write_roo_task_fixture( + source_home.path(), + client, + &stale_id, + "stale-model", + "stale-agent", + ); + let (active_ui, active_history) = write_roo_task_fixture( + source_home.path(), + client, + &active_id, + "active-model", + "active-agent", + ); + for path in [&stale_ui, &stale_history, &active_ui, &active_history] { + let file = std::fs::OpenOptions::new().write(true).open(path).unwrap(); + let Ok(()) = file.set_modified(stale_time) else { + return; + }; + } + let file = std::fs::OpenOptions::new() + .write(true) + .open(&active_history) + .unwrap(); + let Ok(()) = file.set_modified(fresh_time) else { + return; + }; + } + + let clients: Vec = family + .into_iter() + .map(|client| client.as_str().to_string()) + .collect(); + let parsed = with_isolated_tokscale_cache(cache_home.path(), || { parse_local_clients(LocalParseOptions { - home_dir: Some(home), + home_dir: Some(source_home.path().to_str().unwrap().to_string()), use_env_roots: false, clients: Some(clients), since: None, until: None, year: None, scanner_settings: scanner::ScannerSettings::default(), - modified_after: Some(future_ms), + modified_after: Some(1_700_043_200_000), }) .unwrap() }); - assert_eq!(counted.counts.get(ClientId::Hermes), 2); - assert_eq!(counted.messages.len(), 1); - assert_eq!(counted.messages[0].session_id, "hermes-auto-profile"); - assert_eq!(counted.messages[0].input, 100); - assert_eq!(counted.messages[0].output, 25); - } - - #[test] - fn test_modified_after_never_prunes_hermes_dbs_from_extra_scan_paths() { - // SQLite WAL writes may leave the main db file's mtime untouched, so - // `modified_after` must not prune Hermes/Zed dbs even when they come - // from user scan roots (the `files` lanes) rather than the default - // single-db path. A threshold in the future would prune any mtime. - let temp_dir = tempfile::TempDir::new().unwrap(); - let profile_dir = temp_dir.path().join("external-hermes/director_planning"); - std::fs::create_dir_all(&profile_dir).unwrap(); - let profile_db = profile_dir.join("state.db"); - let conn = create_hermes_sqlite_db(&profile_db); - insert_hermes_session( - &conn, - "hermes-wal-session", - "claude-sonnet-4", - 1, - 50, - 10, - 0.03, - ); - drop(conn); - - let future_ms = std::time::SystemTime::now() - .duration_since(std::time::UNIX_EPOCH) - .unwrap() - .as_millis() as u64 - + 3_600_000; - let mut extra_scan_paths = std::collections::BTreeMap::new(); - extra_scan_paths.insert("hermes".to_string(), vec![profile_dir]); - let parsed = parse_local_clients(LocalParseOptions { - home_dir: Some(temp_dir.path().to_str().unwrap().to_string()), - use_env_roots: false, - clients: Some(vec!["hermes".to_string()]), - since: None, - until: None, - year: None, - scanner_settings: scanner::ScannerSettings { - extra_scan_paths, - ..Default::default() - }, - modified_after: Some(future_ms), - }) - .unwrap(); - assert_eq!(parsed.counts.get(ClientId::Hermes), 1); - assert_eq!(parsed.messages.len(), 1); - assert_eq!(parsed.messages[0].session_id, "hermes-wal-session"); + assert_eq!(parsed.messages.len(), 3); + for client in family { + assert_eq!(parsed.counts.get(client), 1); + assert!(parsed.messages.iter().any(|message| { + message.client == client.as_str() + && message.session_id == format!("{}-active", client.as_str()) + && message.model_id == "active-model" + })); + } } + #[cfg(unix)] #[test] - fn test_modified_after_never_prunes_antigravity_cli_dbs() { - // Antigravity CLI conversation `.db` files arrive via the generic - // `files` lane (a `*.db` glob), but they are SQLite — WAL writes may - // leave the main db mtime untouched, so they must be exempt from mtime - // pruning like Hermes/Zed. A plain-file client with the same old mtime - // is still pruned (control). + fn test_modified_after_roo_history_stat_failure_keeps_source() { let temp_dir = tempfile::TempDir::new().unwrap(); - let cli_db = temp_dir.path().join("conv.db"); - std::fs::File::create(&cli_db).unwrap(); - let claude_log = temp_dir.path().join("session.jsonl"); - std::fs::File::create(&claude_log).unwrap(); + let ui_messages = temp_dir.path().join("ui_messages.json"); + std::fs::write(&ui_messages, b"[]").unwrap(); + let history = sessions::roocode::history_path_for_ui_messages(&ui_messages); + std::os::unix::fs::symlink("api_conversation_history.json", &history).unwrap(); let mut scan_result = scanner::ScanResult::default(); scan_result - .get_mut(ClientId::AntigravityCli) - .push(cli_db.clone()); - scan_result - .get_mut(ClientId::Claude) - .push(claude_log.clone()); - - // A threshold in the future would prune any real on-disk mtime. + .get_mut(ClientId::RooCode) + .push(ui_messages.clone()); let future_ms = std::time::SystemTime::now() .duration_since(std::time::UNIX_EPOCH) .unwrap() .as_millis() as u64 + 3_600_000; - crate::prune_scan_result_by_mtime(&mut scan_result, future_ms); + prune_scan_result_by_mtime(&mut scan_result, future_ms); assert_eq!( - scan_result.get(ClientId::AntigravityCli), - std::slice::from_ref(&cli_db), - "Antigravity CLI .db (a WAL-mode SQLite source) must survive mtime pruning" - ); - assert!( - scan_result.get(ClientId::Claude).is_empty(), - "a plain-file client's stale log is still pruned" + scan_result.get(ClientId::RooCode), + std::slice::from_ref(&ui_messages), + "an unreadable history sibling must fail open during pruning" ); } - /// Write a minimal Antigravity CLI conversation DB (one priced - /// `gen_metadata` row carrying `response_id`). The `trajectory_metadata_blob` - /// table is omitted on purpose — the parser tolerates its absence and falls - /// back to the file mtime for the timestamp. - fn write_antigravity_cli_db( - conversations_dir: &std::path::Path, - file_stem: &str, - response_id: &str, - ) { - fn encode_varint(mut value: u64) -> Vec { - let mut out = Vec::new(); - loop { - let mut byte = (value & 0x7f) as u8; - value >>= 7; - if value != 0 { - byte |= 0x80; - } - out.push(byte); - if value == 0 { - break; - } - } - out - } - fn enc_varint(field: u64, value: u64) -> Vec { - let mut out = encode_varint(field << 3); - out.extend(encode_varint(value)); - out - } - fn enc_len(field: u64, payload: &[u8]) -> Vec { - let mut out = encode_varint((field << 3) | 2); - out.extend(encode_varint(payload.len() as u64)); - out.extend_from_slice(payload); - out - } + struct M13RelatedFixture { + droid_source: PathBuf, + droid_related: PathBuf, + kimi_source: PathBuf, + kimi_related: PathBuf, + kiro_source: PathBuf, + kiro_related: PathBuf, + } + + fn write_m13_primary_fixtures(home: &Path) -> M13RelatedFixture { + let droid_dir = home.join(".factory/sessions"); + std::fs::create_dir_all(&droid_dir).unwrap(); + let droid_source = droid_dir.join("droid-session.settings.json"); + std::fs::write( + &droid_source, + r#"{"providerLock":"anthropic","providerLockTimestamp":"2026-01-01T00:00:00Z","tokenUsage":{"inputTokens":100,"outputTokens":20,"cacheCreationTokens":5,"cacheReadTokens":10,"thinkingTokens":2}}"#, + ) + .unwrap(); + let droid_related = droid_dir.join("droid-session.jsonl"); + + let kimi_session_dir = home.join(".kimi/sessions/group-1/kimi-session"); + std::fs::create_dir_all(&kimi_session_dir).unwrap(); + let kimi_source = kimi_session_dir.join("wire.jsonl"); + std::fs::write( + &kimi_source, + r#"{"timestamp":1767225600.0,"message":{"type":"StatusUpdate","payload":{"token_usage":{"input_other":50,"output":5,"input_cache_read":3,"input_cache_creation":2},"message_id":"kimi-turn"}}}"#, + ) + .unwrap(); + let kimi_related = home.join(".kimi/config.json"); + + let kiro_dir = home.join(".kiro/sessions/cli"); + std::fs::create_dir_all(&kiro_dir).unwrap(); + let kiro_source = kiro_dir.join("kiro-session.json"); + std::fs::write( + &kiro_source, + r#"{"session_id":"kiro-session","cwd":"/tmp/m13-project","session_state":{"rts_model_state":{"model_info":{"model_id":"kiro-model","context_window_tokens":1000}},"conversation_metadata":{"user_turn_metadatas":[{"input_token_count":0,"output_token_count":0,"end_timestamp":1767225601,"total_request_count":1,"message_ids":["kiro-turn"],"context_usage_percentage":10.0}]}}}"#, + ) + .unwrap(); + let kiro_related = kiro_source.with_extension("jsonl"); - let mut usage = Vec::new(); - usage.extend(enc_varint(2, 500)); // new input - usage.extend(enc_varint(9, 300)); // output - usage.extend(enc_len(11, response_id.as_bytes())); // responseId - let mut chat_model = Vec::new(); - chat_model.extend(enc_len(4, &usage)); - chat_model.extend(enc_len(19, b"gemini-3-flash-a")); - let gen_blob = enc_len(1, &chat_model); + M13RelatedFixture { + droid_source, + droid_related, + kimi_source, + kimi_related, + kiro_source, + kiro_related, + } + } - std::fs::create_dir_all(conversations_dir).unwrap(); - let path = conversations_dir.join(format!("{file_stem}.db")); - let conn = rusqlite::Connection::open(&path).unwrap(); - conn.execute_batch("CREATE TABLE gen_metadata (idx integer, data blob, size integer);") - .unwrap(); - conn.execute( - "INSERT INTO gen_metadata (idx, data, size) VALUES (0, ?1, 0)", - rusqlite::params![gen_blob], + fn write_m13_related_fixtures(paths: &M13RelatedFixture) { + std::fs::write( + &paths.droid_related, + r#"{"message":"Model: Claude-Opus-4.5-[Anthropic]"}"#, + ) + .unwrap(); + std::fs::write(&paths.kimi_related, r#"{"model":"kimi-new-model"}"#).unwrap(); + std::fs::write( + &paths.kiro_related, + concat!( + r#"{"version":"v1","kind":"Prompt","data":{"message_id":"kiro-prompt","content":[{"kind":"text","data":"prompt body"}],"meta":{"timestamp":1767225600.0}}}"#, + "\n", + r#"{"version":"v1","kind":"AssistantMessage","data":{"message_id":"kiro-turn","content":[{"kind":"text","data":"abcdefghijklmnop"}]}}"#, + "\n" + ), ) .unwrap(); } - // Two independent Antigravity CLI conversation DBs that reuse the same - // responseId must both survive the streaming report path. responseIds are - // unique only within a conversation, so the cross-file dedup gate is - // namespaced by session; with a bare-responseId key (the pre-fix behaviour) - // the second conversation is silently dropped and this fails (count == 1). + fn set_m13_fixture_mtimes( + paths: &M13RelatedFixture, + primary_time: std::time::SystemTime, + related_time: std::time::SystemTime, + ) -> bool { + for path in [&paths.droid_source, &paths.kimi_source, &paths.kiro_source] { + let file = std::fs::OpenOptions::new().write(true).open(path).unwrap(); + if file.set_modified(primary_time).is_err() { + return false; + } + } + for path in [ + &paths.droid_related, + &paths.kimi_related, + &paths.kiro_related, + ] { + let file = std::fs::OpenOptions::new().write(true).open(path).unwrap(); + if file.set_modified(related_time).is_err() { + return false; + } + } + true + } + #[test] #[serial_test::serial] - fn test_streaming_antigravity_cli_keeps_colliding_response_ids_across_conversations() { - let cache_home = tempfile::TempDir::new().unwrap(); + fn test_m13_related_sources_refresh_materialized_and_streaming_caches() { let source_home = tempfile::TempDir::new().unwrap(); - let _env = EnvGuard::set(&[ - ("HOME", cache_home.path().as_os_str()), - ("TOKSCALE_CONFIG_DIR", cache_home.path().as_os_str()), - ]); + let materialized_cache = tempfile::TempDir::new().unwrap(); + let streaming_cache = tempfile::TempDir::new().unwrap(); + let paths = write_m13_primary_fixtures(source_home.path()); + let clients = vec!["droid".to_string(), "kimi".to_string(), "kiro".to_string()]; - { - let conversations_dir = source_home - .path() - .join(".gemini/antigravity-cli/conversations"); - write_antigravity_cli_db(&conversations_dir, "conv-aaa", "SHARED"); - write_antigravity_cli_db(&conversations_dir, "conv-bbb", "SHARED"); + let mut litellm = HashMap::new(); + for (model, input_rate, output_rate) in [ + ("claude-unknown", 0.001, 0.002), + ("claude-opus-4-5", 0.01, 0.02), + ("kimi-for-coding", 0.001, 0.002), + ("kimi-new-model", 0.01, 0.02), + ("kiro-model", 0.003, 0.004), + ] { + litellm.insert( + model.to_string(), + pricing::ModelPricing { + input_cost_per_token: Some(input_rate), + output_cost_per_token: Some(output_rate), + ..Default::default() + }, + ); + } + let pricing = pricing::PricingService::new(litellm, HashMap::new()); + let home = source_home.path().to_str().unwrap().to_string(); + let run_materialized = || { + with_isolated_tokscale_cache(materialized_cache.path(), || { + parse_all_messages_with_pricing_with_env_strategy( + &home, + &clients, + Some(&pricing), + false, + &scanner::ScannerSettings::default(), + ) + }) + }; + let run_streaming = || { + with_isolated_tokscale_cache(streaming_cache.path(), || { + let mut messages = Vec::new(); + scan_messages_streaming( + &home, + &clients, + Some(&pricing), + false, + &scanner::ScannerSettings::default(), + &|_: &UnifiedMessage| true, + &mut |message: &UnifiedMessage| messages.push(message.clone()), + ); + messages + }) + }; - let mut sessions: Vec = Vec::new(); - scan_messages_streaming( - source_home.path().to_str().unwrap(), - &["antigravity-cli".to_string()], - None, - false, - &scanner::ScannerSettings::default(), - &|_m: &UnifiedMessage| true, - &mut |m: &UnifiedMessage| sessions.push(m.session_id.clone()), + let materialized_before = run_materialized(); + let streaming_before = run_streaming(); + assert_eq!(materialized_before.len(), 3); + assert_eq!(streaming_before.len(), 3); + for (client, model) in [ + ("droid", "claude-unknown"), + ("kimi", "kimi-for-coding"), + ("kiro", "kiro-model"), + ] { + assert_eq!( + materialized_before + .iter() + .find(|message| message.client == client) + .unwrap() + .model_id, + model + ); + assert_eq!( + streaming_before + .iter() + .find(|message| message.client == client) + .unwrap() + .model_id, + model ); + } - sessions.sort(); + write_m13_related_fixtures(&paths); + + let materialized_after = run_materialized(); + let streaming_after = run_streaming(); + assert_eq!(materialized_after.len(), 3); + assert_eq!(streaming_after.len(), 3); + for client in ["droid", "kimi", "kiro"] { + let before = materialized_before + .iter() + .find(|message| message.client == client) + .unwrap(); + let materialized = materialized_after + .iter() + .find(|message| message.client == client) + .unwrap(); + let streaming = streaming_after + .iter() + .find(|message| message.client == client) + .unwrap(); + + let expected_model = match client { + "droid" => "claude-opus-4-5", + "kimi" => "kimi-new-model", + "kiro" => "kiro-model", + _ => unreachable!(), + }; + assert_eq!(materialized.model_id, expected_model); + assert!( + materialized.cost > before.cost, + "{client} related-source creation must refresh derived cost" + ); + if client == "kiro" { + assert_eq!(before.tokens.output, 0); + assert_eq!(materialized.tokens.output, 4); + assert_eq!(materialized.timestamp, 1_767_225_600_000); + assert_eq!(materialized.duration_ms, Some(1_000)); + } + + assert_eq!(streaming.client, materialized.client); + assert_eq!(streaming.session_id, materialized.session_id); + assert_eq!(streaming.model_id, materialized.model_id); + assert_eq!(streaming.provider_id, materialized.provider_id); + assert_eq!(streaming.workspace_key, materialized.workspace_key); + assert_eq!(streaming.workspace_label, materialized.workspace_label); + assert_eq!(streaming.timestamp, materialized.timestamp); + assert_eq!(streaming.duration_ms, materialized.duration_ms); + assert_eq!(streaming.message_count, materialized.message_count); + assert_eq!(streaming.tokens.input, materialized.tokens.input); + assert_eq!(streaming.tokens.output, materialized.tokens.output); + assert_eq!(streaming.tokens.cache_read, materialized.tokens.cache_read); assert_eq!( - sessions, - vec!["conv-aaa".to_string(), "conv-bbb".to_string()], - "both conversations reusing responseId \"SHARED\" must survive" + streaming.tokens.cache_write, + materialized.tokens.cache_write + ); + assert_eq!(streaming.tokens.reasoning, materialized.tokens.reasoning); + assert!((streaming.cost - materialized.cost).abs() < 1e-9); + } + } + + #[test] + fn test_latest_source_mtime_ms_probes_m13_related_sources() { + let source_home = tempfile::TempDir::new().unwrap(); + let paths = write_m13_primary_fixtures(source_home.path()); + write_m13_related_fixtures(&paths); + let stale_time = + std::time::SystemTime::UNIX_EPOCH + std::time::Duration::from_secs(1_700_000_000); + let fresh_time = + std::time::SystemTime::UNIX_EPOCH + std::time::Duration::from_secs(1_700_086_400); + if !set_m13_fixture_mtimes(&paths, stale_time, fresh_time) { + return; + } + + for client in [ClientId::Droid, ClientId::Kimi, ClientId::Kiro] { + let token = latest_source_mtime_ms(&LocalParseOptions { + home_dir: Some(source_home.path().to_str().unwrap().to_string()), + use_env_roots: false, + clients: Some(vec![client.as_str().to_string()]), + since: None, + until: None, + year: None, + scanner_settings: scanner::ScannerSettings::default(), + modified_after: None, + }) + .unwrap(); + assert_eq!( + token, + 1_700_086_400_000, + "{} change token must include its parser dependency", + client.as_str() ); } } - // jcode (`~/.jcode/sessions/session_*.json`) must be discovered by the - // generic scanner (EnvVar JCODE_HOME / .jcode root, `session_*.json` glob) - // and flow through the streaming lane with its authoritative per-message - // token_usage. #[test] #[serial_test::serial] - fn test_streaming_jcode_flows_through_lane() { - let cache_home = tempfile::TempDir::new().unwrap(); + fn test_modified_after_keeps_m13_sources_with_fresh_dependencies() { let source_home = tempfile::TempDir::new().unwrap(); - let _env = EnvGuard::set(&[ - ("HOME", cache_home.path().as_os_str()), - ("TOKSCALE_CONFIG_DIR", cache_home.path().as_os_str()), - ]); - - { - let sessions_dir = source_home.path().join(".jcode/sessions"); - std::fs::create_dir_all(&sessions_dir).unwrap(); - std::fs::write( - sessions_dir.join("session_test.json"), - r#"{"id":"session_test","provider_key":"cliproxyapi","model":"claude-sonnet-4","working_dir":"/x","messages":[{"id":"u1","role":"user","timestamp":"2026-06-16T12:00:00Z"},{"id":"a1","role":"assistant","timestamp":"2026-06-16T12:00:01Z","token_usage":{"input_tokens":1200,"output_tokens":300}}]}"#, - ) - .unwrap(); + let cache_home = tempfile::TempDir::new().unwrap(); + let paths = write_m13_primary_fixtures(source_home.path()); + write_m13_related_fixtures(&paths); + let stale_time = + std::time::SystemTime::UNIX_EPOCH + std::time::Duration::from_secs(1_700_000_000); + let fresh_time = + std::time::SystemTime::UNIX_EPOCH + std::time::Duration::from_secs(1_700_086_400); + if !set_m13_fixture_mtimes(&paths, stale_time, fresh_time) { + return; + } - let mut input_sum = 0i64; - let mut count = 0usize; - scan_messages_streaming( - source_home.path().to_str().unwrap(), - &["jcode".to_string()], - None, - false, - &scanner::ScannerSettings::default(), - &|_m: &UnifiedMessage| true, - &mut |m: &UnifiedMessage| { - input_sum += m.tokens.input; - count += 1; - }, - ); + let parsed = with_isolated_tokscale_cache(cache_home.path(), || { + parse_local_clients(LocalParseOptions { + home_dir: Some(source_home.path().to_str().unwrap().to_string()), + use_env_roots: false, + clients: Some(vec![ + "droid".to_string(), + "kimi".to_string(), + "kiro".to_string(), + ]), + since: None, + until: None, + year: None, + scanner_settings: scanner::ScannerSettings::default(), + modified_after: Some(1_700_043_200_000), + }) + .unwrap() + }); - assert_eq!(count, 1, "the jcode assistant message must flow through the streaming lane"); - assert_eq!(input_sum, 1200); - } + assert_eq!(parsed.messages.len(), 3); + assert_eq!(parsed.counts.get(ClientId::Droid), 1); + assert_eq!(parsed.counts.get(ClientId::Kimi), 1); + assert_eq!(parsed.counts.get(ClientId::Kiro), 1); + assert!(parsed + .messages + .iter() + .any(|message| { message.client == "droid" && message.model_id == "claude-opus-4-5" })); + assert!(parsed + .messages + .iter() + .any(|message| { message.client == "kimi" && message.model_id == "kimi-new-model" })); + assert!(parsed.messages.iter().any(|message| { + message.client == "kiro" && message.output == 4 && message.duration_ms == Some(1_000) + })); } - // micode (`$XDG_DATA_HOME/micode/*.db`, WAL-mode SQLite) must be discovered - // via the generic `*.db` glob and flow through the streaming lane, keeping - // its authoritative per-message cost intact (MiMo models are unpriced, so - // apply_pricing leaves the embedded cost alone). #[test] - #[serial_test::serial] - fn test_streaming_micode_flows_with_authoritative_cost() { - let cache_home = tempfile::TempDir::new().unwrap(); + fn test_modified_after_prunes_stale_m13_sources_without_dependencies() { let source_home = tempfile::TempDir::new().unwrap(); - let _env = EnvGuard::set(&[ - ("HOME", cache_home.path().as_os_str()), - ("TOKSCALE_CONFIG_DIR", cache_home.path().as_os_str()), - ]); - - { - let micode_dir = source_home.path().join(".local/share/mimocode"); - std::fs::create_dir_all(&micode_dir).unwrap(); - let db_path = micode_dir.join("test.db"); - { - let conn = rusqlite::Connection::open(&db_path).unwrap(); - conn.execute_batch( - "CREATE TABLE message (id TEXT PRIMARY KEY, session_id TEXT NOT NULL, data TEXT NOT NULL);", - ) - .unwrap(); - conn.execute( - "INSERT INTO message (id, session_id, data) VALUES (?1, ?2, ?3)", - rusqlite::params![ - "msg_001", - "ses_001", - r#"{"role":"assistant","modelID":"mimo-v2.5-pro","providerID":"mimo","cost":0.05,"tokens":{"input":1000,"output":500},"time":{"created":1700000000000.0,"completed":1700000001000.0}}"# - ], - ) - .unwrap(); - } + let paths = write_m13_primary_fixtures(source_home.path()); + let mut scan_result = scanner::ScanResult::default(); + scan_result + .get_mut(ClientId::Droid) + .push(paths.droid_source); + scan_result.get_mut(ClientId::Kimi).push(paths.kimi_source); + scan_result.get_mut(ClientId::Kiro).push(paths.kiro_source); - let mut cost_sum = 0.0f64; - let mut count = 0usize; - scan_messages_streaming( - source_home.path().to_str().unwrap(), - &["micode".to_string()], - None, - false, - &scanner::ScannerSettings::default(), - &|_m: &UnifiedMessage| true, - &mut |m: &UnifiedMessage| { - cost_sum += m.cost; - count += 1; - }, - ); + let future_ms = std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .unwrap() + .as_millis() as u64 + + 3_600_000; + prune_scan_result_by_mtime(&mut scan_result, future_ms); - assert_eq!(count, 1, "the micode assistant message must flow through the streaming lane"); - assert!( - (cost_sum - 0.05).abs() < 1e-9, - "authoritative micode cost must survive pricing (got {cost_sum})" - ); - } + assert!(scan_result.get(ClientId::Droid).is_empty()); + assert!(scan_result.get(ClientId::Kimi).is_empty()); + assert!(scan_result.get(ClientId::Kiro).is_empty()); } - // #742 Part 2: the micode lane is cost-guarded so MiMo Code's authoritative - // embedded cost is never overwritten by a recomputed tokens*rate when the - // model resolves to a price. reprice_lane_message(.., guard=true) reprices - // only when the embedded cost is absent (<= 0.0); guard=false is the old - // unconditional behavior that this fix replaces for micode. + #[cfg(unix)] #[test] - fn test_reprice_lane_message_guards_authoritative_micode_cost() { - // A pricing service that WOULD recompute a large cost for the MiMo model - // (1000*0.001 + 500*0.002 = 2.0, versus the embedded 0.05). - let mut litellm = HashMap::new(); - litellm.insert( - "mimo-v2.5-pro".into(), - pricing::ModelPricing { - input_cost_per_token: Some(0.001), - output_cost_per_token: Some(0.002), - ..Default::default() - }, - ); - let pricing = pricing::PricingService::new(litellm, HashMap::new()); - let recomputed = 1000.0 * 0.001 + 500.0 * 0.002; // 2.0 + fn test_modified_after_m13_dependency_stat_failures_keep_sources() { + let source_home = tempfile::TempDir::new().unwrap(); + let paths = write_m13_primary_fixtures(source_home.path()); + for related in [ + &paths.droid_related, + &paths.kimi_related, + &paths.kiro_related, + ] { + std::os::unix::fs::symlink(related.file_name().unwrap(), related).unwrap(); + } - let make = |embedded_cost: f64| { - UnifiedMessage::new( - "micode", - "mimo-v2.5-pro", - "mimo", - "ses", - 1_700_000_000_000, - TokenBreakdown { - input: 1000, - output: 500, - cache_read: 0, - cache_write: 0, - reasoning: 0, - }, - embedded_cost, - ) - }; + let mut scan_result = scanner::ScanResult::default(); + scan_result + .get_mut(ClientId::Droid) + .push(paths.droid_source.clone()); + scan_result + .get_mut(ClientId::Kimi) + .push(paths.kimi_source.clone()); + scan_result + .get_mut(ClientId::Kiro) + .push(paths.kiro_source.clone()); + let future_ms = std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .unwrap() + .as_millis() as u64 + + 3_600_000; + prune_scan_result_by_mtime(&mut scan_result, future_ms); - // guard=true + embedded cost present -> authoritative cost survives. - let mut guarded = make(0.05); - reprice_lane_message(&mut guarded, Some(&pricing), true); - assert!( - (guarded.cost - 0.05).abs() < 1e-9, - "cost-guarded reprice must keep the embedded 0.05, got {}", - guarded.cost + assert_eq!( + scan_result.get(ClientId::Droid), + std::slice::from_ref(&paths.droid_source) ); - - // guard=false (old behavior) -> unconditionally overwritten by the recompute. - let mut unguarded = make(0.05); - reprice_lane_message(&mut unguarded, Some(&pricing), false); - assert!( - (unguarded.cost - recomputed).abs() < 1e-9, - "unguarded reprice overwrites the embedded cost with {recomputed}, got {}", - unguarded.cost + assert_eq!( + scan_result.get(ClientId::Kimi), + std::slice::from_ref(&paths.kimi_source) ); - - // guard=true + embedded cost absent (<= 0.0) -> still repriced (fallback). - let mut absent = make(0.0); - reprice_lane_message(&mut absent, Some(&pricing), true); - assert!( - (absent.cost - recomputed).abs() < 1e-9, - "a missing embedded cost must still be priced, got {}", - absent.cost + assert_eq!( + scan_result.get(ClientId::Kiro), + std::slice::from_ref(&paths.kiro_source) ); } - fn roo_task_root(home: &Path, client: ClientId) -> PathBuf { - let relative = match client { - ClientId::RooCode => ".config/Code/User/globalStorage/rooveterinaryinc.roo-cline/tasks", - ClientId::KiloCode => ".config/Code/User/globalStorage/kilocode.kilo-code/tasks", - ClientId::Cline => ".config/Code/User/globalStorage/saoudrizwan.claude-dev/tasks", - _ => panic!("not a Roo-family client: {client:?}"), - }; - home.join(relative) - } + // micode `.db` is WAL-mode SQLite reached via the generic `*.db` glob, so it + // must be exempt from mtime pruning (a WAL-only write leaves the main db's + // mtime untouched) — same treatment as Antigravity CLI / Hermes / Zed. + #[test] + fn test_modified_after_never_prunes_micode_dbs() { + let temp_dir = tempfile::TempDir::new().unwrap(); + let micode_db = temp_dir.path().join("micode.db"); + std::fs::File::create(&micode_db).unwrap(); + let claude_log = temp_dir.path().join("session.jsonl"); + std::fs::File::create(&claude_log).unwrap(); + + let mut scan_result = scanner::ScanResult::default(); + scan_result + .get_mut(ClientId::MiMoCode) + .push(micode_db.clone()); + scan_result + .get_mut(ClientId::Claude) + .push(claude_log.clone()); - fn write_roo_history_fixture(history: &Path, model: &str, agent: &str) { - std::fs::write( - history, - format!( - "{model}{agent}" - ), - ) - .unwrap(); - } + let future_ms = std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .unwrap() + .as_millis() as u64 + + 3_600_000; + crate::prune_scan_result_by_mtime(&mut scan_result, future_ms); - fn write_roo_task_fixture( - home: &Path, - client: ClientId, - task_id: &str, - model: &str, - agent: &str, - ) -> (PathBuf, PathBuf) { - let task_dir = roo_task_root(home, client).join(task_id); - std::fs::create_dir_all(&task_dir).unwrap(); - let ui_messages = task_dir.join("ui_messages.json"); - std::fs::write( - &ui_messages, - r#"[{"type":"say","say":"api_req_started","ts":"2026-06-25T10:00:00Z","text":"{\"cost\":0.125,\"tokensIn\":100,\"tokensOut\":25,\"cacheReads\":10,\"cacheWrites\":5,\"apiProtocol\":\"anthropic\"}"}]"#, - ) - .unwrap(); - let history = sessions::roocode::history_path_for_ui_messages(&ui_messages); - write_roo_history_fixture(&history, model, agent); - (ui_messages, history) + assert_eq!( + scan_result.get(ClientId::MiMoCode), + std::slice::from_ref(&micode_db), + "micode .db (a WAL-mode SQLite source) must survive mtime pruning" + ); + assert!( + scan_result.get(ClientId::Claude).is_empty(), + "a plain-file client's stale log is still pruned" + ); } + // gjc (`$GJC_CODING_AGENT_DIR/sessions/*.jsonl`) must be discovered via the + // EnvVar fallback root (`.gjc/agent`) + `*.jsonl` glob and flow through the + // streaming lane, keeping its authoritative embedded `usage.cost.total` + // (A1). With pricing absent the guard's reprice branch is a no-op; the + // materialized path mirrors upstream's proven reprice-when-absent guard. #[test] #[serial_test::serial] - fn test_roo_family_history_rewrite_refreshes_materialized_and_streaming_caches() { + fn test_streaming_gjc_flows_with_authoritative_cost() { + let cache_home = tempfile::TempDir::new().unwrap(); let source_home = tempfile::TempDir::new().unwrap(); - let materialized_cache = tempfile::TempDir::new().unwrap(); - let streaming_cache = tempfile::TempDir::new().unwrap(); - let family = [ClientId::RooCode, ClientId::KiloCode, ClientId::Cline]; - let mut litellm = HashMap::new(); - litellm.insert( - "old-model".to_string(), - pricing::ModelPricing { - input_cost_per_token: Some(0.001), - output_cost_per_token: Some(0.002), - ..Default::default() - }, - ); - litellm.insert( - "new-model-with-longer-id".to_string(), - pricing::ModelPricing { - input_cost_per_token: Some(0.01), - output_cost_per_token: Some(0.02), - ..Default::default() - }, - ); - let pricing = pricing::PricingService::new(litellm, HashMap::new()); - let mut histories = Vec::new(); - for client in family { - let task_id = format!("{}-task", client.as_str()); - let (_, history) = write_roo_task_fixture( - source_home.path(), - client, - &task_id, - "old-model", - "old-agent", - ); - histories.push(history); - } + let _env = EnvGuard::set(&[ + ("HOME", cache_home.path().as_os_str()), + ("TOKSCALE_CONFIG_DIR", cache_home.path().as_os_str()), + ]); - let home = source_home.path().to_str().unwrap().to_string(); - let clients: Vec = family - .into_iter() - .map(|client| client.as_str().to_string()) - .collect(); - let run_materialized = || { - with_isolated_tokscale_cache(materialized_cache.path(), || { - parse_all_messages_with_pricing_with_env_strategy( - &home, - &clients, - Some(&pricing), - false, - &scanner::ScannerSettings::default(), - ) - }) - }; - let run_streaming = || { - with_isolated_tokscale_cache(streaming_cache.path(), || { - let mut messages = Vec::new(); - scan_messages_streaming( - &home, - &clients, - Some(&pricing), - false, - &scanner::ScannerSettings::default(), - &|_: &UnifiedMessage| true, - &mut |message: &UnifiedMessage| messages.push(message.clone()), - ); - messages - }) - }; + { + let gjc_dir = source_home.path().join(".gjc/agent/sessions"); + std::fs::create_dir_all(&gjc_dir).unwrap(); + std::fs::write( + gjc_dir.join("test.jsonl"), + "{\"type\":\"session\",\"id\":\"gjc_ses_001\",\"cwd\":\"/work/pi\"}\n{\"type\":\"message\",\"id\":\"msg_001\",\"message\":{\"role\":\"assistant\",\"model\":\"claude-sonnet-4\",\"provider\":\"anthropic\",\"timestamp\":1767225601000,\"usage\":{\"input\":100,\"output\":50,\"cost\":{\"total\":0.3}}}}\n", + ) + .unwrap(); - let materialized_before = run_materialized(); - let streaming_before = run_streaming(); - for (lane, messages) in [ - ("materialized", &materialized_before), - ("streaming", &streaming_before), - ] { - assert_eq!(messages.len(), 3, "{lane} seed message count"); - assert!(messages.iter().all(|message| { - message.model_id == "old-model" && message.agent.as_deref() == Some("old-agent") - })); - } + let mut cost_sum = 0.0f64; + let mut count = 0usize; + scan_messages_streaming( + source_home.path().to_str().unwrap(), + &["gjc".to_string()], + None, + false, + &scanner::ScannerSettings::default(), + &|_m: &UnifiedMessage| true, + &mut |m: &UnifiedMessage| { + cost_sum += m.cost; + count += 1; + }, + ); - for history in &histories { - write_roo_history_fixture(history, "new-model-with-longer-id", "new-agent"); + assert_eq!(count, 1, "the gjc assistant message must flow through the streaming lane"); + assert!( + (cost_sum - 0.3).abs() < 1e-9, + "authoritative gjc cost must reach the sink (got {cost_sum})" + ); } + } - let materialized_after = run_materialized(); - let streaming_after = run_streaming(); - assert_eq!(materialized_after.len(), 3); - assert_eq!(streaming_after.len(), 3); - for client in family { - let client_name = client.as_str(); - let materialized = materialized_after - .iter() - .find(|message| message.client == client_name) - .unwrap(); - let streaming = streaming_after - .iter() - .find(|message| message.client == client_name) - .unwrap(); - let materialized_before = materialized_before - .iter() - .find(|message| message.client == client_name) - .unwrap(); + // jcode's `session_*.json` snapshot is a file-lane source whose sibling + // `.journal.jsonl` is appended between snapshot rewrites without touching + // the snapshot mtime, so it must be exempt from mtime pruning like the WAL + // db lanes — otherwise an active session with a stale snapshot is dropped + // and its recent journal turns vanish from the live tail. + #[test] + fn test_modified_after_never_prunes_jcode_sessions() { + let temp_dir = tempfile::TempDir::new().unwrap(); + let jcode_snapshot = temp_dir.path().join("session_x.json"); + std::fs::File::create(&jcode_snapshot).unwrap(); + let claude_log = temp_dir.path().join("session.jsonl"); + std::fs::File::create(&claude_log).unwrap(); - assert_eq!(materialized.session_id, format!("{client_name}-task")); - assert_eq!(materialized.model_id, "new-model-with-longer-id"); - assert_eq!(materialized.agent.as_deref(), Some("new-agent")); - assert_eq!(materialized.tokens.input, 100); - assert_eq!(materialized.tokens.output, 25); - assert_eq!(materialized.tokens.cache_read, 10); - assert_eq!(materialized.tokens.cache_write, 5); - assert!( - materialized.cost > materialized_before.cost, - "{client_name} history model rewrite must refresh derived pricing" - ); + let mut scan_result = scanner::ScanResult::default(); + scan_result + .get_mut(ClientId::Jcode) + .push(jcode_snapshot.clone()); + scan_result + .get_mut(ClientId::Claude) + .push(claude_log.clone()); - assert_eq!(streaming.client, materialized.client); - assert_eq!(streaming.session_id, materialized.session_id); - assert_eq!(streaming.model_id, materialized.model_id); - assert_eq!(streaming.agent, materialized.agent); - assert_eq!(streaming.tokens.input, materialized.tokens.input); - assert_eq!(streaming.tokens.output, materialized.tokens.output); - assert_eq!(streaming.tokens.cache_read, materialized.tokens.cache_read); - assert_eq!( - streaming.tokens.cache_write, - materialized.tokens.cache_write - ); - assert!((streaming.cost - materialized.cost).abs() < 1e-9); + let future_ms = std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .unwrap() + .as_millis() as u64 + + 3_600_000; + crate::prune_scan_result_by_mtime(&mut scan_result, future_ms); + + assert_eq!( + scan_result.get(ClientId::Jcode), + std::slice::from_ref(&jcode_snapshot), + "jcode snapshot (its journal sibling can change without it) must survive mtime pruning" + ); + assert!( + scan_result.get(ClientId::Claude).is_empty(), + "a plain-file client's stale log is still pruned" + ); + } + + #[test] + fn test_modified_after_prunes_grok_by_updates_or_signals_mtime() { + let temp_dir = tempfile::TempDir::new().unwrap(); + let stale_dir = temp_dir.path().join("stale"); + let active_dir = temp_dir.path().join("active"); + std::fs::create_dir_all(&stale_dir).unwrap(); + std::fs::create_dir_all(&active_dir).unwrap(); + + let stale_updates = stale_dir.join("updates.jsonl"); + let stale_signals = stale_dir.join("signals.json"); + let active_updates = active_dir.join("updates.jsonl"); + let active_signals = active_dir.join("signals.json"); + for path in [&stale_updates, &stale_signals, &active_updates, &active_signals] { + std::fs::File::create(path).unwrap(); } - } - fn latest_source_mtime_probes_roo_history(client: ClientId) { - let source_home = tempfile::TempDir::new().unwrap(); - let (ui_messages, history) = write_roo_task_fixture( - source_home.path(), - client, - "mtime-task", - "test-model", - "test-agent", - ); let stale_time = std::time::SystemTime::UNIX_EPOCH + std::time::Duration::from_secs(1_700_000_000); - let fresh_time = + let active_time = std::time::SystemTime::UNIX_EPOCH + std::time::Duration::from_secs(1_700_086_400); - for (path, time) in [(&ui_messages, stale_time), (&history, fresh_time)] { + for path in [&stale_updates, &stale_signals, &active_updates] { let file = std::fs::OpenOptions::new().write(true).open(path).unwrap(); - let Ok(()) = file.set_modified(time) else { + let Ok(()) = file.set_modified(stale_time) else { return; }; } + let file = std::fs::OpenOptions::new() + .write(true) + .open(&active_signals) + .unwrap(); + let Ok(()) = file.set_modified(active_time) else { + return; + }; - let token = latest_source_mtime_ms(&LocalParseOptions { - home_dir: Some(source_home.path().to_str().unwrap().to_string()), - use_env_roots: false, - clients: Some(vec![client.as_str().to_string()]), - since: None, - until: None, - year: None, - scanner_settings: scanner::ScannerSettings::default(), - modified_after: None, - }) - .unwrap(); + let mut scan_result = scanner::ScanResult::default(); + scan_result + .get_mut(ClientId::Grok) + .extend([stale_updates.clone(), active_updates.clone()]); + + crate::prune_scan_result_by_mtime(&mut scan_result, 1_700_043_200_000); assert_eq!( - token, - 1_700_086_400_000, - "{} change token must include the history sibling", - client.as_str() + scan_result.get(ClientId::Grok), + std::slice::from_ref(&active_updates), + "stale Grok sessions should be pruned while a fresh signals sibling keeps its session" ); } #[test] - fn test_latest_source_mtime_ms_probes_roocode_history() { - latest_source_mtime_probes_roo_history(ClientId::RooCode); - } - - #[test] - fn test_latest_source_mtime_ms_probes_kilocode_history() { - latest_source_mtime_probes_roo_history(ClientId::KiloCode); - } + fn test_modified_after_retains_grok_authority_sources() { + let temp_dir = tempfile::TempDir::new().unwrap(); + let stale_dir = temp_dir.path().join("stale"); + let active_dir = temp_dir.path().join("active"); + let logs_dir = temp_dir.path().join("logs"); + for dir in [&stale_dir, &active_dir, &logs_dir] { + std::fs::create_dir_all(dir).unwrap(); + } - #[test] - fn test_latest_source_mtime_ms_probes_cline_history() { - latest_source_mtime_probes_roo_history(ClientId::Cline); - } + let stale_updates = stale_dir.join("updates.jsonl"); + let active_updates = active_dir.join("updates.jsonl"); + let unified = logs_dir.join("unified.jsonl"); + for path in [&stale_updates, &active_updates, &unified] { + std::fs::File::create(path).unwrap(); + } - #[test] - #[serial_test::serial] - fn test_modified_after_keeps_roo_family_with_fresh_history() { - let source_home = tempfile::TempDir::new().unwrap(); - let cache_home = tempfile::TempDir::new().unwrap(); - let family = [ClientId::RooCode, ClientId::KiloCode, ClientId::Cline]; let stale_time = std::time::SystemTime::UNIX_EPOCH + std::time::Duration::from_secs(1_700_000_000); - let fresh_time = + let active_time = std::time::SystemTime::UNIX_EPOCH + std::time::Duration::from_secs(1_700_086_400); - - for client in family { - let stale_id = format!("{}-stale", client.as_str()); - let active_id = format!("{}-active", client.as_str()); - let (stale_ui, stale_history) = write_roo_task_fixture( - source_home.path(), - client, - &stale_id, - "stale-model", - "stale-agent", - ); - let (active_ui, active_history) = write_roo_task_fixture( - source_home.path(), - client, - &active_id, - "active-model", - "active-agent", - ); - for path in [&stale_ui, &stale_history, &active_ui, &active_history] { - let file = std::fs::OpenOptions::new().write(true).open(path).unwrap(); - let Ok(()) = file.set_modified(stale_time) else { - return; - }; - } - let file = std::fs::OpenOptions::new() - .write(true) - .open(&active_history) - .unwrap(); - let Ok(()) = file.set_modified(fresh_time) else { + for path in [&stale_updates, &unified] { + let file = std::fs::OpenOptions::new().write(true).open(path).unwrap(); + let Ok(()) = file.set_modified(stale_time) else { return; }; } + let file = std::fs::OpenOptions::new() + .write(true) + .open(&active_updates) + .unwrap(); + let Ok(()) = file.set_modified(active_time) else { + return; + }; - let clients: Vec = family - .into_iter() - .map(|client| client.as_str().to_string()) - .collect(); - let parsed = with_isolated_tokscale_cache(cache_home.path(), || { - parse_local_clients(LocalParseOptions { - home_dir: Some(source_home.path().to_str().unwrap().to_string()), - use_env_roots: false, - clients: Some(clients), - since: None, - until: None, - year: None, - scanner_settings: scanner::ScannerSettings::default(), - modified_after: Some(1_700_043_200_000), - }) - .unwrap() - }); - - assert_eq!(parsed.messages.len(), 3); - for client in family { - assert_eq!(parsed.counts.get(client), 1); - assert!(parsed.messages.iter().any(|message| { - message.client == client.as_str() - && message.session_id == format!("{}-active", client.as_str()) - && message.model_id == "active-model" - })); - } - } - - #[cfg(unix)] - #[test] - fn test_modified_after_roo_history_stat_failure_keeps_source() { - let temp_dir = tempfile::TempDir::new().unwrap(); - let ui_messages = temp_dir.path().join("ui_messages.json"); - std::fs::write(&ui_messages, b"[]").unwrap(); - let history = sessions::roocode::history_path_for_ui_messages(&ui_messages); - std::os::unix::fs::symlink("api_conversation_history.json", &history).unwrap(); + let mut legacy_fresh = scanner::ScanResult::default(); + legacy_fresh.get_mut(ClientId::Grok).extend([ + stale_updates.clone(), + active_updates.clone(), + unified.clone(), + ]); + crate::prune_scan_result_by_mtime(&mut legacy_fresh, 1_700_043_200_000); + assert_eq!( + legacy_fresh.get(ClientId::Grok), + &[active_updates.clone(), unified.clone()], + "a stale unified authority source must survive while a legacy source is fresh" + ); - let mut scan_result = scanner::ScanResult::default(); - scan_result - .get_mut(ClientId::RooCode) - .push(ui_messages.clone()); - let future_ms = std::time::SystemTime::now() - .duration_since(std::time::UNIX_EPOCH) - .unwrap() - .as_millis() as u64 - + 3_600_000; - prune_scan_result_by_mtime(&mut scan_result, future_ms); + let file = std::fs::OpenOptions::new() + .write(true) + .open(&unified) + .unwrap(); + let Ok(()) = file.set_modified(active_time) else { + return; + }; + let file = std::fs::OpenOptions::new() + .write(true) + .open(&active_updates) + .unwrap(); + let Ok(()) = file.set_modified(stale_time) else { + return; + }; + let mut unified_fresh = scanner::ScanResult::default(); + unified_fresh.get_mut(ClientId::Grok).extend([ + stale_updates.clone(), + active_updates.clone(), + unified.clone(), + ]); + crate::prune_scan_result_by_mtime(&mut unified_fresh, 1_700_043_200_000); assert_eq!( - scan_result.get(ClientId::RooCode), - std::slice::from_ref(&ui_messages), - "an unreadable history sibling must fail open during pruning" + unified_fresh.get(ClientId::Grok), + &[stale_updates, active_updates, unified], + "a fresh unified source needs the legacy cohort for workspace attribution" ); } - struct M13RelatedFixture { - droid_source: PathBuf, - droid_related: PathBuf, - kimi_source: PathBuf, - kimi_related: PathBuf, - kiro_source: PathBuf, - kiro_related: PathBuf, - } + // The pruning helper must treat a fresh write to *any* metadata sibling the + // parser reads (not just signals.json) as source activity — otherwise a + // summary.json- or events.jsonl-only write (a late-arriving model id) lets an + // otherwise-stale session be pruned and its fresh model never re-parsed. + fn prune_grok_keeps_session_with_fresh_sibling(fresh_sibling: &str) { + let temp_dir = tempfile::TempDir::new().unwrap(); + let stale_dir = temp_dir.path().join("stale"); + let active_dir = temp_dir.path().join("active"); + std::fs::create_dir_all(&stale_dir).unwrap(); + std::fs::create_dir_all(&active_dir).unwrap(); - fn write_m13_primary_fixtures(home: &Path) -> M13RelatedFixture { - let droid_dir = home.join(".factory/sessions"); - std::fs::create_dir_all(&droid_dir).unwrap(); - let droid_source = droid_dir.join("droid-session.settings.json"); - std::fs::write( - &droid_source, - r#"{"providerLock":"anthropic","providerLockTimestamp":"2026-01-01T00:00:00Z","tokenUsage":{"inputTokens":100,"outputTokens":20,"cacheCreationTokens":5,"cacheReadTokens":10,"thinkingTokens":2}}"#, - ) - .unwrap(); - let droid_related = droid_dir.join("droid-session.jsonl"); + let stale_updates = stale_dir.join("updates.jsonl"); + let active_updates = active_dir.join("updates.jsonl"); + let active_fresh = active_dir.join(fresh_sibling); + // Every session file starts stale; only the one fresh sibling moves. + let mut all_paths = vec![stale_updates.clone(), active_updates.clone()]; + for name in message_cache::GROK_METADATA_SIBLINGS { + all_paths.push(stale_dir.join(name)); + all_paths.push(active_dir.join(name)); + } + for path in &all_paths { + std::fs::File::create(path).unwrap(); + } - let kimi_session_dir = home.join(".kimi/sessions/group-1/kimi-session"); - std::fs::create_dir_all(&kimi_session_dir).unwrap(); - let kimi_source = kimi_session_dir.join("wire.jsonl"); - std::fs::write( - &kimi_source, - r#"{"timestamp":1767225600.0,"message":{"type":"StatusUpdate","payload":{"token_usage":{"input_other":50,"output":5,"input_cache_read":3,"input_cache_creation":2},"message_id":"kimi-turn"}}}"#, - ) - .unwrap(); - let kimi_related = home.join(".kimi/config.json"); + let stale_time = + std::time::SystemTime::UNIX_EPOCH + std::time::Duration::from_secs(1_700_000_000); + let active_time = + std::time::SystemTime::UNIX_EPOCH + std::time::Duration::from_secs(1_700_086_400); + for path in &all_paths { + let file = std::fs::OpenOptions::new().write(true).open(path).unwrap(); + let Ok(()) = file.set_modified(stale_time) else { + return; + }; + } + let file = std::fs::OpenOptions::new() + .write(true) + .open(&active_fresh) + .unwrap(); + let Ok(()) = file.set_modified(active_time) else { + return; + }; + + let mut scan_result = scanner::ScanResult::default(); + scan_result + .get_mut(ClientId::Grok) + .extend([stale_updates.clone(), active_updates.clone()]); - let kiro_dir = home.join(".kiro/sessions/cli"); - std::fs::create_dir_all(&kiro_dir).unwrap(); - let kiro_source = kiro_dir.join("kiro-session.json"); - std::fs::write( - &kiro_source, - r#"{"session_id":"kiro-session","cwd":"/tmp/m13-project","session_state":{"rts_model_state":{"model_info":{"model_id":"kiro-model","context_window_tokens":1000}},"conversation_metadata":{"user_turn_metadatas":[{"input_token_count":0,"output_token_count":0,"end_timestamp":1767225601,"total_request_count":1,"message_ids":["kiro-turn"],"context_usage_percentage":10.0}]}}}"#, - ) - .unwrap(); - let kiro_related = kiro_source.with_extension("jsonl"); + crate::prune_scan_result_by_mtime(&mut scan_result, 1_700_043_200_000); - M13RelatedFixture { - droid_source, - droid_related, - kimi_source, - kimi_related, - kiro_source, - kiro_related, - } + assert_eq!( + scan_result.get(ClientId::Grok), + std::slice::from_ref(&active_updates), + "a fresh {fresh_sibling} sibling must keep its otherwise-stale Grok session" + ); } - fn write_m13_related_fixtures(paths: &M13RelatedFixture) { - std::fs::write( - &paths.droid_related, - r#"{"message":"Model: Claude-Opus-4.5-[Anthropic]"}"#, - ) - .unwrap(); - std::fs::write(&paths.kimi_related, r#"{"model":"kimi-new-model"}"#).unwrap(); - std::fs::write( - &paths.kiro_related, - concat!( - r#"{"version":"v1","kind":"Prompt","data":{"message_id":"kiro-prompt","content":[{"kind":"text","data":"prompt body"}],"meta":{"timestamp":1767225600.0}}}"#, - "\n", - r#"{"version":"v1","kind":"AssistantMessage","data":{"message_id":"kiro-turn","content":[{"kind":"text","data":"abcdefghijklmnop"}]}}"#, - "\n" - ), - ) - .unwrap(); + #[test] + fn test_modified_after_prunes_grok_keeps_session_with_fresh_summary() { + prune_grok_keeps_session_with_fresh_sibling("summary.json"); } - fn set_m13_fixture_mtimes( - paths: &M13RelatedFixture, - primary_time: std::time::SystemTime, - related_time: std::time::SystemTime, - ) -> bool { - for path in [&paths.droid_source, &paths.kimi_source, &paths.kiro_source] { - let file = std::fs::OpenOptions::new().write(true).open(path).unwrap(); - if file.set_modified(primary_time).is_err() { - return false; - } - } - for path in [ - &paths.droid_related, - &paths.kimi_related, - &paths.kiro_related, - ] { - let file = std::fs::OpenOptions::new().write(true).open(path).unwrap(); - if file.set_modified(related_time).is_err() { - return false; - } - } - true + #[test] + fn test_modified_after_prunes_grok_keeps_session_with_fresh_events() { + prune_grok_keeps_session_with_fresh_sibling("events.jsonl"); } - #[test] - #[serial_test::serial] - fn test_m13_related_sources_refresh_materialized_and_streaming_caches() { + // The live-tail change token must move when Grok rewrites *any* metadata + // sibling the parser reads, even though updates.jsonl is unchanged; otherwise + // UsageTail short-circuits and the session stays pinned to its fallback model. + fn latest_source_mtime_ms_probes_grok_sibling(fresh_sibling: &str) { let source_home = tempfile::TempDir::new().unwrap(); - let materialized_cache = tempfile::TempDir::new().unwrap(); - let streaming_cache = tempfile::TempDir::new().unwrap(); - let paths = write_m13_primary_fixtures(source_home.path()); - let clients = vec!["droid".to_string(), "kimi".to_string(), "kiro".to_string()]; - - let mut litellm = HashMap::new(); - for (model, input_rate, output_rate) in [ - ("claude-unknown", 0.001, 0.002), - ("claude-opus-4-5", 0.01, 0.02), - ("kimi-for-coding", 0.001, 0.002), - ("kimi-new-model", 0.01, 0.02), - ("kiro-model", 0.003, 0.004), - ] { - litellm.insert( - model.to_string(), - pricing::ModelPricing { - input_cost_per_token: Some(input_rate), - output_cost_per_token: Some(output_rate), - ..Default::default() - }, - ); + let session_dir = source_home + .path() + .join(".grok/sessions/%2Ftmp%2Fproject/session-uuid-1"); + std::fs::create_dir_all(&session_dir).unwrap(); + let updates = session_dir.join("updates.jsonl"); + std::fs::write(&updates, b"{\"totalTokens\":1}\n").unwrap(); + // Every sibling exists and is stale; only the target sibling is newer. + for name in message_cache::GROK_METADATA_SIBLINGS { + std::fs::write(session_dir.join(name), b"{}").unwrap(); } - let pricing = pricing::PricingService::new(litellm, HashMap::new()); - let home = source_home.path().to_str().unwrap().to_string(); - let run_materialized = || { - with_isolated_tokscale_cache(materialized_cache.path(), || { - parse_all_messages_with_pricing_with_env_strategy( - &home, - &clients, - Some(&pricing), - false, - &scanner::ScannerSettings::default(), - ) - }) - }; - let run_streaming = || { - with_isolated_tokscale_cache(streaming_cache.path(), || { - let mut messages = Vec::new(); - scan_messages_streaming( - &home, - &clients, - Some(&pricing), - false, - &scanner::ScannerSettings::default(), - &|_: &UnifiedMessage| true, - &mut |message: &UnifiedMessage| messages.push(message.clone()), - ); - messages - }) - }; - let materialized_before = run_materialized(); - let streaming_before = run_streaming(); - assert_eq!(materialized_before.len(), 3); - assert_eq!(streaming_before.len(), 3); - for (client, model) in [ - ("droid", "claude-unknown"), - ("kimi", "kimi-for-coding"), - ("kiro", "kiro-model"), - ] { - assert_eq!( - materialized_before - .iter() - .find(|message| message.client == client) - .unwrap() - .model_id, - model - ); - assert_eq!( - streaming_before - .iter() - .find(|message| message.client == client) - .unwrap() - .model_id, - model - ); + let stale_time = + std::time::SystemTime::UNIX_EPOCH + std::time::Duration::from_secs(1_700_000_000); + let fresh_time = + std::time::SystemTime::UNIX_EPOCH + std::time::Duration::from_secs(1_700_086_400); + for name in std::iter::once("updates.jsonl").chain(message_cache::GROK_METADATA_SIBLINGS) { + let f = std::fs::OpenOptions::new() + .write(true) + .open(session_dir.join(name)) + .unwrap(); + let Ok(()) = f.set_modified(stale_time) else { + return; + }; } + let f = std::fs::OpenOptions::new() + .write(true) + .open(session_dir.join(fresh_sibling)) + .unwrap(); + let Ok(()) = f.set_modified(fresh_time) else { + return; + }; - write_m13_related_fixtures(&paths); + let options = LocalParseOptions { + home_dir: Some(source_home.path().to_str().unwrap().to_string()), + use_env_roots: false, + clients: Some(vec!["grok".to_string()]), + since: None, + until: None, + year: None, + scanner_settings: scanner::ScannerSettings::default(), + modified_after: None, + }; + let token = crate::latest_source_mtime_ms(&options).unwrap(); - let materialized_after = run_materialized(); - let streaming_after = run_streaming(); - assert_eq!(materialized_after.len(), 3); - assert_eq!(streaming_after.len(), 3); - for client in ["droid", "kimi", "kiro"] { - let before = materialized_before - .iter() - .find(|message| message.client == client) - .unwrap(); - let materialized = materialized_after - .iter() - .find(|message| message.client == client) - .unwrap(); - let streaming = streaming_after - .iter() - .find(|message| message.client == client) - .unwrap(); + assert_eq!( + token, 1_700_086_400_000, + "the change token must reflect the fresh {fresh_sibling} mtime, not just updates.jsonl" + ); + } - let expected_model = match client { - "droid" => "claude-opus-4-5", - "kimi" => "kimi-new-model", - "kiro" => "kiro-model", - _ => unreachable!(), - }; - assert_eq!(materialized.model_id, expected_model); - assert!( - materialized.cost > before.cost, - "{client} related-source creation must refresh derived cost" - ); - if client == "kiro" { - assert_eq!(before.tokens.output, 0); - assert_eq!(materialized.tokens.output, 4); - assert_eq!(materialized.timestamp, 1_767_225_600_000); - assert_eq!(materialized.duration_ms, Some(1_000)); - } + #[test] + #[serial_test::serial] + fn test_latest_source_mtime_ms_probes_grok_summary() { + latest_source_mtime_ms_probes_grok_sibling("summary.json"); + } - assert_eq!(streaming.client, materialized.client); - assert_eq!(streaming.session_id, materialized.session_id); - assert_eq!(streaming.model_id, materialized.model_id); - assert_eq!(streaming.provider_id, materialized.provider_id); - assert_eq!(streaming.workspace_key, materialized.workspace_key); - assert_eq!(streaming.workspace_label, materialized.workspace_label); - assert_eq!(streaming.timestamp, materialized.timestamp); - assert_eq!(streaming.duration_ms, materialized.duration_ms); - assert_eq!(streaming.message_count, materialized.message_count); - assert_eq!(streaming.tokens.input, materialized.tokens.input); - assert_eq!(streaming.tokens.output, materialized.tokens.output); - assert_eq!(streaming.tokens.cache_read, materialized.tokens.cache_read); - assert_eq!( - streaming.tokens.cache_write, - materialized.tokens.cache_write - ); - assert_eq!(streaming.tokens.reasoning, materialized.tokens.reasoning); - assert!((streaming.cost - materialized.cost).abs() < 1e-9); - } + #[test] + #[serial_test::serial] + fn test_latest_source_mtime_ms_probes_grok_events() { + latest_source_mtime_ms_probes_grok_sibling("events.jsonl"); } #[test] - fn test_latest_source_mtime_ms_probes_m13_related_sources() { + #[serial_test::serial] + fn grok_streaming_cache_hits_refresh_derived_date_before_filtering() { let source_home = tempfile::TempDir::new().unwrap(); - let paths = write_m13_primary_fixtures(source_home.path()); - write_m13_related_fixtures(&paths); - let stale_time = - std::time::SystemTime::UNIX_EPOCH + std::time::Duration::from_secs(1_700_000_000); - let fresh_time = - std::time::SystemTime::UNIX_EPOCH + std::time::Duration::from_secs(1_700_086_400); - if !set_m13_fixture_mtimes(&paths, stale_time, fresh_time) { - return; - } + let cache_home = tempfile::TempDir::new().unwrap(); + let _env = EnvGuard::set(&[ + ("HOME", cache_home.path().as_os_str()), + ("TOKSCALE_CONFIG_DIR", cache_home.path().as_os_str()), + ("TOKSCALE_PRICING_CACHE_ONLY", std::ffi::OsStr::new("1")), + ]); + let logs_dir = source_home.path().join(".grok/logs"); + std::fs::create_dir_all(&logs_dir).unwrap(); + let unified = logs_dir.join("unified.jsonl"); + std::fs::write( + &unified, + "{\"ts\":\"2023-11-14T22:13:20Z\",\"pid\":7,\"sid\":\"cached\",\"msg\":\"shell.turn.inference_done\",\"ctx\":{\"loop_index\":1,\"prompt_tokens\":10,\"completion_tokens\":2}}\n", + ) + .unwrap(); - for client in [ClientId::Droid, ClientId::Kimi, ClientId::Kiro] { - let token = latest_source_mtime_ms(&LocalParseOptions { - home_dir: Some(source_home.path().to_str().unwrap().to_string()), - use_env_roots: false, - clients: Some(vec![client.as_str().to_string()]), - since: None, - until: None, - year: None, - scanner_settings: scanner::ScannerSettings::default(), - modified_after: None, - }) - .unwrap(); - assert_eq!( - token, - 1_700_086_400_000, - "{} change token must include its parser dependency", - client.as_str() - ); - } + let mut cached_messages = sessions::grok::parse_grok_unified_log_file(&unified); + assert_eq!(cached_messages.len(), 1); + let expected_date = cached_messages[0].date.clone(); + cached_messages[0].date = "stale-cached-date".to_string(); + let fingerprint = message_cache::SourceFingerprint::from_grok_path(&unified).unwrap(); + let mut cache = message_cache::SourceMessageCache::load(); + cache.insert(message_cache::CachedSourceEntry::new( + message_cache::CacheIdentity::for_client(ClientId::Grok), + &unified, + fingerprint, + cached_messages, + Vec::new(), + None, + )); + cache.save_if_dirty(); + + let clients = vec!["grok".to_string()]; + let mut streamed = Vec::new(); + scan_messages_streaming( + source_home.path().to_str().unwrap(), + &clients, + None, + false, + &scanner::ScannerSettings::default(), + &|message| message.date == expected_date, + &mut |message| streamed.push(message.clone()), + ); + + assert_eq!(streamed.len(), 1); + assert_eq!(streamed[0].date, expected_date); } #[test] #[serial_test::serial] - fn test_modified_after_keeps_m13_sources_with_fresh_dependencies() { + fn grok_materialized_reprices_after_legacy_model_carry_over() { let source_home = tempfile::TempDir::new().unwrap(); let cache_home = tempfile::TempDir::new().unwrap(); - let paths = write_m13_primary_fixtures(source_home.path()); - write_m13_related_fixtures(&paths); - let stale_time = - std::time::SystemTime::UNIX_EPOCH + std::time::Duration::from_secs(1_700_000_000); - let fresh_time = - std::time::SystemTime::UNIX_EPOCH + std::time::Duration::from_secs(1_700_086_400); - if !set_m13_fixture_mtimes(&paths, stale_time, fresh_time) { - return; - } + let _env = EnvGuard::set(&[ + ("HOME", cache_home.path().as_os_str()), + ("TOKSCALE_CONFIG_DIR", cache_home.path().as_os_str()), + ]); + let covered_dir = source_home + .path() + .join(".grok/sessions/%2Ftmp%2Fproject/covered"); + std::fs::create_dir_all(&covered_dir).unwrap(); + std::fs::write( + covered_dir.join("updates.jsonl"), + concat!( + "{\"method\":\"session/update\",\"params\":{\"sessionId\":\"covered\",\"update\":{\"sessionUpdate\":\"user_message_chunk\",\"_meta\":{\"modelId\":\"grok-build\"}},\"_meta\":{\"agentTimestampMs\":1700000000000}}}\n", + "{\"method\":\"session/update\",\"params\":{\"sessionId\":\"covered\",\"update\":{\"sessionUpdate\":\"agent_message_chunk\"},\"_meta\":{\"totalTokens\":999,\"agentTimestampMs\":1700000001000}}}\n", + ), + ) + .unwrap(); + let logs_dir = source_home.path().join(".grok/logs"); + std::fs::create_dir_all(&logs_dir).unwrap(); + std::fs::write( + logs_dir.join("unified.jsonl"), + "{\"ts\":\"2023-11-14T22:13:20Z\",\"pid\":7,\"sid\":\"covered\",\"msg\":\"shell.turn.inference_done\",\"ctx\":{\"loop_index\":1,\"prompt_tokens\":100,\"cached_prompt_tokens\":60,\"completion_tokens\":25,\"reasoning_tokens\":5}}\n", + ) + .unwrap(); - let parsed = with_isolated_tokscale_cache(cache_home.path(), || { - parse_local_clients(LocalParseOptions { - home_dir: Some(source_home.path().to_str().unwrap().to_string()), - use_env_roots: false, - clients: Some(vec![ - "droid".to_string(), - "kimi".to_string(), - "kiro".to_string(), - ]), - since: None, - until: None, - year: None, - scanner_settings: scanner::ScannerSettings::default(), - modified_after: Some(1_700_043_200_000), - }) - .unwrap() - }); + let mut litellm = HashMap::new(); + litellm.insert( + "grok-build".to_string(), + pricing::ModelPricing { + input_cost_per_token: Some(0.001), + output_cost_per_token: Some(0.002), + cache_read_input_token_cost: Some(0.0005), + ..Default::default() + }, + ); + let pricing = pricing::PricingService::new(litellm, HashMap::new()); + let clients = vec!["grok".to_string()]; + let scanner_settings = scanner::ScannerSettings::default(); + let materialized = || { + parse_all_messages_with_pricing_with_env_strategy( + source_home.path().to_str().unwrap(), + &clients, + Some(&pricing), + false, + &scanner_settings, + ) + }; + let streamed = || { + let mut messages = Vec::new(); + scan_messages_streaming( + source_home.path().to_str().unwrap(), + &clients, + Some(&pricing), + false, + &scanner_settings, + &|_| true, + &mut |message| messages.push(message.clone()), + ); + messages + }; - assert_eq!(parsed.messages.len(), 3); - assert_eq!(parsed.counts.get(ClientId::Droid), 1); - assert_eq!(parsed.counts.get(ClientId::Kimi), 1); - assert_eq!(parsed.counts.get(ClientId::Kiro), 1); - assert!(parsed - .messages - .iter() - .any(|message| { message.client == "droid" && message.model_id == "claude-opus-4-5" })); - assert!(parsed - .messages - .iter() - .any(|message| { message.client == "kimi" && message.model_id == "kimi-new-model" })); - assert!(parsed.messages.iter().any(|message| { - message.client == "kiro" && message.output == 4 && message.duration_ms == Some(1_000) - })); + let cold = materialized(); + assert_eq!(cold.len(), 1); + assert_eq!(cold[0].model_id, "grok-build"); + let expected_cost = pricing.calculate_cost_with_provider( + &cold[0].model_id, + Some(&cold[0].provider_id), + &cold[0].tokens, + ); + assert!(expected_cost > 0.0); + assert_eq!(cold[0].cost, expected_cost); + assert_eq!(cold[0].cost_source, CostSource::Estimated); + assert_eq!(streamed(), cold); + assert_eq!(materialized(), cold, "warm materialized cache must reprice"); + assert_eq!(streamed(), cold, "warm streaming cache must stay in parity"); } #[test] - fn test_modified_after_prunes_stale_m13_sources_without_dependencies() { + #[serial_test::serial] + fn m17_grok_unified_precedence_tracks_source_lifecycle_across_all_lanes() { let source_home = tempfile::TempDir::new().unwrap(); - let paths = write_m13_primary_fixtures(source_home.path()); - let mut scan_result = scanner::ScanResult::default(); - scan_result - .get_mut(ClientId::Droid) - .push(paths.droid_source); - scan_result.get_mut(ClientId::Kimi).push(paths.kimi_source); - scan_result.get_mut(ClientId::Kiro).push(paths.kiro_source); + let cache_home = tempfile::TempDir::new().unwrap(); + let _env = EnvGuard::set(&[ + ("HOME", cache_home.path().as_os_str()), + ("TOKSCALE_CONFIG_DIR", cache_home.path().as_os_str()), + ("TOKSCALE_PRICING_CACHE_ONLY", std::ffi::OsStr::new("1")), + ]); + let clients = vec!["grok".to_string()]; + let scanner_settings = scanner::ScannerSettings::default(); - let future_ms = std::time::SystemTime::now() - .duration_since(std::time::UNIX_EPOCH) - .unwrap() - .as_millis() as u64 - + 3_600_000; - prune_scan_result_by_mtime(&mut scan_result, future_ms); + let covered_dir = source_home + .path() + .join(".grok/sessions/%2Ftmp%2Fproject/covered"); + let legacy_only_dir = source_home + .path() + .join(".grok/sessions/%2Ftmp%2Fproject/legacy-only"); + std::fs::create_dir_all(&covered_dir).unwrap(); + std::fs::create_dir_all(&legacy_only_dir).unwrap(); + let covered_updates = covered_dir.join("updates.jsonl"); + let legacy_only_updates = legacy_only_dir.join("updates.jsonl"); + std::fs::write( + &covered_updates, + concat!( + "{\"method\":\"session/update\",\"params\":{\"sessionId\":\"covered\",\"update\":{\"sessionUpdate\":\"user_message_chunk\",\"_meta\":{\"modelId\":\"grok-build\"}},\"_meta\":{\"agentTimestampMs\":1700000000000}}}\n", + "{\"method\":\"session/update\",\"params\":{\"sessionId\":\"covered\",\"update\":{\"sessionUpdate\":\"agent_message_chunk\"},\"_meta\":{\"totalTokens\":999,\"agentTimestampMs\":1700000001000}}}\n", + ), + ) + .unwrap(); + std::fs::write( + &legacy_only_updates, + concat!( + "{\"method\":\"session/update\",\"params\":{\"sessionId\":\"legacy-only\",\"update\":{\"sessionUpdate\":\"user_message_chunk\",\"_meta\":{\"modelId\":\"grok-build\"}},\"_meta\":{\"agentTimestampMs\":1700000010000}}}\n", + "{\"method\":\"session/update\",\"params\":{\"sessionId\":\"legacy-only\",\"update\":{\"sessionUpdate\":\"agent_message_chunk\"},\"_meta\":{\"totalTokens\":17,\"agentTimestampMs\":1700000011000}}}\n", + ), + ) + .unwrap(); + let legacy_time = + std::time::SystemTime::UNIX_EPOCH + std::time::Duration::from_secs(1_700_086_400); + for path in [&covered_updates, &legacy_only_updates] { + let file = std::fs::OpenOptions::new().write(true).open(path).unwrap(); + if file.set_modified(legacy_time).is_err() { + return; + } + } - assert!(scan_result.get(ClientId::Droid).is_empty()); - assert!(scan_result.get(ClientId::Kimi).is_empty()); - assert!(scan_result.get(ClientId::Kiro).is_empty()); - } + let local_options = || LocalParseOptions { + home_dir: Some(source_home.path().to_string_lossy().into_owned()), + use_env_roots: false, + clients: Some(clients.clone()), + since: None, + until: None, + year: None, + scanner_settings: scanner_settings.clone(), + modified_after: None, + }; + let report_options = || ReportOptions { + home_dir: Some(source_home.path().to_string_lossy().into_owned()), + use_env_roots: false, + clients: Some(clients.clone()), + scanner_settings: scanner_settings.clone(), + ..Default::default() + }; + let sorted = |mut messages: Vec| { + messages.sort_by(|left, right| left.dedup_key.cmp(&right.dedup_key)); + messages + }; + let materialized = || { + sorted(parse_all_messages_with_pricing_with_env_strategy( + source_home.path().to_str().unwrap(), + &clients, + None, + false, + &scanner_settings, + )) + }; + let streamed = || { + let mut messages = Vec::new(); + scan_messages_streaming( + source_home.path().to_str().unwrap(), + &clients, + None, + false, + &scanner_settings, + &|_| true, + &mut |message| messages.push(message.clone()), + ); + sorted(messages) + }; - #[cfg(unix)] - #[test] - fn test_modified_after_m13_dependency_stat_failures_keep_sources() { - let source_home = tempfile::TempDir::new().unwrap(); - let paths = write_m13_primary_fixtures(source_home.path()); - for related in [ - &paths.droid_related, - &paths.kimi_related, - &paths.kiro_related, - ] { - std::os::unix::fs::symlink(related.file_name().unwrap(), related).unwrap(); + let legacy = materialized(); + assert_eq!(legacy.len(), 2); + assert_eq!( + legacy + .iter() + .map(|message| message.tokens.total()) + .sum::(), + 1_016 + ); + assert_eq!(materialized(), legacy, "legacy cache hits must stay stable"); + let before_unified = latest_source_mtime_ms(&local_options()).unwrap(); + let legacy_change_token = local_source_change_token(&local_options()).unwrap(); + assert_eq!(before_unified, 1_700_086_400_000); + + let logs_dir = source_home.path().join(".grok/logs"); + std::fs::create_dir_all(&logs_dir).unwrap(); + let unified = logs_dir.join("unified.jsonl"); + std::fs::write( + &unified, + "{\"ts\":\"2023-11-14T22:13:20Z\",\"pid\":7,\"sid\":\"covered\",\"msg\":\"shell.turn.inference_done\",\"ctx\":{\"loop_index\":1,\"prompt_tokens\":100,\"cached_prompt_tokens\":60,\"completion_tokens\":25,\"reasoning_tokens\":5}}\n", + ) + .unwrap(); + let unified_time = + std::time::SystemTime::UNIX_EPOCH + std::time::Duration::from_secs(1_700_000_000); + let unified_file = std::fs::OpenOptions::new() + .write(true) + .open(&unified) + .unwrap(); + if unified_file.set_modified(unified_time).is_err() { + return; } + assert_eq!( + latest_source_mtime_ms(&local_options()).unwrap(), + before_unified, + "the newer legacy mtime deliberately masks unified topology changes" + ); + let selected_change_token = local_source_change_token(&local_options()).unwrap(); + assert_ne!(selected_change_token, legacy_change_token); - let mut scan_result = scanner::ScanResult::default(); - scan_result - .get_mut(ClientId::Droid) - .push(paths.droid_source.clone()); - scan_result - .get_mut(ClientId::Kimi) - .push(paths.kimi_source.clone()); - scan_result - .get_mut(ClientId::Kiro) - .push(paths.kiro_source.clone()); - let future_ms = std::time::SystemTime::now() - .duration_since(std::time::UNIX_EPOCH) - .unwrap() - .as_millis() as u64 - + 3_600_000; - prune_scan_result_by_mtime(&mut scan_result, future_ms); + let selected = materialized(); + assert_eq!(selected.len(), 2); + assert_eq!( + selected + .iter() + .find(|message| message.session_id == "covered") + .unwrap() + .model_id, + "grok-build" + ); + let selected_tokens = + selected + .iter() + .fold(TokenBreakdown::default(), |mut total, message| { + total.input += message.tokens.input; + total.output += message.tokens.output; + total.cache_read += message.tokens.cache_read; + total.cache_write += message.tokens.cache_write; + total.reasoning += message.tokens.reasoning; + total + }); + assert_eq!( + selected_tokens, + TokenBreakdown { + input: 57, + output: 20, + cache_read: 60, + cache_write: 0, + reasoning: 5, + } + ); + assert!(selected.iter().any(|message| { + message.session_id == "covered" + && message + .dedup_key + .as_deref() + .is_some_and(|key| key.starts_with("grok-unified:")) + })); + assert!(selected + .iter() + .any(|message| message.session_id == "legacy-only")); + assert_eq!( + materialized(), + selected, + "selected cache hits must stay stable" + ); + assert_eq!(streamed(), selected); + let counted = parse_local_clients(local_options()).unwrap(); + assert_eq!(counted.counts.get(ClientId::Grok), 2); + assert_eq!(counted.messages.len(), 2); assert_eq!( - scan_result.get(ClientId::Droid), - std::slice::from_ref(&paths.droid_source) + counted + .messages + .iter() + .map(|message| message.input) + .sum::(), + 57 ); assert_eq!( - scan_result.get(ClientId::Kimi), - std::slice::from_ref(&paths.kimi_source) + counted + .messages + .iter() + .map(|message| message.cache_read) + .sum::(), + 60 ); + + let runtime = tokio::runtime::Builder::new_current_thread() + .enable_all() + .build() + .unwrap(); + let assert_reports = |expected: (i32, i64, i64, i64, i64, i64)| { + let (messages, input, output, cache_read, cache_write, reasoning) = expected; + let model = runtime + .block_on(get_model_report(report_options())) + .unwrap(); + let monthly = runtime + .block_on(get_monthly_report(report_options())) + .unwrap(); + let hourly = runtime + .block_on(get_hourly_report(report_options())) + .unwrap(); + let agents = runtime + .block_on(get_agents_report(report_options())) + .unwrap(); + + assert_eq!( + ( + model.total_messages, + model.total_input, + model.total_output, + model.total_cache_read, + model.total_cache_write, + model + .entries + .iter() + .map(|entry| entry.reasoning) + .sum::(), + ), + expected + ); + assert_eq!( + monthly + .entries + .iter() + .fold((0, 0, 0, 0, 0), |totals, entry| ( + totals.0 + entry.message_count, + totals.1 + entry.input, + totals.2 + entry.output, + totals.3 + entry.cache_read, + totals.4 + entry.cache_write, + ),), + (messages, input, output, cache_read, cache_write) + ); + assert_eq!( + hourly + .entries + .iter() + .fold((0, 0, 0, 0, 0, 0), |totals, entry| ( + totals.0 + entry.message_count, + totals.1 + entry.input, + totals.2 + entry.output, + totals.3 + entry.cache_read, + totals.4 + entry.cache_write, + totals.5 + entry.reasoning, + ),), + expected + ); + assert_eq!( + agents + .entries + .iter() + .fold((0, 0, 0, 0, 0), |totals, entry| ( + totals.0 + entry.input, + totals.1 + entry.output, + totals.2 + entry.cache_read, + totals.3 + entry.cache_write, + totals.4 + entry.reasoning, + ),), + (input, output, cache_read, cache_write, reasoning) + ); + assert_eq!(agents.total_messages, messages); + }; + assert_reports((2, 57, 20, 60, 0, 5)); + + std::fs::remove_file(&unified).unwrap(); assert_eq!( - scan_result.get(ClientId::Kiro), - std::slice::from_ref(&paths.kiro_source) + latest_source_mtime_ms(&local_options()).unwrap(), + before_unified ); + assert_eq!( + local_source_change_token(&local_options()).unwrap(), + legacy_change_token, + "deleting a non-max authority source must still invalidate consumers" + ); + let restored = materialized(); + assert_eq!(restored, legacy); + assert_eq!(streamed(), legacy); + let restored_count = parse_local_clients(local_options()).unwrap(); + assert_eq!(restored_count.counts.get(ClientId::Grok), 2); + assert_reports((2, 1_016, 0, 0, 0, 0)); + assert!(message_cache::SourceMessageCache::load() + .get( + message_cache::CacheIdentity::for_client(ClientId::Grok), + &unified, + ) + .is_none()); } - // micode `.db` is WAL-mode SQLite reached via the generic `*.db` glob, so it - // must be exempt from mtime pruning (a WAL-only write leaves the main db's - // mtime untouched) — same treatment as Antigravity CLI / Hermes / Zed. #[test] - fn test_modified_after_never_prunes_micode_dbs() { - let temp_dir = tempfile::TempDir::new().unwrap(); - let micode_db = temp_dir.path().join("micode.db"); - std::fs::File::create(&micode_db).unwrap(); - let claude_log = temp_dir.path().join("session.jsonl"); - std::fs::File::create(&claude_log).unwrap(); + fn test_latest_source_mtime_ms_probes_auto_discovered_hermes_profile_wal() { + let source_home = tempfile::TempDir::new().unwrap(); + let profile_dir = source_home.path().join(".hermes/profiles/research"); + std::fs::create_dir_all(&profile_dir).unwrap(); + let db = profile_dir.join("state.db"); + let wal = profile_dir.join("state.db-wal"); + std::fs::File::create(&db).unwrap(); + std::fs::File::create(&wal).unwrap(); - let mut scan_result = scanner::ScanResult::default(); - scan_result - .get_mut(ClientId::MiMoCode) - .push(micode_db.clone()); - scan_result - .get_mut(ClientId::Claude) - .push(claude_log.clone()); + let db_time = + std::time::SystemTime::UNIX_EPOCH + std::time::Duration::from_secs(1_700_000_000); + let wal_time = + std::time::SystemTime::UNIX_EPOCH + std::time::Duration::from_secs(1_700_086_400); + let db_file = std::fs::OpenOptions::new().write(true).open(&db).unwrap(); + let Ok(()) = db_file.set_modified(db_time) else { + return; + }; + drop(db_file); + let wal_file = std::fs::OpenOptions::new().write(true).open(&wal).unwrap(); + let Ok(()) = wal_file.set_modified(wal_time) else { + return; + }; + drop(wal_file); - let future_ms = std::time::SystemTime::now() - .duration_since(std::time::UNIX_EPOCH) - .unwrap() - .as_millis() as u64 - + 3_600_000; - crate::prune_scan_result_by_mtime(&mut scan_result, future_ms); + let options = LocalParseOptions { + home_dir: Some(source_home.path().to_str().unwrap().to_string()), + use_env_roots: false, + clients: Some(vec!["hermes".to_string()]), + since: None, + until: None, + year: None, + scanner_settings: scanner::ScannerSettings::default(), + modified_after: None, + }; + let token = crate::latest_source_mtime_ms(&options).unwrap(); assert_eq!( - scan_result.get(ClientId::MiMoCode), - std::slice::from_ref(&micode_db), - "micode .db (a WAL-mode SQLite source) must survive mtime pruning" - ); - assert!( - scan_result.get(ClientId::Claude).is_empty(), - "a plain-file client's stale log is still pruned" + token, 1_700_086_400_000, + "the change token must include an auto-discovered profile WAL" ); } - // gjc (`$GJC_CODING_AGENT_DIR/sessions/*.jsonl`) must be discovered via the - // EnvVar fallback root (`.gjc/agent`) + `*.jsonl` glob and flow through the - // streaming lane, keeping its authoritative embedded `usage.cost.total` - // (A1). With pricing absent the guard's reprice branch is a no-op; the - // materialized path mirrors upstream's proven reprice-when-absent guard. + // The live-tail change token must move when jcode appends to the sibling + // `.journal.jsonl` even though the snapshot mtime is unchanged; otherwise + // UsageTail short-circuits and never reflects the new turn. #[test] #[serial_test::serial] - fn test_streaming_gjc_flows_with_authoritative_cost() { - let cache_home = tempfile::TempDir::new().unwrap(); + fn test_latest_source_mtime_ms_probes_jcode_journal() { let source_home = tempfile::TempDir::new().unwrap(); - let _env = EnvGuard::set(&[ - ("HOME", cache_home.path().as_os_str()), - ("TOKSCALE_CONFIG_DIR", cache_home.path().as_os_str()), - ]); + let sessions_dir = source_home.path().join(".jcode/sessions"); + std::fs::create_dir_all(&sessions_dir).unwrap(); + let snapshot = sessions_dir.join("session_x.json"); + std::fs::write(&snapshot, br#"{"id":"session_x","messages":[]}"#).unwrap(); + let journal = sessions_dir.join("session_x.journal.jsonl"); + std::fs::write(&journal, b"{\"append_messages\":[]}\n").unwrap(); - { - let gjc_dir = source_home.path().join(".gjc/agent/sessions"); - std::fs::create_dir_all(&gjc_dir).unwrap(); - std::fs::write( - gjc_dir.join("test.jsonl"), - "{\"type\":\"session\",\"id\":\"gjc_ses_001\",\"cwd\":\"/work/pi\"}\n{\"type\":\"message\",\"id\":\"msg_001\",\"message\":{\"role\":\"assistant\",\"model\":\"claude-sonnet-4\",\"provider\":\"anthropic\",\"timestamp\":1767225601000,\"usage\":{\"input\":100,\"output\":50,\"cost\":{\"total\":0.3}}}}\n", - ) - .unwrap(); + // Snapshot old, journal strictly newer — the journal-only append the + // probe must catch. Skip gracefully if the FS rejects set_modified. + let snapshot_time = + std::time::SystemTime::UNIX_EPOCH + std::time::Duration::from_secs(1_700_000_000); + let journal_time = + std::time::SystemTime::UNIX_EPOCH + std::time::Duration::from_secs(1_700_086_400); + let sf = std::fs::OpenOptions::new().write(true).open(&snapshot).unwrap(); + let Ok(()) = sf.set_modified(snapshot_time) else { + return; + }; + drop(sf); + let jf = std::fs::OpenOptions::new().write(true).open(&journal).unwrap(); + let Ok(()) = jf.set_modified(journal_time) else { + return; + }; + drop(jf); - let mut cost_sum = 0.0f64; - let mut count = 0usize; - scan_messages_streaming( - source_home.path().to_str().unwrap(), - &["gjc".to_string()], - None, - false, - &scanner::ScannerSettings::default(), - &|_m: &UnifiedMessage| true, - &mut |m: &UnifiedMessage| { - cost_sum += m.cost; - count += 1; - }, - ); + let options = LocalParseOptions { + home_dir: Some(source_home.path().to_str().unwrap().to_string()), + use_env_roots: false, + clients: Some(vec!["jcode".to_string()]), + since: None, + until: None, + year: None, + scanner_settings: scanner::ScannerSettings::default(), + modified_after: None, + }; + let token = crate::latest_source_mtime_ms(&options).unwrap(); - assert_eq!(count, 1, "the gjc assistant message must flow through the streaming lane"); - assert!( - (cost_sum - 0.3).abs() < 1e-9, - "authoritative gjc cost must reach the sink (got {cost_sum})" - ); - } + // The newer journal mtime must dominate; without the journal probe the + // token would stop at the older snapshot mtime (1_700_000_000_000). + assert_eq!( + token, 1_700_086_400_000, + "the change token must reflect the jcode journal mtime, not just the snapshot" + ); } - // jcode's `session_*.json` snapshot is a file-lane source whose sibling - // `.journal.jsonl` is appended between snapshot rewrites without touching - // the snapshot mtime, so it must be exempt from mtime pruning like the WAL - // db lanes — otherwise an active session with a stale snapshot is dropped - // and its recent journal turns vanish from the live tail. #[test] - fn test_modified_after_never_prunes_jcode_sessions() { + fn test_parse_local_clients_honors_scanner_extra_scan_paths_for_zed_threads_db() { let temp_dir = tempfile::TempDir::new().unwrap(); - let jcode_snapshot = temp_dir.path().join("session_x.json"); - std::fs::File::create(&jcode_snapshot).unwrap(); - let claude_log = temp_dir.path().join("session.jsonl"); - std::fs::File::create(&claude_log).unwrap(); + let extra_threads_dir = temp_dir.path().join("custom-zed/threads"); + std::fs::create_dir_all(&extra_threads_dir).unwrap(); + let threads_db = extra_threads_dir.join("threads.db"); + let conn = create_zed_sqlite_db(&threads_db); + insert_zed_thread(&conn, "zed-extra-thread", "claude-sonnet-4-5"); + drop(conn); - let mut scan_result = scanner::ScanResult::default(); - scan_result - .get_mut(ClientId::Jcode) - .push(jcode_snapshot.clone()); - scan_result - .get_mut(ClientId::Claude) - .push(claude_log.clone()); + let parsed_default = parse_local_clients(LocalParseOptions { + home_dir: Some(temp_dir.path().to_str().unwrap().to_string()), + use_env_roots: false, + clients: Some(vec!["zed".to_string()]), + since: None, + until: None, + year: None, + scanner_settings: scanner::ScannerSettings::default(), + modified_after: None, + }) + .unwrap(); + assert_eq!(parsed_default.counts.get(ClientId::Zed), 0); + assert!(parsed_default.messages.is_empty()); - let future_ms = std::time::SystemTime::now() - .duration_since(std::time::UNIX_EPOCH) - .unwrap() - .as_millis() as u64 - + 3_600_000; - crate::prune_scan_result_by_mtime(&mut scan_result, future_ms); + let mut extra_scan_paths = std::collections::BTreeMap::new(); + extra_scan_paths.insert("zed".to_string(), vec![extra_threads_dir]); + let parsed_with_settings = parse_local_clients(LocalParseOptions { + home_dir: Some(temp_dir.path().to_str().unwrap().to_string()), + use_env_roots: false, + clients: Some(vec!["zed".to_string()]), + since: None, + until: None, + year: None, + scanner_settings: scanner::ScannerSettings { + extra_scan_paths, + ..Default::default() + }, + modified_after: None, + }) + .unwrap(); + assert_eq!(parsed_with_settings.counts.get(ClientId::Zed), 1); + assert_eq!(parsed_with_settings.messages.len(), 1); + assert_eq!(parsed_with_settings.messages[0].client, "zed"); assert_eq!( - scan_result.get(ClientId::Jcode), - std::slice::from_ref(&jcode_snapshot), - "jcode snapshot (its journal sibling can change without it) must survive mtime pruning" + parsed_with_settings.messages[0].session_id, + "zed-extra-thread" ); - assert!( - scan_result.get(ClientId::Claude).is_empty(), - "a plain-file client's stale log is still pruned" + assert_eq!( + parsed_with_settings.messages[0].model_id, + "claude-sonnet-4-5" ); + assert_eq!(parsed_with_settings.messages[0].input, 42); + assert_eq!(parsed_with_settings.messages[0].output, 7); } #[test] - fn test_modified_after_prunes_grok_by_updates_or_signals_mtime() { + fn test_parse_local_clients_dedups_zed_threads_across_default_and_extra_dbs() { let temp_dir = tempfile::TempDir::new().unwrap(); - let stale_dir = temp_dir.path().join("stale"); - let active_dir = temp_dir.path().join("active"); - std::fs::create_dir_all(&stale_dir).unwrap(); - std::fs::create_dir_all(&active_dir).unwrap(); - - let stale_updates = stale_dir.join("updates.jsonl"); - let stale_signals = stale_dir.join("signals.json"); - let active_updates = active_dir.join("updates.jsonl"); - let active_signals = active_dir.join("signals.json"); - for path in [&stale_updates, &stale_signals, &active_updates, &active_signals] { - std::fs::File::create(path).unwrap(); - } - - let stale_time = - std::time::SystemTime::UNIX_EPOCH + std::time::Duration::from_secs(1_700_000_000); - let active_time = - std::time::SystemTime::UNIX_EPOCH + std::time::Duration::from_secs(1_700_086_400); - for path in [&stale_updates, &stale_signals, &active_updates] { - let file = std::fs::OpenOptions::new().write(true).open(path).unwrap(); - let Ok(()) = file.set_modified(stale_time) else { - return; - }; - } - let file = std::fs::OpenOptions::new() - .write(true) - .open(&active_signals) - .unwrap(); - let Ok(()) = file.set_modified(active_time) else { - return; - }; - let mut scan_result = scanner::ScanResult::default(); - scan_result - .get_mut(ClientId::Grok) - .extend([stale_updates.clone(), active_updates.clone()]); + // Place threads.db at the default platform path so the scanner finds it + // as `zed_db` AND we also pass it via extraScanPaths. + let default_threads_dir = temp_dir.path().join(".local/share/zed/threads"); + std::fs::create_dir_all(&default_threads_dir).unwrap(); + let default_db = default_threads_dir.join("threads.db"); + let conn = create_zed_sqlite_db(&default_db); + insert_zed_thread(&conn, "shared-zed-thread", "claude-sonnet-4-5"); + drop(conn); - crate::prune_scan_result_by_mtime(&mut scan_result, 1_700_043_200_000); + // Point extraScanPaths.zed at the same directory — dedup should prevent + // the thread from appearing twice. + let mut extra_scan_paths = std::collections::BTreeMap::new(); + extra_scan_paths.insert("zed".to_string(), vec![default_threads_dir.clone()]); + let parsed = parse_local_clients(LocalParseOptions { + home_dir: Some(temp_dir.path().to_str().unwrap().to_string()), + use_env_roots: false, + clients: Some(vec!["zed".to_string()]), + since: None, + until: None, + year: None, + scanner_settings: scanner::ScannerSettings { + extra_scan_paths, + ..Default::default() + }, + modified_after: None, + }) + .unwrap(); - assert_eq!( - scan_result.get(ClientId::Grok), - std::slice::from_ref(&active_updates), - "stale Grok sessions should be pruned while a fresh signals sibling keeps its session" - ); + // Should see exactly 1 message, not 2 (deduped by canonicalize). + assert_eq!(parsed.counts.get(ClientId::Zed), 1); + assert_eq!(parsed.messages.len(), 1); + assert_eq!(parsed.messages[0].session_id, "shared-zed-thread"); } - // The pruning helper must treat a fresh write to *any* metadata sibling the - // parser reads (not just signals.json) as source activity — otherwise a - // summary.json- or events.jsonl-only write (a late-arriving model id) lets an - // otherwise-stale session be pruned and its fresh model never re-parsed. - fn prune_grok_keeps_session_with_fresh_sibling(fresh_sibling: &str) { + #[test] + fn test_parse_local_clients_zed_extra_scan_paths_nonexistent_dir_is_silent() { let temp_dir = tempfile::TempDir::new().unwrap(); - let stale_dir = temp_dir.path().join("stale"); - let active_dir = temp_dir.path().join("active"); - std::fs::create_dir_all(&stale_dir).unwrap(); - std::fs::create_dir_all(&active_dir).unwrap(); - let stale_updates = stale_dir.join("updates.jsonl"); - let active_updates = active_dir.join("updates.jsonl"); - let active_fresh = active_dir.join(fresh_sibling); - // Every session file starts stale; only the one fresh sibling moves. - let mut all_paths = vec![stale_updates.clone(), active_updates.clone()]; - for name in message_cache::GROK_METADATA_SIBLINGS { - all_paths.push(stale_dir.join(name)); - all_paths.push(active_dir.join(name)); - } - for path in &all_paths { - std::fs::File::create(path).unwrap(); - } + let mut extra_scan_paths = std::collections::BTreeMap::new(); + extra_scan_paths.insert( + "zed".to_string(), + vec![temp_dir.path().join("does/not/exist")], + ); + let parsed = parse_local_clients(LocalParseOptions { + home_dir: Some(temp_dir.path().to_str().unwrap().to_string()), + use_env_roots: false, + clients: Some(vec!["zed".to_string()]), + since: None, + until: None, + year: None, + scanner_settings: scanner::ScannerSettings { + extra_scan_paths, + ..Default::default() + }, + modified_after: None, + }) + .unwrap(); - let stale_time = - std::time::SystemTime::UNIX_EPOCH + std::time::Duration::from_secs(1_700_000_000); - let active_time = - std::time::SystemTime::UNIX_EPOCH + std::time::Duration::from_secs(1_700_086_400); - for path in &all_paths { - let file = std::fs::OpenOptions::new().write(true).open(path).unwrap(); - let Ok(()) = file.set_modified(stale_time) else { - return; - }; - } - let file = std::fs::OpenOptions::new() - .write(true) - .open(&active_fresh) - .unwrap(); - let Ok(()) = file.set_modified(active_time) else { - return; - }; + assert_eq!(parsed.counts.get(ClientId::Zed), 0); + assert!(parsed.messages.is_empty()); + } - let mut scan_result = scanner::ScanResult::default(); - scan_result - .get_mut(ClientId::Grok) - .extend([stale_updates.clone(), active_updates.clone()]); + #[test] + fn test_parse_local_clients_dedups_default_and_auto_discovered_hermes_profile() { + let temp_dir = tempfile::TempDir::new().unwrap(); - crate::prune_scan_result_by_mtime(&mut scan_result, 1_700_043_200_000); + let default_dir = temp_dir.path().join(".hermes"); + std::fs::create_dir_all(&default_dir).unwrap(); + let default_db = default_dir.join("state.db"); + let default_conn = create_hermes_sqlite_db(&default_db); + insert_hermes_session( + &default_conn, + "shared-hermes-session", + "claude-sonnet-4", + 2, + 100, + 25, + 0.07, + ); + drop(default_conn); - assert_eq!( - scan_result.get(ClientId::Grok), - std::slice::from_ref(&active_updates), - "a fresh {fresh_sibling} sibling must keep its otherwise-stale Grok session" + let profile_dir = temp_dir.path().join(".hermes/profiles/director_planning"); + std::fs::create_dir_all(&profile_dir).unwrap(); + let profile_db = profile_dir.join("state.db"); + let profile_conn = create_hermes_sqlite_db(&profile_db); + insert_hermes_session( + &profile_conn, + "shared-hermes-session", + "claude-sonnet-4", + 9, + 999, + 999, + 9.99, ); - } + insert_hermes_session( + &profile_conn, + "profile-only-session", + "claude-sonnet-4", + 1, + 30, + 3, + 0.02, + ); + drop(profile_conn); - #[test] - fn test_modified_after_prunes_grok_keeps_session_with_fresh_summary() { - prune_grok_keeps_session_with_fresh_sibling("summary.json"); + let parsed = parse_local_clients(LocalParseOptions { + home_dir: Some(temp_dir.path().to_str().unwrap().to_string()), + use_env_roots: false, + clients: Some(vec!["hermes".to_string()]), + since: None, + until: None, + year: None, + scanner_settings: scanner::ScannerSettings::default(), + modified_after: None, + }) + .unwrap(); + + assert_eq!(parsed.counts.get(ClientId::Hermes), 3); + assert_eq!(parsed.messages.len(), 2); + let shared = parsed + .messages + .iter() + .find(|message| message.session_id == "shared-hermes-session") + .unwrap(); + assert_eq!(shared.input, 100); + assert_eq!(shared.output, 25); + assert!(parsed + .messages + .iter() + .any(|message| message.session_id == "profile-only-session")); } #[test] - fn test_modified_after_prunes_grok_keeps_session_with_fresh_events() { - prune_grok_keeps_session_with_fresh_sibling("events.jsonl"); - } + fn test_parse_local_clients_claude_filter_ignores_scanner_settings_opencode_db_paths() { + // Regression guard for the scanner client-filter bypass: even + // when `scanner.opencodeDbPaths` pins an external opencode db, + // a `--clients claude` request must NOT pull in OpenCode rows. + // Before the fix, the merge ran outside the OpenCode-enabled + // guard so user-pinned dbs leaked through both `messages` and + // `counts` (the latter is computed before the message-level + // client filter, so even the post-filter pipeline could not + // hide a leaked count). + let temp_dir = tempfile::TempDir::new().unwrap(); - // The live-tail change token must move when Grok rewrites *any* metadata - // sibling the parser reads, even though updates.jsonl is unchanged; otherwise - // UsageTail short-circuits and the session stays pinned to its fallback model. - fn latest_source_mtime_ms_probes_grok_sibling(fresh_sibling: &str) { - let source_home = tempfile::TempDir::new().unwrap(); - let session_dir = source_home - .path() - .join(".grok/sessions/%2Ftmp%2Fproject/session-uuid-1"); - std::fs::create_dir_all(&session_dir).unwrap(); - let updates = session_dir.join("updates.jsonl"); - std::fs::write(&updates, b"{\"totalTokens\":1}\n").unwrap(); - // Every sibling exists and is stale; only the target sibling is newer. - for name in message_cache::GROK_METADATA_SIBLINGS { - std::fs::write(session_dir.join(name), b"{}").unwrap(); - } + // Claude session: one assistant message, the only thing the + // filter should accept. + let claude_dir = temp_dir.path().join(".claude/projects/myproject"); + std::fs::create_dir_all(&claude_dir).unwrap(); + std::fs::write( + claude_dir.join("conversation.jsonl"), + r#"{"type":"assistant","timestamp":"2024-12-01T10:00:00.000Z","requestId":"req_001","message":{"id":"msg_001","model":"claude-3-5-sonnet","usage":{"input_tokens":100,"output_tokens":50}}} +"#, + ) + .unwrap(); - let stale_time = - std::time::SystemTime::UNIX_EPOCH + std::time::Duration::from_secs(1_700_000_000); - let fresh_time = - std::time::SystemTime::UNIX_EPOCH + std::time::Duration::from_secs(1_700_086_400); - for name in std::iter::once("updates.jsonl").chain(message_cache::GROK_METADATA_SIBLINGS) { - let f = std::fs::OpenOptions::new() - .write(true) - .open(session_dir.join(name)) - .unwrap(); - let Ok(()) = f.set_modified(stale_time) else { - return; - }; - } - let f = std::fs::OpenOptions::new() - .write(true) - .open(session_dir.join(fresh_sibling)) - .unwrap(); - let Ok(()) = f.set_modified(fresh_time) else { - return; - }; + // External opencode.db that the user has pinned via + // scanner.opencodeDbPaths. Without the fix, this would leak + // into the Claude-only result. + let outside_dir = temp_dir.path().join("elsewhere"); + std::fs::create_dir_all(&outside_dir).unwrap(); + let external_db = outside_dir.join("opencode.db"); + let conn = rusqlite::Connection::open(&external_db).unwrap(); + conn.execute_batch( + "PRAGMA journal_mode=WAL; + CREATE TABLE message ( + id TEXT PRIMARY KEY, + session_id TEXT NOT NULL, + data TEXT NOT NULL + );", + ) + .unwrap(); + conn.execute( + "INSERT INTO message (id, session_id, data) VALUES (?1, ?2, ?3)", + rusqlite::params![ + "leaked-opencode", + "should-not-show-up", + r#"{ + "role": "assistant", + "modelID": "claude-sonnet-4", + "providerID": "anthropic", + "tokens": { "input": 9999, "output": 9999, "reasoning": 0, "cache": { "read": 0, "write": 0 } }, + "time": { "created": 1700000000000.0 } + }"# + ], + ) + .unwrap(); + drop(conn); - let options = LocalParseOptions { - home_dir: Some(source_home.path().to_str().unwrap().to_string()), + let parsed = parse_local_clients(LocalParseOptions { + home_dir: Some(temp_dir.path().to_str().unwrap().to_string()), use_env_roots: false, - clients: Some(vec!["grok".to_string()]), + clients: Some(vec!["claude".to_string()]), since: None, until: None, year: None, - scanner_settings: scanner::ScannerSettings::default(), + scanner_settings: scanner::ScannerSettings { + opencode_db_paths: vec![external_db.clone()], + ..Default::default() + }, modified_after: None, - }; - let token = crate::latest_source_mtime_ms(&options).unwrap(); + }) + .unwrap(); assert_eq!( - token, 1_700_086_400_000, - "the change token must reflect the fresh {fresh_sibling} mtime, not just updates.jsonl" + parsed.counts.get(ClientId::OpenCode), + 0, + "OpenCode count must stay zero under a Claude-only filter even \ + when scanner.opencodeDbPaths is set" + ); + assert_eq!( + parsed.counts.get(ClientId::Claude), + 1, + "Claude message must still be counted" + ); + assert_eq!(parsed.messages.len(), 1); + assert_eq!(parsed.messages[0].client, "claude"); + assert!( + parsed.messages.iter().all(|m| m.client != "opencode"), + "no OpenCode messages may leak into a Claude-only result, got {:?}", + parsed.messages ); } #[test] - #[serial_test::serial] - fn test_latest_source_mtime_ms_probes_grok_summary() { - latest_source_mtime_ms_probes_grok_sibling("summary.json"); - } - - #[test] - #[serial_test::serial] - fn test_latest_source_mtime_ms_probes_grok_events() { - latest_source_mtime_ms_probes_grok_sibling("events.jsonl"); - } - - #[test] - fn test_latest_source_mtime_ms_probes_auto_discovered_hermes_profile_wal() { - let source_home = tempfile::TempDir::new().unwrap(); - let profile_dir = source_home.path().join(".hermes/profiles/research"); - std::fs::create_dir_all(&profile_dir).unwrap(); - let db = profile_dir.join("state.db"); - let wal = profile_dir.join("state.db-wal"); - std::fs::File::create(&db).unwrap(); - std::fs::File::create(&wal).unwrap(); - - let db_time = - std::time::SystemTime::UNIX_EPOCH + std::time::Duration::from_secs(1_700_000_000); - let wal_time = - std::time::SystemTime::UNIX_EPOCH + std::time::Duration::from_secs(1_700_086_400); - let db_file = std::fs::OpenOptions::new().write(true).open(&db).unwrap(); - let Ok(()) = db_file.set_modified(db_time) else { - return; - }; - drop(db_file); - let wal_file = std::fs::OpenOptions::new().write(true).open(&wal).unwrap(); - let Ok(()) = wal_file.set_modified(wal_time) else { - return; - }; - drop(wal_file); + fn test_parse_local_clients_claude_transcripts_count_only_usage_metadata() { + let temp_dir = tempfile::TempDir::new().unwrap(); + let transcripts_dir = temp_dir.path().join(".claude/transcripts"); + std::fs::create_dir_all(&transcripts_dir).unwrap(); + std::fs::write( + transcripts_dir.join("ses_123456789012345678901234567.jsonl"), + r#"{"type":"user","timestamp":"2026-04-01T10:00:00.000Z","message":{"content":"Wrapped prompt"}} +{"type":"assistant","timestamp":"2026-04-01T10:00:01.000Z","requestId":"req_wrapper","message":{"id":"msg_wrapper","model":"claude-sonnet-4","usage":{"input_tokens":123,"output_tokens":45,"cache_read_input_tokens":67,"cache_creation_input_tokens":8}}} +"#, + ) + .unwrap(); + std::fs::write( + transcripts_dir.join("ses_765432109876543210987654321.jsonl"), + r#"{"type":"user","timestamp":"2026-04-01T10:00:00.000Z","message":{"content":"Wrapped prompt"}} +{"type":"tool_use","timestamp":"2026-04-01T10:00:01.000Z","message":{"content":"Run tool"}} +{"type":"tool_result","timestamp":"2026-04-01T10:00:02.000Z","message":{"content":"Tool result"}} +"#, + ) + .unwrap(); - let options = LocalParseOptions { - home_dir: Some(source_home.path().to_str().unwrap().to_string()), + let parsed = parse_local_clients(LocalParseOptions { + home_dir: Some(temp_dir.path().to_str().unwrap().to_string()), use_env_roots: false, - clients: Some(vec!["hermes".to_string()]), + clients: Some(vec!["claude".to_string()]), since: None, until: None, year: None, scanner_settings: scanner::ScannerSettings::default(), modified_after: None, - }; - let token = crate::latest_source_mtime_ms(&options).unwrap(); + }) + .unwrap(); + assert_eq!(parsed.counts.get(ClientId::Claude), 1); + assert_eq!(parsed.messages.len(), 1); + assert_eq!(parsed.messages[0].client, "claude"); assert_eq!( - token, 1_700_086_400_000, - "the change token must include an auto-discovered profile WAL" + parsed.messages[0].session_id, + "ses_123456789012345678901234567" ); + assert_eq!(parsed.messages[0].model_id, "claude-sonnet-4"); + assert_eq!(parsed.messages[0].input, 123); + assert_eq!(parsed.messages[0].output, 45); + assert_eq!(parsed.messages[0].cache_read, 67); + assert_eq!(parsed.messages[0].cache_write, 8); } - // The live-tail change token must move when jcode appends to the sibling - // `.journal.jsonl` even though the snapshot mtime is unchanged; otherwise - // UsageTail short-circuits and never reflects the new turn. #[test] #[serial_test::serial] - fn test_latest_source_mtime_ms_probes_jcode_journal() { + fn test_parse_all_messages_refreshes_cc_mirror_provider_when_variant_metadata_changes() { + let cache_home = tempfile::TempDir::new().unwrap(); let source_home = tempfile::TempDir::new().unwrap(); - let sessions_dir = source_home.path().join(".jcode/sessions"); - std::fs::create_dir_all(&sessions_dir).unwrap(); - let snapshot = sessions_dir.join("session_x.json"); - std::fs::write(&snapshot, br#"{"id":"session_x","messages":[]}"#).unwrap(); - let journal = sessions_dir.join("session_x.journal.jsonl"); - std::fs::write(&journal, b"{\"append_messages\":[]}\n").unwrap(); + let _env = EnvGuard::set(&[ + ("HOME", cache_home.path().as_os_str()), + ("TOKSCALE_CONFIG_DIR", cache_home.path().as_os_str()), + ]); - // Snapshot old, journal strictly newer — the journal-only append the - // probe must catch. Skip gracefully if the FS rejects set_modified. - let snapshot_time = - std::time::SystemTime::UNIX_EPOCH + std::time::Duration::from_secs(1_700_000_000); - let journal_time = - std::time::SystemTime::UNIX_EPOCH + std::time::Duration::from_secs(1_700_086_400); - let sf = std::fs::OpenOptions::new().write(true).open(&snapshot).unwrap(); - let Ok(()) = sf.set_modified(snapshot_time) else { - return; - }; - drop(sf); - let jf = std::fs::OpenOptions::new().write(true).open(&journal).unwrap(); - let Ok(()) = jf.set_modified(journal_time) else { - return; - }; - drop(jf); + { + let variant_dir = source_home.path().join(".cc-mirror/kimi-code"); + let config_dir = source_home.path().join("mirror-configs/kimi-code"); + let project_dir = config_dir.join("projects/project-one"); + std::fs::create_dir_all(&project_dir).unwrap(); + std::fs::create_dir_all(&variant_dir).unwrap(); + let variant_path = variant_dir.join("variant.json"); + std::fs::write( + &variant_path, + serde_json::json!({ + "name": "kimi-code", + "provider": "kimi", + "configDir": config_dir, + }) + .to_string(), + ) + .unwrap(); + let session_path = project_dir.join("session.jsonl"); + std::fs::write( + &session_path, + r#"{"type":"assistant","timestamp":"2024-12-01T10:00:00.000Z","requestId":"req_001","message":{"id":"msg_001","model":"claude-3-5-sonnet","usage":{"input_tokens":100,"output_tokens":50}}} +"#, + ) + .unwrap(); - let options = LocalParseOptions { - home_dir: Some(source_home.path().to_str().unwrap().to_string()), - use_env_roots: false, - clients: Some(vec!["jcode".to_string()]), - since: None, - until: None, - year: None, - scanner_settings: scanner::ScannerSettings::default(), - modified_after: None, - }; - let token = crate::latest_source_mtime_ms(&options).unwrap(); + let first_messages = parse_all_messages_with_pricing( + source_home.path().to_str().unwrap(), + &["claude".to_string()], + None, + ); + assert_eq!(first_messages.len(), 1); + assert_eq!(first_messages[0].client, "cc-mirror/kimi-code"); + assert_eq!(first_messages[0].provider_id, "kimi"); - // The newer journal mtime must dominate; without the journal probe the - // token would stop at the older snapshot mtime (1_700_000_000_000). - assert_eq!( - token, 1_700_086_400_000, - "the change token must reflect the jcode journal mtime, not just the snapshot" - ); + std::fs::write( + &variant_path, + serde_json::json!({ + "name": "kimi-code", + "provider": "minimax", + "configDir": config_dir, + }) + .to_string(), + ) + .unwrap(); + + let refreshed_messages = parse_all_messages_with_pricing( + source_home.path().to_str().unwrap(), + &["claude".to_string()], + None, + ); + assert_eq!(refreshed_messages.len(), 1); + assert_eq!(refreshed_messages[0].client, "cc-mirror/kimi-code"); + assert_eq!(refreshed_messages[0].provider_id, "minimax"); + } } #[test] - fn test_parse_local_clients_honors_scanner_extra_scan_paths_for_zed_threads_db() { - let temp_dir = tempfile::TempDir::new().unwrap(); - let windows_threads_dir = temp_dir.path().join("AppData/Local/Zed/threads"); - std::fs::create_dir_all(&windows_threads_dir).unwrap(); - let threads_db = windows_threads_dir.join("threads.db"); - let conn = create_zed_sqlite_db(&threads_db); - insert_zed_thread(&conn, "zed-extra-thread", "claude-sonnet-4-5"); - drop(conn); + #[serial_test::serial] + fn test_parse_all_messages_keeps_normal_claude_when_cc_mirror_points_at_claude_config() { + let cache_home = tempfile::TempDir::new().unwrap(); + let source_home = tempfile::TempDir::new().unwrap(); + let _env = EnvGuard::set(&[ + ("HOME", cache_home.path().as_os_str()), + ("TOKSCALE_CONFIG_DIR", cache_home.path().as_os_str()), + ]); - let parsed_default = parse_local_clients(LocalParseOptions { - home_dir: Some(temp_dir.path().to_str().unwrap().to_string()), - use_env_roots: false, - clients: Some(vec!["zed".to_string()]), - since: None, - until: None, - year: None, - scanner_settings: scanner::ScannerSettings::default(), - modified_after: None, - }) - .unwrap(); - assert_eq!(parsed_default.counts.get(ClientId::Zed), 0); - assert!(parsed_default.messages.is_empty()); + { + let claude_dir = source_home.path().join(".claude"); + let project_dir = claude_dir.join("projects/project-one"); + std::fs::create_dir_all(&project_dir).unwrap(); + let session_path = project_dir.join("session.jsonl"); + std::fs::write( + &session_path, + r#"{"type":"assistant","timestamp":"2024-12-01T10:00:00.000Z","requestId":"req_001","message":{"id":"msg_001","model":"claude-3-5-sonnet","usage":{"input_tokens":100,"output_tokens":50}}} +"#, + ) + .unwrap(); - let mut extra_scan_paths = std::collections::BTreeMap::new(); - extra_scan_paths.insert("zed".to_string(), vec![windows_threads_dir]); - let parsed_with_settings = parse_local_clients(LocalParseOptions { - home_dir: Some(temp_dir.path().to_str().unwrap().to_string()), - use_env_roots: false, - clients: Some(vec!["zed".to_string()]), - since: None, - until: None, - year: None, - scanner_settings: scanner::ScannerSettings { - extra_scan_paths, - ..Default::default() - }, - modified_after: None, - }) - .unwrap(); + let variant_dir = source_home.path().join(".cc-mirror/plain-mirror"); + std::fs::create_dir_all(&variant_dir).unwrap(); + std::fs::write( + variant_dir.join("variant.json"), + serde_json::json!({ + "name": "plain-mirror", + "provider": "mirror", + "configDir": claude_dir, + }) + .to_string(), + ) + .unwrap(); - assert_eq!(parsed_with_settings.counts.get(ClientId::Zed), 1); - assert_eq!(parsed_with_settings.messages.len(), 1); - assert_eq!(parsed_with_settings.messages[0].client, "zed"); - assert_eq!( - parsed_with_settings.messages[0].session_id, - "zed-extra-thread" - ); - assert_eq!( - parsed_with_settings.messages[0].model_id, - "claude-sonnet-4-5" - ); - assert_eq!(parsed_with_settings.messages[0].input, 42); - assert_eq!(parsed_with_settings.messages[0].output, 7); + let messages = parse_all_messages_with_pricing( + source_home.path().to_str().unwrap(), + &["claude".to_string()], + None, + ); + assert_eq!(messages.len(), 1); + assert_eq!(messages[0].client, "claude"); + } } #[test] - fn test_parse_local_clients_dedups_zed_threads_across_default_and_extra_dbs() { + fn test_parse_local_clients_amp_partial_ledger_recovers_message_fallback_day() { + use chrono::TimeZone; + let temp_dir = tempfile::TempDir::new().unwrap(); + let amp_dir = temp_dir.path().join(".local/share/amp/threads"); + std::fs::create_dir_all(&_dir).unwrap(); - // Place threads.db at the default platform path so the scanner finds it - // as `zed_db` AND we also pass it via extraScanPaths. - let default_threads_dir = temp_dir.path().join(".local/share/zed/threads"); - std::fs::create_dir_all(&default_threads_dir).unwrap(); - let default_db = default_threads_dir.join("threads.db"); - let conn = create_zed_sqlite_db(&default_db); - insert_zed_thread(&conn, "shared-zed-thread", "claude-sonnet-4-5"); - drop(conn); + let thread_created = chrono::DateTime::parse_from_rfc3339("2026-04-04T12:00:00Z") + .unwrap() + .timestamp_millis(); + let ledger_timestamp = chrono::DateTime::parse_from_rfc3339("2026-04-08T12:00:00Z") + .unwrap() + .timestamp_millis(); + + let thread = format!( + r#"{{ + "id": "thread-amp-gap", + "created": {thread_created}, + "usageLedger": {{ + "events": [ + {{ + "timestamp": "2026-04-08T12:00:00Z", + "model": "claude-sonnet-4-0", + "credits": 0.75, + "tokens": {{ "input": 100, "output": 20 }} + }} + ] + }}, + "messages": [ + {{ + "role": "assistant", + "messageId": 1, + "usage": {{ + "model": "claude-sonnet-4-0", + "inputTokens": 100, + "outputTokens": 20, + "credits": 0.75 + }} + }}, + {{ + "role": "assistant", + "messageId": 2, + "usage": {{ + "model": "claude-sonnet-4-0", + "inputTokens": 50, + "outputTokens": 10, + "credits": 0.40 + }} + }} + ] + }}"# + ); + std::fs::write(amp_dir.join("T-thread-amp-gap.json"), thread).unwrap(); - // Point extraScanPaths.zed at the same directory — dedup should prevent - // the thread from appearing twice. - let mut extra_scan_paths = std::collections::BTreeMap::new(); - extra_scan_paths.insert("zed".to_string(), vec![default_threads_dir.clone()]); let parsed = parse_local_clients(LocalParseOptions { home_dir: Some(temp_dir.path().to_str().unwrap().to_string()), use_env_roots: false, - clients: Some(vec!["zed".to_string()]), + clients: Some(vec!["amp".to_string()]), since: None, until: None, year: None, - scanner_settings: scanner::ScannerSettings { - extra_scan_paths, - ..Default::default() - }, + scanner_settings: scanner::ScannerSettings::default(), modified_after: None, }) .unwrap(); - // Should see exactly 1 message, not 2 (deduped by canonicalize). - assert_eq!(parsed.counts.get(ClientId::Zed), 1); - assert_eq!(parsed.messages.len(), 1); - assert_eq!(parsed.messages[0].session_id, "shared-zed-thread"); + assert_eq!(parsed.counts.get(ClientId::Amp), 2); + assert_eq!(parsed.messages.len(), 2); + + let dates: HashSet = parsed.messages.iter().map(|msg| msg.date.clone()).collect(); + let local_date = |timestamp_ms: i64| { + chrono::Local + .timestamp_millis_opt(timestamp_ms) + .single() + .unwrap() + .format("%Y-%m-%d") + .to_string() + }; + assert!(dates.contains(&local_date(thread_created + 2000))); + assert!(dates.contains(&local_date(ledger_timestamp))); } - #[test] - fn test_parse_local_clients_zed_extra_scan_paths_nonexistent_dir_is_silent() { - let temp_dir = tempfile::TempDir::new().unwrap(); + // ========================================================================= + // fold_messages_streaming parity tests (RED — fold_messages_streaming not yet impl) + // ========================================================================= - let mut extra_scan_paths = std::collections::BTreeMap::new(); - extra_scan_paths.insert( - "zed".to_string(), - vec![temp_dir.path().join("does/not/exist")], - ); - let parsed = parse_local_clients(LocalParseOptions { - home_dir: Some(temp_dir.path().to_str().unwrap().to_string()), - use_env_roots: false, - clients: Some(vec!["zed".to_string()]), - since: None, - until: None, - year: None, - scanner_settings: scanner::ScannerSettings { - extra_scan_paths, - ..Default::default() + /// Deterministic UnifiedMessage fixture helper shared with parity tests. + /// Uses no real JSONL files; all fields are constructed inline. + #[cfg(test)] + #[allow(clippy::too_many_arguments)] + fn parity_msg( + date: &str, + client: &str, + model: &str, + session_id: &str, + dedup_key: Option<&str>, + timestamp_ms: i64, + input: i64, + output: i64, + cost: f64, + ) -> crate::sessions::UnifiedMessage { + use crate::TokenBreakdown; + crate::sessions::UnifiedMessage { + client: client.to_string(), + model_id: model.to_string(), + provider_id: "anthropic".to_string(), + session_id: session_id.to_string(), + workspace_key: None, + workspace_label: None, + timestamp: timestamp_ms, + date: date.to_string(), + tokens: TokenBreakdown { + input, + output, + cache_read: 0, + cache_write: 0, + reasoning: 0, }, - modified_after: None, - }) - .unwrap(); + cost, + cost_source: crate::CostSource::Unknown, + duration_ms: None, + message_count: 1, + agent: None, + dedup_key: dedup_key.map(|s| s.to_string()), + dedup_aliases: Vec::new(), + is_turn_start: false, + } + } + + // A/B parity: fold_messages_streaming output == aggregate_by_date output + // for the same deterministic fixture (no dedup_keys, no trae). + #[test] + fn test_fold_messages_streaming_parity_with_aggregate_by_date_no_dedup() { + let messages = vec![ + parity_msg("2025-06-01", "claude", "claude-sonnet-4-5", "s1", None, + 1_748_000_000_000, 100, 50, 0.01), + parity_msg("2025-06-01", "opencode", "gpt-4o", "s2", None, + 1_748_000_001_000, 200, 100, 0.02), + parity_msg("2025-06-02", "codex", "gpt-5", "s3", None, + 1_748_086_400_000, 400, 200, 0.04), + ]; - assert_eq!(parsed.counts.get(ClientId::Zed), 0); - assert!(parsed.messages.is_empty()); + // Reference: existing aggregate_by_date (clone-based) + let reference = crate::aggregator::aggregate_by_date(messages.clone()); + + // Subject: new streaming path + let streaming = fold_messages_streaming(&messages); + + assert_eq!( + reference.len(), streaming.len(), + "parity: day bucket count must match" + ); + for (ref_day, stream_day) in reference.iter().zip(streaming.iter()) { + assert_eq!(ref_day.date, stream_day.date, "parity: date must match"); + assert_eq!( + ref_day.totals.tokens, stream_day.totals.tokens, + "parity: tokens must match for date {}", ref_day.date + ); + assert!( + (ref_day.totals.cost - stream_day.totals.cost).abs() < 1e-9, + "parity: cost must match for date {}", ref_day.date + ); + assert_eq!( + ref_day.totals.messages, stream_day.totals.messages, + "parity: message_count must match for date {}", ref_day.date + ); + } } + // A/B parity with cross-file dedup: fold_messages_streaming must apply + // the same dedup_key filtering as the existing pipeline does via seen_keys. #[test] - fn test_parse_local_clients_dedups_default_and_auto_discovered_hermes_profile() { - let temp_dir = tempfile::TempDir::new().unwrap(); + fn test_fold_messages_streaming_parity_cross_file_dedup() { + // Construct messages that include a duplicated dedup_key pair. + // The existing pipeline filters duplicates via the seen_keys HashSet. + // fold_messages_streaming must produce the same counts. + let unique = parity_msg("2025-06-10", "claude", "claude-sonnet-4-5", "u1", + Some("unique-key-1"), 1_749_000_000_000, 300, 150, 0.06); + let dup_first = parity_msg("2025-06-10", "claude", "claude-sonnet-4-5", "d1", + Some("dup-key-shared"), 1_749_000_001_000, 200, 100, 0.04); + let dup_second = parity_msg("2025-06-10", "claude", "claude-haiku-4-5", "d2", + Some("dup-key-shared"), 1_749_000_002_000, 200, 100, 0.04); - let default_dir = temp_dir.path().join(".hermes"); - std::fs::create_dir_all(&default_dir).unwrap(); - let default_db = default_dir.join("state.db"); - let default_conn = create_hermes_sqlite_db(&default_db); - insert_hermes_session( - &default_conn, - "shared-hermes-session", - "claude-sonnet-4", - 2, - 100, - 25, - 0.07, - ); - drop(default_conn); + // The reference pipeline keeps only the first occurrence of dup-key-shared + // (seen_keys.insert returns false on second) — 2 messages total. + let all_msgs = vec![unique.clone(), dup_first.clone(), dup_second.clone()]; - let profile_dir = temp_dir.path().join(".hermes/profiles/director_planning"); - std::fs::create_dir_all(&profile_dir).unwrap(); - let profile_db = profile_dir.join("state.db"); - let profile_conn = create_hermes_sqlite_db(&profile_db); - insert_hermes_session( - &profile_conn, - "shared-hermes-session", - "claude-sonnet-4", - 9, - 999, - 999, - 9.99, + let streaming = fold_messages_streaming(&all_msgs); + + assert_eq!(streaming.len(), 1, "all on same date -> 1 bucket"); + assert_eq!( + streaming[0].totals.messages, 2, + "parity: duplicate dedup_key must reduce count from 3 to 2" ); - insert_hermes_session( - &profile_conn, - "profile-only-session", - "claude-sonnet-4", - 1, - 30, - 3, - 0.02, + assert!( + (streaming[0].totals.cost - 0.10).abs() < 1e-9, + "parity: cost must exclude the duplicate message" ); - drop(profile_conn); - - let parsed = parse_local_clients(LocalParseOptions { - home_dir: Some(temp_dir.path().to_str().unwrap().to_string()), - use_env_roots: false, - clients: Some(vec!["hermes".to_string()]), - since: None, - until: None, - year: None, - scanner_settings: scanner::ScannerSettings::default(), - modified_after: None, - }) - .unwrap(); - - assert_eq!(parsed.counts.get(ClientId::Hermes), 3); - assert_eq!(parsed.messages.len(), 2); - let shared = parsed - .messages - .iter() - .find(|message| message.session_id == "shared-hermes-session") - .unwrap(); - assert_eq!(shared.input, 100); - assert_eq!(shared.output, 25); - assert!(parsed - .messages - .iter() - .any(|message| message.session_id == "profile-only-session")); } - #[test] - fn test_parse_local_clients_claude_filter_ignores_scanner_settings_opencode_db_paths() { - // Regression guard for the scanner client-filter bypass: even - // when `scanner.opencodeDbPaths` pins an external opencode db, - // a `--clients claude` request must NOT pull in OpenCode rows. - // Before the fix, the merge ran outside the OpenCode-enabled - // guard so user-pinned dbs leaked through both `messages` and - // `counts` (the latter is computed before the message-level - // client filter, so even the post-filter pipeline could not - // hide a leaked count). - let temp_dir = tempfile::TempDir::new().unwrap(); + // ========================================================================= + // Phase 2 RED tests: build_graph_result_from_messages streaming entry-point + // ========================================================================= + // + // `build_graph_result_from_messages` does NOT exist yet. These tests + // define the observable contract that the GREEN implementation must satisfy: + // - accept a `&[UnifiedMessage]` slice and an optional `since` date string + // - apply the `since` post-parse filter (date string prefix comparison, + // same semantics as `filter_messages_for_report`) + // - drive aggregation via `StreamingAggregator` (zero-clone fold path) + // - return a `GraphResult` whose per-day tokens/cost match the reference + // `aggregate_by_date` pipeline exactly (zero tolerance) + // + // All three tests will produce a **compile error** until the function is + // declared in lib.rs, which is the required RED state. - // Claude session: one assistant message, the only thing the - // filter should accept. - let claude_dir = temp_dir.path().join(".claude/projects/myproject"); - std::fs::create_dir_all(&claude_dir).unwrap(); - std::fs::write( - claude_dir.join("conversation.jsonl"), - r#"{"type":"assistant","timestamp":"2024-12-01T10:00:00.000Z","requestId":"req_001","message":{"id":"msg_001","model":"claude-3-5-sonnet","usage":{"input_tokens":100,"output_tokens":50}}} -"#, - ) - .unwrap(); + use crate::GraphResult; - // External opencode.db that the user has pinned via - // scanner.opencodeDbPaths. Without the fix, this would leak - // into the Claude-only result. - let outside_dir = temp_dir.path().join("elsewhere"); - std::fs::create_dir_all(&outside_dir).unwrap(); - let external_db = outside_dir.join("opencode.db"); - let conn = rusqlite::Connection::open(&external_db).unwrap(); - conn.execute_batch( - "PRAGMA journal_mode=WAL; - CREATE TABLE message ( - id TEXT PRIMARY KEY, - session_id TEXT NOT NULL, - data TEXT NOT NULL - );", - ) - .unwrap(); - conn.execute( - "INSERT INTO message (id, session_id, data) VALUES (?1, ?2, ?3)", - rusqlite::params![ - "leaked-opencode", - "should-not-show-up", - r#"{ - "role": "assistant", - "modelID": "claude-sonnet-4", - "providerID": "anthropic", - "tokens": { "input": 9999, "output": 9999, "reasoning": 0, "cache": { "read": 0, "write": 0 } }, - "time": { "created": 1700000000000.0 } - }"# - ], - ) - .unwrap(); - drop(conn); + /// Multi-client, multi-day fixture: streaming path total tokens and cost + /// per daily bucket must match hand-computed expected values. + /// + /// Hardcoded expected values — calculation: + /// + /// 2025-06-01 (no dedup): + /// s1: input=500, output=250 → tokens=750, cost=0.05 + /// s2: input=300, output=150 → tokens=450, cost=0.03 + /// TOTAL: tokens=1200, cost=0.08, messages=2 + /// + /// 2025-06-02 (trae session dedup — same session_id="trae-sess"): + /// trae-k1: ts=1_748_822_400_000, input=100, output=50 → tokens=150, cost=0.01 + /// trae-k2: ts=1_748_822_500_000, input=200, output=100 → tokens=300, cost=0.02 + /// StreamingAggregator keeps latest timestamp → trae-k2 wins + /// TOTAL: tokens=300, cost=0.02, messages=1 + /// + /// 2025-06-03 (cross-file dedup_key — both carry "dup-phase2"): + /// d1: input=400, output=200 → tokens=600, cost=0.04 (first seen — kept) + /// d2: input=400, output=200 → tokens=600, cost=0.04 (same dedup_key — dropped) + /// TOTAL: tokens=600, cost=0.04, messages=1 + #[test] + fn test_build_graph_result_from_messages_matches_aggregate_by_date() { + let messages = vec![ + // Day 2025-06-01: two clients, no dedup + parity_msg("2025-06-01", "claude", "claude-sonnet-4-5", "s1", None, + 1_748_736_000_000, 500, 250, 0.05), + parity_msg("2025-06-01", "opencode", "gpt-4o", "s2", None, + 1_748_736_001_000, 300, 150, 0.03), + // Day 2025-06-02: trae dedup by session_id — two entries same session, keep latest + parity_msg("2025-06-02", "trae", "gpt-5.2", "trae-sess", Some("trae-k1"), + 1_748_822_400_000, 100, 50, 0.01), + parity_msg("2025-06-02", "trae", "gpt-5.2", "trae-sess", Some("trae-k2"), + 1_748_822_500_000, 200, 100, 0.02), // newer timestamp -> wins + // Day 2025-06-03: cross-file dedup pair — same dedup_key, second dropped + parity_msg("2025-06-03", "claude", "claude-haiku-4-5", "d1", + Some("dup-phase2"), 1_748_908_800_000, 400, 200, 0.04), + parity_msg("2025-06-03", "claude", "claude-haiku-4-5", "d2", + Some("dup-phase2"), 1_748_908_801_000, 400, 200, 0.04), // same dedup_key -> discarded + ]; - let parsed = parse_local_clients(LocalParseOptions { - home_dir: Some(temp_dir.path().to_str().unwrap().to_string()), - use_env_roots: false, - clients: Some(vec!["claude".to_string()]), - since: None, - until: None, - year: None, - scanner_settings: scanner::ScannerSettings { - opencode_db_paths: vec![external_db.clone()], - ..Default::default() - }, - modified_after: None, - }) - .unwrap(); + // Subject: new streaming entry-point + let result: GraphResult = + crate::build_graph_result_from_messages(&messages, None); + // Verify bucket count: 3 distinct dates assert_eq!( - parsed.counts.get(ClientId::OpenCode), - 0, - "OpenCode count must stay zero under a Claude-only filter even \ - when scanner.opencodeDbPaths is set" + result.contributions.len(), 3, + "phase2 streaming: must produce exactly 3 daily buckets" + ); + + // Locate each day bucket by date (sort order: ascending) + let day1 = result.contributions.iter().find(|c| c.date == "2025-06-01") + .expect("phase2: 2025-06-01 bucket must exist"); + let day2 = result.contributions.iter().find(|c| c.date == "2025-06-02") + .expect("phase2: 2025-06-02 bucket must exist"); + let day3 = result.contributions.iter().find(|c| c.date == "2025-06-03") + .expect("phase2: 2025-06-03 bucket must exist"); + + // 2025-06-01: s1 (750) + s2 (450) = 1200 tokens, 0.05+0.03=0.08 cost, 2 messages + assert_eq!(day1.totals.tokens, 1200, + "2025-06-01: tokens must be 750+450=1200"); + assert!( + (day1.totals.cost - 0.08).abs() < 1e-9, + "2025-06-01: cost must be 0.05+0.03=0.08" ); - assert_eq!( - parsed.counts.get(ClientId::Claude), - 1, - "Claude message must still be counted" + assert_eq!(day1.totals.messages, 2, + "2025-06-01: both non-trae non-dedup messages must be counted"); + + // 2025-06-02: trae session dedup — trae-k2 wins (larger timestamp) + // trae-k2: input=200, output=100 -> tokens=300, cost=0.02 + assert_eq!(day2.totals.tokens, 300, + "2025-06-02: trae dedup — only winner (trae-k2, tokens=300) counted"); + assert!( + (day2.totals.cost - 0.02).abs() < 1e-9, + "2025-06-02: trae dedup — cost must be 0.02 (trae-k2 only)" ); - assert_eq!(parsed.messages.len(), 1); - assert_eq!(parsed.messages[0].client, "claude"); + assert_eq!(day2.totals.messages, 1, + "2025-06-02: trae dedup collapses 2 entries to 1 per session_id"); + + // 2025-06-03: cross-file dedup — d1 kept, d2 dropped (same dedup_key) + // d1: input=400, output=200 -> tokens=600, cost=0.04 + assert_eq!(day3.totals.tokens, 600, + "2025-06-03: cross-file dedup — only d1 (tokens=600) counted, d2 dropped"); assert!( - parsed.messages.iter().all(|m| m.client != "opencode"), - "no OpenCode messages may leak into a Claude-only result, got {:?}", - parsed.messages + (day3.totals.cost - 0.04).abs() < 1e-9, + "2025-06-03: cross-file dedup — cost must be 0.04 (d1 only)" ); + assert_eq!(day3.totals.messages, 1, + "2025-06-03: duplicate dedup_key dropped, 1 message retained"); } + /// `since` filter semantics: same fixture with `since = "2025-06-02"` must + /// produce only the 2025-06-02 and 2025-06-03 buckets, with their + /// token/cost totals matching a manually filtered reference. #[test] - fn test_parse_local_clients_claude_transcripts_count_only_usage_metadata() { - let temp_dir = tempfile::TempDir::new().unwrap(); - let transcripts_dir = temp_dir.path().join(".claude/transcripts"); - std::fs::create_dir_all(&transcripts_dir).unwrap(); - std::fs::write( - transcripts_dir.join("ses_123456789012345678901234567.jsonl"), - r#"{"type":"user","timestamp":"2026-04-01T10:00:00.000Z","message":{"content":"Wrapped prompt"}} -{"type":"assistant","timestamp":"2026-04-01T10:00:01.000Z","requestId":"req_wrapper","message":{"id":"msg_wrapper","model":"claude-sonnet-4","usage":{"input_tokens":123,"output_tokens":45,"cache_read_input_tokens":67,"cache_creation_input_tokens":8}}} -"#, - ) - .unwrap(); - std::fs::write( - transcripts_dir.join("ses_765432109876543210987654321.jsonl"), - r#"{"type":"user","timestamp":"2026-04-01T10:00:00.000Z","message":{"content":"Wrapped prompt"}} -{"type":"tool_use","timestamp":"2026-04-01T10:00:01.000Z","message":{"content":"Run tool"}} -{"type":"tool_result","timestamp":"2026-04-01T10:00:02.000Z","message":{"content":"Tool result"}} -"#, - ) - .unwrap(); + fn test_build_graph_result_from_messages_since_filter_excludes_earlier_dates() { + let messages = vec![ + parity_msg("2025-06-01", "claude", "claude-sonnet-4-5", "s1", None, + 1_748_736_000_000, 500, 250, 0.05), + parity_msg("2025-06-01", "opencode", "gpt-4o", "s2", None, + 1_748_736_001_000, 300, 150, 0.03), + parity_msg("2025-06-02", "codex", "gpt-5", "s3", None, + 1_748_822_400_000, 400, 200, 0.04), + parity_msg("2025-06-03", "claude", "claude-haiku-4-5", "s4", None, + 1_748_908_800_000, 200, 100, 0.02), + ]; - let parsed = parse_local_clients(LocalParseOptions { - home_dir: Some(temp_dir.path().to_str().unwrap().to_string()), - use_env_roots: false, - clients: Some(vec!["claude".to_string()]), - since: None, - until: None, - year: None, - scanner_settings: scanner::ScannerSettings::default(), - modified_after: None, - }) - .unwrap(); + // Subject: streaming entry with since = "2025-06-02" + // (function does not exist yet -> RED compile error) + let result: GraphResult = + crate::build_graph_result_from_messages(&messages, Some("2025-06-02")); - assert_eq!(parsed.counts.get(ClientId::Claude), 1); - assert_eq!(parsed.messages.len(), 1); - assert_eq!(parsed.messages[0].client, "claude"); + // Only 2025-06-02 and 2025-06-03 must be present assert_eq!( - parsed.messages[0].session_id, - "ses_123456789012345678901234567" + result.contributions.len(), 2, + "since filter: must exclude 2025-06-01, leaving 2 buckets" + ); + + let dates: Vec<&str> = result.contributions.iter().map(|c| c.date.as_str()).collect(); + assert!(dates.contains(&"2025-06-02"), + "since filter: 2025-06-02 bucket must be present"); + assert!(dates.contains(&"2025-06-03"), + "since filter: 2025-06-03 bucket must be present"); + assert!(!dates.contains(&"2025-06-01"), + "since filter: 2025-06-01 bucket must be absent"); + + // 2025-06-02 token total: input 400 + output 200 = 600 + let day2 = result.contributions.iter().find(|c| c.date == "2025-06-02").unwrap(); + assert_eq!(day2.totals.tokens, 600, + "since filter: 2025-06-02 token total must be 600"); + assert!( + (day2.totals.cost - 0.04).abs() < 1e-9, + "since filter: 2025-06-02 cost must be 0.04" ); - assert_eq!(parsed.messages[0].model_id, "claude-sonnet-4"); - assert_eq!(parsed.messages[0].input, 123); - assert_eq!(parsed.messages[0].output, 45); - assert_eq!(parsed.messages[0].cache_read, 67); - assert_eq!(parsed.messages[0].cache_write, 8); } + /// Trae dedup in streaming path: two messages for the same trae session + /// (same `session_id`, different `dedup_key`, later timestamp wins) must + /// produce exactly ONE message worth of tokens/cost in the daily bucket. #[test] - #[serial_test::serial] - fn test_parse_all_messages_refreshes_cc_mirror_provider_when_variant_metadata_changes() { - let cache_home = tempfile::TempDir::new().unwrap(); - let source_home = tempfile::TempDir::new().unwrap(); - let _env = EnvGuard::set(&[ - ("HOME", cache_home.path().as_os_str()), - ("TOKSCALE_CONFIG_DIR", cache_home.path().as_os_str()), - ]); - - { - let variant_dir = source_home.path().join(".cc-mirror/kimi-code"); - let config_dir = source_home.path().join("mirror-configs/kimi-code"); - let project_dir = config_dir.join("projects/project-one"); - std::fs::create_dir_all(&project_dir).unwrap(); - std::fs::create_dir_all(&variant_dir).unwrap(); - let variant_path = variant_dir.join("variant.json"); - std::fs::write( - &variant_path, - serde_json::json!({ - "name": "kimi-code", - "provider": "kimi", - "configDir": config_dir, - }) - .to_string(), - ) - .unwrap(); - let session_path = project_dir.join("session.jsonl"); - std::fs::write( - &session_path, - r#"{"type":"assistant","timestamp":"2024-12-01T10:00:00.000Z","requestId":"req_001","message":{"id":"msg_001","model":"claude-3-5-sonnet","usage":{"input_tokens":100,"output_tokens":50}}} -"#, - ) - .unwrap(); + fn test_build_graph_result_from_messages_trae_session_dedup_keeps_latest() { + let messages = vec![ + // Earlier trae message (should be dropped) + parity_msg("2025-06-10", "trae", "gpt-5.2", "trae-sess-a", Some("trae-early"), + 1_749_513_600_000, 100, 50, 0.01), + // Later trae message for same session_id (should win) + parity_msg("2025-06-10", "trae", "gpt-5.2", "trae-sess-a", Some("trae-late"), + 1_749_513_700_000, 300, 150, 0.03), + // Non-trae message (should be included as-is) + parity_msg("2025-06-10", "claude", "claude-sonnet-4-5", "c1", None, + 1_749_513_800_000, 200, 100, 0.02), + ]; - let first_messages = parse_all_messages_with_pricing( - source_home.path().to_str().unwrap(), - &["claude".to_string()], - None, - ); - assert_eq!(first_messages.len(), 1); - assert_eq!(first_messages[0].client, "cc-mirror/kimi-code"); - assert_eq!(first_messages[0].provider_id, "kimi"); + // Subject: streaming entry (does not exist yet -> RED compile error) + let result: GraphResult = + crate::build_graph_result_from_messages(&messages, None); - std::fs::write( - &variant_path, - serde_json::json!({ - "name": "kimi-code", - "provider": "minimax", - "configDir": config_dir, - }) - .to_string(), - ) - .unwrap(); + assert_eq!(result.contributions.len(), 1, + "trae dedup: all messages on same date -> 1 bucket"); - let refreshed_messages = parse_all_messages_with_pricing( - source_home.path().to_str().unwrap(), - &["claude".to_string()], - None, - ); - assert_eq!(refreshed_messages.len(), 1); - assert_eq!(refreshed_messages[0].client, "cc-mirror/kimi-code"); - assert_eq!(refreshed_messages[0].provider_id, "minimax"); - } + let day = &result.contributions[0]; + // Kept messages: trae-late (tokens=450) + claude (tokens=300) = 750 total tokens + assert_eq!(day.totals.tokens, 750, + "trae dedup: token total must reflect only the winning trae entry (450) + claude (300)"); + assert!( + (day.totals.cost - 0.05).abs() < 1e-9, + "trae dedup: cost must be 0.03 (latest trae) + 0.02 (claude) = 0.05" + ); + assert_eq!(day.totals.messages, 2, + "trae dedup: message count must be 2 (1 trae winner + 1 claude)"); } #[test] - #[serial_test::serial] - fn test_parse_all_messages_keeps_normal_claude_when_cc_mirror_points_at_claude_config() { - let cache_home = tempfile::TempDir::new().unwrap(); - let source_home = tempfile::TempDir::new().unwrap(); - let _env = EnvGuard::set(&[ - ("HOME", cache_home.path().as_os_str()), - ("TOKSCALE_CONFIG_DIR", cache_home.path().as_os_str()), - ]); - - { - let claude_dir = source_home.path().join(".claude"); - let project_dir = claude_dir.join("projects/project-one"); - std::fs::create_dir_all(&project_dir).unwrap(); - let session_path = project_dir.join("session.jsonl"); - std::fs::write( - &session_path, - r#"{"type":"assistant","timestamp":"2024-12-01T10:00:00.000Z","requestId":"req_001","message":{"id":"msg_001","model":"claude-3-5-sonnet","usage":{"input_tokens":100,"output_tokens":50}}} -"#, - ) - .unwrap(); - - let variant_dir = source_home.path().join(".cc-mirror/plain-mirror"); - std::fs::create_dir_all(&variant_dir).unwrap(); - std::fs::write( - variant_dir.join("variant.json"), - serde_json::json!({ - "name": "plain-mirror", - "provider": "mirror", - "configDir": claude_dir, - }) - .to_string(), + fn model_aggregation_saturates_overflowing_token_folds() { + // token_total_saturates_on_overlarge_buckets (see positive_token_total's + // callers) covers a single message's grand total; the per-field + // CROSS-MESSAGE fold in aggregate_model_usage_entries must saturate too. + // An antigravity-cli row can carry an i64::MAX bucket after the + // untrusted-varint clamp (sessions/antigravity_cli.rs to_i64), so two + // such rows folded into one model group with plain `+=` overflow (debug + // panic / release wrap) before the already-saturating grand total runs. + let make = || { + UnifiedMessage::new( + "antigravity-cli", + "gemini-3-pro", + "antigravity", + "session-overflow", + 1_733_011_200_000, + TokenBreakdown { + input: i64::MAX, + output: 0, + cache_read: i64::MAX, + cache_write: 0, + reasoning: 0, + }, + 0.0, ) - .unwrap(); + }; - let messages = parse_all_messages_with_pricing( - source_home.path().to_str().unwrap(), - &["claude".to_string()], - None, - ); - assert_eq!(messages.len(), 1); - assert_eq!(messages[0].client, "claude"); - } + let entries = aggregate_model_usage_entries(vec![make(), make()], &GroupBy::Model); + assert_eq!(entries.len(), 1); + assert_eq!(entries[0].input, i64::MAX); + assert_eq!(entries[0].cache_read, i64::MAX); } #[test] - fn test_parse_local_clients_amp_partial_ledger_recovers_message_fallback_day() { - use chrono::TimeZone; + fn model_report_totals_saturate_across_groups() { + // aggregate_model_usage_entries saturates each entry's fields, so an + // entry can be i64::MAX. get_model_report sums the entries into the + // report-level totals via model_report_token_totals; two saturated + // entries (two distinct models) must not overflow that sum either. + let make = |model: &str| { + UnifiedMessage::new( + "antigravity-cli", + model, + "antigravity", + "session-overflow", + 1_733_011_200_000, + TokenBreakdown { + input: i64::MAX, + output: 0, + cache_read: i64::MAX, + cache_write: 0, + reasoning: 0, + }, + 0.0, + ) + }; - let temp_dir = tempfile::TempDir::new().unwrap(); - let amp_dir = temp_dir.path().join(".local/share/amp/threads"); - std::fs::create_dir_all(&_dir).unwrap(); + let entries = aggregate_model_usage_entries( + vec![make("gemini-3-pro"), make("claude-opus-4-6")], + &GroupBy::Model, + ); + assert_eq!(entries.len(), 2); + let (total_input, _total_output, total_cache_read, _total_cache_write) = + super::model_report_token_totals(&entries); + assert_eq!(total_input, i64::MAX); + assert_eq!(total_cache_read, i64::MAX); + } - let thread_created = chrono::DateTime::parse_from_rfc3339("2026-04-04T12:00:00Z") - .unwrap() - .timestamp_millis(); - let ledger_timestamp = chrono::DateTime::parse_from_rfc3339("2026-04-08T12:00:00Z") - .unwrap() - .timestamp_millis(); + fn m15a_global_root(home: &Path) -> PathBuf { + home.join("Library/Application Support/Kiro/User/globalStorage/kiro.kiroagent") + } - let thread = format!( + fn write_m15a_snapshot(home: &Path, body: &str) -> PathBuf { + let path = m15a_global_root(home).join("workspace-a/conversation.chat"); + std::fs::create_dir_all(path.parent().unwrap()).unwrap(); + std::fs::write(path, body).unwrap(); + m15a_global_root(home).join("workspace-a/conversation.chat") + } + + fn write_m15a_execution(home: &Path, status: &str, start_time: &str) -> PathBuf { + let path = m15a_global_root(home).join("workspace-a/execution-store/execution"); + std::fs::create_dir_all(path.parent().unwrap()).unwrap(); + std::fs::write( + &path, + format!( + r#"{{ + "executionId": "exec-1", + "chatSessionId": "chat-1", + "status": "{status}", + "startTime": {start_time}, + "endTime": 1770983427500, + "completionOptions": {{"modelId": "claude-sonnet-4-5"}}, + "context": {{"messages": [{{"entries": [{{"type": "text", "text": "execution input"}}]}}]}}, + "actions": [{{"actionType": "say", "output": "execution output"}}] + }}"# + ), + ) + .unwrap(); + path + } + + fn m15a_snapshot_body(execution_id: &str, prompt: &str, response: &str) -> String { + format!( r#"{{ - "id": "thread-amp-gap", - "created": {thread_created}, - "usageLedger": {{ - "events": [ - {{ - "timestamp": "2026-04-08T12:00:00Z", - "model": "claude-sonnet-4-0", - "credits": 0.75, - "tokens": {{ "input": 100, "output": 20 }} - }} - ] - }}, + "executionId": "{execution_id}", + "model": "claude-sonnet-4-5", "messages": [ - {{ - "role": "assistant", - "messageId": 1, - "usage": {{ - "model": "claude-sonnet-4-0", - "inputTokens": 100, - "outputTokens": 20, - "credits": 0.75 - }} - }}, - {{ - "role": "assistant", - "messageId": 2, - "usage": {{ - "model": "claude-sonnet-4-0", - "inputTokens": 50, - "outputTokens": 10, - "credits": 0.40 - }} - }} + {{"role": "user", "content": "{prompt}"}}, + {{"role": "assistant", "content": "{response}"}} ] }}"# - ); - std::fs::write(amp_dir.join("T-thread-amp-gap.json"), thread).unwrap(); + ) + } - let parsed = parse_local_clients(LocalParseOptions { - home_dir: Some(temp_dir.path().to_str().unwrap().to_string()), + fn m15a_local_options(home: &Path) -> LocalParseOptions { + LocalParseOptions { + home_dir: Some(home.to_string_lossy().into_owned()), use_env_roots: false, - clients: Some(vec!["amp".to_string()]), - since: None, - until: None, - year: None, - scanner_settings: scanner::ScannerSettings::default(), - modified_after: None, - }) - .unwrap(); - - assert_eq!(parsed.counts.get(ClientId::Amp), 2); - assert_eq!(parsed.messages.len(), 2); - - let dates: HashSet = parsed.messages.iter().map(|msg| msg.date.clone()).collect(); - let local_date = |timestamp_ms: i64| { - chrono::Local - .timestamp_millis_opt(timestamp_ms) - .single() - .unwrap() - .format("%Y-%m-%d") - .to_string() - }; - assert!(dates.contains(&local_date(thread_created + 2000))); - assert!(dates.contains(&local_date(ledger_timestamp))); + clients: Some(vec!["kiro".to_string()]), + ..Default::default() + } } - // ========================================================================= - // fold_messages_streaming parity tests (RED — fold_messages_streaming not yet impl) - // ========================================================================= - - /// Deterministic UnifiedMessage fixture helper shared with parity tests. - /// Uses no real JSONL files; all fields are constructed inline. - #[cfg(test)] - #[allow(clippy::too_many_arguments)] - fn parity_msg( - date: &str, - client: &str, - model: &str, - session_id: &str, - dedup_key: Option<&str>, - timestamp_ms: i64, - input: i64, - output: i64, - cost: f64, - ) -> crate::sessions::UnifiedMessage { - use crate::TokenBreakdown; - crate::sessions::UnifiedMessage { - client: client.to_string(), - model_id: model.to_string(), - provider_id: "anthropic".to_string(), - session_id: session_id.to_string(), - workspace_key: None, - workspace_label: None, - timestamp: timestamp_ms, - date: date.to_string(), - tokens: TokenBreakdown { - input, - output, - cache_read: 0, - cache_write: 0, - reasoning: 0, - }, - cost, - cost_source: crate::CostSource::Unknown, - duration_ms: None, - message_count: 1, - agent: None, - dedup_key: dedup_key.map(|s| s.to_string()), - is_turn_start: false, + fn m15a_report_options(home: &Path) -> ReportOptions { + ReportOptions { + home_dir: Some(home.to_string_lossy().into_owned()), + use_env_roots: false, + clients: Some(vec!["kiro".to_string()]), + ..Default::default() } } - // A/B parity: fold_messages_streaming output == aggregate_by_date output - // for the same deterministic fixture (no dedup_keys, no trae). + fn m15b_ide_paths(home: &Path) -> (PathBuf, PathBuf) { + let dir = home.join(".kiro/sessions/workspace-a/sess_m15b"); + (dir.join("session.json"), dir.join("messages.jsonl")) + } + #[test] - fn test_fold_messages_streaming_parity_with_aggregate_by_date_no_dedup() { - let messages = vec![ - parity_msg("2025-06-01", "claude", "claude-sonnet-4-5", "s1", None, - 1_748_000_000_000, 100, 50, 0.01), - parity_msg("2025-06-01", "opencode", "gpt-4o", "s2", None, - 1_748_000_001_000, 200, 100, 0.02), - parity_msg("2025-06-02", "codex", "gpt-5", "s3", None, - 1_748_086_400_000, 400, 200, 0.04), - ]; + #[serial_test::serial] + fn m15b_ide_sibling_changes_reach_cache_all_lanes_and_reports() { + let source_home = tempfile::TempDir::new().unwrap(); + let cache_home = tempfile::TempDir::new().unwrap(); + let _env = EnvGuard::set(&[ + ("HOME", cache_home.path().as_os_str()), + ("TOKSCALE_CONFIG_DIR", cache_home.path().as_os_str()), + ("TOKSCALE_PRICING_CACHE_ONLY", std::ffi::OsStr::new("1")), + ]); + let (session, messages) = m15b_ide_paths(source_home.path()); + std::fs::create_dir_all(session.parent().unwrap()).unwrap(); + std::fs::write( + &session, + r#"{ + "id":"sess_m15b", + "modelId":"claude-opus-4.6", + "workspacePaths":["/tmp/m15b-project"] + }"#, + ) + .unwrap(); - // Reference: existing aggregate_by_date (clone-based) - let reference = crate::aggregator::aggregate_by_date(messages.clone()); + let clients = ["kiro".to_string()]; + let scanner_settings = scanner::ScannerSettings::default(); + let scan_result = scan_test_sources(source_home.path(), &clients, &scanner_settings); + let cache_session = scanner_spelling(scan_result.get(ClientId::Kiro), &session); + let parse_materialized = || { + let mut parsed = parse_all_messages_with_pricing_with_env_strategy( + source_home.path().to_str().unwrap(), + &clients, + None, + false, + &scanner_settings, + ); + parsed.sort_by(|left, right| left.dedup_key.cmp(&right.dedup_key)); + parsed + }; - // Subject: new streaming path - let streaming = fold_messages_streaming(&messages); + let source_fingerprint = message_cache::SourceFingerprint::from_path(&session).unwrap(); + let missing_sidecar_fingerprint = + message_cache::SourceFingerprint::from_kiro_path(&session).unwrap(); + let before = latest_source_mtime_ms(&m15a_local_options(source_home.path())).unwrap(); + assert!(parse_materialized().is_empty()); + std::thread::sleep(std::time::Duration::from_millis(20)); + std::fs::write( + &messages, + concat!( + "{\"timestamp\":\"2026-06-20T10:00:00Z\",\"payload\":{\"type\":\"user\",\"content\":\"hello\"}}\n", + "{\"payload\":{\"type\":\"session_metadata\",\"key\":\"contextUsage\",\"value\":{\"usagePercentage\":10.0}}}\n", + "{\"payload\":{\"type\":\"assistant\",\"content\":\"answer\"}}\n", + "{\"payload\":{\"type\":\"usage_summary\",\"elapsedTime\":1000}}\n", + "{\"timestamp\":\"2026-06-20T10:00:01Z\",\"payload\":{\"type\":\"turn_end\"}}\n", + ), + ) + .unwrap(); + let first_sidecar_fingerprint = + message_cache::SourceFingerprint::from_kiro_path(&session).unwrap(); + assert_ne!(missing_sidecar_fingerprint, first_sidecar_fingerprint); assert_eq!( - reference.len(), streaming.len(), - "parity: day bucket count must match" - ); - for (ref_day, stream_day) in reference.iter().zip(streaming.iter()) { - assert_eq!(ref_day.date, stream_day.date, "parity: date must match"); - assert_eq!( - ref_day.totals.tokens, stream_day.totals.tokens, - "parity: tokens must match for date {}", ref_day.date - ); - assert!( - (ref_day.totals.cost - stream_day.totals.cost).abs() < 1e-9, - "parity: cost must match for date {}", ref_day.date - ); - assert_eq!( - ref_day.totals.messages, stream_day.totals.messages, - "parity: message_count must match for date {}", ref_day.date - ); - } - } - - // A/B parity with cross-file dedup: fold_messages_streaming must apply - // the same dedup_key filtering as the existing pipeline does via seen_keys. - #[test] - fn test_fold_messages_streaming_parity_cross_file_dedup() { - // Construct messages that include a duplicated dedup_key pair. - // The existing pipeline filters duplicates via the seen_keys HashSet. - // fold_messages_streaming must produce the same counts. - let unique = parity_msg("2025-06-10", "claude", "claude-sonnet-4-5", "u1", - Some("unique-key-1"), 1_749_000_000_000, 300, 150, 0.06); - let dup_first = parity_msg("2025-06-10", "claude", "claude-sonnet-4-5", "d1", - Some("dup-key-shared"), 1_749_000_001_000, 200, 100, 0.04); - let dup_second = parity_msg("2025-06-10", "claude", "claude-haiku-4-5", "d2", - Some("dup-key-shared"), 1_749_000_002_000, 200, 100, 0.04); + source_fingerprint, + message_cache::SourceFingerprint::from_path(&session).unwrap(), + "messages.jsonl must invalidate through the related-file fingerprint" + ); + let after_first = latest_source_mtime_ms(&m15a_local_options(source_home.path())).unwrap(); + assert!(after_first > before); + + let first = parse_materialized(); + assert_eq!(first.len(), 1); + assert_eq!(first[0].tokens.input, 20_000); + assert_eq!(first[0].dedup_key.as_deref(), Some("sess_m15b:ide:0")); + assert_eq!(first[0].model_id, "claude-opus-4.6"); + assert_eq!(first[0].workspace_key.as_deref(), Some("/tmp/m15b-project")); + assert_eq!( + parse_materialized(), + first, + "the first warm hit must be stable" + ); - // The reference pipeline keeps only the first occurrence of dup-key-shared - // (seen_keys.insert returns false on second) — 2 messages total. - let all_msgs = vec![unique.clone(), dup_first.clone(), dup_second.clone()]; + std::thread::sleep(std::time::Duration::from_millis(20)); + std::fs::write( + &messages, + concat!( + "{\"timestamp\":\"2026-06-20T10:00:00Z\",\"payload\":{\"type\":\"user\",\"content\":\"hello\"}}\n", + "{\"payload\":{\"type\":\"session_metadata\",\"key\":\"contextUsage\",\"value\":{\"usagePercentage\":10.0}}}\n", + "{\"payload\":{\"type\":\"assistant\",\"content\":\"answer\"}}\n", + "{\"payload\":{\"type\":\"usage_summary\",\"elapsedTime\":1000}}\n", + "{\"timestamp\":\"2026-06-20T10:00:01Z\",\"payload\":{\"type\":\"turn_end\"}}\n", + "{\"timestamp\":\"2026-06-20T10:01:00Z\",\"payload\":{\"type\":\"user\",\"content\":\"next\"}}\n", + "{\"payload\":{\"type\":\"session_metadata\",\"key\":\"contextUsage\",\"value\":{\"usagePercentage\":20.0}}}\n", + "{\"payload\":{\"type\":\"assistant\",\"content\":\"done\"}}\n", + "{\"payload\":{\"type\":\"usage_summary\",\"elapsedTime\":1000}}\n", + "{\"timestamp\":\"2026-06-20T10:01:01Z\",\"payload\":{\"type\":\"turn_end\"}}\n", + ), + ) + .unwrap(); + let updated_sidecar_fingerprint = + message_cache::SourceFingerprint::from_kiro_path(&session).unwrap(); + assert_ne!(first_sidecar_fingerprint, updated_sidecar_fingerprint); + let after_second = latest_source_mtime_ms(&m15a_local_options(source_home.path())).unwrap(); + assert!(after_second > after_first); + + let updated = parse_materialized(); + assert_eq!(updated.len(), 2); + assert_eq!( + updated + .iter() + .map(|message| message.tokens.input) + .sum::(), + 60_000 + ); + assert_eq!( + parse_materialized(), + updated, + "the rebuilt cache must stay warm-complete" + ); + assert_eq!( + message_cache::SourceMessageCache::load() + .get( + message_cache::CacheIdentity::for_client(ClientId::Kiro), + &cache_session, + ) + .unwrap() + .messages + .len(), + 2 + ); - let streaming = fold_messages_streaming(&all_msgs); + let mut streamed = Vec::new(); + scan_messages_streaming( + source_home.path().to_str().unwrap(), + &clients, + None, + false, + &scanner_settings, + &|_| true, + &mut |message| streamed.push(message.clone()), + ); + streamed.sort_by(|left, right| left.dedup_key.cmp(&right.dedup_key)); + assert_eq!(streamed, updated); - assert_eq!(streaming.len(), 1, "all on same date -> 1 bucket"); + let counted = parse_local_clients(m15a_local_options(source_home.path())).unwrap(); + assert_eq!(counted.counts.get(ClientId::Kiro), 2); + assert_eq!(counted.messages.len(), 2); assert_eq!( - streaming[0].totals.messages, 2, - "parity: duplicate dedup_key must reduce count from 3 to 2" + counted + .messages + .iter() + .map(|message| message.input) + .sum::(), + 60_000 ); - assert!( - (streaming[0].totals.cost - 0.10).abs() < 1e-9, - "parity: cost must exclude the duplicate message" + let pruned = parse_local_clients(LocalParseOptions { + modified_after: Some(after_first + 1), + ..m15a_local_options(source_home.path()) + }) + .unwrap(); + assert_eq!(pruned.counts.get(ClientId::Kiro), 2); + + let runtime = tokio::runtime::Builder::new_current_thread() + .enable_all() + .build() + .unwrap(); + let report_options = m15a_report_options(source_home.path()); + let model = runtime + .block_on(get_model_report(report_options.clone())) + .unwrap(); + let monthly = runtime + .block_on(get_monthly_report(report_options.clone())) + .unwrap(); + let hourly = runtime + .block_on(get_hourly_report(report_options.clone())) + .unwrap(); + let agents = runtime.block_on(get_agents_report(report_options)).unwrap(); + assert_eq!(model.total_messages, 2); + assert_eq!(model.total_input, 60_000); + assert_eq!( + monthly + .entries + .iter() + .map(|entry| entry.message_count) + .sum::(), + 2 + ); + assert_eq!( + hourly + .entries + .iter() + .map(|entry| entry.message_count) + .sum::(), + 2 + ); + assert_eq!(agents.total_messages, 2); + assert_eq!( + agents.entries.iter().map(|entry| entry.input).sum::(), + 60_000 ); } - // ========================================================================= - // Phase 2 RED tests: build_graph_result_from_messages streaming entry-point - // ========================================================================= - // - // `build_graph_result_from_messages` does NOT exist yet. These tests - // define the observable contract that the GREEN implementation must satisfy: - // - accept a `&[UnifiedMessage]` slice and an optional `since` date string - // - apply the `since` post-parse filter (date string prefix comparison, - // same semantics as `filter_messages_for_report`) - // - drive aggregation via `StreamingAggregator` (zero-clone fold path) - // - return a `GraphResult` whose per-day tokens/cost match the reference - // `aggregate_by_date` pipeline exactly (zero tolerance) - // - // All three tests will produce a **compile error** until the function is - // declared in lib.rs, which is the required RED state. - - use crate::GraphResult; - - /// Multi-client, multi-day fixture: streaming path total tokens and cost - /// per daily bucket must match hand-computed expected values. - /// - /// Hardcoded expected values — calculation: - /// - /// 2025-06-01 (no dedup): - /// s1: input=500, output=250 → tokens=750, cost=0.05 - /// s2: input=300, output=150 → tokens=450, cost=0.03 - /// TOTAL: tokens=1200, cost=0.08, messages=2 - /// - /// 2025-06-02 (trae session dedup — same session_id="trae-sess"): - /// trae-k1: ts=1_748_822_400_000, input=100, output=50 → tokens=150, cost=0.01 - /// trae-k2: ts=1_748_822_500_000, input=200, output=100 → tokens=300, cost=0.02 - /// StreamingAggregator keeps latest timestamp → trae-k2 wins - /// TOTAL: tokens=300, cost=0.02, messages=1 - /// - /// 2025-06-03 (cross-file dedup_key — both carry "dup-phase2"): - /// d1: input=400, output=200 → tokens=600, cost=0.04 (first seen — kept) - /// d2: input=400, output=200 → tokens=600, cost=0.04 (same dedup_key — dropped) - /// TOTAL: tokens=600, cost=0.04, messages=1 + #[cfg(target_os = "macos")] #[test] - fn test_build_graph_result_from_messages_matches_aggregate_by_date() { - let messages = vec![ - // Day 2025-06-01: two clients, no dedup - parity_msg("2025-06-01", "claude", "claude-sonnet-4-5", "s1", None, - 1_748_736_000_000, 500, 250, 0.05), - parity_msg("2025-06-01", "opencode", "gpt-4o", "s2", None, - 1_748_736_001_000, 300, 150, 0.03), - // Day 2025-06-02: trae dedup by session_id — two entries same session, keep latest - parity_msg("2025-06-02", "trae", "gpt-5.2", "trae-sess", Some("trae-k1"), - 1_748_822_400_000, 100, 50, 0.01), - parity_msg("2025-06-02", "trae", "gpt-5.2", "trae-sess", Some("trae-k2"), - 1_748_822_500_000, 200, 100, 0.02), // newer timestamp -> wins - // Day 2025-06-03: cross-file dedup pair — same dedup_key, second dropped - parity_msg("2025-06-03", "claude", "claude-haiku-4-5", "d1", - Some("dup-phase2"), 1_748_908_800_000, 400, 200, 0.04), - parity_msg("2025-06-03", "claude", "claude-haiku-4-5", "d2", - Some("dup-phase2"), 1_748_908_801_000, 400, 200, 0.04), // same dedup_key -> discarded - ]; + #[serial_test::serial] + fn m15a_cli_keys_cannot_seed_or_collide_with_globalstorage() { + let source_home = tempfile::TempDir::new().unwrap(); + let cache_home = tempfile::TempDir::new().unwrap(); + let _env = EnvGuard::set(&[ + ("HOME", cache_home.path().as_os_str()), + ("TOKSCALE_CONFIG_DIR", cache_home.path().as_os_str()), + ("TOKSCALE_PRICING_CACHE_ONLY", std::ffi::OsStr::new("1")), + ]); + let cli_dir = source_home.path().join(".kiro/sessions/cli"); + std::fs::create_dir_all(&cli_dir).unwrap(); + std::fs::write( + cli_dir.join("cli.json"), + r#"{"session_id":"execution","cwd":"workspace-a","session_state":{"rts_model_state":{"model_info":{"model_id":"cli-model"}},"conversation_metadata":{"user_turn_metadatas":[{"input_token_count":1}]}}}"#, + ) + .unwrap(); + std::fs::write(cli_dir.join("cli.jsonl"), "").unwrap(); + std::fs::write( + cli_dir.join("cli-collision.json"), + r#"{"session_id":"workspace-a/conversation:globalstorage:exec","cwd":"workspace-a","session_state":{"rts_model_state":{"model_info":{"model_id":"cli-model"}},"conversation_metadata":{"user_turn_metadatas":[{"input_token_count":1}]}}}"#, + ) + .unwrap(); + std::fs::write(cli_dir.join("cli-collision.jsonl"), "").unwrap(); + write_m15a_snapshot(source_home.path(), &m15a_snapshot_body("0", "ABCD", "")); + let execution_zero = + m15a_global_root(source_home.path()).join("workspace-a/execution-store/execution-zero"); + std::fs::create_dir_all(execution_zero.parent().unwrap()).unwrap(); + std::fs::write( + execution_zero, + r#"{ + "executionId": "0", + "chatSessionId": "conversation", + "status": "succeed", + "startTime": 1770983426, + "endTime": 1770983427500, + "completionOptions": {"modelId": "claude-sonnet-4-5"}, + "context": {"messages": [{"entries": [{"type": "text", "text": "execution input"}]}]}, + "actions": [{"actionType": "say", "output": "execution output"}] + }"#, + ) + .unwrap(); - // Subject: new streaming entry-point - let result: GraphResult = - crate::build_graph_result_from_messages(&messages, None); + let expected = vec![ + "conversation", + "execution", + "workspace-a/conversation:globalstorage:exec", + ]; + let mut materialized: Vec<_> = parse_all_messages_with_pricing_with_env_strategy( + source_home.path().to_str().unwrap(), + &["kiro".to_string()], + None, + false, + &scanner::ScannerSettings::default(), + ) + .into_iter() + .map(|message| message.session_id) + .collect(); + materialized.sort(); + assert_eq!(materialized, expected); - // Verify bucket count: 3 distinct dates - assert_eq!( - result.contributions.len(), 3, - "phase2 streaming: must produce exactly 3 daily buckets" + let mut streamed = Vec::new(); + scan_messages_streaming( + source_home.path().to_str().unwrap(), + &["kiro".to_string()], + None, + false, + &scanner::ScannerSettings::default(), + &|_| true, + &mut |message| streamed.push(message.session_id.clone()), ); + streamed.sort(); + assert_eq!(streamed, expected); - // Locate each day bucket by date (sort order: ascending) - let day1 = result.contributions.iter().find(|c| c.date == "2025-06-01") - .expect("phase2: 2025-06-01 bucket must exist"); - let day2 = result.contributions.iter().find(|c| c.date == "2025-06-02") - .expect("phase2: 2025-06-02 bucket must exist"); - let day3 = result.contributions.iter().find(|c| c.date == "2025-06-03") - .expect("phase2: 2025-06-03 bucket must exist"); + let mut counted: Vec<_> = parse_local_clients(m15a_local_options(source_home.path())) + .unwrap() + .messages + .into_iter() + .map(|message| message.session_id) + .collect(); + counted.sort(); + assert_eq!(counted, expected); + } - // 2025-06-01: s1 (750) + s2 (450) = 1200 tokens, 0.05+0.03=0.08 cost, 2 messages - assert_eq!(day1.totals.tokens, 1200, - "2025-06-01: tokens must be 750+450=1200"); - assert!( - (day1.totals.cost - 0.08).abs() < 1e-9, - "2025-06-01: cost must be 0.05+0.03=0.08" - ); - assert_eq!(day1.totals.messages, 2, - "2025-06-01: both non-trae non-dedup messages must be counted"); + #[cfg(target_os = "macos")] + #[test] + #[serial_test::serial] + fn m15a_duplicate_snapshot_extensions_are_exact_once() { + let source_home = tempfile::TempDir::new().unwrap(); + let cache_home = tempfile::TempDir::new().unwrap(); + let _env = EnvGuard::set(&[ + ("HOME", cache_home.path().as_os_str()), + ("TOKSCALE_CONFIG_DIR", cache_home.path().as_os_str()), + ("TOKSCALE_PRICING_CACHE_ONLY", std::ffi::OsStr::new("1")), + ]); + let body = m15a_snapshot_body("unused", "ABCD", "WXYZ"); + let chat = write_m15a_snapshot(source_home.path(), &body); + std::fs::write(chat.with_extension("json"), body).unwrap(); - // 2025-06-02: trae session dedup — trae-k2 wins (larger timestamp) - // trae-k2: input=200, output=100 -> tokens=300, cost=0.02 - assert_eq!(day2.totals.tokens, 300, - "2025-06-02: trae dedup — only winner (trae-k2, tokens=300) counted"); - assert!( - (day2.totals.cost - 0.02).abs() < 1e-9, - "2025-06-02: trae dedup — cost must be 0.02 (trae-k2 only)" + let materialized = parse_all_messages_with_pricing_with_env_strategy( + source_home.path().to_str().unwrap(), + &["kiro".to_string()], + None, + false, + &scanner::ScannerSettings::default(), ); - assert_eq!(day2.totals.messages, 1, - "2025-06-02: trae dedup collapses 2 entries to 1 per session_id"); + assert_eq!(materialized.len(), 1); - // 2025-06-03: cross-file dedup — d1 kept, d2 dropped (same dedup_key) - // d1: input=400, output=200 -> tokens=600, cost=0.04 - assert_eq!(day3.totals.tokens, 600, - "2025-06-03: cross-file dedup — only d1 (tokens=600) counted, d2 dropped"); - assert!( - (day3.totals.cost - 0.04).abs() < 1e-9, - "2025-06-03: cross-file dedup — cost must be 0.04 (d1 only)" + let mut streamed = Vec::new(); + scan_messages_streaming( + source_home.path().to_str().unwrap(), + &["kiro".to_string()], + None, + false, + &scanner::ScannerSettings::default(), + &|_| true, + &mut |message| streamed.push(message.clone()), ); - assert_eq!(day3.totals.messages, 1, - "2025-06-03: duplicate dedup_key dropped, 1 message retained"); + assert_eq!(streamed.len(), 1); + + let counted = parse_local_clients(m15a_local_options(source_home.path())).unwrap(); + assert_eq!(counted.counts.get(ClientId::Kiro), 1); + assert_eq!(counted.messages.len(), 1); } - /// `since` filter semantics: same fixture with `since = "2025-06-02"` must - /// produce only the 2025-06-02 and 2025-06-03 buckets, with their - /// token/cost totals matching a manually filtered reference. + #[cfg(target_os = "macos")] #[test] - fn test_build_graph_result_from_messages_since_filter_excludes_earlier_dates() { - let messages = vec![ - parity_msg("2025-06-01", "claude", "claude-sonnet-4-5", "s1", None, - 1_748_736_000_000, 500, 250, 0.05), - parity_msg("2025-06-01", "opencode", "gpt-4o", "s2", None, - 1_748_736_001_000, 300, 150, 0.03), - parity_msg("2025-06-02", "codex", "gpt-5", "s3", None, - 1_748_822_400_000, 400, 200, 0.04), - parity_msg("2025-06-03", "claude", "claude-haiku-4-5", "s4", None, - 1_748_908_800_000, 200, 100, 0.02), - ]; + #[serial_test::serial] + fn m15a_materialized_streaming_count_and_report_parity() { + let source_home = tempfile::TempDir::new().unwrap(); + let cache_home = tempfile::TempDir::new().unwrap(); + let _env = EnvGuard::set(&[ + ("HOME", cache_home.path().as_os_str()), + ("TOKSCALE_CONFIG_DIR", cache_home.path().as_os_str()), + ("TOKSCALE_PRICING_CACHE_ONLY", std::ffi::OsStr::new("1")), + ]); + let snapshot = write_m15a_snapshot( + source_home.path(), + &m15a_snapshot_body("exec-1", "snapshot input", "snapshot output"), + ); + let execution = write_m15a_execution(source_home.path(), "succeed", "1770983426"); + let mut pricing_data = HashMap::new(); + pricing_data.insert( + "claude-sonnet-4-5".to_string(), + pricing::ModelPricing { + input_cost_per_token: Some(1.0), + output_cost_per_token: Some(1.0), + ..Default::default() + }, + ); + let pricing_service = pricing::PricingService::new(pricing_data, HashMap::new()); + + let materialized = parse_all_messages_with_pricing_with_env_strategy( + source_home.path().to_str().unwrap(), + &["kiro".to_string()], + Some(&pricing_service), + false, + &scanner::ScannerSettings::default(), + ); + assert_eq!(materialized.len(), 1); + assert_eq!( + materialized[0].dedup_key.as_deref(), + Some("execution:exec-1") + ); + assert!(materialized[0].cost > 0.0); + + // Both source entries are raw and independently cached even though the + // snapshot is suppressed in the merged result. + let cache = message_cache::SourceMessageCache::load(); + assert!(cache + .get( + message_cache::CacheIdentity::for_client(ClientId::Kiro), + &snapshot + ) + .is_some_and(|entry| entry.messages[0] + .dedup_key + .as_deref() + .is_some_and(|key| key.ends_with(":globalstorage:exec:exec-1")))); + assert!(cache + .get( + message_cache::CacheIdentity::for_client(ClientId::Kiro), + &execution + ) + .is_some_and( + |entry| entry.messages[0].dedup_key.as_deref() == Some("execution:exec-1") + )); + assert!(cache + .get( + message_cache::CacheIdentity::for_client(ClientId::Kiro), + &snapshot, + ) + .is_some_and(|entry| entry.messages[0].cost == 0.0)); + assert!(cache + .get( + message_cache::CacheIdentity::for_client(ClientId::Kiro), + &execution, + ) + .is_some_and(|entry| entry.messages[0].cost == 0.0)); - // Subject: streaming entry with since = "2025-06-02" - // (function does not exist yet -> RED compile error) - let result: GraphResult = - crate::build_graph_result_from_messages(&messages, Some("2025-06-02")); + let mut streamed = Vec::new(); + scan_messages_streaming( + source_home.path().to_str().unwrap(), + &["kiro".to_string()], + Some(&pricing_service), + false, + &scanner::ScannerSettings::default(), + &|_| true, + &mut |message| streamed.push(message.clone()), + ); + assert_eq!(streamed.len(), 1); + assert_eq!(streamed[0].dedup_key.as_deref(), Some("execution:exec-1")); + assert!(streamed[0].cost > 0.0); + assert_eq!(materialized[0].cost, streamed[0].cost); + assert_eq!( + (materialized[0].tokens.input, materialized[0].tokens.output), + (streamed[0].tokens.input, streamed[0].tokens.output) + ); + assert_eq!(materialized[0].model_id, "claude-sonnet-4-5"); + assert_eq!(materialized[0].session_id, "chat-1"); + assert_eq!( + materialized[0].workspace_key.as_deref(), + Some("workspace-a") + ); + assert_eq!( + materialized[0].workspace_label.as_deref(), + Some("workspace-a") + ); + assert_eq!(materialized[0].timestamp, 1_770_983_426_000); + assert_eq!(materialized[0].duration_ms, Some(1_500)); + assert_eq!(materialized[0].message_count, 1); + assert_eq!( + ( + streamed[0].model_id.as_str(), + streamed[0].session_id.as_str(), + streamed[0].workspace_key.as_deref(), + streamed[0].workspace_label.as_deref(), + streamed[0].timestamp, + streamed[0].duration_ms, + streamed[0].message_count, + ), + ( + materialized[0].model_id.as_str(), + materialized[0].session_id.as_str(), + materialized[0].workspace_key.as_deref(), + materialized[0].workspace_label.as_deref(), + materialized[0].timestamp, + materialized[0].duration_ms, + materialized[0].message_count, + ) + ); - // Only 2025-06-02 and 2025-06-03 must be present + let counted = parse_local_clients(m15a_local_options(source_home.path())).unwrap(); + assert_eq!(counted.counts.get(ClientId::Kiro), 1); + assert_eq!(counted.messages.len(), 1); + assert_eq!(counted.messages[0].model_id, materialized[0].model_id); + assert_eq!(counted.messages[0].session_id, materialized[0].session_id); assert_eq!( - result.contributions.len(), 2, - "since filter: must exclude 2025-06-01, leaving 2 buckets" + counted.messages[0].workspace_key, + materialized[0].workspace_key + ); + assert_eq!( + counted.messages[0].workspace_label, + materialized[0].workspace_label + ); + assert_eq!(counted.messages[0].timestamp, materialized[0].timestamp); + assert_eq!(counted.messages[0].duration_ms, materialized[0].duration_ms); + assert_eq!( + counted.messages[0].message_count, + materialized[0].message_count ); - let dates: Vec<&str> = result.contributions.iter().map(|c| c.date.as_str()).collect(); - assert!(dates.contains(&"2025-06-02"), - "since filter: 2025-06-02 bucket must be present"); - assert!(dates.contains(&"2025-06-03"), - "since filter: 2025-06-03 bucket must be present"); - assert!(!dates.contains(&"2025-06-01"), - "since filter: 2025-06-01 bucket must be absent"); + let runtime = tokio::runtime::Builder::new_current_thread() + .enable_all() + .build() + .unwrap(); + let options = m15a_report_options(source_home.path()); + let model = runtime.block_on(get_model_report(options.clone())).unwrap(); + let monthly = runtime + .block_on(get_monthly_report(options.clone())) + .unwrap(); + let hourly = runtime + .block_on(get_hourly_report(options.clone())) + .unwrap(); + let agents = runtime + .block_on(get_agents_report(options.clone())) + .unwrap(); + let mut session_options = options.clone(); + session_options.group_by = GroupBy::Session; + let session_model = runtime.block_on(get_model_report(session_options)).unwrap(); + let mut workspace_options = options; + workspace_options.group_by = GroupBy::WorkspaceModel; + let workspace_model = runtime + .block_on(get_model_report(workspace_options)) + .unwrap(); - // 2025-06-02 token total: input 400 + output 200 = 600 - let day2 = result.contributions.iter().find(|c| c.date == "2025-06-02").unwrap(); - assert_eq!(day2.totals.tokens, 600, - "since filter: 2025-06-02 token total must be 600"); - assert!( - (day2.totals.cost - 0.04).abs() < 1e-9, - "since filter: 2025-06-02 cost must be 0.04" + assert_eq!(model.total_messages, 1); + assert_eq!(model.entries.len(), 1); + assert_eq!(model.entries[0].model, materialized[0].model_id); + assert_eq!( + model.entries[0].message_count, + materialized[0].message_count + ); + assert_eq!(session_model.entries.len(), 1); + assert_eq!( + session_model.entries[0].session_id.as_deref(), + Some("chat-1") + ); + assert_eq!(workspace_model.entries.len(), 1); + assert_eq!( + workspace_model.entries[0].workspace_key.as_deref(), + Some("workspace-a") ); + assert_eq!( + workspace_model.entries[0].workspace_label.as_deref(), + Some("workspace-a") + ); + assert_eq!( + monthly + .entries + .iter() + .map(|entry| entry.message_count) + .sum::(), + 1 + ); + assert_eq!( + hourly + .entries + .iter() + .map(|entry| entry.message_count) + .sum::(), + 1 + ); + assert_eq!(agents.total_messages, 1); + assert_eq!(model.total_input, monthly.entries[0].input); + assert_eq!(model.total_input, hourly.entries[0].input); + assert_eq!(model.total_input, agents.entries[0].input); + assert_eq!(model.total_output, monthly.entries[0].output); + assert_eq!(model.total_output, hourly.entries[0].output); + assert_eq!(model.total_output, agents.entries[0].output); } - /// Trae dedup in streaming path: two messages for the same trae session - /// (same `session_id`, different `dedup_key`, later timestamp wins) must - /// produce exactly ONE message worth of tokens/cost in the daily bucket. + #[cfg(target_os = "macos")] #[test] - fn test_build_graph_result_from_messages_trae_session_dedup_keeps_latest() { - let messages = vec![ - // Earlier trae message (should be dropped) - parity_msg("2025-06-10", "trae", "gpt-5.2", "trae-sess-a", Some("trae-early"), - 1_749_513_600_000, 100, 50, 0.01), - // Later trae message for same session_id (should win) - parity_msg("2025-06-10", "trae", "gpt-5.2", "trae-sess-a", Some("trae-late"), - 1_749_513_700_000, 300, 150, 0.03), - // Non-trae message (should be included as-is) - parity_msg("2025-06-10", "claude", "claude-sonnet-4-5", "c1", None, - 1_749_513_800_000, 200, 100, 0.02), - ]; + #[serial_test::serial] + fn m15a_warm_cache_mixed_hits_reapply_suppression_and_restore_snapshot() { + let source_home = tempfile::TempDir::new().unwrap(); + let cache_home = tempfile::TempDir::new().unwrap(); + let _env = EnvGuard::set(&[ + ("HOME", cache_home.path().as_os_str()), + ("TOKSCALE_CONFIG_DIR", cache_home.path().as_os_str()), + ("TOKSCALE_PRICING_CACHE_ONLY", std::ffi::OsStr::new("1")), + ]); + let snapshot = write_m15a_snapshot( + source_home.path(), + &m15a_snapshot_body("exec-1", "snapshot input", "snapshot output"), + ); + let execution = write_m15a_execution(source_home.path(), "succeed", "1770983426"); + let home = source_home.path().to_str().unwrap(); + let clients = ["kiro".to_string()]; - // Subject: streaming entry (does not exist yet -> RED compile error) - let result: GraphResult = - crate::build_graph_result_from_messages(&messages, None); + let first = parse_all_messages_with_pricing_with_env_strategy( + home, + &clients, + None, + false, + &scanner::ScannerSettings::default(), + ); + assert_eq!(first.len(), 1); + assert_eq!(first[0].dedup_key.as_deref(), Some("execution:exec-1")); - assert_eq!(result.contributions.len(), 1, - "trae dedup: all messages on same date -> 1 bucket"); + // Snapshot hit + execution miss: only the newly parsed execution wins. + std::thread::sleep(std::time::Duration::from_millis(5)); + write_m15a_execution(source_home.path(), "succeed", "1770983426.5"); + let mut streamed = Vec::new(); + scan_messages_streaming( + home, + &clients, + None, + false, + &scanner::ScannerSettings::default(), + &|_| true, + &mut |message| streamed.push(message.clone()), + ); + assert_eq!(streamed.len(), 1); + assert_eq!(streamed[0].dedup_key.as_deref(), Some("execution:exec-1")); - let day = &result.contributions[0]; - // Kept messages: trae-late (tokens=450) + claude (tokens=300) = 750 total tokens - assert_eq!(day.totals.tokens, 750, - "trae dedup: token total must reflect only the winning trae entry (450) + claude (300)"); - assert!( - (day.totals.cost - 0.05).abs() < 1e-9, - "trae dedup: cost must be 0.03 (latest trae) + 0.02 (claude) = 0.05" + // Execution hit + snapshot miss: the snapshot change is still suppressed. + std::thread::sleep(std::time::Duration::from_millis(5)); + write_m15a_snapshot( + source_home.path(), + &m15a_snapshot_body( + "exec-1", + "changed snapshot input", + "changed snapshot output", + ), ); - assert_eq!(day.totals.messages, 2, - "trae dedup: message count must be 2 (1 trae winner + 1 claude)"); + let second = parse_all_messages_with_pricing_with_env_strategy( + home, + &clients, + None, + false, + &scanner::ScannerSettings::default(), + ); + assert_eq!(second.len(), 1); + assert_eq!(second[0].dedup_key.as_deref(), Some("execution:exec-1")); + + // A successful execution rewritten as failed removes its stale cache + // entry and exposes the raw cached snapshot. + std::thread::sleep(std::time::Duration::from_millis(5)); + write_m15a_execution(source_home.path(), "failed", "1770983426"); + let mut failed = Vec::new(); + scan_messages_streaming( + home, + &clients, + None, + false, + &scanner::ScannerSettings::default(), + &|_| true, + &mut |message| failed.push(message.clone()), + ); + assert_eq!(failed.len(), 1); + assert!(failed[0] + .dedup_key + .as_deref() + .is_some_and(|key| key.ends_with(":globalstorage:exec:exec-1"))); + + // Restore a successful execution after the failed rewrite. This proves + // a cached successful execution can become authoritative again. + std::thread::sleep(std::time::Duration::from_millis(5)); + write_m15a_execution(source_home.path(), "succeed", "1770983426"); + let restored_execution = parse_all_messages_with_pricing_with_env_strategy( + home, + &clients, + None, + false, + &scanner::ScannerSettings::default(), + ); + assert_eq!(restored_execution.len(), 1); + assert_eq!( + restored_execution[0].dedup_key.as_deref(), + Some("execution:exec-1") + ); + assert!(message_cache::SourceMessageCache::load() + .get( + message_cache::CacheIdentity::for_client(ClientId::Kiro), + &execution, + ) + .is_some()); + + // Removing that cached successful execution must expose the raw cached + // snapshot on the other (streaming) lane. + std::fs::remove_file(&execution).unwrap(); + let mut restored = Vec::new(); + scan_messages_streaming( + home, + &clients, + None, + false, + &scanner::ScannerSettings::default(), + &|_| true, + &mut |message| restored.push(message.clone()), + ); + assert_eq!(restored.len(), 1); + assert!(restored[0] + .dedup_key + .as_deref() + .is_some_and(|key| key.ends_with(":globalstorage:exec:exec-1"))); + assert!(message_cache::SourceMessageCache::load() + .get( + message_cache::CacheIdentity::for_client(ClientId::Kiro), + &snapshot, + ) + .is_some()); } + #[cfg(target_os = "macos")] #[test] - fn model_aggregation_saturates_overflowing_token_folds() { - // token_total_saturates_on_overlarge_buckets (see positive_token_total's - // callers) covers a single message's grand total; the per-field - // CROSS-MESSAGE fold in aggregate_model_usage_entries must saturate too. - // An antigravity-cli row can carry an i64::MAX bucket after the - // untrusted-varint clamp (sessions/antigravity_cli.rs to_i64), so two - // such rows folded into one model group with plain `+=` overflow (debug - // panic / release wrap) before the already-saturating grand total runs. - let make = || { - UnifiedMessage::new( - "antigravity-cli", - "gemini-3-pro", - "antigravity", - "session-overflow", - 1_733_011_200_000, - TokenBreakdown { - input: i64::MAX, - output: 0, - cache_read: i64::MAX, - cache_write: 0, - reasoning: 0, - }, - 0.0, - ) - }; + #[serial_test::serial] + fn m15a_suppression_precedes_report_date_filter() { + let source_home = tempfile::TempDir::new().unwrap(); + let cache_home = tempfile::TempDir::new().unwrap(); + let _env = EnvGuard::set(&[ + ("HOME", cache_home.path().as_os_str()), + ("TOKSCALE_CONFIG_DIR", cache_home.path().as_os_str()), + ("TOKSCALE_PRICING_CACHE_ONLY", std::ffi::OsStr::new("1")), + ]); + let snapshot = write_m15a_snapshot( + source_home.path(), + &m15a_snapshot_body("exec-future", "snapshot input", "snapshot output"), + ); + let execution = m15a_global_root(source_home.path()) + .join("workspace-a/execution-store/execution-future"); + std::fs::create_dir_all(execution.parent().unwrap()).unwrap(); + std::fs::write( + &execution, + r#"{"executionId":"exec-future","chatSessionId":"chat-future","status":"succeed","startTime":4102444800000,"endTime":4102444801000,"actions":[{"actionType":"say","output":"future answer"}],"input":{"data":{"messages":[{"content":"future question"}]}}}"#, + ) + .unwrap(); + let snapshot_date = sessions::kiro::parse_kiro_file(&snapshot)[0].date.clone(); - let entries = aggregate_model_usage_entries(vec![make(), make()], &GroupBy::Model); - assert_eq!(entries.len(), 1); - assert_eq!(entries[0].input, i64::MAX); - assert_eq!(entries[0].cache_read, i64::MAX); + let mut options = m15a_report_options(source_home.path()); + options.until = Some(snapshot_date); + let runtime = tokio::runtime::Builder::new_current_thread() + .enable_all() + .build() + .unwrap(); + let report = runtime.block_on(get_model_report(options)).unwrap(); + assert!(report.entries.is_empty()); } + #[cfg(target_os = "macos")] #[test] - fn model_report_totals_saturate_across_groups() { - // aggregate_model_usage_entries saturates each entry's fields, so an - // entry can be i64::MAX. get_model_report sums the entries into the - // report-level totals via model_report_token_totals; two saturated - // entries (two distinct models) must not overflow that sum either. - let make = |model: &str| { - UnifiedMessage::new( - "antigravity-cli", - model, - "antigravity", - "session-overflow", - 1_733_011_200_000, - TokenBreakdown { - input: i64::MAX, - output: 0, - cache_read: i64::MAX, - cache_write: 0, - reasoning: 0, - }, - 0.0, - ) - }; + #[serial_test::serial] + fn m15a_globalstorage_mtime_pruning_and_stat_failure_fail_open() { + let source_home = tempfile::TempDir::new().unwrap(); + let cache_home = tempfile::TempDir::new().unwrap(); + let _env = EnvGuard::set(&[ + ("HOME", cache_home.path().as_os_str()), + ("TOKSCALE_CONFIG_DIR", cache_home.path().as_os_str()), + ]); + let snapshot = write_m15a_snapshot( + source_home.path(), + &m15a_snapshot_body("exec-1", "initial input", "initial output"), + ); + let options = m15a_local_options(source_home.path()); + let before = latest_source_mtime_ms(&options).unwrap(); + std::thread::sleep(std::time::Duration::from_millis(20)); + write_m15a_snapshot( + source_home.path(), + &m15a_snapshot_body("exec-1", "rewritten input", "rewritten output"), + ); + let after = latest_source_mtime_ms(&options).unwrap(); + assert!(after > before, "globalStorage primary mtime must advance"); - let entries = aggregate_model_usage_entries( - vec![make("gemini-3-pro"), make("claude-opus-4-6")], - &GroupBy::Model, + let parsed = parse_local_clients(LocalParseOptions { + modified_after: Some(before + 1), + ..options.clone() + }) + .unwrap(); + assert_eq!(parsed.counts.get(ClientId::Kiro), 1); + + let execution = write_m15a_execution(source_home.path(), "succeed", "1770983426"); + let execution_mtime = super::kiro_source_mtime_ms(&execution).unwrap(); + std::thread::sleep(std::time::Duration::from_millis(20)); + write_m15a_snapshot( + source_home.path(), + &m15a_snapshot_body("exec-1", "newer snapshot", "must stay suppressed"), ); - assert_eq!(entries.len(), 2); - let (total_input, _total_output, total_cache_read, _total_cache_write) = - super::model_report_token_totals(&entries); - assert_eq!(total_input, i64::MAX); - assert_eq!(total_cache_read, i64::MAX); + let snapshot_mtime = super::kiro_source_mtime_ms(&snapshot).unwrap(); + assert!(snapshot_mtime > execution_mtime); + let snapshot_date = sessions::kiro::parse_kiro_file(&snapshot)[0].date.clone(); + let parsed = parse_local_clients(LocalParseOptions { + modified_after: Some(execution_mtime + 1), + since: Some(snapshot_date), + ..options.clone() + }) + .unwrap(); + assert_eq!(parsed.counts.get(ClientId::Kiro), 1); + assert!( + parsed.messages.is_empty(), + "mtime pruning must retain the older execution until suppression" + ); + + let missing = source_home.path().join("missing.chat"); + let mut scan = scanner::ScanResult::default(); + scan.get_mut(ClientId::Kiro).push(missing); + prune_scan_result_by_mtime(&mut scan, u64::MAX); + assert_eq!(scan.get(ClientId::Kiro).len(), 1); + + // Keep the fixture path live for the cache/source identity assertion. + assert!(message_cache::SourceFingerprint::from_kiro_path(&snapshot).is_some()); } } diff --git a/vendor/tokscale-core/src/message_cache.rs b/vendor/tokscale-core/src/message_cache.rs index e503e0d..c403eb8 100644 --- a/vendor/tokscale-core/src/message_cache.rs +++ b/vendor/tokscale-core/src/message_cache.rs @@ -1,3 +1,4 @@ +use crate::clients::ClientId; use crate::sessions::codex::CodexParseState; use crate::UnifiedMessage; use bincode::Options; @@ -8,100 +9,31 @@ use std::ffi::OsString; use std::fs::{self, File, OpenOptions}; use std::io::{BufReader, BufWriter, Read, Seek, SeekFrom, Write}; use std::path::{Path, PathBuf}; -use std::sync::{Arc, LazyLock, Mutex}; +use std::sync::{Mutex, OnceLock}; use std::time::UNIX_EPOCH; -// 19: Codex fork-replay parsing now skips replayed parent usage, scopes the -// token_count dedup key to the fork parent, and keeps user-fork turns after -// repeated child session_meta rows. Cached messages store their dedup_key and -// older entries can be empty, so they must be reparsed. -// 20 (#760 vendoring): session parsers now canonicalize provider ids -// (fireworks->fireworks_ai, vertex/gemini/azure aliases), recover -// missing-provider messages (gjc/pi), infer codex provider from the model, and -// emit stable dedup keys (qwen/mux) — all of which change cached parser output, -// so stale entries must be reparsed. (Our own schema counter; do not mirror -// upstream's number.) -// 21 (M5b: vendoring our own upstream-merged fixes): codex same-millisecond -// fork-replay gate (#735), micode row-id-fallback dedup namespaced by db (#742), -// and gjc header-less session fallback by file name (#743) all change cached -// dedup keys / codex token output, so stale entries must be reparsed. (#741 -// roo-family sibling fingerprinting also invalidates naturally, and #737's -// antigravity-cli timestamps converge with our dropped local patch.) -// 22 (M6: follow-up correctness + attribution): jcode journal-wins over stale -// snapshot + tz-less timestamps (#754), micode epoch seconds/ms normalization -// (#747), fable->anthropic provider inference (#762), and copilot per-message -// agent attribution (#724/#751) all change cached parser output, so stale -// entries must be reparsed. (Our own schema counter; do not mirror upstream's.) -// 23 (PR #30: drop synthetic claude rows): claudecode parsing now discards -// assistant turns whose model is the `` placeholder (locally -// fabricated by Claude Code, zero-token, no real cost), so a session cached -// before this change still carries the phantom zero-token row until reparsed. -// (Our own schema counter; do not mirror upstream's number.) -// 24 (M8-B1: cost provenance contract): UnifiedMessage now serializes -// cost_source, so schema-23 bincode entries must be rebuilt with the new field. -// 25 (M10-B: Jcode journal corrections that only replace a snapshotted message -// are now turn-neutral, so a following brand-new journal turn is no longer robbed -// of its is_turn_start; schema-24 caches carry the under-counted turn flags, so -// invalidate them.) -// 26 (M10-C: Pi session_info now supplies subagent attribution. Non-empty -// schema-25 caches replay Pi messages with no agent metadata, so invalidate them. -// This is TokenBar's own counter, replacing upstream #856's per-client -// parser_version rather than importing it.) -// 27 (M10-D: Claude workflow journals are now excluded, and deep nested workflow -// transcripts recover Tier-2 parent-session attribution. Schema-26 caches may -// replay usage-shaped workflow journals or stale/generic deep Tier-2 agent -// attribution, so invalidate them. The #856 parser_version/shard-cache -// architecture remains excluded.) -// 28 (M10-E: Copilot trace fallback now resolves the first root invoke_agent -// span by walking the parentSpanId hierarchy, so nested sub-agent invokes -// exported first no longer become the trace default. Schema-27 caches can carry -// the wrong nested fallback and must be rebuilt. The final #834 resolver covers -// #821's reversed-export-order semantic.) -// 29 (M14: Codex token durations now advance from the last accepted token -// snapshot instead of repeatedly measuring from the turn start. Schema-28 -// caches can replay overlapping duration_ms values or resume an incremental -// parse without the new cursor, so unchanged sources must be rebuilt.) -const CACHE_SCHEMA_VERSION: u32 = 29; -const CACHE_FILENAME: &str = "source-message-cache.bin"; +// CACHE_FORMAT_VERSION changes only when the serialized storage layout or a +// cross-client type such as UnifiedMessage changes incompatibly. Parser-only +// changes belong in parser_version() so one client cannot evict every other +// client's cached transcripts. +// 2: Related-file fingerprints now retain their paths and whether they were +// absent when cached. Claude sidechain parent candidates can therefore be +// revalidated without reparsing the sidechain on every warm scan, while a +// later-created parent transcript still invalidates the entry. +const CACHE_FORMAT_VERSION: u32 = 2; +// V2 intentionally starts cold and leaves source-message-cache.bin untouched: +// the monolith did not record a trustworthy parser owner for migration. +const CACHE_SHARD_DIRNAME: &str = "source-message-cache-v2"; const CACHE_LOCK_FILENAME: &str = "source-message-cache.lock"; -const MAX_CACHE_FILE_BYTES: u64 = 256 * 1024 * 1024; +const CACHE_SHARD_COUNT: usize = 256; +const MAX_CACHE_SHARD_BYTES: u64 = 256 * 1024 * 1024; const FINGERPRINT_SAMPLE_BYTES: usize = 4096; const FINGERPRINT_SAMPLE_POINTS: usize = 5; const HASH_BUFFER_BYTES: usize = 64 * 1024; -/// Process-level memo keyed by (path, size, modified_ns). On a cache hit the -/// caller skips `compute_sample_hashes` + `hash_prefix` (the expensive I/O + -/// SHA-256 work). On lock poison or stat failure the caller falls back to the -/// full recompute path — FAIL-LOUD, never returns stale/empty data. -struct HashMemoEntry { - size: u64, - modified_ns: u64, - sample_hashes: Vec, - content_hash: [u8; 32], -} - -static HASH_MEMO: LazyLock>> = - LazyLock::new(|| Mutex::new(HashMap::new())); - -/// Process-level memo for the deserialized store. On a cache hit (`path + -/// size + modified_ns` all match) `load()` skips the bincode deserialize. -/// Path is compared so a test-env config-dir switch never returns data from -/// a different store file. -struct StoreMemo { - cache_file: PathBuf, - file_size: u64, - file_modified_ns: u64, - entries: HashMap>, -} - -static STORE_MEMO: LazyLock>> = - LazyLock::new(|| Mutex::new(None)); - -/// Convert a `Duration` (typically elapsed since UNIX_EPOCH) to nanoseconds -/// as `u64`. -/// u64 holds ~584 years of ns since epoch — truncation is theoretical. -fn duration_to_nanos(d: std::time::Duration) -> u64 { - d.as_nanos() as u64 +#[cfg(test)] +thread_local! { + static FULL_HASH_CALLS: std::cell::Cell = const { std::cell::Cell::new(0) }; } fn cache_dir() -> Option { @@ -115,28 +47,14 @@ fn cache_dir() -> Option { } } -fn cache_path() -> Option { - Some(cache_dir()?.join(CACHE_FILENAME)) +fn cache_shard_dir() -> Option { + Some(cache_dir()?.join(CACHE_SHARD_DIRNAME)) } fn cache_lock_path() -> Option { Some(cache_dir()?.join(CACHE_LOCK_FILENAME)) } -fn legacy_cache_paths() -> Vec { - if crate::paths::is_config_dir_overridden() { - return Vec::new(); - } - - [ - crate::paths::legacy_dirs_cache_dir().map(|d| d.join(CACHE_FILENAME)), - crate::paths::legacy_dot_cache_tokscale_dir().map(|d| d.join(CACHE_FILENAME)), - ] - .into_iter() - .flatten() - .collect() -} - fn fallback_cache_dir() -> Option { std::env::var_os("XDG_RUNTIME_DIR") .map(PathBuf::from) @@ -179,6 +97,19 @@ fn ensure_cache_dir(dir: &Path) -> std::io::Result<()> { Ok(()) } +fn warn_cache_failure_once(context: &'static str, path: &Path, error: &impl std::fmt::Display) { + tracing::warn!(path = %path.display(), %error, %context, "source message cache failure"); + + // Most non-TUI commands (including `submit`) do not install a tracing + // subscriber. Surface persistence failures directly once per process so a + // permanently cold cache can never fail silently again. + static WARNED_CONTEXTS: OnceLock>> = OnceLock::new(); + let warned = WARNED_CONTEXTS.get_or_init(|| Mutex::new(HashSet::new())); + if warned.lock().is_ok_and(|mut warned| warned.insert(context)) { + eprintln!("tokscale: warning: {context} ({}): {error}", path.display()); + } +} + #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] pub(crate) struct FileSampleHash { pub offset: u64, @@ -198,97 +129,73 @@ pub(crate) struct SourceFingerprint { #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] pub(crate) struct RelatedFileFingerprint { pub suffix: String, + pub path: CachedPath, + pub exists: bool, pub size: u64, pub modified_ns: u64, pub sample_hashes: Vec, pub content_hash: [u8; 32], } -/// Metadata siblings a Grok `updates.jsonl` session's parse depends on, in the -/// same session directory. `parse_grok_updates_file` reconciles totals from -/// `signals.json` and `read_metadata` additionally reads `summary.json` and -/// `events.jsonl` for the model id. The fingerprint (`from_grok_path`) and every -/// mtime change probe in `lib.rs` (`latest_source_mtime_ms`, `grok_source_mtime_ms`) -/// must watch the *same* set, or a sibling-only write goes unnoticed and the -/// cache serves a stale (fallback-model) session. +/// Metadata siblings a Grok `updates.jsonl` session depends on. +/// Keep this list aligned with the parser fingerprint and live-tail probes. pub(crate) const GROK_METADATA_SIBLINGS: [&str; 3] = ["signals.json", "summary.json", "events.jsonl"]; +#[derive(Debug, Clone, PartialEq, Eq)] +pub(crate) enum FingerprintStatus { + /// Size and nanosecond mtime still match for the source and every parser + /// sidecar, and their bounded samples still match. No full-file SHA-256 was + /// computed, so a warm scan reads at most 20 KiB per watched file. + Unchanged, + /// Metadata changed, so a complete fingerprint was rebuilt to distinguish + /// a real content change from a metadata-only touch. + Changed(SourceFingerprint), +} + impl SourceFingerprint { pub(crate) fn from_path(path: &Path) -> Option { Self::from_path_with_related(path, std::iter::empty()) } + pub(crate) fn from_path_samples_only(path: &Path) -> Option { + Self::from_path_with_related_mode(path, std::iter::empty(), ContentHashMode::SamplesOnly) + } + pub(crate) fn from_sqlite_path(path: &Path) -> Option { let related_paths = ["-wal"] .into_iter() .map(|suffix| (suffix.to_string(), append_path_suffix(path, suffix))); - Self::from_path_with_related(path, related_paths) + Self::from_path_with_related_mode(path, related_paths, ContentHashMode::SamplesOnly) } - /// Fingerprint for a jcode session snapshot plus its sibling - /// `.journal.jsonl` append-log. jcode appends new turns to the - /// journal between snapshot rewrites, so a journal-only write leaves the - /// snapshot's size/mtime unchanged; including the journal in the fingerprint - /// invalidates the cache instead of serving stale (missing-turn) data. - pub(crate) fn from_jcode_path(path: &Path) -> Option { - let related_paths = - std::iter::once((".journal.jsonl".to_string(), jcode_journal_path(path))); - Self::from_path_with_related(path, related_paths) - } - - /// Fingerprint a Droid settings snapshot together with the fallback JSONL - /// that supplies its model when the snapshot omits one. - pub(crate) fn from_droid_path(path: &Path) -> Option { - let Some(jsonl) = crate::sessions::droid::droid_jsonl_path(path) else { - return Self::from_path(path); - }; - let related_paths = std::iter::once(("session.jsonl".to_string(), jsonl)); - Self::from_path_with_related(path, related_paths) + /// Fingerprint for Copilot Desktop's SQLite database, WAL, and dynamic + /// `session-state/*/events.jsonl` dependencies. Any unreadable dependency + /// fails open to a cache miss instead of serving a DB-only stale entry. + pub(crate) fn from_copilot_desktop_path(path: &Path) -> Option { + let related_paths = copilot_desktop_related_paths(path)?; + Self::from_path_with_related_mode(path, related_paths, ContentHashMode::SamplesOnly) } - /// Fingerprint a legacy Kimi wire log together with the shared config that - /// supplies its model. - pub(crate) fn from_kimi_path(path: &Path) -> Option { - let Some(config) = crate::sessions::kimi::kimi_config_path(path) else { - return Self::from_path(path); - }; - let related_paths = std::iter::once(("config.json".to_string(), config)); - Self::from_path_with_related(path, related_paths) + pub(crate) fn check_copilot_desktop_path_samples_only( + path: &Path, + cached: Option<&Self>, + ) -> Option { + let related_paths = copilot_desktop_related_paths(path)?; + Self::check_path_with_related_mode( + path, + related_paths, + cached, + ContentHashMode::SamplesOnly, + ) } - /// Fingerprint a Kiro CLI session header together with its same-stem - /// message sidecar. - pub(crate) fn from_kiro_path(path: &Path) -> Option { - let Some(messages) = crate::sessions::kiro::kiro_related_messages_path(path) else { - return Self::from_path(path); - }; - let related_paths = std::iter::once(("messages.jsonl".to_string(), messages)); - Self::from_path_with_related(path, related_paths) - } - - /// Fingerprint for a Grok `updates.jsonl` session and every sibling - /// `read_metadata` consults. `parse_grok_updates_file` reconciles session - /// totals from `signals.json` (compaction), and `read_metadata` additionally - /// reads `summary.json` (model id + `updated_at`/`created_at` timestamp) and - /// `events.jsonl` (model id, session id, `ts` timestamp) — so a sibling that - /// is written or rewritten after the last `updates.jsonl` write must still - /// invalidate the cache. An `updates.jsonl`-only (or `signals.json`-only) - /// fingerprint would leave a session pinned to its fallback model forever - /// when a late-arriving `summary.json`/`events.jsonl` is the only carrier of - /// the real model id. - /// - /// LOCAL DIVERGENCE from upstream junhoyeo/tokscale, which fingerprints only - /// `signals.json` here even though its `read_metadata` reads all three — the - /// same class of gap as our reported #741 (roo history sibling). Must be - /// re-applied on any re-vendor of `message_cache.rs`; candidate to report - /// upstream. See vendor/README.md. - pub(crate) fn from_grok_path(path: &Path) -> Option { - let session_dir = path.parent().unwrap_or_else(|| Path::new(".")); - let related_paths = GROK_METADATA_SIBLINGS - .into_iter() - .map(|name| (name.to_string(), session_dir.join(name))); - Self::from_path_with_related(path, related_paths) + pub(crate) fn from_jcode_path(path: &Path) -> Option { + let related_paths = std::iter::once(( + ".journal.jsonl".to_string(), + crate::sessions::jcode::jcode_journal_path(path), + )); + Self::from_path_with_related_mode(path, related_paths, ContentHashMode::SamplesOnly) } /// Fingerprint for a Roo-family task (`ui_messages.json`) and its sibling @@ -299,12 +206,13 @@ impl SourceFingerprint { pub(crate) fn from_roo_path(path: &Path) -> Option { let history = crate::sessions::roocode::history_path_for_ui_messages(path); let related_paths = std::iter::once(("api_conversation_history.json".to_string(), history)); - Self::from_path_with_related(path, related_paths) + Self::from_path_with_related_mode(path, related_paths, ContentHashMode::SamplesOnly) } /// Fingerprint for a Claude Code JSONL file that may have a sibling `.meta.json` /// sidecar. When the sidecar appears or changes (e.g. after a Claude Code upgrade), /// the fingerprint changes and the cache invalidates. + #[cfg(test)] pub(crate) fn from_claude_code_path_with_home( path: &Path, home_dir: Option<&Path>, @@ -328,20 +236,327 @@ impl SourceFingerprint { related.push((format!("parent-session-{index}.jsonl"), parent_path)); } - Self::from_path_with_related(path, related) + Self::from_path_with_related_mode(path, related, ContentHashMode::SamplesOnly) + } + + /// Fingerprint for a Grok `updates.jsonl` session and every sibling read by + /// its parser for rollup and session metadata. + pub(crate) fn from_grok_path(path: &Path) -> Option { + if path.file_name().and_then(|name| name.to_str()) == Some("unified.jsonl") { + return Self::from_path_samples_only(path); + } + let parent = path.parent().unwrap_or_else(|| Path::new(".")); + let related_paths = ["signals.json", "summary.json", "events.jsonl"] + .into_iter() + .map(|name| (name.to_string(), parent.join(name))); + Self::from_path_with_related_mode(path, related_paths, ContentHashMode::SamplesOnly) + } + + /// Fingerprint for a Kiro source file. IDE sessions consume a sibling + /// `messages.jsonl`, while CLI `*.json` headers consume same-stem `*.jsonl`. + /// Global-storage and `.chat` snapshots are self-contained. + pub(crate) fn from_kiro_path(path: &Path) -> Option { + let Some(messages) = crate::sessions::kiro::kiro_related_messages_path(path) else { + return Self::from_path_samples_only(path); + }; + let related_paths = std::iter::once(("messages.jsonl".to_string(), messages)); + Self::from_path_with_related_mode(path, related_paths, ContentHashMode::SamplesOnly) + } + + pub(crate) fn from_droid_path(path: &Path) -> Option { + let Some(jsonl) = crate::sessions::droid::droid_jsonl_path(path) else { + return Self::from_path_samples_only(path); + }; + let related_paths = std::iter::once(("session.jsonl".to_string(), jsonl)); + Self::from_path_with_related_mode(path, related_paths, ContentHashMode::SamplesOnly) + } + + pub(crate) fn from_kimi_path(path: &Path) -> Option { + if crate::sessions::kimi::is_kimi_code_path(path) { + return Self::from_path_samples_only(path); + } + let Some(config) = crate::sessions::kimi::kimi_config_path(path) else { + return Self::from_path_samples_only(path); + }; + let related_paths = std::iter::once(("config.json".to_string(), config)); + Self::from_path_with_related_mode(path, related_paths, ContentHashMode::SamplesOnly) + } + + pub(crate) fn check_path(path: &Path, cached: Option<&Self>) -> Option { + Self::check_path_with_related(path, std::iter::empty(), cached) + } + + /// Check a non-Codex source without rebuilding its write-only whole-file + /// hash when metadata or samples changed. Codex uses `check_path` because + /// its incremental resume state compares the full content hash; generic + /// parsers only need the bounded samples for invalidation. + pub(crate) fn check_path_samples_only( + path: &Path, + cached: Option<&Self>, + ) -> Option { + Self::check_path_with_related_mode( + path, + std::iter::empty(), + cached, + ContentHashMode::SamplesOnly, + ) + } + + pub(crate) fn check_sqlite_path( + path: &Path, + cached: Option<&Self>, + ) -> Option { + let related_paths = ["-wal"] + .into_iter() + .map(|suffix| (suffix.to_string(), append_path_suffix(path, suffix))); + // SQLite databases can be tens of GB; skip the whole-file content hash + // (size + mtime + samples detect changes, and no SQLite source reads + // content_hash). See ContentHashMode. + Self::check_path_with_related_mode( + path, + related_paths, + cached, + ContentHashMode::SamplesOnly, + ) + } + + pub(crate) fn check_jcode_path_samples_only( + path: &Path, + cached: Option<&Self>, + ) -> Option { + Self::check_jcode_path_with_mode(path, cached, ContentHashMode::SamplesOnly) + } + + fn check_jcode_path_with_mode( + path: &Path, + cached: Option<&Self>, + mode: ContentHashMode, + ) -> Option { + let related_paths = std::iter::once(( + ".journal.jsonl".to_string(), + crate::sessions::jcode::jcode_journal_path(path), + )); + Self::check_path_with_related_mode(path, related_paths, cached, mode) + } + + pub(crate) fn check_roo_path_samples_only( + path: &Path, + cached: Option<&Self>, + ) -> Option { + Self::check_roo_path_with_mode(path, cached, ContentHashMode::SamplesOnly) + } + + fn check_roo_path_with_mode( + path: &Path, + cached: Option<&Self>, + mode: ContentHashMode, + ) -> Option { + let history = crate::sessions::roocode::history_path_for_ui_messages(path); + let related_paths = std::iter::once(("api_conversation_history.json".to_string(), history)); + Self::check_path_with_related_mode(path, related_paths, cached, mode) + } + + pub(crate) fn check_claude_code_path_with_home_samples_only( + path: &Path, + cached: Option<&Self>, + home_dir: Option<&Path>, + ) -> Option { + Self::check_claude_code_path_with_home_mode( + path, + cached, + home_dir, + ContentHashMode::SamplesOnly, + ) + } + + fn check_claude_code_path_with_home_mode( + path: &Path, + cached: Option<&Self>, + home_dir: Option<&Path>, + mode: ContentHashMode, + ) -> Option { + let mut related = Vec::new(); + + if let Some(stem) = path.file_stem().and_then(|s| s.to_str()) { + let meta_filename = format!("{}.meta.json", stem); + related.push((".meta.json".to_string(), path.with_file_name(meta_filename))); + } + + if let Some(variant_path) = crate::cc_mirror::variant_file_for_session_path(path, home_dir) + { + related.push(("cc-mirror/variant.json".to_string(), variant_path)); + } + + let primary_matches = + cached.and_then(|fingerprint| primary_fingerprint_matches(path, fingerprint)); + let parent_paths = cached + .filter(|_| primary_matches == Some(true)) + .map(cached_claude_parent_paths) + .unwrap_or_else(|| { + crate::sessions::claudecode::parent_session_paths_for_cache(path) + .into_iter() + .enumerate() + .map(|(index, parent_path)| { + (format!("parent-session-{index}.jsonl"), parent_path) + }) + .collect() + }); + related.extend(parent_paths); + + Self::check_path_with_related_mode_and_primary(path, related, cached, mode, primary_matches) + } + + pub(crate) fn check_grok_path_samples_only( + path: &Path, + cached: Option<&Self>, + ) -> Option { + Self::check_grok_path_with_mode(path, cached, ContentHashMode::SamplesOnly) + } + + fn check_grok_path_with_mode( + path: &Path, + cached: Option<&Self>, + mode: ContentHashMode, + ) -> Option { + if path.file_name().and_then(|name| name.to_str()) == Some("unified.jsonl") { + return Self::check_path_with_related_mode(path, std::iter::empty(), cached, mode); + } + let parent = path.parent().unwrap_or_else(|| Path::new(".")); + let related_paths = ["signals.json", "summary.json", "events.jsonl"] + .into_iter() + .map(|name| (name.to_string(), parent.join(name))); + Self::check_path_with_related_mode(path, related_paths, cached, mode) + } + + pub(crate) fn check_kiro_path_samples_only( + path: &Path, + cached: Option<&Self>, + ) -> Option { + Self::check_kiro_path_with_mode(path, cached, ContentHashMode::SamplesOnly) + } + + fn check_kiro_path_with_mode( + path: &Path, + cached: Option<&Self>, + mode: ContentHashMode, + ) -> Option { + let Some(messages) = crate::sessions::kiro::kiro_related_messages_path(path) else { + return Self::check_path_with_related_mode(path, std::iter::empty(), cached, mode); + }; + let related_paths = std::iter::once(("messages.jsonl".to_string(), messages)); + Self::check_path_with_related_mode(path, related_paths, cached, mode) + } + + pub(crate) fn check_droid_path_samples_only( + path: &Path, + cached: Option<&Self>, + ) -> Option { + Self::check_droid_path_with_mode(path, cached, ContentHashMode::SamplesOnly) + } + + fn check_droid_path_with_mode( + path: &Path, + cached: Option<&Self>, + mode: ContentHashMode, + ) -> Option { + let Some(jsonl) = crate::sessions::droid::droid_jsonl_path(path) else { + return Self::check_path_with_related_mode(path, std::iter::empty(), cached, mode); + }; + let related_paths = std::iter::once(("session.jsonl".to_string(), jsonl)); + Self::check_path_with_related_mode(path, related_paths, cached, mode) + } + + pub(crate) fn check_kimi_path_samples_only( + path: &Path, + cached: Option<&Self>, + ) -> Option { + Self::check_kimi_path_with_mode(path, cached, ContentHashMode::SamplesOnly) + } + + fn check_kimi_path_with_mode( + path: &Path, + cached: Option<&Self>, + mode: ContentHashMode, + ) -> Option { + if crate::sessions::kimi::is_kimi_code_path(path) { + return Self::check_path_with_related_mode(path, std::iter::empty(), cached, mode); + } + let Some(config) = crate::sessions::kimi::kimi_config_path(path) else { + return Self::check_path_with_related_mode(path, std::iter::empty(), cached, mode); + }; + let related_paths = std::iter::once(("config.json".to_string(), config)); + Self::check_path_with_related_mode(path, related_paths, cached, mode) + } + + fn check_path_with_related( + path: &Path, + related_paths: I, + cached: Option<&Self>, + ) -> Option + where + I: IntoIterator, + { + Self::check_path_with_related_mode(path, related_paths, cached, ContentHashMode::Full) + } + + fn check_path_with_related_mode( + path: &Path, + related_paths: I, + cached: Option<&Self>, + mode: ContentHashMode, + ) -> Option + where + I: IntoIterator, + { + Self::check_path_with_related_mode_and_primary(path, related_paths, cached, mode, None) + } + + fn check_path_with_related_mode_and_primary( + path: &Path, + related_paths: I, + cached: Option<&Self>, + mode: ContentHashMode, + primary_matches: Option, + ) -> Option + where + I: IntoIterator, + { + let related_paths: Vec<(String, PathBuf)> = related_paths.into_iter().collect(); + let cache_hit = cached.is_some_and(|fingerprint| { + primary_matches + .unwrap_or_else(|| primary_fingerprint_matches(path, fingerprint).unwrap_or(false)) + && related_fingerprint_metadata_matches(&related_paths, fingerprint) + .unwrap_or(false) + }); + if cache_hit { + return Some(FingerprintStatus::Unchanged); + } + + Self::from_path_with_related_mode(path, related_paths, mode).map(FingerprintStatus::Changed) } fn from_path_with_related(path: &Path, related_paths: I) -> Option where I: IntoIterator, { - let (size, modified_ns, sample_hashes, content_hash) = file_fingerprint_parts(path)?; + Self::from_path_with_related_mode(path, related_paths, ContentHashMode::Full) + } + + fn from_path_with_related_mode( + path: &Path, + related_paths: I, + mode: ContentHashMode, + ) -> Option + where + I: IntoIterator, + { + let (size, modified_ns, sample_hashes, content_hash) = file_fingerprint_parts(path, mode)?; let mut related_files: Vec = related_paths .into_iter() - .filter_map(|(suffix, related_path)| { - RelatedFileFingerprint::from_path(suffix, &related_path) + .map(|(suffix, related_path)| { + RelatedFileFingerprint::from_path(suffix, &related_path, mode) }) - .collect(); + .collect::>()?; related_files.sort_by(|left, right| left.suffix.cmp(&right.suffix)); Some(Self { @@ -354,17 +569,131 @@ impl SourceFingerprint { } } +fn copilot_desktop_related_paths(path: &Path) -> Option> { + let root = path.parent().unwrap_or_else(|| Path::new(".")); + File::open(path).ok()?; + let events = crate::sessions::copilot_desktop::session_state_event_paths(path).ok()?; + let mut related_paths = Vec::new(); + + let wal_path = append_path_suffix(path, "-wal"); + match fs::metadata(&wal_path) { + Ok(metadata) if metadata.is_file() => { + File::open(&wal_path).ok()?; + related_paths.push(("-wal".to_string(), wal_path)); + } + Ok(_) => {} + Err(error) if error.kind() == std::io::ErrorKind::NotFound => {} + Err(_) => return None, + } + + for event in events { + File::open(&event).ok()?; + let label = event + .strip_prefix(root) + .unwrap_or(&event) + .to_string_lossy() + .into_owned(); + related_paths.push((label, event)); + } + Some(related_paths) +} + impl RelatedFileFingerprint { - fn from_path(suffix: String, path: &Path) -> Option { - let (size, modified_ns, sample_hashes, content_hash) = file_fingerprint_parts(path)?; - Some(Self { - suffix, - size, - modified_ns, - sample_hashes, - content_hash, - }) + fn from_path(suffix: String, path: &Path, mode: ContentHashMode) -> Option { + let cached_path = CachedPath::from_path(path); + match path.metadata() { + Ok(_) => { + let (size, modified_ns, sample_hashes, content_hash) = + file_fingerprint_parts(path, mode)?; + Some(Self { + suffix, + path: cached_path, + exists: true, + size, + modified_ns, + sample_hashes, + content_hash, + }) + } + Err(error) if error.kind() == std::io::ErrorKind::NotFound => Some(Self { + suffix, + path: cached_path, + exists: false, + size: 0, + modified_ns: 0, + sample_hashes: Vec::new(), + content_hash: [0; 32], + }), + Err(_) => None, + } + } +} + +fn cached_claude_parent_paths(cached: &SourceFingerprint) -> Vec<(String, PathBuf)> { + cached + .related_files + .iter() + .filter(|related| related.suffix.starts_with("parent-session-")) + .map(|related| (related.suffix.clone(), related.path.to_path_buf())) + .collect() +} + +fn primary_fingerprint_matches(path: &Path, cached: &SourceFingerprint) -> Option { + let (size, modified_ns) = metadata_signature(path).ok()?; + if size != cached.size || modified_ns != cached.modified_ns { + return Some(false); + } + Some(compute_sample_hashes(path, size)? == cached.sample_hashes) +} + +fn metadata_signature(path: &Path) -> std::io::Result<(u64, u64)> { + let metadata = path.metadata()?; + let modified_ns = metadata + .modified()? + .duration_since(UNIX_EPOCH) + .map_err(std::io::Error::other)? + .as_nanos() as u64; + Ok((metadata.len(), modified_ns)) +} + +fn related_fingerprint_metadata_matches( + related_paths: &[(String, PathBuf)], + cached: &SourceFingerprint, +) -> Option { + if cached.related_files.len() != related_paths.len() { + return Some(false); + } + + for (suffix, related_path) in related_paths { + let Some(related) = cached + .related_files + .iter() + .find(|related| related.suffix == *suffix) + else { + return Some(false); + }; + if related.path != CachedPath::from_path(related_path) { + return Some(false); + } + match metadata_signature(related_path) { + Ok((size, modified_ns)) => { + if !related.exists || related.size != size || related.modified_ns != modified_ns { + return Some(false); + } + if compute_sample_hashes(related_path, size)? != related.sample_hashes { + return Some(false); + } + } + Err(error) if error.kind() == std::io::ErrorKind::NotFound => { + if related.exists { + return Some(false); + } + } + Err(_) => return None, + } } + + Some(true) } #[derive(Debug, Clone, Serialize, Deserialize)] @@ -393,6 +722,10 @@ impl CachedPath { PathBuf::from(OsString::from_vec(self.0.clone())) } + + fn update_digest(&self, hasher: &mut Sha256) { + hasher.update(&self.0); + } } #[cfg(windows)] @@ -413,6 +746,12 @@ impl CachedPath { PathBuf::from(OsString::from_wide(&self.0)) } + + fn update_digest(&self, hasher: &mut Sha256) { + for code_unit in &self.0 { + hasher.update(code_unit.to_le_bytes()); + } + } } #[cfg(not(any(unix, windows)))] @@ -428,10 +767,103 @@ impl CachedPath { pub(crate) fn to_path_buf(&self) -> PathBuf { PathBuf::from(&self.0) } + + fn update_digest(&self, hasher: &mut Sha256) { + hasher.update(self.0.as_bytes()); + } +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] +pub(crate) struct CacheIdentity { + namespace: &'static str, + parser_version: u32, +} + +impl CacheIdentity { + pub(crate) fn for_client(client: ClientId) -> Self { + Self { + namespace: client.as_str(), + parser_version: parser_version(client), + } + } + + pub(crate) const fn synthetic() -> Self { + Self { + namespace: "synthetic", + parser_version: 1, + } + } + + fn current_for_namespace(namespace: &str) -> Option { + if namespace == "synthetic" { + return Some(Self::synthetic()); + } + ClientId::from_str(namespace).map(Self::for_client) + } + + fn all() -> impl Iterator { + ClientId::iter() + .map(Self::for_client) + .chain(std::iter::once(Self::synthetic())) + } +} + +fn parser_version(client: ClientId) -> u32 { + match client { + // These clients accumulated parser-only invalidations under the old + // global schema. Their independent counters start from those histories + // so future changes have an obvious local version to increment. + ClientId::Codex => 4, + ClientId::Jcode => 4, + ClientId::Copilot => 4, + _ => 1, + } +} + +#[derive(Debug, Clone, PartialEq, Eq, Hash)] +pub(crate) struct CacheKey { + namespace: String, + path: CachedPath, +} + +impl CacheKey { + fn new(identity: CacheIdentity, path: &Path) -> Self { + Self { + namespace: identity.namespace.to_string(), + path: CachedPath::from_path(path), + } + } + + fn from_entry(entry: &CachedSourceEntry) -> Self { + Self { + namespace: entry.parser_namespace.clone(), + path: entry.path.clone(), + } + } + + fn shard(&self) -> CacheShardKey { + let mut hasher = Sha256::new(); + hasher.update(self.namespace.as_bytes()); + hasher.update([0]); + self.path.update_digest(&mut hasher); + let digest = hasher.finalize(); + CacheShardKey { + namespace: self.namespace.clone(), + index: usize::from(digest[0]) % CACHE_SHARD_COUNT, + } + } +} + +#[derive(Debug, Clone, PartialEq, Eq, Hash)] +struct CacheShardKey { + namespace: String, + index: usize, } #[derive(Debug, Clone, Serialize, Deserialize)] pub(crate) struct CachedSourceEntry { + parser_namespace: String, + parser_version: u32, pub path: CachedPath, pub fingerprint: SourceFingerprint, pub messages: Vec, @@ -441,6 +873,7 @@ pub(crate) struct CachedSourceEntry { impl CachedSourceEntry { pub(crate) fn new( + identity: CacheIdentity, path: &Path, fingerprint: SourceFingerprint, messages: Vec, @@ -448,6 +881,8 @@ impl CachedSourceEntry { codex_incremental: Option, ) -> Self { Self { + parser_namespace: identity.namespace.to_string(), + parser_version: identity.parser_version, path: CachedPath::from_path(path), fingerprint, messages, @@ -455,254 +890,200 @@ impl CachedSourceEntry { codex_incremental, } } -} -#[derive(Debug, Default, Serialize, Deserialize)] -struct CachedSourceStore { - schema_version: u32, - entries: Vec, -} - -#[derive(Default)] -pub(crate) struct SourceMessageCache { - pub entries: HashMap>, - dirty: bool, - dirty_keys: HashSet, - deleted_paths: HashSet, -} - -/// Return the memo's entry map if `path + size + mtime_ns` all match. -/// Returns `None` on lock poison, stat failure, or any mismatch. -fn store_memo_entries_if_current( - path: &Path, -) -> Option>> { - let guard = STORE_MEMO.lock().ok()?; - let memo = guard.as_ref()?; - if memo.cache_file != path { - return None; - } - let meta = fs::metadata(path).ok()?; - let size = meta.len(); - let modified_ns = meta - .modified() - .ok() - .and_then(|t| t.duration_since(UNIX_EPOCH).ok()) - .map(duration_to_nanos)?; - if size != memo.file_size || modified_ns != memo.file_modified_ns { - return None; + fn identity_is_current(&self) -> bool { + CacheIdentity::current_for_namespace(&self.parser_namespace) + .is_some_and(|identity| identity.parser_version == self.parser_version) } - Some(memo.entries.clone()) } -/// Write or replace the STORE_MEMO entry. Silently skips on lock poison or -/// unreadable mtime — a miss on the next `load()` will just re-read disk. -fn store_memo_update( - path: &Path, - meta: &std::fs::Metadata, - entries: &HashMap>, -) { - let size = meta.len(); - let Some(modified_ns) = meta - .modified() - .ok() - .and_then(|t| t.duration_since(UNIX_EPOCH).ok()) - .map(duration_to_nanos) - else { - return; - }; - if let Ok(mut guard) = STORE_MEMO.lock() { - *guard = Some(StoreMemo { - cache_file: path.to_path_buf(), - file_size: size, - file_modified_ns: modified_ns, - entries: entries.clone(), - }); - } +/// The envelope is deliberately independent from CachedSourceEntry's binary +/// layout. A parser version can therefore be checked before its payload is +/// deserialized, so (for example) a CodexParseState layout change cannot make +/// Claude's independently sharded cache unreadable. +#[derive(Debug, Serialize, Deserialize)] +struct CachedShardEnvelope { + format_version: u32, + parser_namespace: String, + parser_version: u32, + payload: Vec, } -/// Load from legacy cache paths without memoizing (path would differ from -/// the canonical path so the memo would never be a valid hit). -fn load_from_legacy_paths() -> SourceMessageCache { - legacy_cache_paths() - .into_iter() - .find_map(|p| read_store_from_path(&p)) - .map(|store| SourceMessageCache { - entries: store - .entries - .into_iter() - .map(|e| (e.path.clone(), Arc::new(e))) - .collect(), - dirty: false, - dirty_keys: HashSet::new(), - deleted_paths: HashSet::new(), - }) - .unwrap_or_default() +#[derive(Debug, Clone)] +enum DeletionReason { + Invalidated(SourceFingerprint), + Missing, } -/// Build the `merged_entries` base for `save_if_dirty`. Tries STORE_MEMO -/// first to skip re-reading the on-disk store; falls back to disk on a miss. -fn save_merge_base( - final_path: &Path, -) -> HashMap> { - if let Some(entries) = store_memo_entries_if_current(final_path) { - return entries; - } - read_store_from_path(final_path) - .map(|store| { - store - .entries - .into_iter() - .map(|e| (e.path.clone(), Arc::new(e))) - .collect() - }) - .unwrap_or_default() +#[derive(Default)] +pub(crate) struct SourceMessageCache { + pub entries: HashMap, + dirty: bool, + dirty_keys: HashSet, + deleted_keys: HashMap, + rewrite_shards: HashSet, } impl SourceMessageCache { pub(crate) fn load() -> Self { - let Some(path) = cache_path() else { + Self::load_with_limit(MAX_CACHE_SHARD_BYTES) + } + + fn load_with_limit(max_shard_bytes: u64) -> Self { + let Some(shard_root) = cache_shard_dir() else { return Self::default(); }; let Some(lock_path) = cache_lock_path() else { return Self::default(); }; - if let Some(lock_dir) = lock_path.parent() { - if ensure_cache_dir(lock_dir).is_err() { - return Self::default(); - } - } - - // Check STORE_MEMO before acquiring the file lock. A hit means the - // on-disk store is almost certainly unchanged — the stat is taken - // outside the file lock, so a concurrent atomic-rename writer can - // theoretically race it; (path, size, mtime_ns) collision makes a - // false hit vanishingly unlikely. Path is part of the key so a - // test-env config-dir switch never returns data from a different file. - if let Some(entries) = store_memo_entries_if_current(&path) { - return Self { - entries, - dirty: false, - dirty_keys: HashSet::new(), - deleted_paths: HashSet::new(), - }; + if let Err(error) = ensure_cache_dir(&shard_root) { + warn_cache_failure_once( + "source message cache directory is unavailable", + &shard_root, + &error, + ); + return Self::default(); } - let lock_file = match OpenOptions::new() .read(true) .write(true) .create(true) .truncate(false) - .open(lock_path) + .open(&lock_path) { Ok(file) => file, - Err(_) => return Self::default(), + Err(error) => { + warn_cache_failure_once( + "source message cache lock is unavailable", + &lock_path, + &error, + ); + return Self::default(); + } }; - if fs2::FileExt::lock_shared(&lock_file).is_err() { + if let Err(error) = fs2::FileExt::lock_shared(&lock_file) { + warn_cache_failure_once("source message cache lock failed", &lock_path, &error); return Self::default(); } - // Stat inside the shared lock for a consistent size+mtime pair. - let locked_meta = fs::metadata(&path).ok(); + let mut cache = Self::default(); + for identity in CacheIdentity::all() { + let parser_dir = shard_root.join(identity.namespace); + let read_dir = match fs::read_dir(&parser_dir) { + Ok(read_dir) => read_dir, + Err(error) if error.kind() == std::io::ErrorKind::NotFound => continue, + Err(error) => { + warn_cache_failure_once( + "source message cache parser directory is unreadable", + &parser_dir, + &error, + ); + continue; + } + }; - let store = match read_store_from_path_status(&path) { - CacheReadStatus::Loaded(store) => Some(store), - CacheReadStatus::Missing => { - // Legacy paths: load but do not memoize (path mismatch). - return load_from_legacy_paths(); + for dir_entry in read_dir.filter_map(Result::ok) { + let Some(index) = parse_shard_filename(&dir_entry.file_name()) else { + continue; + }; + let shard_key = CacheShardKey { + namespace: identity.namespace.to_string(), + index, + }; + let path = dir_entry.path(); + match read_shard_with_limit(&path, identity, max_shard_bytes) { + ShardReadStatus::Loaded(entries) => { + for entry in entries { + let key = CacheKey::from_entry(&entry); + if key.shard() == shard_key && entry.identity_is_current() { + cache.entries.insert(key, entry); + } else { + cache.rewrite_shards.insert(shard_key.clone()); + } + } + } + ShardReadStatus::Missing => {} + ShardReadStatus::Stale => { + cache.rewrite_shards.insert(shard_key); + } + ShardReadStatus::Invalid(error) => { + warn_cache_failure_once( + "source message cache shard is invalid", + &path, + &error, + ); + cache.rewrite_shards.insert(shard_key); + } + } } - CacheReadStatus::Invalid => None, - }; - let Some(store) = store else { - return Self::default(); - }; - - let entries: HashMap> = store - .entries - .into_iter() - .map(|entry| (entry.path.clone(), Arc::new(entry))) - .collect(); - - // Update STORE_MEMO so the next load() on an unchanged file is free. - if let Some(meta) = locked_meta.as_ref() { - store_memo_update(&path, meta, &entries); } - Self { - entries, - dirty: false, - dirty_keys: HashSet::new(), - deleted_paths: HashSet::new(), - } + cache.dirty = !cache.rewrite_shards.is_empty(); + cache } pub(crate) fn insert(&mut self, entry: CachedSourceEntry) { - let key = entry.path.clone(); - self.entries.insert(key.clone(), Arc::new(entry)); - self.deleted_paths.remove(&key); + let key = CacheKey::from_entry(&entry); + self.entries.insert(key.clone(), entry); + self.deleted_keys.remove(&key); self.dirty_keys.insert(key); self.dirty = true; } - pub(crate) fn get(&self, path: &Path) -> Option<&CachedSourceEntry> { - let key = CachedPath::from_path(path); - self.entries.get(&key).map(|a| a.as_ref()) + pub(crate) fn get(&self, identity: CacheIdentity, path: &Path) -> Option<&CachedSourceEntry> { + let key = CacheKey::new(identity, path); + self.entries.get(&key).filter(|entry| { + entry.parser_namespace == identity.namespace + && entry.parser_version == identity.parser_version + }) } - pub(crate) fn remove(&mut self, path: &Path) { - let key = CachedPath::from_path(path); - if self.entries.remove(&key).is_some() { + pub(crate) fn remove(&mut self, identity: CacheIdentity, path: &Path) { + let key = CacheKey::new(identity, path); + if let Some(entry) = self.entries.remove(&key) { self.dirty_keys.remove(&key); - self.deleted_paths.insert(key); + self.deleted_keys + .insert(key, DeletionReason::Invalidated(entry.fingerprint)); self.dirty = true; } - // Evict from HASH_MEMO so a deleted session file leaves no stale memo - // entry. Lock poison → silently skip; next access recomputes. - if let Ok(mut memo) = HASH_MEMO.lock() { - memo.remove(path); - } } pub(crate) fn prune_missing_files(&mut self) { - let removed_paths: Vec = self + let removed_keys: Vec = self .entries .keys() - .filter(|path| !path.to_path_buf().exists()) + .filter(|key| !key.path.to_path_buf().exists()) .cloned() .collect(); - if removed_paths.is_empty() { - return; - } - for path in &removed_paths { - let path_buf = path.to_path_buf(); - self.entries.remove(path); - self.dirty_keys.remove(path); - self.deleted_paths.insert(path.clone()); - // Evict from HASH_MEMO so deleted session files leave no stale - // memo entries. Lock poison → silently skip; next access - // recomputes. - if let Ok(mut memo) = HASH_MEMO.lock() { - memo.remove(&path_buf); - } + for key in removed_keys { + self.entries.remove(&key); + self.dirty_keys.remove(&key); + self.deleted_keys.insert(key, DeletionReason::Missing); + self.dirty = true; } - self.dirty = true; } pub(crate) fn save_if_dirty(&mut self) { + self.save_if_dirty_with_limit(MAX_CACHE_SHARD_BYTES); + } + + fn save_if_dirty_with_limit(&mut self, max_shard_bytes: u64) { if !self.dirty { return; } - let Some(dir) = cache_dir() else { + let Some(shard_root) = cache_shard_dir() else { return; }; - if ensure_cache_dir(&dir).is_err() { + if let Err(error) = ensure_cache_dir(&shard_root) { + warn_cache_failure_once( + "source message cache directory is unavailable", + &shard_root, + &error, + ); return; } - - let Some(final_path) = cache_path() else { - return; - }; let Some(lock_path) = cache_lock_path() else { return; }; @@ -711,138 +1092,264 @@ impl SourceMessageCache { .write(true) .create(true) .truncate(false) - .open(lock_path) + .open(&lock_path) { Ok(file) => file, - Err(_) => return, + Err(error) => { + warn_cache_failure_once( + "source message cache lock is unavailable", + &lock_path, + &error, + ); + return; + } }; - if fs2::FileExt::lock_exclusive(&lock_file).is_err() { + if let Err(error) = fs2::FileExt::lock_exclusive(&lock_file) { + warn_cache_failure_once("source message cache lock failed", &lock_path, &error); return; } - // Build the merged entries. Try STORE_MEMO first to skip re-reading - // the on-disk store; fall back to disk read on a miss. - let mut merged_entries: HashMap> = - save_merge_base(&final_path); - - for path in &self.deleted_paths { - if !path.to_path_buf().exists() { - merged_entries.remove(path); - } + // Bucket dirty and deleted keys by shard up front. CacheKey::shard() + // computes a SHA-256 digest, so grouping once keeps hashing at O(keys). + // The previous per-shard `.filter(|k| k.shard() == shard_key)` recomputed + // that digest for every key on every shard — O(shards * keys) — which + // dominated cold-cache builds (hundreds of shards * tens of thousands of + // files re-hashed). + let mut dirty_by_shard: HashMap> = HashMap::new(); + for key in &self.dirty_keys { + dirty_by_shard + .entry(key.shard()) + .or_default() + .push(key.clone()); } - for path in &self.dirty_keys { - if let Some(arc) = self.entries.get(path) { - merged_entries.insert(path.clone(), Arc::clone(arc)); - } + let mut deleted_by_shard: HashMap> = + HashMap::new(); + for (key, reason) in &self.deleted_keys { + deleted_by_shard + .entry(key.shard()) + .or_default() + .push((key.clone(), reason.clone())); } - let store = CachedSourceStore { - schema_version: CACHE_SCHEMA_VERSION, - entries: merged_entries.values().map(|e| (**e).clone()).collect(), - }; - - let nanos = std::time::SystemTime::now() - .duration_since(std::time::UNIX_EPOCH) - .map(duration_to_nanos) - .unwrap_or(0); - let tmp_path = dir.join(format!( - ".{}.{}.{:x}.tmp", - CACHE_FILENAME, - std::process::id(), - nanos - )); + let mut affected_shards = self.rewrite_shards.clone(); + affected_shards.extend(dirty_by_shard.keys().cloned()); + affected_shards.extend(deleted_by_shard.keys().cloned()); - // INVARIANT: All cache writes use atomic temp-file rename. NEVER delete - // the canonical cache file before writing — a partial save or process - // crash between delete and rename would lose the cache. The temp-file - // pattern makes corruption-on-crash impossible. - let write_result = (|| -> std::io::Result<()> { - let file = File::create(&tmp_path)?; - let mut writer = BufWriter::new(file); - bincode::options() - .with_limit(MAX_CACHE_FILE_BYTES) - .serialize_into(&mut writer, &store) - .map_err(std::io::Error::other)?; - writer.flush()?; - writer.get_ref().sync_all()?; - drop(writer); - crate::fs_atomic::replace_file(&tmp_path, &final_path)?; - let final_file = OpenOptions::new() - .read(true) - .write(true) - .open(&final_path)?; - final_file.sync_all()?; - Ok(()) - })(); - - if write_result.is_err() { - let _ = fs::remove_file(&tmp_path); - return; - } + let mut successful_shards = HashSet::new(); + for shard_key in affected_shards { + let Some(identity) = CacheIdentity::current_for_namespace(&shard_key.namespace) else { + continue; + }; + let parser_dir = shard_root.join(identity.namespace); + if let Err(error) = ensure_cache_dir(&parser_dir) { + warn_cache_failure_once( + "source message cache parser directory is unavailable", + &parser_dir, + &error, + ); + continue; + } + let final_path = shard_path(&shard_root, &shard_key); + + let mut merged_entries: HashMap = + match read_shard_with_limit(&final_path, identity, max_shard_bytes) { + ShardReadStatus::Loaded(entries) => entries + .into_iter() + .filter(|entry| entry.identity_is_current()) + .map(|entry| (CacheKey::from_entry(&entry), entry)) + .filter(|(key, _)| key.shard() == shard_key) + .collect(), + ShardReadStatus::Missing | ShardReadStatus::Stale => HashMap::new(), + ShardReadStatus::Invalid(error) => { + warn_cache_failure_once( + "source message cache shard is invalid", + &final_path, + &error, + ); + HashMap::new() + } + }; + + if let Some(deleted) = deleted_by_shard.get(&shard_key) { + for (key, reason) in deleted { + let should_remove = match reason { + DeletionReason::Missing => !key.path.to_path_buf().exists(), + DeletionReason::Invalidated(expected) => merged_entries + .get(key) + .is_some_and(|entry| entry.fingerprint == *expected), + }; + if should_remove { + merged_entries.remove(key); + } + } + } + if let Some(dirty) = dirty_by_shard.get(&shard_key) { + for key in dirty { + if let Some(entry) = self.entries.get(key) { + merged_entries.insert(key.clone(), entry.clone()); + } + } + } - // Stat final_path while exclusive lock is still held, then refresh - // STORE_MEMO so the next load() skips the disk read (race-safe). - if let Ok(meta) = fs::metadata(&final_path) { - store_memo_update(&final_path, &meta, &merged_entries); + let mut entries: Vec = merged_entries.into_values().collect(); + entries.sort_by_key(|left| left.path.to_path_buf()); + match write_shard_with_limit(&final_path, identity, &entries, max_shard_bytes) { + Ok(()) => { + successful_shards.insert(shard_key); + } + Err(error) => { + warn_cache_failure_once( + "source message cache shard could not be saved; future scans may remain cold", + &final_path, + &error, + ); + } + } } - self.entries = merged_entries; - self.dirty = false; - self.dirty_keys.clear(); - self.deleted_paths.clear(); + self.dirty_keys + .retain(|key| !successful_shards.contains(&key.shard())); + self.deleted_keys + .retain(|key, _| !successful_shards.contains(&key.shard())); + self.rewrite_shards + .retain(|shard| !successful_shards.contains(shard)); + self.dirty = !(self.dirty_keys.is_empty() + && self.deleted_keys.is_empty() + && self.rewrite_shards.is_empty()); } } -fn read_store_from_path(path: &Path) -> Option { - let file = File::open(path).ok()?; - let metadata = file.metadata().ok()?; - if metadata.len() > MAX_CACHE_FILE_BYTES { - return None; - } +fn shard_filename(index: usize) -> String { + format!("shard-{index:02x}.bin") +} - let reader = BufReader::new(file); - let store: CachedSourceStore = bincode::options() - .with_limit(MAX_CACHE_FILE_BYTES) - .deserialize_from(reader) - .ok()?; - if store.schema_version != CACHE_SCHEMA_VERSION { - return None; - } - Some(store) +fn parse_shard_filename(filename: &std::ffi::OsStr) -> Option { + let filename = filename.to_str()?; + let encoded = filename.strip_prefix("shard-")?.strip_suffix(".bin")?; + let index = usize::from_str_radix(encoded, 16).ok()?; + (index < CACHE_SHARD_COUNT).then_some(index) +} + +fn shard_path(root: &Path, key: &CacheShardKey) -> PathBuf { + root.join(&key.namespace).join(shard_filename(key.index)) } -enum CacheReadStatus { +enum ShardReadStatus { Missing, - Invalid, - Loaded(CachedSourceStore), + Stale, + Invalid(String), + Loaded(Vec), +} + +#[cfg(test)] +fn read_shard(path: &Path, identity: CacheIdentity) -> ShardReadStatus { + read_shard_with_limit(path, identity, MAX_CACHE_SHARD_BYTES) } -fn read_store_from_path_status(path: &Path) -> CacheReadStatus { +fn read_shard_with_limit( + path: &Path, + identity: CacheIdentity, + max_shard_bytes: u64, +) -> ShardReadStatus { let file = match File::open(path) { Ok(file) => file, - Err(err) if err.kind() == std::io::ErrorKind::NotFound => return CacheReadStatus::Missing, - Err(_) => return CacheReadStatus::Invalid, + Err(error) if error.kind() == std::io::ErrorKind::NotFound => { + return ShardReadStatus::Missing + } + Err(error) => return ShardReadStatus::Invalid(error.to_string()), }; let metadata = match file.metadata() { Ok(metadata) => metadata, - Err(_) => return CacheReadStatus::Invalid, + Err(error) => return ShardReadStatus::Invalid(error.to_string()), }; - if metadata.len() > MAX_CACHE_FILE_BYTES { - return CacheReadStatus::Invalid; + if metadata.len() > max_shard_bytes { + return ShardReadStatus::Invalid(format!( + "{} bytes exceeds the {}-byte shard limit", + metadata.len(), + max_shard_bytes + )); } - let reader = BufReader::new(file); - let store: CachedSourceStore = match bincode::options() - .with_limit(MAX_CACHE_FILE_BYTES) - .deserialize_from(reader) + let envelope: CachedShardEnvelope = match bincode::options() + .with_limit(max_shard_bytes) + .deserialize_from(BufReader::new(file)) { - Ok(store) => store, - Err(_) => return CacheReadStatus::Invalid, + Ok(envelope) => envelope, + Err(error) => return ShardReadStatus::Invalid(error.to_string()), }; - if store.schema_version != CACHE_SCHEMA_VERSION { - return CacheReadStatus::Invalid; + if envelope.format_version != CACHE_FORMAT_VERSION { + return ShardReadStatus::Stale; + } + if envelope.parser_namespace != identity.namespace + || envelope.parser_version != identity.parser_version + { + return ShardReadStatus::Stale; + } + + match bincode::options() + .with_limit(max_shard_bytes) + .deserialize(&envelope.payload) + { + Ok(entries) => ShardReadStatus::Loaded(entries), + Err(error) => ShardReadStatus::Invalid(error.to_string()), } - CacheReadStatus::Loaded(store) +} + +fn write_shard_with_limit( + final_path: &Path, + identity: CacheIdentity, + entries: &[CachedSourceEntry], + max_shard_bytes: u64, +) -> std::io::Result<()> { + let payload = bincode::options() + .with_limit(max_shard_bytes) + .serialize(entries) + .map_err(std::io::Error::other)?; + let envelope = CachedShardEnvelope { + format_version: CACHE_FORMAT_VERSION, + parser_namespace: identity.namespace.to_string(), + parser_version: identity.parser_version, + payload, + }; + let parent = final_path + .parent() + .ok_or_else(|| std::io::Error::other("cache shard has no parent directory"))?; + let nanos = std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .map(|duration| duration.as_nanos() as u64) + .unwrap_or(0); + let tmp_path = parent.join(format!( + ".{}.{}.{nanos:x}.tmp", + final_path + .file_name() + .and_then(|name| name.to_str()) + .unwrap_or("source-message-cache"), + std::process::id(), + )); + + // INVARIANT: shard writes use atomic temp-file replacement. Never remove + // the canonical shard before the replacement is completely serialized and + // fsynced, or one failed large shard write could destroy its last good copy. + let write_result = (|| -> std::io::Result<()> { + let file = File::create(&tmp_path)?; + let mut writer = BufWriter::new(file); + bincode::options() + .with_limit(max_shard_bytes) + .serialize_into(&mut writer, &envelope) + .map_err(std::io::Error::other)?; + writer.flush()?; + writer.get_ref().sync_all()?; + drop(writer); + crate::fs_atomic::replace_file(&tmp_path, final_path)?; + let final_file = OpenOptions::new().read(true).write(true).open(final_path)?; + final_file.sync_all()?; + Ok(()) + })(); + + if write_result.is_err() { + let _ = fs::remove_file(&tmp_path); + } + write_result } fn read_sample_hash(file: &mut File, offset: u64, len: usize) -> Option { @@ -914,43 +1421,36 @@ fn hash_bytes(bytes: &[u8]) -> u64 { hash } -fn file_fingerprint_parts(path: &Path) -> Option<(u64, u64, Vec, [u8; 32])> { +/// Whether a fingerprint carries a whole-file `content_hash`. +/// +/// Validation uses size + mtime + samples ([`primary_fingerprint_matches`] and +/// [`related_fingerprint_metadata_matches`]) for every source. Only Codex reads +/// `content_hash` for incremental resume; +/// generic parsers and SQLite sources store a zero sentinel so changed or cold +/// files do not pay for a second whole-file hash that cannot affect parsing. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +enum ContentHashMode { + Full, + SamplesOnly, +} + +fn file_fingerprint_parts( + path: &Path, + mode: ContentHashMode, +) -> Option<(u64, u64, Vec, [u8; 32])> { let metadata = path.metadata().ok()?; let size = metadata.len(); let modified_ns = metadata .modified() .ok()? .duration_since(UNIX_EPOCH) - .ok() - .map(duration_to_nanos)?; - - // Check HASH_MEMO: if (path, size, modified_ns) all match, skip the - // expensive compute_sample_hashes + hash_prefix (file I/O + SHA-256). - // Lock poison or stat failure → fall through to full recompute (FAIL-LOUD). - if let Ok(memo) = HASH_MEMO.lock() { - if let Some(entry) = memo.get(path) { - if entry.size == size && entry.modified_ns == modified_ns { - return Some((size, modified_ns, entry.sample_hashes.clone(), entry.content_hash)); - } - } - } - + .ok()? + .as_nanos() as u64; let sample_hashes = compute_sample_hashes(path, size)?; - let content_hash = hash_prefix(path, size)?; - - // Update HASH_MEMO. Silently skip on lock poison — next call recomputes. - if let Ok(mut memo) = HASH_MEMO.lock() { - memo.insert( - path.to_path_buf(), - HashMemoEntry { - size, - modified_ns, - sample_hashes: sample_hashes.clone(), - content_hash, - }, - ); - } - + let content_hash = match mode { + ContentHashMode::Full => hash_prefix(path, size)?, + ContentHashMode::SamplesOnly => [0_u8; 32], + }; Some((size, modified_ns, sample_hashes, content_hash)) } @@ -960,22 +1460,15 @@ fn append_path_suffix(path: &Path, suffix: &str) -> PathBuf { PathBuf::from(os) } -/// Sibling journal path for a jcode session snapshot: `session_x.json` -> -/// `session_x.journal.jsonl` (replacing the `.json` suffix; falling back to a -/// plain append when the name has no `.json` suffix). `pub(crate)` so the -/// live-tail change-token probe can stat the journal alongside the snapshot. +/// Compatibility wrapper for the local live-tail mtime probe. pub(crate) fn jcode_journal_path(path: &Path) -> PathBuf { - let Some(file_name) = path.file_name().and_then(|name| name.to_str()) else { - return append_path_suffix(path, ".journal.jsonl"); - }; - let journal_name = file_name - .strip_suffix(".json") - .map(|stem| format!("{stem}.journal.jsonl")) - .unwrap_or_else(|| format!("{file_name}.journal.jsonl")); - path.with_file_name(journal_name) + crate::sessions::jcode::jcode_journal_path(path) } fn hash_prefix(path: &Path, len: u64) -> Option<[u8; 32]> { + #[cfg(test)] + FULL_HASH_CALLS.with(|calls| calls.set(calls.get() + 1)); + let mut file = File::open(path).ok()?; let mut hasher = Sha256::new(); let mut remaining = len; @@ -994,6 +1487,12 @@ fn hash_prefix(path: &Path, len: u64) -> Option<[u8; 32]> { Some(hasher.finalize().into()) } +#[cfg(test)] +fn full_hash_call_count() -> usize { + FULL_HASH_CALLS.with(std::cell::Cell::get) +} + +#[cfg(test)] pub(crate) fn build_codex_incremental_cache( path: &Path, consumed_offset: u64, @@ -1012,6 +1511,29 @@ pub(crate) fn build_codex_incremental_cache( }) } +/// Build Codex incremental state when the caller already hashed the complete +/// consumed prefix. Full-file Codex fingerprints are also the prefix hash when +/// `consumed_offset` equals the current file size, so accepting that digest +/// avoids a second read of the transcript. +pub(crate) fn build_codex_incremental_cache_with_prefix_hash( + path: &Path, + consumed_offset: u64, + state: CodexParseState, + prefix_hash: [u8; 32], +) -> Option { + let ends_with_newline = consumed_offset == 0 || file_ends_with_newline(path, consumed_offset); + if !ends_with_newline { + return None; + } + + Some(CodexIncrementalCache { + state, + consumed_offset, + ends_with_newline, + prefix_hash, + }) +} + fn file_ends_with_newline(path: &Path, size: u64) -> bool { if size == 0 { return true; @@ -1060,40 +1582,6 @@ mod tests { use std::io::Write; use tempfile::{NamedTempFile, TempDir}; - #[test] - fn from_jcode_path_invalidates_on_journal_only_change() { - // jcode appends new turns to `.journal.jsonl` between snapshot - // rewrites. The cache fingerprint must include the journal, or a - // journal-only write (snapshot byte-identical) serves stale data. - let dir = TempDir::new().unwrap(); - let snapshot = dir.path().join("session_test.json"); - std::fs::write(&snapshot, br#"{"id":"session_test","messages":[]}"#).unwrap(); - let journal = dir.path().join("session_test.journal.jsonl"); - std::fs::write(&journal, b"{\"append_messages\":[]}\n").unwrap(); - - let jcode_before = SourceFingerprint::from_jcode_path(&snapshot).unwrap(); - let plain_before = SourceFingerprint::from_path(&snapshot).unwrap(); - - // Append to the journal only; leave the snapshot byte-identical. - std::fs::write( - &journal, - b"{\"append_messages\":[]}\n{\"append_messages\":[{\"id\":\"x\"}]}\n", - ) - .unwrap(); - - let jcode_after = SourceFingerprint::from_jcode_path(&snapshot).unwrap(); - let plain_after = SourceFingerprint::from_path(&snapshot).unwrap(); - - assert_ne!( - jcode_before, jcode_after, - "a journal-only change must alter the jcode fingerprint" - ); - assert_eq!( - plain_before, plain_after, - "from_path ignores the journal sibling (control)" - ); - } - #[test] fn from_roo_path_invalidates_on_history_only_change() { // parse_roo_kilo_file reads model/agent from the sibling @@ -1125,284 +1613,311 @@ mod tests { ); } - #[test] - fn from_roo_path_invalidates_when_history_appears() { - let dir = TempDir::new().unwrap(); - let ui = dir.path().join("ui_messages.json"); - std::fs::write(&ui, b"[]").unwrap(); + fn restore_env_var(key: &str, value: Option>) { + unsafe { + match value { + Some(value) => std::env::set_var(key, value), + None => std::env::remove_var(key), + } + } + } - let before = SourceFingerprint::from_roo_path(&ui).unwrap(); - std::fs::write( - crate::sessions::roocode::history_path_for_ui_messages(&ui), - b"gpt-5", - ) - .unwrap(); - let after = SourceFingerprint::from_roo_path(&ui).unwrap(); + /// Pin every env var the cache resolvers consult so the test stays + /// inside `temp_home`. CI runners can leak `XDG_CONFIG_HOME` / + /// `XDG_CACHE_HOME` from the host, which would resolve cache shards outside + /// the sandbox. `dirs::config_dir()` uses Windows Known Folders instead of + /// HOME/XDG, so `TOKSCALE_CONFIG_DIR` must also be pinned explicitly there. + /// Returns the previous values so the caller can restore. + fn sandbox_cache_env( + temp_home: &std::path::Path, + ) -> ( + Option, + Option, + Option, + Option, + ) { + let prev_home = std::env::var_os("HOME"); + let prev_xdg_config = std::env::var_os("XDG_CONFIG_HOME"); + let prev_xdg_cache = std::env::var_os("XDG_CACHE_HOME"); + let prev_override = std::env::var_os("TOKSCALE_CONFIG_DIR"); + unsafe { + std::env::set_var("HOME", temp_home); + std::env::set_var("XDG_CONFIG_HOME", temp_home.join(".config")); + std::env::set_var("XDG_CACHE_HOME", temp_home.join(".cache")); + std::env::set_var("TOKSCALE_CONFIG_DIR", temp_home); + } + (prev_home, prev_xdg_config, prev_xdg_cache, prev_override) + } - assert_ne!( - before, after, - "creating the optional history sibling must invalidate the roo fingerprint" - ); + fn restore_cache_env( + prev: ( + Option, + Option, + Option, + Option, + ), + ) { + restore_env_var("HOME", prev.0); + restore_env_var("XDG_CONFIG_HOME", prev.1); + restore_env_var("XDG_CACHE_HOME", prev.2); + restore_env_var("TOKSCALE_CONFIG_DIR", prev.3); } - #[test] - fn from_droid_path_invalidates_when_fallback_jsonl_appears() { - let dir = TempDir::new().unwrap(); - let settings = dir.path().join("session.settings.json"); - std::fs::write(&settings, br#"{"tokenUsage":{"inputTokens":1}}"#).unwrap(); + fn write_temp_file(content: &[u8]) -> NamedTempFile { + let mut file = NamedTempFile::new().unwrap(); + file.write_all(content).unwrap(); + file.flush().unwrap(); + file + } - let before = SourceFingerprint::from_droid_path(&settings).unwrap(); - let plain_before = SourceFingerprint::from_path(&settings).unwrap(); - let jsonl = crate::sessions::droid::droid_jsonl_path(&settings).unwrap(); - assert_eq!(jsonl, dir.path().join("session.jsonl")); - assert!(before.related_files.is_empty()); + fn test_entry(identity: CacheIdentity, path: &Path, session_id: &str) -> CachedSourceEntry { + CachedSourceEntry::new( + identity, + path, + SourceFingerprint::from_path(path).unwrap(), + vec![UnifiedMessage::new( + identity.namespace, + "gpt-5", + "provider", + session_id, + 1, + TokenBreakdown { + input: 1, + output: 2, + cache_read: 3, + cache_write: 0, + reasoning: 0, + }, + 0.0, + )], + Vec::new(), + None, + ) + } - std::fs::write(&jsonl, b"Model: Claude Sonnet 4\n").unwrap(); + fn write_sources_in_distinct_shards( + dir: &TempDir, + identity: CacheIdentity, + ) -> (PathBuf, PathBuf) { + let first = dir.path().join("source-0.jsonl"); + std::fs::write(&first, b"source-0\n").unwrap(); + let first_shard = CacheKey::new(identity, &first).shard(); - let after = SourceFingerprint::from_droid_path(&settings).unwrap(); - let plain_after = SourceFingerprint::from_path(&settings).unwrap(); - assert_ne!(before, after); - assert_eq!(plain_before, plain_after); - assert_eq!(after.related_files[0].suffix, "session.jsonl"); + for index in 1..=CACHE_SHARD_COUNT * 2 { + let candidate = dir.path().join(format!("source-{index}.jsonl")); + std::fs::write(&candidate, format!("source-{index}\n")).unwrap(); + if CacheKey::new(identity, &candidate).shard() != first_shard { + return (first, candidate); + } + } - std::fs::write(&jsonl, b"Model: Claude Opus 4.5 Thinking\n").unwrap(); - let rewritten = SourceFingerprint::from_droid_path(&settings).unwrap(); - assert_ne!(after, rewritten); - assert_eq!( - plain_after, - SourceFingerprint::from_path(&settings).unwrap() - ); + panic!("failed to find paths in distinct cache shards"); } - #[test] - fn from_kimi_path_invalidates_when_config_appears() { - let dir = TempDir::new().unwrap(); - let wire = dir.path().join(".kimi/sessions/group/session/wire.jsonl"); - std::fs::create_dir_all(wire.parent().unwrap()).unwrap(); - std::fs::write(&wire, b"usage\n").unwrap(); - - let before = SourceFingerprint::from_kimi_path(&wire).unwrap(); - let plain_before = SourceFingerprint::from_path(&wire).unwrap(); - let config = crate::sessions::kimi::kimi_config_path(&wire).unwrap(); - assert_eq!(config, dir.path().join(".kimi/config.json")); - assert!(before.related_files.is_empty()); - - std::fs::write(&config, br#"{"model":"kimi-k2"}"#).unwrap(); + fn write_sources_in_same_shard(dir: &TempDir, identity: CacheIdentity) -> (PathBuf, PathBuf) { + let mut paths_by_shard = HashMap::new(); + for index in 0..=CACHE_SHARD_COUNT * 4 { + let candidate = dir.path().join(format!("source-{index}.jsonl")); + std::fs::write(&candidate, format!("source-{index}\n")).unwrap(); + let shard = CacheKey::new(identity, &candidate).shard(); + if let Some(first) = paths_by_shard.insert(shard, candidate.clone()) { + return (first, candidate); + } + } - let after = SourceFingerprint::from_kimi_path(&wire).unwrap(); - let plain_after = SourceFingerprint::from_path(&wire).unwrap(); - assert_ne!(before, after); - assert_eq!(plain_before, plain_after); - assert_eq!(after.related_files[0].suffix, "config.json"); + panic!("failed to find paths in the same cache shard"); + } - std::fs::write(&config, br#"{"model":"kimi-k2-thinking"}"#).unwrap(); - let rewritten = SourceFingerprint::from_kimi_path(&wire).unwrap(); - assert_ne!(after, rewritten); - assert_eq!(plain_after, SourceFingerprint::from_path(&wire).unwrap()); + fn cache_shard_path(identity: CacheIdentity, path: &Path) -> PathBuf { + let root = cache_shard_dir().unwrap(); + shard_path(&root, &CacheKey::new(identity, path).shard()) } #[test] - fn from_kiro_path_invalidates_when_messages_sidecar_appears() { - let dir = TempDir::new().unwrap(); - let session = dir.path().join("session.json"); - std::fs::write(&session, b"{}").unwrap(); - - let before = SourceFingerprint::from_kiro_path(&session).unwrap(); - let plain_before = SourceFingerprint::from_path(&session).unwrap(); - let messages = crate::sessions::kiro::kiro_related_messages_path(&session).unwrap(); - assert_eq!(messages, dir.path().join("session.jsonl")); - assert!(before.related_files.is_empty()); - - std::fs::write(&messages, b"message\n").unwrap(); + fn test_codex_prefix_matches_appended_file() { + let file = write_temp_file(b"line-1\nline-2\n"); + let fingerprint = SourceFingerprint::from_path(file.path()).unwrap(); + let incremental_cache = build_codex_incremental_cache( + file.path(), + fingerprint.size, + CodexParseState::default(), + ) + .unwrap(); - let after = SourceFingerprint::from_kiro_path(&session).unwrap(); - let plain_after = SourceFingerprint::from_path(&session).unwrap(); - assert_ne!(before, after); - assert_eq!(plain_before, plain_after); - assert_eq!(after.related_files[0].suffix, "messages.jsonl"); + let mut reopened = file.reopen().unwrap(); + reopened.seek(SeekFrom::End(0)).unwrap(); + reopened.write_all(b"line-3\n").unwrap(); + reopened.flush().unwrap(); - std::fs::write(&messages, b"rewritten message sidecar\n").unwrap(); - let rewritten = SourceFingerprint::from_kiro_path(&session).unwrap(); - assert_ne!(after, rewritten); - assert_eq!(plain_after, SourceFingerprint::from_path(&session).unwrap()); + assert!(codex_prefix_matches(file.path(), &incremental_cache,)); } #[test] - fn from_grok_path_invalidates_on_summary_only_change() { - // read_metadata reads the model id (and timestamp) from the sibling - // summary.json when updates.jsonl carries no per-turn model, so a - // summary-only rewrite (updates.jsonl + signals.json byte-identical) must - // change the fingerprint or the cache keeps a session pinned to its - // fallback model. Local divergence: upstream fingerprints only - // signals.json here. - let dir = TempDir::new().unwrap(); - let updates = dir.path().join("updates.jsonl"); - std::fs::write(&updates, b"{\"totalTokens\":10}\n").unwrap(); - let summary = dir.path().join("summary.json"); - std::fs::write(&summary, br#"{"current_model_id":"grok-4"}"#).unwrap(); - - let grok_before = SourceFingerprint::from_grok_path(&updates).unwrap(); - let plain_before = SourceFingerprint::from_path(&updates).unwrap(); - - // Rewrite summary.json only; leave updates.jsonl byte-identical. - std::fs::write(&summary, br#"{"current_model_id":"grok-4-fast"}"#).unwrap(); + fn test_codex_incremental_cache_reuses_full_hash() { + let file = write_temp_file(b"line-1\nline-2\n"); + let fingerprint = SourceFingerprint::from_path(file.path()).unwrap(); + let full_hashes_before = full_hash_call_count(); - let grok_after = SourceFingerprint::from_grok_path(&updates).unwrap(); - let plain_after = SourceFingerprint::from_path(&updates).unwrap(); + let incremental_cache = build_codex_incremental_cache_with_prefix_hash( + file.path(), + fingerprint.size, + CodexParseState::default(), + fingerprint.content_hash, + ) + .unwrap(); - assert_ne!( - grok_before, grok_after, - "a summary-only change must alter the grok fingerprint" - ); assert_eq!( - plain_before, plain_after, - "from_path ignores the summary sibling (control)" + full_hash_call_count(), + full_hashes_before, + "a supplied Codex fingerprint must avoid a second whole-file SHA-256" ); + assert_eq!(incremental_cache.prefix_hash, fingerprint.content_hash); + assert!(incremental_cache.ends_with_newline); } #[test] - fn from_grok_path_invalidates_on_events_only_change() { - // events.jsonl is the other metadata sibling read_metadata consults for - // the model id / session id / timestamp, so an events-only rewrite must - // invalidate the cache too. - let dir = TempDir::new().unwrap(); - let updates = dir.path().join("updates.jsonl"); - std::fs::write(&updates, b"{\"totalTokens\":10}\n").unwrap(); - let events = dir.path().join("events.jsonl"); - std::fs::write(&events, b"{\"model_id\":\"grok-4\"}\n").unwrap(); - - let before = SourceFingerprint::from_grok_path(&updates).unwrap(); - - std::fs::write(&events, b"{\"model_id\":\"grok-4-fast\"}\n").unwrap(); + fn test_check_path_returns_unchanged_for_matching_metadata_and_samples() { + let file = write_temp_file(&vec![b'a'; 32 * 1024]); + let fingerprint = SourceFingerprint::from_path(file.path()).unwrap(); + let full_hashes_before = full_hash_call_count(); - let after = SourceFingerprint::from_grok_path(&updates).unwrap(); + let status = SourceFingerprint::check_path(file.path(), Some(&fingerprint)).unwrap(); - assert_ne!( - before, after, - "an events-only change must alter the grok fingerprint" + assert!(matches!(status, FingerprintStatus::Unchanged)); + assert_eq!( + full_hash_call_count(), + full_hashes_before, + "an unchanged fingerprint must not compute a full SHA-256" ); } - struct EnvGuard(Vec<(&'static str, Option)>); - - impl EnvGuard { - fn capture(keys: &[&'static str]) -> Self { - Self( - keys.iter() - .map(|key| (*key, std::env::var_os(key))) - .collect(), - ) - } + #[test] + fn test_check_path_returns_changed_when_sample_changes_with_same_metadata() { + let original = vec![b'a'; 32 * 1024]; + let file = write_temp_file(&original); + let fingerprint = SourceFingerprint::from_path(file.path()).unwrap(); + let original_signature = metadata_signature(file.path()).unwrap(); + let original_modified = std::fs::metadata(file.path()).unwrap().modified().unwrap(); - fn set(&mut self, key: &'static str, value: impl AsRef) { - unsafe { std::env::set_var(key, value) }; - } + let mut rewritten = original; + rewritten[0] = b'z'; + std::fs::write(file.path(), rewritten).unwrap(); + File::options() + .write(true) + .open(file.path()) + .unwrap() + .set_times(std::fs::FileTimes::new().set_modified(original_modified)) + .unwrap(); + assert_eq!(metadata_signature(file.path()).unwrap(), original_signature); + let full_hashes_before = full_hash_call_count(); - fn remove(&mut self, key: &'static str) { - unsafe { std::env::remove_var(key) }; - } - } + let status = SourceFingerprint::check_path(file.path(), Some(&fingerprint)).unwrap(); - impl Drop for EnvGuard { - fn drop(&mut self) { - unsafe { - for (key, previous) in self.0.drain(..) { - match previous { - Some(value) => std::env::set_var(key, value), - None => std::env::remove_var(key), - } - } - } - } + let FingerprintStatus::Changed(changed) = status else { + panic!("changed sample must rebuild the full fingerprint"); + }; + assert_ne!(changed, fingerprint); + assert_eq!( + full_hash_call_count(), + full_hashes_before + 1, + "a changed sample must rebuild the full fingerprint" + ); } - /// Pin every env var the cache resolvers consult so the test stays - /// inside `temp_home`, including on Windows where HOME/XDG do not control - /// the platform known folders. The override also keeps canonical cache - /// tests from reading or writing a real profile. - fn sandbox_cache_env(temp_home: &std::path::Path) -> EnvGuard { - let config_dir = temp_home.join(".config"); - let cache_dir = temp_home.join(".cache"); - let mut guard = EnvGuard::capture(&[ - "HOME", - "XDG_CONFIG_HOME", - "XDG_CACHE_HOME", - "TOKSCALE_CONFIG_DIR", - ]); - guard.set("HOME", temp_home); - guard.set("XDG_CONFIG_HOME", &config_dir); - guard.set("XDG_CACHE_HOME", &cache_dir); - guard.set("TOKSCALE_CONFIG_DIR", temp_home); - guard - } - - /// Set legacy roots without the override for the non-Windows migration - /// tests. The returned guard restores every process environment variable - /// on normal return and during unwinding. - fn legacy_cache_env( - temp_home: &std::path::Path, - xdg_cache: Option<&std::path::Path>, - ) -> EnvGuard { - let config_dir = temp_home.join(".config"); - let mut guard = EnvGuard::capture(&[ - "HOME", - "XDG_CONFIG_HOME", - "XDG_CACHE_HOME", - "TOKSCALE_CONFIG_DIR", - ]); - guard.set("HOME", temp_home); - guard.set("XDG_CONFIG_HOME", &config_dir); - match xdg_cache { - Some(path) => guard.set("XDG_CACHE_HOME", path), - None => guard.remove("XDG_CACHE_HOME"), - } - guard.remove("TOKSCALE_CONFIG_DIR"); - guard - } + #[test] + fn test_generic_sources_skip_full_hash() { + let original = vec![b'a'; 64 * 1024]; + let file = write_temp_file(&original); + let fingerprint = SourceFingerprint::from_path(file.path()).unwrap(); + let original_signature = metadata_signature(file.path()).unwrap(); + let original_modified = std::fs::metadata(file.path()).unwrap().modified().unwrap(); + + let mut rewritten = original; + rewritten[0] = b'z'; + std::fs::write(file.path(), rewritten).unwrap(); + File::options() + .write(true) + .open(file.path()) + .unwrap() + .set_times(std::fs::FileTimes::new().set_modified(original_modified)) + .unwrap(); + assert_eq!(metadata_signature(file.path()).unwrap(), original_signature); + + let full_hashes_before = full_hash_call_count(); + let status = + SourceFingerprint::check_path_samples_only(file.path(), Some(&fingerprint)).unwrap(); + let FingerprintStatus::Changed(changed) = status else { + panic!("changed sample must invalidate a generic source"); + }; + assert_eq!( + full_hash_call_count(), + full_hashes_before, + "generic source fingerprints must not compute a whole-file SHA-256" + ); + assert_eq!(changed.content_hash, [0_u8; 32]); - fn restore_cache_env(guard: EnvGuard) { - drop(guard); + let full_hashes_before = full_hash_call_count(); + let cold = SourceFingerprint::check_path_samples_only(file.path(), None).unwrap(); + let FingerprintStatus::Changed(cold) = cold else { + panic!("an uncached generic source must build a fingerprint"); + }; + assert_eq!(full_hash_call_count(), full_hashes_before); + assert_eq!(cold.content_hash, [0_u8; 32]); } #[test] - #[serial_test::serial] - fn test_env_guard_restores_after_unwind() { - const KEY: &str = "TOKSCALE_MESSAGE_CACHE_ENV_GUARD_SELF_CHECK"; - let mut outer = EnvGuard::capture(&[KEY]); - outer.set(KEY, "before"); - let result = std::panic::catch_unwind(|| { - let mut inner = EnvGuard::capture(&[KEY]); - inner.set(KEY, "during"); - panic!("exercise EnvGuard unwinding"); - }); - assert!(result.is_err()); - assert_eq!(std::env::var_os(KEY), Some("before".into())); - } + fn test_sqlite_fingerprint_skips_full_hash() { + let file = write_temp_file(&vec![b'a'; 64 * 1024]); + let full_hashes_before = full_hash_call_count(); - fn write_temp_file(content: &[u8]) -> NamedTempFile { - let mut file = NamedTempFile::new().unwrap(); - file.write_all(content).unwrap(); - file.flush().unwrap(); - file + let fingerprint = SourceFingerprint::from_sqlite_path(file.path()).unwrap(); + + assert_eq!( + full_hash_call_count(), + full_hashes_before, + "a SQLite fingerprint must not compute a whole-file SHA-256" + ); + assert_eq!( + fingerprint.content_hash, [0_u8; 32], + "a SQLite fingerprint stores a zero content_hash sentinel" + ); + assert!( + !fingerprint.sample_hashes.is_empty(), + "samples still guard SQLite change detection" + ); } #[test] - fn test_codex_prefix_matches_appended_file() { - let file = write_temp_file(b"line-1\nline-2\n"); - let fingerprint = SourceFingerprint::from_path(file.path()).unwrap(); - let incremental_cache = build_codex_incremental_cache( - file.path(), - fingerprint.size, - CodexParseState::default(), - ) - .unwrap(); - - let mut reopened = file.reopen().unwrap(); - reopened.seek(SeekFrom::End(0)).unwrap(); - reopened.write_all(b"line-3\n").unwrap(); - reopened.flush().unwrap(); + fn test_sqlite_check_detects_change_without_full_hash() { + let original = vec![b'a'; 64 * 1024]; + let file = write_temp_file(&original); + let fingerprint = SourceFingerprint::from_sqlite_path(file.path()).unwrap(); + + // Unchanged: metadata + samples match, no full hash. + let full_hashes_before = full_hash_call_count(); + let status = SourceFingerprint::check_sqlite_path(file.path(), Some(&fingerprint)).unwrap(); + assert!(matches!(status, FingerprintStatus::Unchanged)); + + // Changed: a same-size rewrite with a rolled-back mtime is still caught + // by the samples, and still without a whole-file hash. + let original_modified = std::fs::metadata(file.path()).unwrap().modified().unwrap(); + let mut rewritten = original; + rewritten[0] = b'z'; + std::fs::write(file.path(), rewritten).unwrap(); + File::options() + .write(true) + .open(file.path()) + .unwrap() + .set_times(std::fs::FileTimes::new().set_modified(original_modified)) + .unwrap(); - assert!(codex_prefix_matches(file.path(), &incremental_cache,)); + let status = SourceFingerprint::check_sqlite_path(file.path(), Some(&fingerprint)).unwrap(); + assert!(matches!(status, FingerprintStatus::Changed(_))); + assert_eq!( + full_hash_call_count(), + full_hashes_before, + "SQLite change detection must never compute a whole-file SHA-256" + ); } #[test] @@ -1410,14 +1925,7 @@ mod tests { let file = write_temp_file(b"aaaa\nbbbb\ncccc\n"); let before = SourceFingerprint::from_path(file.path()).unwrap(); - // Windows can retain the same timestamp for a fast same-size rewrite; - // use a changed-length fixture there instead of sleeping or changing - // the production fingerprint memo semantics. - #[cfg(not(target_os = "windows"))] - let rewritten = b"aaaa\nzzzz\ncccc\n"; - #[cfg(target_os = "windows")] - let rewritten = b"aaaa\nzzzz\ncccc\nchanged-length\n"; - std::fs::write(file.path(), rewritten).unwrap(); + std::fs::write(file.path(), b"aaaa\nzzzz\ncccc\n").unwrap(); let after = SourceFingerprint::from_path(file.path()).unwrap(); assert_ne!(before, after); @@ -1432,11 +1940,6 @@ mod tests { let mut rewritten = original.clone(); rewritten[73 * 1024] = b'z'; - // Keep the unsampled same-size rewrite on Unix. Windows filesystem - // timestamp precision can memoize a rapid rewrite, so make its - // fixture length-changing without adding a sleep. - #[cfg(target_os = "windows")] - rewritten.extend_from_slice(b"changed-length\n"); std::fs::write(file.path(), &rewritten).unwrap(); let after = SourceFingerprint::from_path(file.path()).unwrap(); @@ -1456,7 +1959,7 @@ mod tests { let with_wal = SourceFingerprint::from_sqlite_path(&db_path).unwrap(); assert_ne!(base, with_wal); - std::fs::write(&wal_path, b"wal-2-changed-length").unwrap(); + std::fs::write(&wal_path, b"wal-2").unwrap(); let updated_wal = SourceFingerprint::from_sqlite_path(&db_path).unwrap(); assert_ne!(with_wal, updated_wal); @@ -1467,6 +1970,332 @@ mod tests { assert_eq!(before_shm, with_shm); } + #[test] + fn test_copilot_desktop_dynamic_event_set_adds_and_removes() { + let dir = TempDir::new().unwrap(); + let copilot_root = dir.path().join(".copilot"); + let db_path = copilot_root.join("copilot.db"); + std::fs::create_dir_all(&copilot_root).unwrap(); + std::fs::write(&db_path, b"database\n").unwrap(); + + let without_events = SourceFingerprint::from_copilot_desktop_path(&db_path).unwrap(); + let event_path = copilot_root + .join("session-state/session-1") + .join("events.jsonl"); + std::fs::create_dir_all(event_path.parent().unwrap()).unwrap(); + std::fs::write(&event_path, b"event\n").unwrap(); + assert!(matches!( + SourceFingerprint::check_copilot_desktop_path_samples_only( + &db_path, + Some(&without_events), + ), + Some(FingerprintStatus::Changed(_)) + )); + + let with_events = SourceFingerprint::from_copilot_desktop_path(&db_path).unwrap(); + std::fs::remove_file(&event_path).unwrap(); + assert!(matches!( + SourceFingerprint::check_copilot_desktop_path_samples_only( + &db_path, + Some(&with_events), + ), + Some(FingerprintStatus::Changed(_)) + )); + } + + #[test] + fn test_related_path_identity_and_set_changes_miss_warm_cache() { + let dir = TempDir::new().unwrap(); + let primary_path = dir.path().join("primary.jsonl"); + let first_related = dir.path().join("first.json"); + let second_related = dir.path().join("second.json"); + std::fs::write(&primary_path, b"primary\n").unwrap(); + std::fs::write(&first_related, b"same\n").unwrap(); + std::fs::write(&second_related, b"same\n").unwrap(); + let first_modified = std::fs::metadata(&first_related) + .unwrap() + .modified() + .unwrap(); + File::options() + .write(true) + .open(&second_related) + .unwrap() + .set_modified(first_modified) + .unwrap(); + + let cached = SourceFingerprint::from_path_with_related_mode( + &primary_path, + vec![("dependency".to_string(), first_related.clone())], + ContentHashMode::SamplesOnly, + ) + .unwrap(); + assert!(matches!( + SourceFingerprint::check_path_with_related_mode( + &primary_path, + vec![("dependency".to_string(), first_related.clone())], + Some(&cached), + ContentHashMode::SamplesOnly, + ), + Some(FingerprintStatus::Unchanged) + )); + + assert!(matches!( + SourceFingerprint::check_path_with_related_mode( + &primary_path, + vec![("dependency".to_string(), second_related.clone())], + Some(&cached), + ContentHashMode::SamplesOnly, + ), + Some(FingerprintStatus::Changed(_)) + )); + assert!(matches!( + SourceFingerprint::check_path_with_related_mode( + &primary_path, + vec![ + ("dependency".to_string(), first_related), + ("extra".to_string(), second_related), + ], + Some(&cached), + ContentHashMode::SamplesOnly, + ), + Some(FingerprintStatus::Changed(_)) + )); + } + + #[cfg(unix)] + #[test] + fn test_related_non_not_found_read_failure_returns_none() { + let dir = TempDir::new().unwrap(); + let primary_path = dir.path().join("primary.jsonl"); + let related_path = dir.path().join("dependency.json"); + std::fs::write(&primary_path, b"primary\n").unwrap(); + std::fs::write(&related_path, b"dependency\n").unwrap(); + let mut cached = SourceFingerprint::from_path_with_related_mode( + &primary_path, + vec![("dependency".to_string(), related_path.clone())], + ContentHashMode::SamplesOnly, + ) + .unwrap(); + + std::fs::remove_file(&related_path).unwrap(); + std::fs::create_dir(&related_path).unwrap(); + let (size, modified_ns) = metadata_signature(&related_path).unwrap(); + cached.related_files[0].size = size; + cached.related_files[0].modified_ns = modified_ns; + assert!( + SourceFingerprint::check_path_with_related_mode( + &primary_path, + vec![("dependency".to_string(), related_path)], + Some(&cached), + ContentHashMode::SamplesOnly, + ) + .is_none(), + "a non-NotFound related read failure must fail open to a cold parse" + ); + } + + #[test] + fn test_jcode_fingerprint_tracks_journal_sidecar_changes() { + let dir = TempDir::new().unwrap(); + let session_path = dir.path().join("session_fixture.json"); + std::fs::write(&session_path, br#"{"messages":[]}"#).unwrap(); + + let base = SourceFingerprint::from_jcode_path(&session_path).unwrap(); + + let journal_path = dir.path().join("session_fixture.journal.jsonl"); + std::fs::write( + &journal_path, + br#"{"append_messages":[]} +"#, + ) + .unwrap(); + let with_journal = SourceFingerprint::from_jcode_path(&session_path).unwrap(); + assert_ne!(base, with_journal); + + std::fs::write( + &journal_path, + br#"{"append_messages":[{"id":"assistant_1"}]} +"#, + ) + .unwrap(); + let updated_journal = SourceFingerprint::from_jcode_path(&session_path).unwrap(); + assert_ne!(with_journal, updated_journal); + } + + #[test] + fn test_grok_fingerprint_tracks_signals_sidecar_changes() { + let dir = TempDir::new().unwrap(); + let updates_path = dir.path().join("updates.jsonl"); + std::fs::write(&updates_path, b"update\n").unwrap(); + + let base = SourceFingerprint::from_grok_path(&updates_path).unwrap(); + + let signals_path = dir.path().join("signals.json"); + std::fs::write(&signals_path, br#"{"input":1}"#).unwrap(); + let with_signals = SourceFingerprint::from_grok_path(&updates_path).unwrap(); + assert_ne!(base, with_signals); + + std::fs::write(&signals_path, br#"{"input":2}"#).unwrap(); + let updated_signals = SourceFingerprint::from_grok_path(&updates_path).unwrap(); + assert_ne!(with_signals, updated_signals); + } + + #[test] + fn test_grok_fingerprint_tracks_summary_and_events_sidecar_changes() { + let dir = TempDir::new().unwrap(); + let updates_path = dir.path().join("updates.jsonl"); + std::fs::write(&updates_path, b"update\n").unwrap(); + + let base = SourceFingerprint::from_grok_path(&updates_path).unwrap(); + + let summary_path = dir.path().join("summary.json"); + std::fs::write(&summary_path, br#"{"model":"grok-3"}"#).unwrap(); + let with_summary = SourceFingerprint::from_grok_path(&updates_path).unwrap(); + assert_ne!(base, with_summary); + + std::fs::write(&summary_path, br#"{"model":"grok-4"}"#).unwrap(); + let updated_summary = SourceFingerprint::from_grok_path(&updates_path).unwrap(); + assert_ne!(with_summary, updated_summary); + + let events_path = dir.path().join("events.jsonl"); + std::fs::write(&events_path, b"event-1\n").unwrap(); + let with_events = SourceFingerprint::from_grok_path(&updates_path).unwrap(); + assert_ne!(updated_summary, with_events); + + std::fs::write(&events_path, b"event-2\n").unwrap(); + let updated_events = SourceFingerprint::from_grok_path(&updates_path).unwrap(); + assert_ne!(with_events, updated_events); + } + + #[test] + fn test_kiro_ide_fingerprint_tracks_messages_sidecar_changes() { + let dir = TempDir::new().unwrap(); + let sess_dir = dir.path().join("workspace-a/sess_02f1c107"); + std::fs::create_dir_all(&sess_dir).unwrap(); + let session_path = sess_dir.join("session.json"); + std::fs::write(&session_path, br#"{"schemaVersion":"1.0.0"}"#).unwrap(); + + let base = SourceFingerprint::from_kiro_path(&session_path).unwrap(); + + // messages.jsonl appearing (session.json untouched) must invalidate. + let messages_path = sess_dir.join("messages.jsonl"); + std::fs::write( + &messages_path, + br#"{"role":"user","content":"hello"} +"#, + ) + .unwrap(); + let with_messages = SourceFingerprint::from_kiro_path(&session_path).unwrap(); + assert_ne!(base, with_messages); + + // An append landing after the last session.json write must invalidate. + std::fs::write( + &messages_path, + br#"{"role":"user","content":"hello"} +{"role":"assistant","content":"world"} +"#, + ) + .unwrap(); + let updated_messages = SourceFingerprint::from_kiro_path(&session_path).unwrap(); + assert_ne!(with_messages, updated_messages); + + // A CLI source records its absent same-stem JSONL sidecar so a later + // creation invalidates the cache without reparsing the primary file. + let cli_path = dir.path().join("cli-session.json"); + std::fs::write(&cli_path, b"{}").unwrap(); + let cli_fingerprint = SourceFingerprint::from_kiro_path(&cli_path).unwrap(); + assert!(cli_fingerprint.related_files.iter().any(|related| { + related.suffix == "messages.jsonl" + && related.path.to_path_buf() == dir.path().join("cli-session.jsonl") + && !related.exists + })); + } + + #[test] + fn test_kiro_cli_fingerprint_tracks_same_stem_jsonl_changes() { + let dir = TempDir::new().unwrap(); + let session_path = dir.path().join("cli-session.json"); + std::fs::write(&session_path, br#"{"sessionId":"session-1"}"#).unwrap(); + + let base = SourceFingerprint::from_kiro_path(&session_path).unwrap(); + assert!(matches!( + SourceFingerprint::check_kiro_path_samples_only(&session_path, Some(&base)), + Some(FingerprintStatus::Unchanged) + )); + + let messages_path = dir.path().join("cli-session.jsonl"); + std::fs::write(&messages_path, b"message-1\n").unwrap(); + assert!(matches!( + SourceFingerprint::check_kiro_path_samples_only(&session_path, Some(&base)), + Some(FingerprintStatus::Changed(_)) + )); + let with_messages = SourceFingerprint::from_kiro_path(&session_path).unwrap(); + assert_ne!(base, with_messages); + + std::fs::write(&messages_path, b"message-2\n").unwrap(); + let updated_messages = SourceFingerprint::from_kiro_path(&session_path).unwrap(); + assert_ne!(with_messages, updated_messages); + + std::fs::remove_file(&messages_path).unwrap(); + assert!(matches!( + SourceFingerprint::check_kiro_path_samples_only(&session_path, Some(&updated_messages),), + Some(FingerprintStatus::Changed(_)) + )); + } + + #[test] + fn test_droid_fingerprint_tracks_fallback_jsonl_changes() { + let dir = TempDir::new().unwrap(); + let settings_path = dir.path().join("session.settings.json"); + std::fs::write(&settings_path, br#"{"tokenUsage":{"inputTokens":1}}"#).unwrap(); + + let base = SourceFingerprint::from_droid_path(&settings_path).unwrap(); + + let jsonl_path = dir.path().join("session.jsonl"); + std::fs::write(&jsonl_path, b"Model: Claude Sonnet 4\n").unwrap(); + let with_jsonl = SourceFingerprint::from_droid_path(&settings_path).unwrap(); + assert_ne!(base, with_jsonl); + + std::fs::write(&jsonl_path, b"Model: Claude Opus 4\n").unwrap(); + let updated_jsonl = SourceFingerprint::from_droid_path(&settings_path).unwrap(); + assert_ne!(with_jsonl, updated_jsonl); + } + + #[test] + fn test_kimi_fingerprint_tracks_legacy_config_but_keeps_kimi_code_self_contained() { + let dir = TempDir::new().unwrap(); + let legacy_path = dir.path().join(".kimi/sessions/group/session/wire.jsonl"); + std::fs::create_dir_all(legacy_path.parent().unwrap()).unwrap(); + std::fs::write(&legacy_path, b"usage\n").unwrap(); + + let legacy_base = SourceFingerprint::from_kimi_path(&legacy_path).unwrap(); + let legacy_config = dir.path().join(".kimi/config.json"); + std::fs::write(&legacy_config, br#"{"model":"kimi-k2"}"#).unwrap(); + let legacy_with_config = SourceFingerprint::from_kimi_path(&legacy_path).unwrap(); + assert_ne!(legacy_base, legacy_with_config); + + std::fs::write(&legacy_config, br#"{"model":"kimi-k3"}"#).unwrap(); + let legacy_updated_config = SourceFingerprint::from_kimi_path(&legacy_path).unwrap(); + assert_ne!(legacy_with_config, legacy_updated_config); + + let code_path = dir + .path() + .join(".kimi-code/sessions/workspace/session/agents/main/wire.jsonl"); + std::fs::create_dir_all(code_path.parent().unwrap()).unwrap(); + std::fs::write(&code_path, b"usage.record\n").unwrap(); + let code_base = SourceFingerprint::from_kimi_path(&code_path).unwrap(); + assert_eq!( + code_base, + SourceFingerprint::from_path_samples_only(&code_path).unwrap() + ); + + assert!(crate::sessions::kimi::kimi_config_path(&code_path).is_none()); + let unrelated_config = dir.path().join(".kimi-code/config.json"); + std::fs::write(&unrelated_config, br#"{"model":"unrelated"}"#).unwrap(); + let code_with_config = SourceFingerprint::from_kimi_path(&code_path).unwrap(); + assert_eq!(code_base, code_with_config); + } + #[test] fn test_claude_sidechain_fingerprint_tracks_nested_parent_session_changes() { let dir = TempDir::new().unwrap(); @@ -1498,7 +2327,7 @@ mod tests { SourceFingerprint::from_claude_code_path_with_home(&sidechain_path, None).unwrap(); assert_ne!(base, with_parent); - std::fs::write(&parent_path, b"parent transcript 2 with changed length\n").unwrap(); + std::fs::write(&parent_path, b"parent transcript 2\n").unwrap(); let updated_parent = SourceFingerprint::from_claude_code_path_with_home(&sidechain_path, None).unwrap(); assert_ne!(with_parent, updated_parent); @@ -1531,12 +2360,53 @@ mod tests { SourceFingerprint::from_claude_code_path_with_home(&sidechain_path, None).unwrap(); assert_ne!(base, with_parent); - std::fs::write(&parent_path, b"flat parent 2 with changed length\n").unwrap(); + std::fs::write(&parent_path, b"flat parent 2\n").unwrap(); let updated_parent = SourceFingerprint::from_claude_code_path_with_home(&sidechain_path, None).unwrap(); assert_ne!(with_parent, updated_parent); } + #[test] + fn test_claude_sidechain_warm_check_reuses_cached_parent_dependencies() { + let dir = TempDir::new().unwrap(); + let project_dir = dir.path().join("projects/project-one"); + std::fs::create_dir_all(&project_dir).unwrap(); + let sidechain_path = project_dir.join("agent-child.jsonl"); + let mut sidechain = format!("{}\n", "x".repeat(4096)).repeat(65); + sidechain.push_str(concat!( + r#"{"type":"assistant","isSidechain":true,"sessionId":"flat-parent","agentId":"child","timestamp":"2026-01-01T00:00:00Z","requestId":"req-1","message":{"id":"msg-1","model":"claude-sonnet-4","usage":{"input_tokens":1,"output_tokens":1}}}"#, + "\n" + )); + std::fs::write(&sidechain_path, sidechain).unwrap(); + + let cached = + SourceFingerprint::from_claude_code_path_with_home(&sidechain_path, None).unwrap(); + let parent_path = project_dir.join("flat-parent.jsonl"); + assert!(cached.related_files.iter().any(|related| { + related.suffix == "parent-session-0.jsonl" + && related.path.to_path_buf() == parent_path + && !related.exists + })); + assert!(matches!( + SourceFingerprint::check_claude_code_path_with_home_samples_only( + &sidechain_path, + Some(&cached), + None, + ), + Some(FingerprintStatus::Unchanged) + )); + + std::fs::write(&parent_path, b"parent transcript\n").unwrap(); + assert!(matches!( + SourceFingerprint::check_claude_code_path_with_home_samples_only( + &sidechain_path, + Some(&cached), + None, + ), + Some(FingerprintStatus::Changed(_)) + )); + } + #[test] fn test_claude_code_fingerprint_tracks_meta_sidecar_changes() { let dir = TempDir::new().unwrap(); @@ -1721,720 +2591,448 @@ mod tests { #[test] #[serial_test::serial] - fn test_source_message_cache_round_trip() { + fn test_source_message_cache_round_trips_across_distinct_shards() { let temp_home = TempDir::new().unwrap(); let prev_env = sandbox_cache_env(temp_home.path()); - - let file = write_temp_file(b"{}\n"); - let fingerprint = SourceFingerprint::from_path(file.path()).unwrap(); - let entry = CachedSourceEntry::new( - file.path(), - fingerprint, - vec![UnifiedMessage::new( - "client", - "gpt-5", - "provider", - "session-1", - 1, - TokenBreakdown { - input: 1, - output: 2, - cache_read: 3, - cache_write: 0, - reasoning: 0, - }, - 0.0, - )], - Vec::new(), - None, - ); + let source_dir = TempDir::new().unwrap(); + let identity = CacheIdentity::for_client(ClientId::Claude); + let (path_one, path_two) = write_sources_in_distinct_shards(&source_dir, identity); + let shard_one = cache_shard_path(identity, &path_one); + let shard_two = cache_shard_path(identity, &path_two); + assert_ne!(shard_one, shard_two); let mut cache = SourceMessageCache::default(); - cache.insert(entry); + cache.insert(test_entry(identity, &path_one, "session-1")); + cache.insert(test_entry(identity, &path_two, "session-2")); cache.save_if_dirty(); + assert!(shard_one.is_file()); + assert!(shard_two.is_file()); let loaded = SourceMessageCache::load(); - assert_eq!(loaded.entries.len(), 1); - assert!(loaded.get(file.path()).is_some()); + assert_eq!(loaded.entries.len(), 2); + assert!(loaded.get(identity, &path_one).is_some()); + assert!(loaded.get(identity, &path_two).is_some()); restore_cache_env(prev_env); } #[test] - fn test_prune_missing_files_removes_deleted_entries() { - let file = write_temp_file(b"{}\n"); - let fingerprint = SourceFingerprint::from_path(file.path()).unwrap(); - let path = file.path().to_path_buf(); + #[serial_test::serial] + fn test_aggregate_cache_can_exceed_individual_shard_limit() { + const TEST_SHARD_LIMIT: u64 = 32 * 1024; + + let temp_home = TempDir::new().unwrap(); + let prev_env = sandbox_cache_env(temp_home.path()); + let source_dir = TempDir::new().unwrap(); + let identity = CacheIdentity::for_client(ClientId::Claude); + let (path_one, path_two) = write_sources_in_distinct_shards(&source_dir, identity); + + let mut entry_one = test_entry(identity, &path_one, "session-1"); + entry_one.messages[0].model_id = "a".repeat(20 * 1024); + let mut entry_two = test_entry(identity, &path_two, "session-2"); + entry_two.messages[0].model_id = "b".repeat(20 * 1024); let mut cache = SourceMessageCache::default(); - cache.insert(CachedSourceEntry::new( - &path, - fingerprint, - Vec::new(), - Vec::new(), - None, - )); + cache.insert(entry_one); + cache.insert(entry_two); + cache.save_if_dirty_with_limit(TEST_SHARD_LIMIT); + assert!( + !cache.dirty, + "both independently bounded shards should save" + ); - std::fs::remove_file(&path).unwrap(); - cache.prune_missing_files(); + let shard_one = cache_shard_path(identity, &path_one); + let shard_two = cache_shard_path(identity, &path_two); + let size_one = std::fs::metadata(&shard_one).unwrap().len(); + let size_two = std::fs::metadata(&shard_two).unwrap().len(); + assert!(size_one <= TEST_SHARD_LIMIT); + assert!(size_two <= TEST_SHARD_LIMIT); + assert!(size_one + size_two > TEST_SHARD_LIMIT); - assert!(cache.entries.is_empty()); + let loaded = SourceMessageCache::load(); + assert!(loaded.get(identity, &path_one).is_some()); + assert!(loaded.get(identity, &path_two).is_some()); + + restore_cache_env(prev_env); } #[test] #[serial_test::serial] - fn test_load_ignores_oversized_cache_file() { + fn test_corrupt_shard_does_not_hide_entries_from_other_shards() { let temp_home = TempDir::new().unwrap(); - let _env = sandbox_cache_env(temp_home.path()); + let prev_env = sandbox_cache_env(temp_home.path()); + let source_dir = TempDir::new().unwrap(); + let identity = CacheIdentity::for_client(ClientId::Claude); + let (corrupt_path, valid_path) = write_sources_in_distinct_shards(&source_dir, identity); + + let mut cache = SourceMessageCache::default(); + cache.insert(test_entry(identity, &corrupt_path, "corrupt-session")); + cache.insert(test_entry(identity, &valid_path, "valid-session")); + cache.save_if_dirty(); - let cache_file = cache_path().unwrap(); - ensure_cache_dir(cache_file.parent().unwrap()).unwrap(); - let file = File::create(&cache_file).unwrap(); - file.set_len(MAX_CACHE_FILE_BYTES + 1).unwrap(); + let corrupt_shard = cache_shard_path(identity, &corrupt_path); + std::fs::write(&corrupt_shard, b"not a bincode shard").unwrap(); + assert!(matches!( + read_shard(&corrupt_shard, identity), + ShardReadStatus::Invalid(_) + )); let loaded = SourceMessageCache::load(); - assert!(loaded.entries.is_empty()); + assert!(loaded.get(identity, &corrupt_path).is_none()); + assert_eq!( + loaded.get(identity, &valid_path).unwrap().messages[0].session_id, + "valid-session" + ); + assert!( + loaded.dirty, + "the corrupt shard should be scheduled for rewrite" + ); + + restore_cache_env(prev_env); } #[test] #[serial_test::serial] - fn test_schema_24_cache_is_stale_and_rebuilt_as_current_schema() { + fn test_stale_parser_shard_is_skipped_before_decoding_garbage_payload() { let temp_home = TempDir::new().unwrap(); let prev_env = sandbox_cache_env(temp_home.path()); + let source = write_temp_file(b"claude\n"); + let claude = CacheIdentity::for_client(ClientId::Claude); + let codex = CacheIdentity::for_client(ClientId::Codex); - { - let source = write_temp_file(b"schema-migration\n"); - let source_fingerprint = SourceFingerprint::from_path(source.path()).unwrap(); - let mut stale_message = UnifiedMessage::new( - "client", - "model", - "provider", - "stale-session", - 1, - TokenBreakdown::default(), - 0.0, - ); - stale_message.is_turn_start = true; - let stale_entry = CachedSourceEntry::new( - source.path(), - source_fingerprint.clone(), - vec![stale_message], - Vec::new(), - None, - ); - let cache_file = cache_path().unwrap(); - ensure_cache_dir(cache_file.parent().unwrap()).unwrap(); - let stale_store = CachedSourceStore { - schema_version: 24, - entries: vec![stale_entry], - }; - - let writer = BufWriter::new(File::create(&cache_file).unwrap()); - bincode::options() - .serialize_into(writer, &stale_store) - .unwrap(); + let mut seed = SourceMessageCache::default(); + seed.insert(test_entry(claude, source.path(), "claude-session")); + seed.save_if_dirty(); - let mut loaded = SourceMessageCache::load(); - assert!(loaded.entries.is_empty(), "schema-24 entries must be stale"); - assert_eq!( - SourceFingerprint::from_path(source.path()).unwrap(), - source_fingerprint, - "the stale cache entry and rebuilt parse have the same source fingerprint" - ); + let stale_key = CacheShardKey { + namespace: codex.namespace.to_string(), + index: 0, + }; + let stale_path = shard_path(&cache_shard_dir().unwrap(), &stale_key); + ensure_cache_dir(stale_path.parent().unwrap()).unwrap(); + let stale_envelope = CachedShardEnvelope { + format_version: CACHE_FORMAT_VERSION, + parser_namespace: codex.namespace.to_string(), + parser_version: codex.parser_version.saturating_sub(1), + payload: b"deliberately invalid entry payload".to_vec(), + }; + let mut writer = BufWriter::new(File::create(&stale_path).unwrap()); + bincode::options() + .serialize_into(&mut writer, &stale_envelope) + .unwrap(); + writer.flush().unwrap(); + drop(writer); - let rebuilt_message = UnifiedMessage::new( - "client", - "model", - "provider", - "rebuilt-session", - 2, - TokenBreakdown::default(), - 0.0, - ); - loaded.insert(CachedSourceEntry::new( - source.path(), - source_fingerprint, - vec![rebuilt_message], - Vec::new(), - None, - )); - loaded.save_if_dirty(); + assert!(matches!( + read_shard(&stale_path, codex), + ShardReadStatus::Stale + )); + let mut loaded = SourceMessageCache::load(); + assert_eq!(loaded.entries.len(), 1); + assert!(loaded.get(claude, source.path()).is_some()); + assert!(loaded.rewrite_shards.contains(&stale_key)); - let rebuilt = read_store_from_path(&cache_file).unwrap(); - assert_eq!(rebuilt.schema_version, CACHE_SCHEMA_VERSION); - assert_eq!(rebuilt.entries.len(), 1); - assert_eq!( - rebuilt.entries[0].messages[0].cost_source, - crate::sessions::CostSource::Unknown - ); - assert_eq!(rebuilt.entries[0].messages[0].session_id, "rebuilt-session"); - assert!( - !rebuilt.entries[0].messages[0].is_turn_start, - "stale schema-24 turn flags must not survive the rebuild" - ); - } + loaded.save_if_dirty(); + assert!(matches!( + read_shard(&stale_path, codex), + ShardReadStatus::Loaded(entries) if entries.is_empty() + )); + assert!(SourceMessageCache::load() + .get(claude, source.path()) + .is_some()); restore_cache_env(prev_env); } #[test] #[serial_test::serial] - fn test_schema_25_pi_cache_is_stale_and_rebuilt_with_agent() { + fn test_copilot_parser_v3_same_fingerprint_rebuilds_merged_duration() { let temp_home = TempDir::new().unwrap(); let prev_env = sandbox_cache_env(temp_home.path()); + let source = write_temp_file( + br#"{"type":"span","traceId":"trace-duplicate","spanId":"span-duplicate","name":"chat gpt-5.4-mini","startTime":[1775934260,0],"endTime":[1775934270,0],"attributes":{"gen_ai.operation.name":"chat","gen_ai.response.model":"gpt-5.4-mini","gen_ai.usage.input_tokens":10,"gen_ai.usage.output_tokens":20}} +{"type":"span","traceId":"trace-duplicate","spanId":"span-duplicate","name":"chat gpt-5.4-mini","startTime":[1775934255,0],"endTime":[1775934261,0],"attributes":{"gen_ai.operation.name":"chat","gen_ai.response.model":"gpt-5.4-mini","gen_ai.usage.input_tokens":30,"gen_ai.usage.output_tokens":50}} +"#, + ); + let current = CacheIdentity::for_client(ClientId::Copilot); + let stale = CacheIdentity { + namespace: current.namespace, + parser_version: 3, + }; + let fingerprint = SourceFingerprint::from_path(source.path()).unwrap(); + let mut stale_message = UnifiedMessage::new( + "copilot", + "gpt-5.4-mini", + "openai", + "trace-duplicate", + 1_775_934_255_000, + TokenBreakdown { + input: 30, + output: 50, + ..Default::default() + }, + 0.0, + ); + stale_message.duration_ms = Some(10_000); + let stale_entry = CachedSourceEntry::new( + stale, + source.path(), + fingerprint.clone(), + vec![stale_message], + Vec::new(), + None, + ); + let payload = bincode::options().serialize(&vec![stale_entry]).unwrap(); + let stale_key = CacheKey::new(current, source.path()).shard(); + let stale_path = shard_path(&cache_shard_dir().unwrap(), &stale_key); + ensure_cache_dir(stale_path.parent().unwrap()).unwrap(); + let stale_envelope = CachedShardEnvelope { + format_version: CACHE_FORMAT_VERSION, + parser_namespace: stale.namespace.to_string(), + parser_version: stale.parser_version, + payload, + }; + let mut writer = BufWriter::new(File::create(&stale_path).unwrap()); + bincode::options() + .serialize_into(&mut writer, &stale_envelope) + .unwrap(); + writer.flush().unwrap(); + drop(writer); - { - let source = write_temp_file(b"pi-schema-migration\n"); - let source_fingerprint = SourceFingerprint::from_path(source.path()).unwrap(); - let stale_message = UnifiedMessage::new( - "pi", - "gpt-5", - "openai", - "pi-session", - 1, - TokenBreakdown::default(), - 0.0, - ); - let stale_entry = CachedSourceEntry::new( - source.path(), - source_fingerprint.clone(), - vec![stale_message], - Vec::new(), - None, - ); - let cache_file = cache_path().unwrap(); - ensure_cache_dir(cache_file.parent().unwrap()).unwrap(); - let stale_store = CachedSourceStore { - schema_version: 25, - entries: vec![stale_entry], - }; - - let writer = BufWriter::new(File::create(&cache_file).unwrap()); - bincode::options() - .serialize_into(writer, &stale_store) - .unwrap(); - - let mut loaded = SourceMessageCache::load(); - assert!( - loaded.entries.is_empty(), - "schema-25 Pi entries must be stale" - ); - assert_eq!( - SourceFingerprint::from_path(source.path()).unwrap(), - source_fingerprint, - "the stale cache entry and rebuilt parse have the same source fingerprint" - ); - - let rebuilt_message = UnifiedMessage::new_with_agent( - "pi", - "gpt-5", - "openai", - "pi-session", - 2, - TokenBreakdown::default(), - 0.0, - Some("go-reviewer".to_string()), - ); - loaded.insert(CachedSourceEntry::new( - source.path(), - source_fingerprint, - vec![rebuilt_message], - Vec::new(), - None, - )); - loaded.save_if_dirty(); - - let rebuilt = read_store_from_path(&cache_file).unwrap(); - assert_eq!(rebuilt.schema_version, CACHE_SCHEMA_VERSION); - assert_eq!(rebuilt.entries.len(), 1); + assert_eq!( + SourceFingerprint::from_path(source.path()).unwrap(), + fingerprint, + "the parser-only invalidation must rebuild an unchanged source" + ); + assert!(matches!( + read_shard(&stale_path, current), + ShardReadStatus::Stale + )); + let mut cache = SourceMessageCache::load(); + assert!(cache.get(current, source.path()).is_none()); - let reloaded = SourceMessageCache::load(); - assert_eq!(reloaded.entries.len(), 1); - assert_eq!( - reloaded.entries.values().next().unwrap().messages[0] - .agent - .as_deref(), - Some("go-reviewer") - ); - } + let rebuilt_messages = crate::sessions::copilot::parse_copilot_file(source.path()); + assert_eq!(rebuilt_messages.len(), 1); + assert_eq!(rebuilt_messages[0].timestamp, 1_775_934_255_000); + assert_eq!(rebuilt_messages[0].duration_ms, Some(15_000)); + cache.insert(CachedSourceEntry::new( + current, + source.path(), + fingerprint, + rebuilt_messages, + Vec::new(), + None, + )); + cache.save_if_dirty(); + let loaded = SourceMessageCache::load(); + assert_eq!( + loaded.get(current, source.path()).unwrap().messages[0].duration_ms, + Some(15_000) + ); restore_cache_env(prev_env); } #[test] #[serial_test::serial] - fn test_schema_26_claude_workflow_cache_is_stale_and_rebuilt_with_parent_agent() { + fn test_format_one_shard_is_stale_without_hiding_format_two_namespace() { let temp_home = TempDir::new().unwrap(); let prev_env = sandbox_cache_env(temp_home.path()); - - { - let project_dir = temp_home.path().join(".claude/projects/project-one"); - let parent_session_id = "claude-schema-parent"; - let parent_path = project_dir.join(format!("{parent_session_id}.jsonl")); - let workflow_dir = project_dir - .join(parent_session_id) - .join("subagents") - .join("workflows") - .join("wf_schema"); - let sidechain_path = workflow_dir.join("agent-deepagent1.jsonl"); - std::fs::create_dir_all(&workflow_dir).unwrap(); - - std::fs::write( - &parent_path, - r#"{"type":"assistant","timestamp":"2024-12-01T10:00:00.000Z","message":{"id":"msg_parent","model":"claude-3-5-sonnet","role":"assistant","content":[{"type":"tool_use","id":"toolu_schema","name":"Agent","input":{"subagent_type":"code-reviewer","prompt":"review"}}],"usage":{"input_tokens":80,"output_tokens":40}}} -{"type":"user","timestamp":"2024-12-01T10:00:01.000Z","message":{"role":"user","content":[{"tool_use_id":"toolu_schema","type":"tool_result","content":[{"type":"text","text":"agentId: deepagent1 (use SendMessage)"}]}]}}"#, - ) - .unwrap(); - std::fs::write( - &sidechain_path, - r#"{"type":"user","isSidechain":true,"sessionId":"claude-schema-parent","agentId":"deepagent1","timestamp":"2024-12-01T10:00:00.500Z","message":{"content":"task"}} -{"type":"assistant","isSidechain":true,"sessionId":"claude-schema-parent","agentId":"deepagent1","timestamp":"2024-12-01T10:00:01.000Z","requestId":"req_schema","message":{"id":"msg_schema","model":"claude-3-5-sonnet","usage":{"input_tokens":300,"output_tokens":120}}}"#, - ) + let source_dir = TempDir::new().unwrap(); + let identity = CacheIdentity::for_client(ClientId::Claude); + let (format_one_path, valid_path) = write_sources_in_distinct_shards(&source_dir, identity); + let codex = CacheIdentity::for_client(ClientId::Codex); + + let mut seed = SourceMessageCache::default(); + seed.insert(test_entry(identity, &valid_path, "format-two-session")); + seed.save_if_dirty(); + + let stale_key = CacheShardKey { + namespace: codex.namespace.to_string(), + index: CacheKey::new(codex, &format_one_path).shard().index, + }; + let stale_path = shard_path(&cache_shard_dir().unwrap(), &stale_key); + ensure_cache_dir(stale_path.parent().unwrap()).unwrap(); + let format_one_envelope = CachedShardEnvelope { + format_version: 1, + parser_namespace: codex.namespace.to_string(), + parser_version: codex.parser_version, + payload: b"deliberately invalid format-1 payload".to_vec(), + }; + let mut writer = BufWriter::new(File::create(&stale_path).unwrap()); + bincode::options() + .serialize_into(&mut writer, &format_one_envelope) .unwrap(); + writer.flush().unwrap(); + drop(writer); - let source_fingerprint = SourceFingerprint::from_claude_code_path_with_home( - &sidechain_path, - Some(temp_home.path()), - ) - .unwrap(); - let stale_message = UnifiedMessage::new( - "claude", - "claude-sonnet-4", - "anthropic", - parent_session_id, - 1, - TokenBreakdown { - input: 300, - output: 120, - cache_read: 0, - cache_write: 0, - reasoning: 0, - }, - 0.0, - ); - assert!(stale_message.agent.is_none()); - let stale_entry = CachedSourceEntry::new( - &sidechain_path, - source_fingerprint.clone(), - vec![stale_message], - Vec::new(), - None, - ); - let cache_file = cache_path().unwrap(); - ensure_cache_dir(cache_file.parent().unwrap()).unwrap(); - let stale_store = CachedSourceStore { - schema_version: 26, - entries: vec![stale_entry], - }; - - let writer = BufWriter::new(File::create(&cache_file).unwrap()); - bincode::options() - .serialize_into(writer, &stale_store) - .unwrap(); - - let mut loaded = SourceMessageCache::load(); - assert!( - loaded.entries.is_empty(), - "schema-26 Claude workflow entries must be stale" - ); - assert_eq!( - SourceFingerprint::from_claude_code_path_with_home( - &sidechain_path, - Some(temp_home.path()) - ) - .unwrap(), - source_fingerprint, - "the stale cache entry and rebuilt parse have the same source fingerprint" - ); - - let rebuilt_messages = crate::sessions::claudecode::parse_claude_file_with_home( - &sidechain_path, - Some(temp_home.path()), - ); - assert_eq!(rebuilt_messages.len(), 1); - assert_eq!( - rebuilt_messages[0].agent.as_deref(), - Some("Code Reviewer"), - "rebuilding must recover Tier-2 attribution from the parent session" - ); - loaded.insert(CachedSourceEntry::new( - &sidechain_path, - source_fingerprint, - rebuilt_messages, - Vec::new(), - None, - )); - loaded.save_if_dirty(); - - let rebuilt = read_store_from_path(&cache_file).unwrap(); - assert_eq!(rebuilt.schema_version, CACHE_SCHEMA_VERSION); - assert_eq!(rebuilt.entries.len(), 1); - assert_eq!( - rebuilt.entries[0].messages[0].agent.as_deref(), - Some("Code Reviewer") - ); - - let reloaded = SourceMessageCache::load(); - assert_eq!(reloaded.entries.len(), 1); - assert_eq!( - reloaded.entries.values().next().unwrap().messages[0] - .agent - .as_deref(), - Some("Code Reviewer") - ); - } + assert!(matches!( + read_shard(&stale_path, codex), + ShardReadStatus::Stale + )); + let mut loaded = SourceMessageCache::load(); + assert!(loaded.get(identity, &valid_path).is_some()); + assert!(loaded.rewrite_shards.contains(&stale_key)); + + loaded.save_if_dirty(); + assert!(matches!( + read_shard(&stale_path, codex), + ShardReadStatus::Loaded(entries) if entries.is_empty() + )); + assert!(SourceMessageCache::load() + .get(identity, &valid_path) + .is_some()); restore_cache_env(prev_env); } #[test] #[serial_test::serial] - fn test_schema_27_copilot_cache_is_stale_and_rebuilt_with_root_agent() { + fn test_explicit_invalidation_of_existing_path_persists() { let temp_home = TempDir::new().unwrap(); let prev_env = sandbox_cache_env(temp_home.path()); + let source = write_temp_file(b"still exists\n"); + let identity = CacheIdentity::for_client(ClientId::Claude); + + let mut seed = SourceMessageCache::default(); + seed.insert(test_entry(identity, source.path(), "session-1")); + seed.save_if_dirty(); + assert!(SourceMessageCache::load() + .get(identity, source.path()) + .is_some()); + + let mut cache = SourceMessageCache::load(); + cache.remove(identity, source.path()); + cache.save_if_dirty(); - { - // The nested invoke_agent is exported before the root invoke_agent, - // while the agentless chat turn is covered by the root span. The - // intermediate tool-task span deliberately has no attributes: the - // pre-hardening resolver dropped its tool-task -> invoke-root edge, - // making the nested invoke look like a root and poisoning the - // schema-27 fallback. - let source = write_temp_file( - br#"{"type":"span","traceId":"trace-nested","spanId":"invoke-sub","parentSpanId":"tool-task","name":"invoke_agent","endTime":[1775934261,0],"attributes":{"gen_ai.operation.name":"invoke_agent","gen_ai.provider.name":"github","gen_ai.request.model":"claude-sonnet-4.6","gen_ai.agent.id":"github.copilot.subagent"}} -{"type":"span","traceId":"trace-nested","spanId":"chat-sub","parentSpanId":"invoke-sub","name":"chat claude-sonnet-4.6","endTime":[1775934262,0],"attributes":{"gen_ai.operation.name":"chat","gen_ai.provider.name":"github","gen_ai.request.model":"claude-sonnet-4.6","gen_ai.response.model":"claude-sonnet-4.6","gen_ai.response.id":"resp-sub","gen_ai.agent.id":"github.copilot.subagent","gen_ai.usage.input_tokens":100,"gen_ai.usage.output_tokens":10}} -{"type":"span","traceId":"trace-nested","spanId":"tool-task","parentSpanId":"invoke-root","name":"execute_tool task","endTime":[1775934263,0]} -{"type":"span","traceId":"trace-nested","spanId":"invoke-root","name":"invoke_agent","endTime":[1775934260,0],"attributes":{"gen_ai.operation.name":"invoke_agent","gen_ai.provider.name":"github","gen_ai.request.model":"claude-sonnet-4.6","gen_ai.agent.id":"github.copilot.default"}} -{"type":"span","traceId":"trace-nested","spanId":"chat-plain","parentSpanId":"invoke-root","name":"chat gpt-5.4-mini","endTime":[1775934264,967317833],"attributes":{"gen_ai.operation.name":"chat","gen_ai.provider.name":"github","gen_ai.request.model":"gpt-5.4-mini","gen_ai.response.model":"gpt-5.4-mini","gen_ai.response.id":"resp-plain","gen_ai.usage.input_tokens":200,"gen_ai.usage.output_tokens":20}}"#, - ); - let source_fingerprint = SourceFingerprint::from_path(source.path()).unwrap(); - let parsed_messages = crate::sessions::copilot::parse_copilot_file(source.path()); - assert_eq!(parsed_messages.len(), 2); - assert_eq!( - parsed_messages - .iter() - .find(|message| message.model_id == "gpt-5.4-mini") - .unwrap() - .agent - .as_deref(), - Some("github.copilot.default") - ); - assert_eq!( - parsed_messages - .iter() - .find(|message| message.model_id == "claude-sonnet-4.6") - .unwrap() - .agent - .as_deref(), - Some("github.copilot.subagent") - ); - - // Reproduce a schema-27 entry produced by the first-invoke resolver: - // the agentless turn incorrectly inherits the nested sub-agent. - let mut stale_messages = parsed_messages.clone(); - let stale_plain = stale_messages - .iter_mut() - .find(|message| message.model_id == "gpt-5.4-mini") - .unwrap(); - stale_plain.agent = Some("github.copilot.subagent".to_string()); - assert_eq!( - stale_plain.agent.as_deref(), - Some("github.copilot.subagent") - ); - let stale_entry = CachedSourceEntry::new( - source.path(), - source_fingerprint.clone(), - stale_messages, - Vec::new(), - None, - ); - let cache_file = cache_path().unwrap(); - ensure_cache_dir(cache_file.parent().unwrap()).unwrap(); - let stale_store = CachedSourceStore { - schema_version: 27, - entries: vec![stale_entry], - }; - - let writer = BufWriter::new(File::create(&cache_file).unwrap()); - bincode::options() - .serialize_into(writer, &stale_store) - .unwrap(); - - let mut loaded = SourceMessageCache::load(); - assert!( - loaded.entries.is_empty(), - "schema-27 Copilot entries must be stale" - ); - assert_eq!( - SourceFingerprint::from_path(source.path()).unwrap(), - source_fingerprint, - "the stale cache entry and rebuilt parse have the same source fingerprint" - ); - - let rebuilt_messages = crate::sessions::copilot::parse_copilot_file(source.path()); - let rebuilt_plain = rebuilt_messages - .iter() - .find(|message| message.model_id == "gpt-5.4-mini") - .unwrap(); - assert_eq!( - rebuilt_plain.agent.as_deref(), - Some("github.copilot.default"), - "reparse must use the root invoke_agent for the agentless turn" - ); - let rebuilt_sub = rebuilt_messages - .iter() - .find(|message| message.model_id == "claude-sonnet-4.6") - .unwrap(); - assert_eq!( - rebuilt_sub.agent.as_deref(), - Some("github.copilot.subagent"), - "per-record nested attribution must remain unchanged" - ); - loaded.insert(CachedSourceEntry::new( - source.path(), - source_fingerprint, - rebuilt_messages, - Vec::new(), - None, - )); - loaded.save_if_dirty(); - - let rebuilt = read_store_from_path(&cache_file).unwrap(); - assert_eq!(rebuilt.schema_version, CACHE_SCHEMA_VERSION); - assert_eq!(rebuilt.entries.len(), 1); - let rebuilt_entry = &rebuilt.entries[0]; - assert_eq!( - rebuilt_entry - .messages - .iter() - .find(|message| message.model_id == "gpt-5.4-mini") - .unwrap() - .agent - .as_deref(), - Some("github.copilot.default") - ); - assert_eq!( - rebuilt_entry - .messages - .iter() - .find(|message| message.model_id == "claude-sonnet-4.6") - .unwrap() - .agent - .as_deref(), - Some("github.copilot.subagent") - ); - - let reloaded = SourceMessageCache::load(); - let reloaded_entry = reloaded.get(source.path()).unwrap(); - assert_eq!( - reloaded_entry - .messages - .iter() - .find(|message| message.model_id == "gpt-5.4-mini") - .unwrap() - .agent - .as_deref(), - Some("github.copilot.default") - ); - assert_eq!( - reloaded_entry - .messages - .iter() - .find(|message| message.model_id == "claude-sonnet-4.6") - .unwrap() - .agent - .as_deref(), - Some("github.copilot.subagent") - ); - } + assert!( + source.path().is_file(), + "invalidation must not remove the source" + ); + assert!(SourceMessageCache::load() + .get(identity, source.path()) + .is_none()); restore_cache_env(prev_env); } #[test] #[serial_test::serial] - fn test_schema_28_codex_cache_is_stale_and_rebuilt_with_non_overlapping_durations() { + fn test_stale_invalidation_preserves_concurrently_refreshed_entry() { let temp_home = TempDir::new().unwrap(); let prev_env = sandbox_cache_env(temp_home.path()); + let source_dir = TempDir::new().unwrap(); + let path = source_dir.path().join("session.jsonl"); + let identity = CacheIdentity::for_client(ClientId::Claude); + std::fs::write(&path, b"old\n").unwrap(); - { - let source = write_temp_file(include_bytes!( - "../tests/fixtures/codex_duration_timing.jsonl" - )); - let source_fingerprint = SourceFingerprint::from_path(source.path()).unwrap(); - let parsed = crate::sessions::codex::parse_codex_file_incremental( - source.path(), - 0, - CodexParseState::default(), - ); - assert!(parsed.parse_succeeded); - assert_eq!( - parsed - .messages - .iter() - .map(|message| message.duration_ms) - .collect::>(), - vec![Some(1_000), Some(4_000), Some(2_000)] - ); + let mut seed = SourceMessageCache::default(); + seed.insert(test_entry(identity, &path, "old-session")); + seed.save_if_dirty(); - let mut stale_messages = parsed.messages.clone(); - stale_messages[1].duration_ms = Some(5_000); - let stale_incremental = - build_codex_incremental_cache(source.path(), parsed.consumed_offset, parsed.state) - .unwrap(); - let stale_entry = CachedSourceEntry::new( - source.path(), - source_fingerprint.clone(), - stale_messages, - Vec::new(), - Some(stale_incremental), - ); - let cache_file = cache_path().unwrap(); - ensure_cache_dir(cache_file.parent().unwrap()).unwrap(); - let stale_store = CachedSourceStore { - schema_version: 28, - entries: vec![stale_entry], - }; + let mut stale_invalidator = SourceMessageCache::load(); + stale_invalidator.remove(identity, &path); - let writer = BufWriter::new(File::create(&cache_file).unwrap()); - bincode::options() - .serialize_into(writer, &stale_store) - .unwrap(); + std::fs::write(&path, b"fresh-content\n").unwrap(); + let mut fresh_writer = SourceMessageCache::load(); + fresh_writer.insert(test_entry(identity, &path, "fresh-session")); + fresh_writer.save_if_dirty(); - let mut loaded = SourceMessageCache::load(); - assert!( - loaded.entries.is_empty(), - "schema-28 Codex durations and incremental state must be stale" - ); - assert_eq!( - SourceFingerprint::from_path(source.path()).unwrap(), - source_fingerprint, - "the stale entry and rebuilt parse have the same source fingerprint" - ); - - let rebuilt = crate::sessions::codex::parse_codex_file_incremental( - source.path(), - 0, - CodexParseState::default(), - ); - assert!(rebuilt.parse_succeeded); - assert_eq!( - rebuilt - .messages - .iter() - .map(|message| message.duration_ms) - .collect::>(), - vec![Some(1_000), Some(4_000), Some(2_000)] - ); - assert!(rebuilt.state.last_accepted_token_timestamp_ms.is_some()); - let rebuilt_incremental = build_codex_incremental_cache( - source.path(), - rebuilt.consumed_offset, - rebuilt.state, - ) - .unwrap(); - loaded.insert(CachedSourceEntry::new( - source.path(), - source_fingerprint, - rebuilt.messages, - rebuilt.fallback_timestamp_indices, - Some(rebuilt_incremental), - )); - loaded.save_if_dirty(); + stale_invalidator.save_if_dirty(); - let persisted = read_store_from_path(&cache_file).unwrap(); - assert_eq!(persisted.schema_version, CACHE_SCHEMA_VERSION); - assert_eq!(persisted.entries.len(), 1); - assert_eq!( - persisted.entries[0] - .messages - .iter() - .map(|message| message.duration_ms) - .collect::>(), - vec![Some(1_000), Some(4_000), Some(2_000)] - ); - assert!(persisted.entries[0] - .codex_incremental - .as_ref() - .unwrap() - .state - .last_accepted_token_timestamp_ms - .is_some()); - } + let loaded = SourceMessageCache::load(); + assert_eq!( + loaded.get(identity, &path).unwrap().messages[0].session_id, + "fresh-session" + ); restore_cache_env(prev_env); } + #[test] + fn test_prune_missing_files_removes_deleted_entries() { + let file = write_temp_file(b"{}\n"); + let path = file.path().to_path_buf(); + let identity = CacheIdentity::for_client(ClientId::Claude); + + let mut cache = SourceMessageCache::default(); + cache.insert(test_entry(identity, &path, "session-1")); + + std::fs::remove_file(&path).unwrap(); + cache.prune_missing_files(); + + assert!(cache.entries.is_empty()); + } + #[test] #[serial_test::serial] fn test_fallback_cache_dir_prefers_runtime_dir() { let runtime_dir = TempDir::new().unwrap(); - let mut _env = EnvGuard::capture(&["XDG_RUNTIME_DIR"]); - _env.set("XDG_RUNTIME_DIR", runtime_dir.path()); + let original_xdg_runtime_dir = std::env::var("XDG_RUNTIME_DIR").ok(); + restore_env_var("XDG_RUNTIME_DIR", Some(runtime_dir.path())); - assert_eq!( - fallback_cache_dir(), - Some(runtime_dir.path().join("tokscale")) - ); + { + assert_eq!( + fallback_cache_dir(), + Some(runtime_dir.path().join("tokscale")) + ); + } + + restore_env_var("XDG_RUNTIME_DIR", original_xdg_runtime_dir); } #[test] #[serial_test::serial] fn test_save_if_dirty_marks_cache_clean() { let temp_home = TempDir::new().unwrap(); - let _env = sandbox_cache_env(temp_home.path()); + let prev_env = sandbox_cache_env(temp_home.path()); let mut cache = SourceMessageCache::default(); assert!(!cache.dirty); - let file = write_temp_file(b"{}\n"); - let fingerprint = SourceFingerprint::from_path(file.path()).unwrap(); - cache.insert(CachedSourceEntry::new( - file.path(), - fingerprint, - Vec::new(), - Vec::new(), - None, - )); - assert!(cache.dirty); + { + let file = write_temp_file(b"{}\n"); + let identity = CacheIdentity::for_client(ClientId::Claude); + cache.insert(test_entry(identity, file.path(), "session-1")); + assert!(cache.dirty); - cache.save_if_dirty(); - assert!(!cache.dirty); + cache.save_if_dirty(); + assert!(!cache.dirty); + } + + restore_cache_env(prev_env); } #[test] #[serial_test::serial] fn test_save_if_dirty_merges_concurrent_writers() { let temp_home = TempDir::new().unwrap(); - let _env = sandbox_cache_env(temp_home.path()); + let prev_env = sandbox_cache_env(temp_home.path()); - let file_one = write_temp_file(b"{\"id\":1}\n"); - let file_two = write_temp_file(b"{\"id\":2}\n"); + { + let source_dir = TempDir::new().unwrap(); + let identity = CacheIdentity::for_client(ClientId::Claude); + let (path_one, path_two) = write_sources_in_same_shard(&source_dir, identity); + assert_eq!( + CacheKey::new(identity, &path_one).shard(), + CacheKey::new(identity, &path_two).shard() + ); - let mut writer_one = SourceMessageCache::load(); - let mut writer_two = SourceMessageCache::load(); + let mut writer_one = SourceMessageCache::load(); + let mut writer_two = SourceMessageCache::load(); - writer_one.insert(CachedSourceEntry::new( - file_one.path(), - SourceFingerprint::from_path(file_one.path()).unwrap(), - Vec::new(), - Vec::new(), - None, - )); - writer_two.insert(CachedSourceEntry::new( - file_two.path(), - SourceFingerprint::from_path(file_two.path()).unwrap(), - Vec::new(), - Vec::new(), - None, - )); + writer_one.insert(test_entry(identity, &path_one, "session-1")); + writer_two.insert(test_entry(identity, &path_two, "session-2")); - writer_one.save_if_dirty(); - writer_two.save_if_dirty(); + writer_one.save_if_dirty(); + writer_two.save_if_dirty(); - let loaded = SourceMessageCache::load(); - assert!(loaded.get(file_one.path()).is_some()); - assert!(loaded.get(file_two.path()).is_some()); + let loaded = SourceMessageCache::load(); + assert!(loaded.get(identity, &path_one).is_some()); + assert!(loaded.get(identity, &path_two).is_some()); + } + + restore_cache_env(prev_env); } #[test] @@ -2447,29 +3045,10 @@ mod tests { let source_dir = TempDir::new().unwrap(); let path = source_dir.path().join("session.jsonl"); std::fs::write(&path, b"{\"id\":\"old\"}\n").unwrap(); + let identity = CacheIdentity::for_client(ClientId::Claude); let mut seed = SourceMessageCache::default(); - seed.insert(CachedSourceEntry::new( - &path, - SourceFingerprint::from_path(&path).unwrap(), - vec![UnifiedMessage::new( - "client", - "gpt-5", - "provider", - "old-session", - 1, - TokenBreakdown { - input: 1, - output: 0, - cache_read: 0, - cache_write: 0, - reasoning: 0, - }, - 0.0, - )], - Vec::new(), - None, - )); + seed.insert(test_entry(identity, &path, "old-session")); seed.save_if_dirty(); let mut stale_deleter = SourceMessageCache::load(); @@ -2478,34 +3057,14 @@ mod tests { std::fs::write(&path, b"{\"id\":\"fresh\"}\n").unwrap(); let mut fresh_writer = SourceMessageCache::load(); - fresh_writer.insert(CachedSourceEntry::new( - &path, - SourceFingerprint::from_path(&path).unwrap(), - vec![UnifiedMessage::new( - "client", - "gpt-5", - "provider", - "fresh-session", - 2, - TokenBreakdown { - input: 2, - output: 0, - cache_read: 0, - cache_write: 0, - reasoning: 0, - }, - 0.0, - )], - Vec::new(), - None, - )); + fresh_writer.insert(test_entry(identity, &path, "fresh-session")); fresh_writer.save_if_dirty(); stale_deleter.save_if_dirty(); let loaded = SourceMessageCache::load(); let entry = loaded - .get(&path) + .get(identity, &path) .expect("recreated source cache entry should survive stale delete"); assert_eq!(entry.messages[0].session_id, "fresh-session"); } @@ -2515,94 +3074,125 @@ mod tests { #[test] #[serial_test::serial] - #[cfg(not(target_os = "windows"))] - fn load_falls_back_to_legacy_dirs_cache_path() { - let temp_home = TempDir::new().unwrap(); - let temp_xdg_cache = TempDir::new().unwrap(); - let _env = legacy_cache_env(temp_home.path(), Some(temp_xdg_cache.path())); - - let source = write_temp_file(b"legacy-dirs\n"); - let entry = CachedSourceEntry::new( - source.path(), - SourceFingerprint::from_path(source.path()).unwrap(), - Vec::new(), - Vec::new(), - None, - ); - - let legacy_path = crate::paths::legacy_dirs_cache_dir() - .unwrap() - .join(CACHE_FILENAME); - ensure_cache_dir(legacy_path.parent().unwrap()).unwrap(); - let store = CachedSourceStore { - schema_version: CACHE_SCHEMA_VERSION, - entries: vec![entry], - }; - let writer = BufWriter::new(File::create(&legacy_path).unwrap()); - bincode::options().serialize_into(writer, &store).unwrap(); + fn test_parser_versions_are_identity_scoped() { + assert_eq!(parser_version(ClientId::Codex), 4); + assert_eq!(parser_version(ClientId::Jcode), 4); + assert_eq!(parser_version(ClientId::Copilot), 4); + assert_eq!(CacheIdentity::synthetic().parser_version, 1); + for client in ClientId::iter() { + if !matches!( + client, + ClientId::Codex | ClientId::Jcode | ClientId::Copilot + ) { + assert_eq!( + parser_version(client), + 1, + "{} parser version", + client.as_str() + ); + } + } + assert!(ClientId::from_str("pi").is_some()); + assert!(CacheIdentity::current_for_namespace("devin").is_none()); + } - let loaded = SourceMessageCache::load(); - assert!(loaded.get(source.path()).is_some()); + #[test] + fn test_cache_key_namespace_prevents_path_collision() { + let dir = TempDir::new().unwrap(); + let path = dir.path().join("same-source.jsonl"); + std::fs::write(&path, b"same\n").unwrap(); + let claude = CacheKey::new(CacheIdentity::for_client(ClientId::Claude), &path); + let codex = CacheKey::new(CacheIdentity::for_client(ClientId::Codex), &path); + assert_ne!(claude, codex); + assert_ne!(claude.namespace, codex.namespace); } #[test] #[serial_test::serial] - #[cfg(not(target_os = "windows"))] - fn load_falls_back_to_legacy_dot_cache_path() { + fn test_legacy_monolith_is_inert_and_untouched() { let temp_home = TempDir::new().unwrap(); - let _env = legacy_cache_env(temp_home.path(), None); - - let source = write_temp_file(b"legacy-dot\n"); - let entry = CachedSourceEntry::new( + let prev_env = sandbox_cache_env(temp_home.path()); + let legacy = cache_dir().unwrap().join("source-message-cache.bin"); + ensure_cache_dir(legacy.parent().unwrap()).unwrap(); + std::fs::write(&legacy, b"legacy sentinel bytes\n").unwrap(); + let before = std::fs::metadata(&legacy).unwrap(); + let bytes = std::fs::read(&legacy).unwrap(); + + let mut cache = SourceMessageCache::load(); + let source = write_temp_file(b"cold-build\n"); + cache.insert(test_entry( + CacheIdentity::for_client(ClientId::Claude), source.path(), - SourceFingerprint::from_path(source.path()).unwrap(), - Vec::new(), - Vec::new(), - None, - ); - - let legacy_path = crate::paths::legacy_dot_cache_tokscale_dir() - .unwrap() - .join(CACHE_FILENAME); - ensure_cache_dir(legacy_path.parent().unwrap()).unwrap(); - let store = CachedSourceStore { - schema_version: CACHE_SCHEMA_VERSION, - entries: vec![entry], - }; - let writer = BufWriter::new(File::create(&legacy_path).unwrap()); - bincode::options().serialize_into(writer, &store).unwrap(); + "cold", + )); + cache.save_if_dirty(); - let loaded = SourceMessageCache::load(); - assert!(loaded.get(source.path()).is_some()); + assert_eq!(std::fs::read(&legacy).unwrap(), bytes); + let after = std::fs::metadata(&legacy).unwrap(); + assert_eq!(before.len(), after.len()); + assert_eq!(before.modified().unwrap(), after.modified().unwrap()); + assert!(cache_shard_dir().unwrap().is_dir()); + assert!(cache_shard_dir().unwrap().join("claude").is_dir()); + restore_cache_env(prev_env); } - #[cfg(windows)] #[test] #[serial_test::serial] - fn legacy_cache_paths_are_ordered_and_override_gated_without_io() { + fn test_oversized_shard_is_isolated_and_rewritten() { + const TEST_LIMIT: u64 = 1024; let temp_home = TempDir::new().unwrap(); - let _env = legacy_cache_env(temp_home.path(), None); - let candidates = legacy_cache_paths(); - assert_eq!(candidates.len(), 2); - assert_eq!( - candidates[0], - dirs::cache_dir() - .expect("Windows exposes a cache directory") - .join("tokscale") - .join(CACHE_FILENAME) - ); - assert_eq!( - candidates[1], - dirs::home_dir() - .expect("Windows exposes a home directory") - .join(".cache") - .join("tokscale") - .join(CACHE_FILENAME) + let prev_env = sandbox_cache_env(temp_home.path()); + let source_dir = TempDir::new().unwrap(); + let identity = CacheIdentity::for_client(ClientId::Claude); + let (oversized_path, valid_path) = write_sources_in_distinct_shards(&source_dir, identity); + let mut seed = SourceMessageCache::default(); + seed.insert(test_entry(identity, &oversized_path, "oversized")); + seed.insert(test_entry(identity, &valid_path, "valid")); + seed.save_if_dirty_with_limit(1024 * 1024); + let oversized_shard = cache_shard_path(identity, &oversized_path); + let mut oversized_bytes = std::fs::read(&oversized_shard).unwrap(); + oversized_bytes.resize((TEST_LIMIT + 1) as usize, 0); + std::fs::write(&oversized_shard, oversized_bytes).unwrap(); + assert!( + std::fs::metadata(cache_shard_path(identity, &valid_path)) + .unwrap() + .len() + <= TEST_LIMIT ); - let mut override_env = EnvGuard::capture(&["TOKSCALE_CONFIG_DIR"]); - override_env.set("TOKSCALE_CONFIG_DIR", temp_home.path()); - assert!(legacy_cache_paths().is_empty()); + // Exercise the production decode path with an injected limit instead + // of relying on the 256 MiB default merely because this fixture is + // larger than an arbitrary test constant. + let mut loaded = SourceMessageCache::load_with_limit(TEST_LIMIT); + assert!(loaded.get(identity, &oversized_path).is_none()); + assert!(loaded.get(identity, &valid_path).is_some()); + assert!(loaded.dirty); + loaded.save_if_dirty_with_limit(1024 * 1024); + assert!(loaded.get(identity, &valid_path).is_some()); + restore_cache_env(prev_env); + } + + #[test] + #[serial_test::serial] + fn test_dirty_shard_only_rewrites_affected_file() { + let temp_home = TempDir::new().unwrap(); + let prev_env = sandbox_cache_env(temp_home.path()); + let source_dir = TempDir::new().unwrap(); + let identity = CacheIdentity::for_client(ClientId::Claude); + let (first, second) = write_sources_in_distinct_shards(&source_dir, identity); + let mut cache = SourceMessageCache::default(); + cache.insert(test_entry(identity, &first, "first")); + cache.insert(test_entry(identity, &second, "second")); + cache.save_if_dirty(); + let first_shard = cache_shard_path(identity, &first); + let second_shard = cache_shard_path(identity, &second); + let first_bytes = std::fs::read(&first_shard).unwrap(); + let second_bytes = std::fs::read(&second_shard).unwrap(); + cache.insert(test_entry(identity, &first, "first-updated")); + cache.save_if_dirty(); + assert_eq!(std::fs::read(&second_shard).unwrap(), second_bytes); + assert_ne!(std::fs::read(&first_shard).unwrap(), first_bytes); + restore_cache_env(prev_env); } #[cfg(unix)] diff --git a/vendor/tokscale-core/src/model_alias.rs b/vendor/tokscale-core/src/model_alias.rs new file mode 100644 index 0000000..01ef024 --- /dev/null +++ b/vendor/tokscale-core/src/model_alias.rs @@ -0,0 +1,558 @@ +//! Config-driven model-name aliasing for **grouping only**. +//! +//! Different supply channels report the same physical model under different +//! name-strings (for example `claude-opus-4-8`, `claude-opus-4-8-cc`, and +//! `anthropic/claude-opus-4-8` are all one model), so usage stats split across +//! multiple rows. A user-configured `{alias: canonical}` map folds those +//! variants into one canonical **display/group** key. +//! +//! The fold runs as the terminal step of [`crate::normalize_model_for_grouping`], +//! so it applies uniformly to local model/monthly/hourly reports and every +//! `GroupBy`. It is **presentation only**: the submit/upload/export/persist path +//! and graph `ClientContribution` keys use [`crate::canonical_model_id`] (the +//! same syntactic normalization *without* the alias fold), so a machine-local +//! alias config can never rewrite the model identity that leaves the machine or +//! fragment history. It is deliberately **not** applied before pricing +//! (per-message cost is computed on the raw model id upstream), so folding can +//! only relabel and merge already-costed buckets and can never change a cost +//! total. It is orthogonal to the static pricing alias table +//! ([`crate::pricing::aliases`]) and to `provider_identity` — it touches only +//! the model dimension for local grouping. +//! +//! TokenBar adaptation vs upstream `9a5aeb65`: the process-wide map is +//! **reloadable** (not load-once). Changing aliases bumps +//! [`model_alias_generation`] and runs every registered +//! [`register_usage_data_invalidation_hook`] so usage-data consumers (reports, +//! later M24 Warp) can refresh. Multi-message report folds take one +//! [`snapshot_grouping_aliases`] at the start and reuse it for every message so +//! a concurrent reload cannot split one report across two alias maps. +//! Message-cache schema stays 31 because aliases are report-time only. + +use serde::{Deserialize, Serialize}; +use std::collections::{BTreeMap, HashMap}; +use std::sync::atomic::{AtomicU64, Ordering}; +use std::sync::{OnceLock, RwLock}; + +/// Upper bound on the number of configured aliases retained. Oversized configs +/// are truncated rather than rejected, mirroring the capacity guard in +/// [`crate::pricing`]'s custom-pricing loader. +const MAX_MODEL_ALIASES: usize = 4096; + +/// On-disk / settings shape of the flat `modelAliases` object +/// (`{ "alias": "canonical" }`). `#[serde(transparent)]` keeps the serialized +/// form a bare map. Deserialization is lossy: a malformed value (a non-object, +/// or an entry whose value is not a string) is skipped instead of failing the +/// whole settings load. +#[derive(Debug, Clone, Default, PartialEq, Eq, Serialize)] +#[serde(transparent)] +pub struct ModelAliasMap { + /// Raw `alias -> canonical` pairs exactly as written in the config. + pub entries: BTreeMap, +} + +impl<'de> Deserialize<'de> for ModelAliasMap { + fn deserialize(deserializer: D) -> Result + where + D: serde::Deserializer<'de>, + { + // Read the node as a generic value first so a malformed `modelAliases` + // (e.g. an array or scalar) degrades to an empty map instead of + // misaligning the parent settings deserializer. Keep only string-valued + // entries; skip anything else. + let value = serde_json::Value::deserialize(deserializer)?; + let entries = match value { + serde_json::Value::Object(object) => object + .into_iter() + .filter_map(|(key, value)| match value { + serde_json::Value::String(canonical) => Some((key, canonical)), + _ => None, + }) + .collect(), + _ => BTreeMap::new(), + }; + Ok(Self { entries }) + } +} + +/// Runtime resolver built from [`ModelAliasMap`]: keys and values are normalized +/// through [`crate::normalize_syntactic`] so lookups match regardless of case, +/// dated suffix, or `.`-vs-`-` spelling, and canonical values land in the same +/// space the grouping key uses. Empty keys/values and self-maps are dropped; the +/// number of entries is capped. `Clone` clones the HashMap so a report fold can +/// hold a stable [`GroupingAliasSnapshot`] while the process-wide map reloads. +#[derive(Debug, Default, Clone)] +struct ModelAliasResolver { + map: HashMap, +} + +impl ModelAliasResolver { + /// Build a resolver from configured aliases. Both sides of each pair are run + /// through [`crate::normalize_syntactic`] exactly once: keys are placed in the + /// same space as incoming (already-normalized) model names, and canonical + /// values are stored pre-normalized. `apply` returns a canonical value + /// verbatim — it is never re-resolved or re-normalized — so the value written + /// here is exactly the label shown in reports. Empty keys/values and + /// self-maps are dropped, and the number of entries is capped. + fn from_config(config: &ModelAliasMap) -> Self { + let mut map = HashMap::new(); + for (raw_alias, raw_canonical) in &config.entries { + if map.len() >= MAX_MODEL_ALIASES { + break; + } + // Store keys under a separator-insensitive match key so matching is + // provider-agnostic (not claude-only): `gpt-5-5` and `gpt-5.5` share + // the key `gpt-5-5`. The stored canonical value keeps its + // `normalize_syntactic` spelling — it is the label shown verbatim. + let alias_norm = crate::normalize_syntactic(raw_alias); + let canonical = crate::normalize_syntactic(raw_canonical); + // Self-map drop compares the *exact* normalized forms, not the match + // keys: `{gpt-5-5: gpt-5.5}` is a real separator relabel that must be + // kept, whereas `{gpt-5.5: gpt-5.5}` is a genuine no-op to drop. + if alias_norm.is_empty() || canonical.is_empty() || alias_norm == canonical { + continue; + } + map.insert(match_key(&alias_norm), canonical); + } + Self { map } + } + + /// Resolve one model name. `name` must already be `normalize_syntactic`'d (it + /// is, since the only caller is [`crate::normalize_model_for_grouping`]). + /// Resolution is single-hop — the canonical value is never re-resolved — so + /// alias chains collapse one step and cycles are structurally impossible. + /// Returns `name` unchanged on a miss. + fn apply(&self, name: String) -> String { + match self.map.get(&match_key(&name)) { + Some(canonical) => canonical.clone(), + None => name, + } + } +} + +/// Reduce an already-`normalize_syntactic`'d model name to a separator- +/// insensitive match key by rewriting every `.` to `-`. This generalizes alias +/// matching beyond claude: `normalize_syntactic` only rewrites `.`→`-` inside +/// *claude* version numbers, so without this a `gpt-5-5` alias would miss +/// `gpt-5.5`. Folding on the match key alone keeps the displayed canonical form +/// (e.g. `gpt-5.5`) untouched for models that were never aliased. +fn match_key(normalized: &str) -> String { + normalized.replace('.', "-") +} + +#[derive(Default)] +struct AliasState { + config: ModelAliasMap, + resolver: ModelAliasResolver, +} + +fn state() -> &'static RwLock { + static STATE: OnceLock> = OnceLock::new(); + STATE.get_or_init(|| RwLock::new(AliasState::default())) +} + +/// Monotonic generation bumped on every successful alias install/clear. +/// Usage-data consumers (and later M24 Warp) can poll this cheaply to detect +/// that report-time grouping input changed without re-reading the full map. +static GENERATION: AtomicU64 = AtomicU64::new(0); + +type InvalidationHook = Box; + +fn hooks() -> &'static RwLock> { + static HOOKS: OnceLock>> = OnceLock::new(); + HOOKS.get_or_init(|| RwLock::new(Vec::new())) +} + +fn install(config: ModelAliasMap) { + let resolver = ModelAliasResolver::from_config(&config); + { + let mut state = state().write().unwrap_or_else(|e| e.into_inner()); + state.config = config; + state.resolver = resolver; + } + GENERATION.fetch_add(1, Ordering::SeqCst); + notify_usage_data_invalidation(); +} + +fn notify_usage_data_invalidation() { + let hooks = hooks().read().unwrap_or_else(|e| e.into_inner()); + for hook in hooks.iter() { + hook(); + } +} + +/// Install (or replace) the process-wide model-alias map used for grouping. +/// +/// Always reloads — later calls replace earlier ones. Bumps +/// [`model_alias_generation`] and fires every registered usage-data invalidation +/// hook so the next report sees the new grouping. Until the first non-empty +/// install (and after [`clear_model_aliases`]), grouping is a strict identity +/// no-op relative to [`crate::canonical_model_id`]. +pub fn set_model_aliases(config: &ModelAliasMap) { + install(config.clone()); +} + +/// Clear all process-wide grouping aliases (identity no-op) and invalidate +/// usage-data consumers. +pub fn clear_model_aliases() { + install(ModelAliasMap::default()); +} + +/// Snapshot of the currently installed raw alias map (config shape, not the +/// normalized resolver). Empty when unset or cleared. +pub fn model_aliases() -> ModelAliasMap { + state() + .read() + .unwrap_or_else(|e| e.into_inner()) + .config + .clone() +} + +/// Process-wide generation for the installed alias map. Starts at 0; increments +/// on every [`set_model_aliases`] / [`clear_model_aliases`]. Independent of +/// message-cache schema (stays 31). +pub fn model_alias_generation() -> u64 { + GENERATION.load(Ordering::SeqCst) +} + +/// Register a process-wide callback invoked whenever grouping aliases change. +/// +/// Hooks are append-only for the process lifetime. Intended for TokenBar +/// usage-data layers (in-process report caches, later M24 Warp) so a settings +/// reload does not require a process restart. Hooks must be cheap and +/// non-reentrant with respect to alias install. +pub fn register_usage_data_invalidation_hook(hook: impl Fn() + Send + Sync + 'static) { + hooks() + .write() + .unwrap_or_else(|e| e.into_inner()) + .push(Box::new(hook)); +} + +/// Apply the installed resolver to an already-`normalize_syntactic`'d name. +/// +/// Reads the process-wide map on every call. Prefer +/// [`snapshot_grouping_aliases`] + [`GroupingAliasSnapshot::fold`] for any +/// multi-message report fold so a mid-fold reload cannot split grouping. +pub(crate) fn apply_global(name: String) -> String { + state() + .read() + .unwrap_or_else(|e| e.into_inner()) + .resolver + .apply(name) +} + +/// Point-in-time clone of the process-wide grouping-alias resolver. +/// +/// Report folds (model / monthly / hourly) take one snapshot at the start and +/// call [`Self::fold`] for every message. A concurrent [`set_model_aliases`] or +/// [`clear_model_aliases`] updates only the live map; the snapshotted HashMap is +/// independent and keeps the whole report on one alias config. +#[derive(Debug, Clone)] +pub struct GroupingAliasSnapshot { + resolver: ModelAliasResolver, +} + +impl GroupingAliasSnapshot { + /// Single-hop alias fold for an already-[`crate::normalize_syntactic`]'d + /// model name. Misses are identity. Canonical values are returned verbatim + /// (never re-resolved), matching [`apply_global`]. + pub fn fold(&self, syntactic_name: String) -> String { + self.resolver.apply(syntactic_name) + } +} + +/// Clone the currently installed grouping-alias resolver for one report fold. +/// +/// Cheap relative to scanning messages: one `RwLock` read and a HashMap clone +/// (capped at [`MAX_MODEL_ALIASES`] entries). Empty/unset aliases yield an +/// identity snapshot. +pub fn snapshot_grouping_aliases() -> GroupingAliasSnapshot { + let resolver = state() + .read() + .unwrap_or_else(|e| e.into_inner()) + .resolver + .clone(); + GroupingAliasSnapshot { resolver } +} + +/// Shared mutex for every test that mutates process-wide alias state. +/// Lives outside `mod tests` so `lib.rs` integration-style unit tests and +/// this module's tests take the **same** guard (Codex: dual locks interleave). +#[cfg(test)] +pub(crate) fn lock_global_alias_tests() -> std::sync::MutexGuard<'static, ()> { + use std::sync::Mutex; + static LOCK: Mutex<()> = Mutex::new(()); + LOCK.lock().unwrap_or_else(|e| e.into_inner()) +} + +#[cfg(test)] +mod tests { + use super::*; + use std::sync::atomic::AtomicUsize; + + fn resolver(pairs: &[(&str, &str)]) -> ModelAliasResolver { + let entries = pairs + .iter() + .map(|(k, v)| ((*k).to_string(), (*v).to_string())) + .collect(); + ModelAliasResolver::from_config(&ModelAliasMap { entries }) + } + + fn alias_map(pairs: &[(&str, &str)]) -> ModelAliasMap { + ModelAliasMap { + entries: pairs + .iter() + .map(|(k, v)| ((*k).to_string(), (*v).to_string())) + .collect(), + } + } + + #[test] + fn folds_three_variants_to_one_canonical() { + let r = resolver(&[ + ("claude-opus-4-8-cc", "claude-opus-4-8"), + ("anthropic/claude-opus-4-8", "claude-opus-4-8"), + ]); + // All three real-world spellings collapse to the canonical name. The + // third needs no map entry: syntactic normalization already lowercases it. + for input in [ + "claude-opus-4-8-cc", + "anthropic/claude-opus-4-8", + "Claude-Opus-4-8", + ] { + assert_eq!( + r.apply(crate::normalize_syntactic(input)), + "claude-opus-4-8", + "input {input} should fold to claude-opus-4-8" + ); + } + } + + #[test] + fn keys_match_case_and_dotted_insensitively() { + // Config key written with upper case and a dotted version still matches + // the normalized input, because both sides run through normalize_syntactic. + let r = resolver(&[("Claude-Opus-4.8-CC", "claude-opus-4-8")]); + assert_eq!( + r.apply(crate::normalize_syntactic("claude-opus-4-8-cc")), + "claude-opus-4-8" + ); + } + + #[test] + fn drops_empty_and_self_maps() { + let r = resolver(&[ + ("", "claude-opus-4-8"), + ("claude-opus-4-8-cc", ""), + ("gpt-5.5", "gpt-5.5"), + ]); + assert!(r.map.is_empty()); + } + + #[test] + fn resolution_is_single_hop() { + // {a: b, b: c} resolves a -> b (not c) and never loops. + let r = resolver(&[("model-a", "model-b"), ("model-b", "model-c")]); + assert_eq!(r.apply("model-a".to_string()), "model-b"); + assert_eq!(r.apply("model-b".to_string()), "model-c"); + } + + #[test] + fn separator_insensitive_match_is_provider_agnostic() { + // Finding A: `normalize_syntactic` rewrites `.`→`-` only for claude, so + // the resolver must fold separators itself for every other provider. The + // regression is when the CONFIGURED alias key and the model string the + // provider actually reports use different separators — the old exact + // HashMap lookup missed and left the variant unfolded. + + // Dashed alias key (`gpt-5-5-cc`), dotted model spelling (`gpt-5.5-cc`): + // must still fold to the canonical `gpt-5.5`. + let dashed_key = resolver(&[("gpt-5-5-cc", "gpt-5.5")]); + assert_eq!( + dashed_key.apply(crate::normalize_syntactic("gpt-5.5-cc")), + "gpt-5.5", + "a dashed alias key must match the dotted model spelling (gpt-5-5 ↔ gpt-5.5)" + ); + + // Mirror: dotted alias key, dashed model spelling. + let dotted_key = resolver(&[("gpt-5.5-cc", "gpt-5.5")]); + assert_eq!( + dotted_key.apply(crate::normalize_syntactic("gpt-5-5-cc")), + "gpt-5.5", + "a dotted alias key must match the dashed model spelling" + ); + } + + #[test] + fn miss_is_identity() { + let r = resolver(&[("claude-opus-4-8-cc", "claude-opus-4-8")]); + assert_eq!(r.apply("gpt-5.5".to_string()), "gpt-5.5"); + } + + #[test] + fn empty_resolver_is_identity() { + let r = ModelAliasResolver::default(); + assert_eq!( + r.apply("claude-opus-4-8-cc".to_string()), + "claude-opus-4-8-cc" + ); + } + + #[test] + fn respects_capacity_cap() { + let entries: BTreeMap = (0..MAX_MODEL_ALIASES + 100) + .map(|i| (format!("alias-{i}"), format!("canonical-{i}"))) + .collect(); + let r = ModelAliasResolver::from_config(&ModelAliasMap { entries }); + assert_eq!(r.map.len(), MAX_MODEL_ALIASES); + } + + #[test] + fn deserialize_is_lossy_over_non_string_values() { + // Non-string values are skipped; string entries survive. + let parsed: ModelAliasMap = + serde_json::from_str(r#"{"a": "b", "n": 5, "arr": ["x"]}"#).unwrap(); + assert_eq!(parsed.entries.len(), 1); + assert_eq!(parsed.entries.get("a").map(String::as_str), Some("b")); + } + + #[test] + fn deserialize_of_non_object_is_empty() { + // A misuse (array/scalar instead of an object) degrades to empty, not error. + assert!(serde_json::from_str::("[]") + .unwrap() + .entries + .is_empty()); + assert!(serde_json::from_str::("\"oops\"") + .unwrap() + .entries + .is_empty()); + } + + #[test] + fn serialize_round_trips_as_flat_map() { + let map = ModelAliasMap { + entries: [( + "claude-opus-4-8-cc".to_string(), + "claude-opus-4-8".to_string(), + )] + .into_iter() + .collect(), + }; + let json = serde_json::to_string(&map).unwrap(); + assert_eq!(json, r#"{"claude-opus-4-8-cc":"claude-opus-4-8"}"#); + assert_eq!(serde_json::from_str::(&json).unwrap(), map); + } + + #[test] + fn reloadable_global_install_and_clear() { + let _guard = lock_global_alias_tests(); + clear_model_aliases(); + let gen0 = model_alias_generation(); + + set_model_aliases(&alias_map(&[("claude-opus-4-8-cc", "claude-opus-4-8")])); + assert_eq!( + apply_global(crate::normalize_syntactic("claude-opus-4-8-cc")), + "claude-opus-4-8" + ); + assert_eq!( + model_aliases() + .entries + .get("claude-opus-4-8-cc") + .map(String::as_str), + Some("claude-opus-4-8") + ); + let gen1 = model_alias_generation(); + assert!(gen1 > gen0); + + // Reload replaces rather than first-wins. + set_model_aliases(&alias_map(&[("gpt-5.5-cc", "gpt-5.5")])); + assert_eq!( + apply_global(crate::normalize_syntactic("claude-opus-4-8-cc")), + "claude-opus-4-8-cc", + "previous alias must be gone after reload" + ); + assert_eq!( + apply_global(crate::normalize_syntactic("gpt-5.5-cc")), + "gpt-5.5" + ); + let gen2 = model_alias_generation(); + assert!(gen2 > gen1); + + clear_model_aliases(); + assert!(model_aliases().entries.is_empty()); + assert_eq!( + apply_global(crate::normalize_syntactic("gpt-5.5-cc")), + "gpt-5.5-cc" + ); + assert!(model_alias_generation() > gen2); + } + + #[test] + fn invalidation_hook_fires_on_set_and_clear() { + let _guard = lock_global_alias_tests(); + clear_model_aliases(); + + static FIRES: AtomicUsize = AtomicUsize::new(0); + // Register once per process; subtract baseline so the assertion is local. + register_usage_data_invalidation_hook(|| { + FIRES.fetch_add(1, Ordering::SeqCst); + }); + let baseline = FIRES.load(Ordering::SeqCst); + + set_model_aliases(&alias_map(&[("a", "b")])); + clear_model_aliases(); + let after = FIRES.load(Ordering::SeqCst); + assert!( + after >= baseline + 2, + "set + clear must each fire the invalidation hook (baseline={baseline}, after={after})" + ); + } + + #[test] + fn grouping_alias_snapshot_stable_across_mid_fold_reload() { + // Codex P2: a multi-message report must not split across two alias maps + // when set_model_aliases runs mid-fold. Snapshot at fold start; mutate + // the live map; prove the snapshot keeps folding with the old config. + let _guard = lock_global_alias_tests(); + clear_model_aliases(); + set_model_aliases(&alias_map(&[("alias-a", "canonical-b")])); + + let snap = snapshot_grouping_aliases(); + // First half of a report fold under the snapshotted map. + assert_eq!( + snap.fold(crate::normalize_syntactic("alias-a")), + "canonical-b" + ); + + // Mid-fold reload replaces the process-wide map. + set_model_aliases(&alias_map(&[("alias-a", "canonical-other")])); + assert_eq!( + apply_global(crate::normalize_syntactic("alias-a")), + "canonical-other", + "live path must see the reloaded map" + ); + // Snapshot remains on the fold-start config for every later message. + assert_eq!( + snap.fold(crate::normalize_syntactic("alias-a")), + "canonical-b", + "report-fold snapshot must ignore mid-fold reload" + ); + assert_eq!( + snap.fold(crate::normalize_syntactic("alias-a")), + "canonical-b", + "second half of the fold must stay consistent with the first" + ); + + clear_model_aliases(); + assert_eq!( + snap.fold(crate::normalize_syntactic("alias-a")), + "canonical-b", + "clear must not poison an already-taken snapshot" + ); + assert_eq!( + apply_global(crate::normalize_syntactic("alias-a")), + "alias-a", + "live path after clear is identity" + ); + } +} diff --git a/vendor/tokscale-core/src/pricing/aliases.rs b/vendor/tokscale-core/src/pricing/aliases.rs index 4b4aabb..9560fbc 100644 --- a/vendor/tokscale-core/src/pricing/aliases.rs +++ b/vendor/tokscale-core/src/pricing/aliases.rs @@ -18,6 +18,28 @@ static MODEL_ALIASES: Lazy> = Lazy::new(|| { m.insert("model_placeholder_m35", "claude-sonnet-4-6"); m.insert("model_placeholder_m36", "gemini-3.1-pro"); m.insert("model_placeholder_m37", "gemini-3.1-pro"); + // Antigravity uses opaque placeholder IDs in IDE metadata and shorter + // responseModel aliases in CLI conversation protobufs. Keep these as + // machine-ID aliases rather than display labels because labels may be + // renamed or localized. + // + // M133/`gemini-3-flash-b`, `gemini-3-flash-a`, and M187/raw + // `gemini-3.5-flash-low` are source-verified exceptions to the obvious + // mapping: M133 and both response aliases are the High tier; the raw + // `gemini-3.5-flash-low` wire value is the Medium tier; M187 is the true + // Low tier with its distinct machine id. + m.insert("model_placeholder_m16", "gemini-3.1-pro"); + m.insert("model_placeholder_m18", "gemini-3-flash-preview"); + m.insert("model_placeholder_m84", "gemini-3-flash-preview"); + m.insert("model_placeholder_m132", "gemini-3.5-flash-high"); + m.insert("model_placeholder_m133", "gemini-3.5-flash-high"); + m.insert("model_placeholder_m187", "gemini-3.5-flash-extra-low"); + m.insert("model_placeholder_m20", "gemini-3.5-flash-medium"); + m.insert("gemini-pro-default", "gemini-3.1-pro"); + m.insert("gemini-pro-agent", "gemini-3.1-pro"); + m.insert("gemini-3-flash-agent", "gemini-3.5-flash-high"); + m.insert("gemini-3-flash-b", "gemini-3.5-flash-high"); + m.insert("gemini-3.5-flash-low", "gemini-3.5-flash-medium"); m.insert("model_placeholder_m47", "gemini-3-flash-preview"); m.insert("model_openai_gpt_oss_120b_medium", "gpt-oss-120b-medium"); m.insert("claude-opus-4-6-thinking", "claude-opus-4-6"); @@ -42,7 +64,7 @@ static MODEL_ALIASES: Lazy> = Lazy::new(|| { m.insert("gemini-3-pro-low", "gemini-3-pro"); m.insert("gemini-3-flash", "gemini-3-flash-preview"); m.insert("gemini-3-flash-c", "gemini-3-flash-preview"); - m.insert("gemini-3-flash-a", "gemini-3-flash-preview"); + m.insert("gemini-3-flash-a", "gemini-3.5-flash-high"); // Synthetic model variants (only where resolver needs help) m.insert("kimi-k2.5-nvfp4", "kimi-k2.5"); // Quantization variant → base model pricing @@ -60,38 +82,32 @@ mod tests { #[test] fn resolves_antigravity_placeholders() { - assert_eq!( - resolve_alias("MODEL_PLACEHOLDER_M26"), - Some("claude-opus-4-6") - ); - assert_eq!( - resolve_alias("model_placeholder_m37"), - Some("gemini-3.1-pro") - ); - assert_eq!( - resolve_alias("MODEL_OPENAI_GPT_OSS_120B_MEDIUM"), - Some("gpt-oss-120b-medium") - ); - assert_eq!( - resolve_alias("gemini-3-flash-c"), - Some("gemini-3-flash-preview") - ); - assert_eq!( - resolve_alias("gemini-3-flash-a"), - Some("gemini-3-flash-preview") - ); - assert_eq!( - resolve_alias("claude-opus-4.6-thinking"), - Some("claude-opus-4-6") - ); - assert_eq!( - resolve_alias("anthropic/claude-4-5-haiku"), - Some("claude-haiku-4-5") - ); - assert_eq!( - resolve_alias("anthropic/claude-4-6-sonnet"), - Some("claude-sonnet-4-6") - ); + let cases = [ + ("MODEL_PLACEHOLDER_M26", "claude-opus-4-6"), + ("model_placeholder_m37", "gemini-3.1-pro"), + ("model_placeholder_m16", "gemini-3.1-pro"), + ("model_placeholder_m18", "gemini-3-flash-preview"), + ("MODEL_PLACEHOLDER_M84", "gemini-3-flash-preview"), + ("model_placeholder_m132", "gemini-3.5-flash-high"), + ("model_placeholder_m133", "gemini-3.5-flash-high"), + ("model_placeholder_m187", "gemini-3.5-flash-extra-low"), + ("model_placeholder_m20", "gemini-3.5-flash-medium"), + ("gemini-pro-default", "gemini-3.1-pro"), + ("gemini-pro-agent", "gemini-3.1-pro"), + ("gemini-3-flash-agent", "gemini-3.5-flash-high"), + ("gemini-3-flash-b", "gemini-3.5-flash-high"), + ("gemini-3.5-flash-low", "gemini-3.5-flash-medium"), + ("MODEL_OPENAI_GPT_OSS_120B_MEDIUM", "gpt-oss-120b-medium"), + ("gemini-3-flash-c", "gemini-3-flash-preview"), + ("gemini-3-flash-a", "gemini-3.5-flash-high"), + ("claude-opus-4.6-thinking", "claude-opus-4-6"), + ("anthropic/claude-4-5-haiku", "claude-haiku-4-5"), + ("anthropic/claude-4-6-sonnet", "claude-sonnet-4-6"), + ]; + + for (raw, expected) in cases { + assert_eq!(resolve_alias(raw), Some(expected), "raw model: {raw}"); + } } #[test] @@ -104,4 +120,16 @@ mod tests { assert_eq!(resolve_alias("k2p5"), Some("kimi-k2-thinking")); assert_eq!(resolve_alias("k2-p5"), Some("kimi-k2-thinking")); } + + #[test] + fn antigravity_low_and_medium_aliases_remain_distinct() { + let low = resolve_alias("model_placeholder_m187").unwrap(); + let medium = resolve_alias("model_placeholder_m20").unwrap(); + let cli_medium = resolve_alias("gemini-3.5-flash-low").unwrap(); + + assert_eq!(low, "gemini-3.5-flash-extra-low"); + assert_eq!(medium, "gemini-3.5-flash-medium"); + assert_ne!(low, medium); + assert_eq!(cli_medium, medium); + } } diff --git a/vendor/tokscale-core/src/pricing/custom.rs b/vendor/tokscale-core/src/pricing/custom.rs index ac6822d..7387626 100644 --- a/vendor/tokscale-core/src/pricing/custom.rs +++ b/vendor/tokscale-core/src/pricing/custom.rs @@ -241,15 +241,22 @@ impl CustomPricing { self.lookup_with_key(model_id).map(|result| result.pricing) } - pub fn lookup_with_key(&self, model_id: &str) -> Option> { + pub(crate) fn lookup_exact_with_key(&self, model_id: &str) -> Option> { let raw_key = model_id.to_lowercase(); - if let Some(pricing) = self.models.get_key_value(&raw_key) { - return Some(CustomLookupResult { + self.models + .get_key_value(&raw_key) + .map(|pricing| CustomLookupResult { matched_key: pricing.0, pricing: pricing.1, - }); + }) + } + + pub fn lookup_with_key(&self, model_id: &str) -> Option> { + if let Some(result) = self.lookup_exact_with_key(model_id) { + return Some(result); } + let raw_key = model_id.to_lowercase(); let normalized_key = normalize_synthetic_model(model_id).to_lowercase(); if normalized_key != raw_key { if let Some(pricing) = self.models.get_key_value(&normalized_key) { diff --git a/vendor/tokscale-core/src/pricing/lookup.rs b/vendor/tokscale-core/src/pricing/lookup.rs index 907529f..9acddd2 100644 --- a/vendor/tokscale-core/src/pricing/lookup.rs +++ b/vendor/tokscale-core/src/pricing/lookup.rs @@ -74,6 +74,17 @@ const TIERED_PRICING_THRESHOLD_200K_TOKENS: f64 = 200_000.0; const TIERED_PRICING_THRESHOLD_256K_TOKENS: f64 = 256_000.0; const TIERED_PRICING_THRESHOLD_272K_TOKENS: f64 = 272_000.0; +// Only these identities document one long-context rate tier for the whole +// request. Other catalog `*_above_*` fields retain marginal semantics. +const FULL_SESSION_LONG_CONTEXT_LITELLM_KEYS: &[&str] = &[ + "gpt-5.4", + "gpt-5.4-2026-03-05", + "gpt-5.4-pro", + "gpt-5.4-pro-2026-03-05", + "gpt-5.5", + "gpt-5.5-2026-04-23", +]; + const MIN_FUZZY_MATCH_LEN: usize = 5; /// Minimum length for a model name candidate after prefix/suffix stripping. @@ -109,6 +120,7 @@ pub struct PricingLookup { litellm: HashMap, openrouter: HashMap, cursor: HashMap, + sakana: HashMap, litellm_keys: Vec, openrouter_keys: Vec, litellm_key_parts: Vec, @@ -117,6 +129,7 @@ pub struct PricingLookup { openrouter_lower: HashMap, openrouter_model_part: HashMap, cursor_lower: HashMap, + sakana_lower: HashMap, lookup_cache: RwLock>>, } @@ -131,6 +144,15 @@ impl PricingLookup { litellm: HashMap, openrouter: HashMap, cursor: HashMap, + ) -> Self { + Self::new_with_sakana(litellm, openrouter, cursor, HashMap::new()) + } + + pub fn new_with_sakana( + litellm: HashMap, + openrouter: HashMap, + cursor: HashMap, + sakana: HashMap, ) -> Self { let mut litellm_keys: Vec = litellm.keys().cloned().collect(); litellm_keys.sort_by_key(|k| std::cmp::Reverse(k.len())); @@ -160,6 +182,11 @@ impl PricingLookup { cursor_lower.insert(key.to_lowercase(), key.clone()); } + let mut sakana_lower = HashMap::with_capacity(sakana.len()); + for key in sakana.keys() { + sakana_lower.insert(key.to_lowercase(), key.clone()); + } + let build_key_parts = |keys: &[String]| -> Vec { keys.iter() .map(|key| { @@ -180,6 +207,7 @@ impl PricingLookup { litellm, openrouter, cursor, + sakana, litellm_keys, openrouter_keys, litellm_key_parts, @@ -188,6 +216,7 @@ impl PricingLookup { openrouter_lower, openrouter_model_part, cursor_lower, + sakana_lower, lookup_cache: RwLock::new(HashMap::with_capacity(64)), } } @@ -201,8 +230,11 @@ impl PricingLookup { model_id: &str, provider_id: Option<&str>, ) -> Option { - let provider_id = normalize_provider_hint(provider_id); - let cache_key = build_lookup_cache_key(model_id, provider_id); + let normalized_provider_id = normalize_provider_hint(provider_id); + if provider_id.is_some() && normalized_provider_id.is_none() { + return self.lookup_with_source_and_provider(model_id, None, provider_id); + } + let cache_key = build_lookup_cache_key(model_id, normalized_provider_id); if let Some(cached) = self .lookup_cache .read() @@ -256,15 +288,92 @@ impl PricingLookup { force_source: Option<&str>, provider_id: Option<&str>, ) -> Option { + self.lookup_with_source_and_provider_and_terminal_custom( + model_id, + force_source, + provider_id, + |_| None, + ) + } + + /// Resolve a model through the normal built-in pipeline while allowing the + /// owning service to add a custom candidate only at the generic terminal + /// fallback stage. This deliberately bypasses `lookup_cache`: the cache is + /// owned by the built-in datasets and must never retain a result supplied by + /// an unrelated custom-pricing closure. + pub(crate) fn lookup_with_source_and_provider_and_terminal_custom( + &self, + model_id: &str, + force_source: Option<&str>, + provider_id: Option<&str>, + terminal_custom: F, + ) -> Option + where + F: Fn(&str) -> Option, + { + let sakana_provider_id = provider_id.map(str::trim); let provider_id = normalize_provider_hint(provider_id); - let canonical = aliases::resolve_alias(model_id).unwrap_or(model_id); - let lower = canonical.to_lowercase(); + let force_source_normalized = force_source.map(str::to_ascii_lowercase); + let force_source = force_source_normalized.as_deref(); + let raw_lower = model_id.to_lowercase(); + + // Helper to perform lookup with the given source constraint. + let do_lookup = |id: &str| match force_source { + Some("litellm") => self.lookup_litellm_only(id, provider_id), + Some("openrouter") => self.lookup_openrouter_only(id, provider_id), + Some("custom") => None, + _ => self.lookup_auto(id, provider_id, sakana_provider_id), + }; + + let force_custom = force_source.is_some_and(|source| source.eq_ignore_ascii_case("custom")); + let has_parenthesized_suffix = raw_lower + .strip_suffix(')') + .and_then(|inner| inner.rsplit_once('(')) + .is_some(); + + // The complete raw id gets first refusal. This must precede tier + // validation and every normalization pass so an explicitly configured + // `(invalid)` or routed key remains authoritative. Forced upstream + // sources never invoke the custom callback. Use the raw id for the + // Claude guard too, so this direct stage does not become an unguarded + // terminal fallback. + let raw_requested_family = claude_family(&raw_lower); + let raw_requested_version = requested_claude_version(&raw_lower); + let raw_unparsed_modern_version = raw_requested_family.is_some() + && raw_requested_version.is_none() + && contains_delimited_modern_major_minor(&raw_lower); + let raw_unsafe_claude_resolution = |result: &LookupResult| { + resolves_unsafe_claude_version( + raw_requested_family, + raw_requested_version.as_deref(), + raw_unparsed_modern_version, + result, + ) + }; + if force_custom || force_source.is_none() { + if let Some(result) = + terminal_custom(&raw_lower).filter(|result| !raw_unsafe_claude_resolution(result)) + { + return Some(result); + } + } + // Built-in exact lookup is limited to parenthesized ids so ordinary + // provider ranking remains unchanged for unsuffixed ids. + if !force_custom && has_parenthesized_suffix { + if let Some(result) = self + .lookup_exact_full(&raw_lower, force_source, provider_id, sakana_provider_id) + .filter(|result| !raw_unsafe_claude_resolution(result)) + { + return Some(result); + } + } // CLIProxyAPI strips `(level)` reasoning-effort suffixes before routing, // so for pricing lookup we resolve to the base model regardless of tier. // Mirrors the dash-suffix path (e.g. `-xhigh`), which is handled by // `try_strip_unknown_suffix` below. - let normalized_owned = strip_parenthesized_reasoning_tier(&lower).map(str::to_owned); + let raw_normalized_owned = + strip_parenthesized_reasoning_tier(&raw_lower).map(str::to_owned); // Guard against silent misresolution: if the input ends with `(...)` // but the contents are not a recognized CLIProxyAPI level, refuse the @@ -272,23 +381,14 @@ impl PricingLookup { // `-` and could match a shorter, unrelated model id by peeling the // parenthesized fragment off (e.g. `gpt-5.2-codex(invalid)` would // strip `-codex(invalid)` and resolve to `gpt-5.2`). - if normalized_owned.is_none() - && lower - .strip_suffix(')') - .and_then(|inner| inner.rsplit_once('(')) - .is_some() - { + if raw_normalized_owned.is_none() && has_parenthesized_suffix { return None; } - let lower_ref: &str = normalized_owned.as_deref().unwrap_or(&lower); - - // Helper to perform lookup with the given source constraint - let do_lookup = |id: &str| match force_source { - Some("litellm") => self.lookup_litellm_only(id, provider_id), - Some("openrouter") => self.lookup_openrouter_only(id, provider_id), - _ => self.lookup_auto(id, provider_id), - }; + let raw_ref = raw_normalized_owned.as_deref().unwrap_or(&raw_lower); + let canonical = aliases::resolve_alias(raw_ref).unwrap_or(raw_ref); + let canonical_lower = canonical.to_lowercase(); + let lower_ref = canonical_lower.as_str(); let requested_family = claude_family(lower_ref); let requested_version = requested_claude_version(lower_ref); let unparsed_modern_version = requested_family.is_some() @@ -302,38 +402,167 @@ impl PricingLookup { result, ) }; + let guarded_custom = |candidate: &str| { + terminal_custom(candidate).filter(|result| !unsafe_claude_resolution(result)) + }; - // 1. Try direct lookup - if let Some(result) = do_lookup(lower_ref) { - if unsafe_claude_resolution(&result) { - return None; + // Full custom keys and full raw model ids are authoritative before any + // fallback. Keep the raw form first so a routed or provider-scoped key + // cannot be stolen by an alias normalization. + if let Some(result) = guarded_custom(raw_ref) { + return Some(result); + } + // Static aliases are a direct lookup stage, not a fuzzy fallback. A + // raw dataset key still gets first refusal when the id is an alias, + // while non-alias ids retain the established provider-ranking pipeline. + if canonical_lower != raw_ref { + if let Some(result) = self + .lookup_exact_full(raw_ref, force_source, provider_id, sakana_provider_id) + .filter(|result| !unsafe_claude_resolution(result)) + { + return Some(result); } + if let Some(result) = guarded_custom(lower_ref) { + return Some(result); + } + } + let routed_terminal = strip_generic_provider_prefix(raw_ref); + let unknown_router_path = routed_terminal.is_some() && !has_known_provider_prefix(raw_ref); + let direct_result = if unknown_router_path { + // An unknown router path gets only a full-key exact lookup here. + // Broad model-part overrides (Cursor/Sakana) must wait until the + // terminal candidate has tried upstream exact/provider-aware + // resolution. Known provider prefixes keep their established + // provider-ranking path. + self.lookup_exact_full(raw_ref, force_source, provider_id, sakana_provider_id) + } else { + do_lookup(lower_ref) + }; + if let Some(result) = direct_result.filter(|result| !unsafe_claude_resolution(result)) { return Some(result); } - if parse_provider_scoped_model_path(lower_ref).is_some() { + // Provider-scoped paths are intentionally fail-closed after their + // provider-aware direct lookup. Never route them through a generic + // terminal or custom fallback. + if parse_provider_scoped_model_path(raw_ref).is_some() + || parse_provider_scoped_model_path(lower_ref).is_some() + { return None; } + let original_allows_sakana = sakana_identity_matches(raw_ref, sakana_provider_id); let guarded_lookup = |candidate: &str| { - do_lookup(candidate).filter(|result| !unsafe_claude_resolution(result)) + do_lookup(candidate) + .filter(|result| !unsafe_claude_resolution(result)) + .filter(|result| { + !result.source.eq_ignore_ascii_case("Sakana") || original_allows_sakana + }) + }; + let guarded_lookup_ref = &guarded_lookup; + let guarded_full_path_lookup = |candidate: &str| { + let result = if unknown_router_path { + self.lookup_exact_full(candidate, force_source, provider_id, sakana_provider_id) + } else { + do_lookup(candidate) + }; + result + .filter(|result| !unsafe_claude_resolution(result)) + .filter(|result| { + !result.source.eq_ignore_ascii_case("Sakana") || original_allows_sakana + }) }; - // 2. Try stripping unknown suffixes (e.g., -thinking, -high, -codex) - if let Some(result) = try_strip_unknown_suffix(lower_ref, guarded_lookup) { + // Preserve the full-path suffix pass before any generic terminal + // fallback. For example, `azure_ai/grok-code-fast-1-high` must first + // try `azure_ai/grok-code-fast-1`, rather than allowing terminal fuzzy + // matching to choose another provider's entry. + if let Some(result) = try_strip_unknown_suffix(raw_ref, guarded_full_path_lookup) { return Some(result); } - // 3. Try stripping unknown prefixes (e.g., antigravity-, myplugin-) - // For each prefix candidate, also try suffix stripping - if let Some(result) = try_strip_unknown_prefix(lower_ref, guarded_lookup) { + // Generic provider-routing prefix fallback. Normalize the terminal + // independently because the outer alias pass cannot see `cx/k2p6`. + // Raw custom keys precede alias-normalized custom keys; built-in lookup + // then sees the same normalized candidate as a direct lookup. + if let Some(terminal) = routed_terminal { + let terminal_lookup = |candidate: &str| { + let normalized = normalize_terminal_candidate(candidate)?; + if let Some(result) = guarded_custom(candidate) { + return Some(result); + } + if normalized.as_str() != candidate { + if let Some(result) = guarded_custom(&normalized) { + return Some(result); + } + } + guarded_lookup(&normalized) + }; + + if let Some(result) = terminal_lookup(terminal) { + return Some(result); + } + if let Some(result) = try_strip_unknown_suffix(terminal, terminal_lookup) { + return Some(result); + } + } + + // Try stripping unknown prefixes (e.g., antigravity-, myplugin-), with + // each candidate also receiving the bounded suffix pass. + if let Some(result) = try_strip_unknown_prefix(raw_ref, guarded_lookup_ref) { return Some(result); } None } - fn lookup_auto(&self, model_id: &str, provider_id: Option<&str>) -> Option { + /// Resolve only the complete raw model id, without model-part, prefix, + /// fuzzy, or suffix fallback. Provider-scoped paths are allowed to use + /// their dedicated provider-aware exact route before the global fail-closed + /// boundary. + fn lookup_exact_full( + &self, + model_id: &str, + force_source: Option<&str>, + provider_id: Option<&str>, + sakana_provider_id: Option<&str>, + ) -> Option { + if parse_provider_scoped_model_path(model_id).is_some() { + return match force_source { + Some("litellm") => self.lookup_provider_scoped_path_litellm(model_id, provider_id), + Some("openrouter") => { + self.lookup_provider_scoped_path_openrouter(model_id, provider_id) + } + Some("custom") => None, + _ => self.lookup_provider_scoped_path(model_id, provider_id), + }; + } + + let exact_litellm = || { + self.exact_match_litellm_full_for_provider(model_id, provider_id) + .or_else(|| self.exact_match_litellm_full(model_id)) + }; + let exact_openrouter = || { + self.exact_match_openrouter_full_for_provider(model_id, provider_id) + .or_else(|| self.exact_match_openrouter_full(model_id)) + }; + + match force_source { + Some("litellm") => exact_litellm(), + Some("openrouter") => exact_openrouter(), + Some("custom") => None, + _ => choose_best_source_result(exact_litellm(), exact_openrouter(), provider_id) + .or_else(|| self.exact_match_cursor_full(model_id)) + .or_else(|| self.exact_match_sakana_full(model_id, sakana_provider_id)), + } + } + + fn lookup_auto( + &self, + model_id: &str, + provider_id: Option<&str>, + sakana_provider_id: Option<&str>, + ) -> Option { if let Some(result) = self.lookup_provider_scoped_path(model_id, provider_id) { return Some(result); } @@ -453,6 +682,15 @@ impl PricingLookup { } } + if let Some(result) = self.exact_match_sakana(model_id, sakana_provider_id) { + return Some(result); + } + if let Some(version_normalized) = normalize_version_separator(model_id) { + if let Some(result) = self.exact_match_sakana(&version_normalized, sakana_provider_id) { + return Some(result); + } + } + if !is_fuzzy_eligible(model_id) { return None; } @@ -687,6 +925,44 @@ impl PricingLookup { lookup_result_if_usable(pricing, "LiteLLM", key) } + fn exact_match_litellm_full(&self, model_id: &str) -> Option { + self.exact_match_litellm(model_id) + } + + fn exact_match_litellm_full_for_provider( + &self, + model_id: &str, + provider_id: Option<&str>, + ) -> Option { + let result = self.exact_match_litellm_full(model_id)?; + if let Some(provider) = provider_id { + if !provider_identity::matches_provider_hint(&result.matched_key, Some(provider)) { + return None; + } + } + Some(result) + } + + fn exact_match_openrouter_full(&self, model_id: &str) -> Option { + let key = self.openrouter_lower.get(model_id)?; + let pricing = self.openrouter.get(key)?; + lookup_result_if_usable(pricing, "OpenRouter", key) + } + + fn exact_match_openrouter_full_for_provider( + &self, + model_id: &str, + provider_id: Option<&str>, + ) -> Option { + let result = self.exact_match_openrouter_full(model_id)?; + if let Some(provider) = provider_id { + if !provider_identity::matches_provider_hint(&result.matched_key, Some(provider)) { + return None; + } + } + Some(result) + } + fn exact_match_openrouter(&self, model_id: &str) -> Option { if let Some(key) = self.openrouter_lower.get(model_id) { if let Some(pricing) = self.openrouter.get(key) { @@ -701,6 +977,11 @@ impl PricingLookup { None } + fn exact_match_cursor_full(&self, model_id: &str) -> Option { + let key = self.cursor_lower.get(model_id)?; + lookup_result_if_usable(self.cursor.get(key)?, "Cursor", key) + } + fn exact_match_cursor(&self, model_id: &str) -> Option { if let Some(key) = self.cursor_lower.get(model_id) { return lookup_result_if_usable(self.cursor.get(key).unwrap(), "Cursor", key); @@ -715,6 +996,39 @@ impl PricingLookup { None } + fn exact_match_sakana_full( + &self, + model_id: &str, + provider_id: Option<&str>, + ) -> Option { + if !sakana_identity_matches(model_id, provider_id) { + return None; + } + let key = self.sakana_lower.get(model_id)?; + lookup_result_if_usable(self.sakana.get(key)?, "Sakana", key) + } + + fn exact_match_sakana( + &self, + model_id: &str, + provider_id: Option<&str>, + ) -> Option { + if !sakana_identity_matches(model_id, provider_id) { + return None; + } + if let Some(key) = self.sakana_lower.get(model_id) { + return lookup_result_if_usable(self.sakana.get(key).unwrap(), "Sakana", key); + } + if let Some(model_part) = model_id.split('/').next_back() { + if model_part != model_id { + if let Some(key) = self.sakana_lower.get(model_part) { + return lookup_result_if_usable(self.sakana.get(key).unwrap(), "Sakana", key); + } + } + } + None + } + fn prefix_match_litellm( &self, model_id: &str, @@ -858,15 +1172,46 @@ impl PricingLookup { None => return 0.0, }; - compute_cost( + compute_cost_for_lookup_result(&result, usage) + } +} + +pub(crate) fn compute_cost_for_lookup_result(result: &LookupResult, usage: &TokenBreakdown) -> f64 { + if uses_full_session_long_context_tier(result) { + return compute_full_session_long_context_cost( &result.pricing, usage.input, usage.output, usage.cache_read, usage.cache_write, usage.reasoning, - ) + ); + } + + compute_cost( + &result.pricing, + usage.input, + usage.output, + usage.cache_read, + usage.cache_write, + usage.reasoning, + ) +} + +fn uses_full_session_long_context_tier(result: &LookupResult) -> bool { + if result.source.eq_ignore_ascii_case("Sakana") { + return result.matched_key.eq_ignore_ascii_case("fugu-ultra"); } + + let terminal_model_id = result + .matched_key + .rsplit('/') + .next() + .unwrap_or(result.matched_key.as_str()); + result.source.eq_ignore_ascii_case("LiteLLM") + && FULL_SESSION_LONG_CONTEXT_LITELLM_KEYS + .iter() + .any(|key| terminal_model_id.eq_ignore_ascii_case(key)) } pub fn compute_cost( @@ -987,6 +1332,58 @@ pub fn compute_cost( input_cost + output_cost + cache_read_cost + cache_write_cost } +/// Apply a provider-documented long-context tier to the whole request. The +/// selector includes input plus cache-read tokens; output and reasoning use the +/// selected tier, while cache-write is priced independently and never selects it. +fn compute_full_session_long_context_cost( + pricing: &ModelPricing, + input: i64, + output: i64, + cache_read: i64, + cache_write: i64, + reasoning: i64, +) -> f64 { + let safe_price = |opt: Option| opt.filter(|v| is_valid_price_value(*v)).unwrap_or(0.0); + let input_clamped = input.max(0) as f64; + let output_clamped = output.max(0).saturating_add(reasoning.max(0)) as f64; + let cache_read_clamped = cache_read.max(0) as f64; + let use_long_context_rates = + input_clamped + cache_read_clamped > TIERED_PRICING_THRESHOLD_272K_TOKENS; + let selected_price = |base: Option, long_context: Option| { + if use_long_context_rates { + safe_price( + long_context + .filter(|value| is_valid_price_value(*value)) + .or(base), + ) + } else { + safe_price(base) + } + }; + + let input_cost = input_clamped + * selected_price( + pricing.input_cost_per_token, + pricing.input_cost_per_token_above_272k_tokens, + ); + let output_cost = output_clamped + * selected_price( + pricing.output_cost_per_token, + pricing.output_cost_per_token_above_272k_tokens, + ); + let cache_read_cost = cache_read_clamped + * selected_price( + pricing.cache_read_input_token_cost, + pricing.cache_read_input_token_cost_above_272k_tokens, + ); + + // Cache-write is an independently reported subset and does not select the + // request's input-context tier. + let cache_write_cost = compute_cost(pricing, 0, 0, 0, cache_write, 0); + + input_cost + output_cost + cache_read_cost + cache_write_cost +} + fn extract_model_family(model_id: &str) -> String { let lower = model_id.to_lowercase(); @@ -1317,6 +1714,60 @@ fn strip_known_provider_prefix(model_id: &str) -> Option<&str> { None } +fn has_known_provider_prefix(model_id: &str) -> bool { + let Some((prefix, _)) = model_id.split_once('/') else { + return false; + }; + PROVIDER_PREFIXES + .iter() + .chain(ORIGINAL_PROVIDER_PREFIXES) + .chain(RESELLER_PROVIDER_PREFIXES) + .any(|candidate| candidate.trim_end_matches('/') == prefix) +} + +/// Generic routing-prefix fallback for ids whose leading segment is not one +/// of the curated `PROVIDER_PREFIXES` (e.g. `cx/gpt-5.5` routed through an +/// `omniroute` proxy, or any other CLI/router-assigned alias). Returns the +/// terminal path segment — the part after the last `/` — when the id actually +/// contains a `/`, so `cx/gpt-5.5` resolves to `gpt-5.5`. +/// +/// This is intentionally unconditional (unlike `strip_known_provider_prefix`, +/// which only recognizes canonical LLM provider names): the caller only +/// invokes it as a fallback AFTER the exact/direct lookup on the full id has +/// already failed, so dataset keys that legitimately keep their prefix (e.g. +/// `anthropic/claude-fable-5`) are resolved by their own exact key first and +/// never reach this fallback. +fn strip_generic_provider_prefix(model_id: &str) -> Option<&str> { + let terminal = model_id.rsplit('/').next()?; + if terminal.is_empty() || terminal == model_id { + return None; + } + Some(terminal) +} + +/// Apply the same bounded reasoning-tier and static-alias normalization used +/// by the outer lookup to a generic terminal candidate. Returning `None` for +/// an invalid parenthesized tier keeps terminal fallback fail-closed instead of +/// peeling the fragment through the dash-suffix helper. +fn normalize_terminal_candidate(candidate: &str) -> Option { + let normalized_owned = strip_parenthesized_reasoning_tier(candidate).map(str::to_owned); + if normalized_owned.is_none() + && candidate + .strip_suffix(')') + .and_then(|inner| inner.rsplit_once('(')) + .is_some() + { + return None; + } + + let normalized = normalized_owned.as_deref().unwrap_or(candidate); + Some( + aliases::resolve_alias(normalized) + .unwrap_or(normalized) + .to_string(), + ) +} + fn is_valid_price_value(value: f64) -> bool { value.is_finite() && value >= 0.0 } @@ -1683,6 +2134,21 @@ fn select_best_match( pick(effective_matches, hint_is_reseller) } +fn sakana_identity_matches(model_id: &str, provider_id: Option<&str>) -> bool { + let only_sakana_tags = |identity: &str| { + let tags = provider_identity::provider_tags(identity); + !tags.is_empty() && tags.iter().all(|tag| tag == "sakana") + }; + + if let Some((path, _)) = model_id.rsplit_once('/') { + if path.split('/').any(|segment| !only_sakana_tags(segment)) { + return false; + } + } + + provider_id.is_none_or(|provider_id| provider_id.split('/').all(only_sakana_tags)) +} + fn model_prefix_matches_provider(model_id: &str, provider_id: Option<&str>) -> bool { let Some(hint) = provider_id else { return true; @@ -1851,10 +2317,13 @@ fn backfill_cache_costs(mut winner: LookupResult, donor: &ModelPricing) -> Looku donor.cache_read_input_token_cost_above_272k_tokens; } } - if p.cache_creation_input_token_cost.is_none() && donor.cache_creation_input_token_cost.is_some() + if p.cache_creation_input_token_cost.is_none() + && donor.cache_creation_input_token_cost.is_some() { p.cache_creation_input_token_cost = donor.cache_creation_input_token_cost; - if p.cache_creation_input_token_cost_above_200k_tokens.is_none() { + if p.cache_creation_input_token_cost_above_200k_tokens + .is_none() + { p.cache_creation_input_token_cost_above_200k_tokens = donor.cache_creation_input_token_cost_above_200k_tokens; } @@ -2843,6 +3312,53 @@ mod tests { ); } + #[test] + fn antigravity_model_aliases_reach_priced_catalog_entries() { + let priced = [ + "gemini-3.1-pro", + "gemini-3-flash-preview", + "gemini-3.5-flash-high", + "gemini-3.5-flash-extra-low", + "gemini-3.5-flash-medium", + ]; + let mut litellm = HashMap::new(); + for model in priced { + litellm.insert( + model.to_string(), + ModelPricing { + input_cost_per_token: Some(0.000002), + output_cost_per_token: Some(0.000012), + cache_read_input_token_cost: Some(0.0000002), + ..Default::default() + }, + ); + } + let lookup = PricingLookup::new(litellm, HashMap::new(), HashMap::new()); + + let cases = [ + ("MODEL_PLACEHOLDER_M16", "gemini-3.1-pro"), + ("MODEL_PLACEHOLDER_M84", "gemini-3-flash-preview"), + ("MODEL_PLACEHOLDER_M133", "gemini-3.5-flash-high"), + ("gemini-3-flash-agent", "gemini-3.5-flash-high"), + ("gemini-3-flash-b", "gemini-3.5-flash-high"), + ("gemini-3-flash-a", "gemini-3.5-flash-high"), + ("MODEL_PLACEHOLDER_M187", "gemini-3.5-flash-extra-low"), + ("MODEL_PLACEHOLDER_M20", "gemini-3.5-flash-medium"), + ]; + + for (raw, expected_key) in cases { + let result = lookup + .lookup(raw) + .unwrap_or_else(|| panic!("unpriced alias: {raw}")); + assert_eq!(result.matched_key, expected_key, "raw model: {raw}"); + assert!(result.pricing.input_cost_per_token.is_some()); + } + + assert!( + lookup.calculate_cost("gemini-3-flash-agent", 1_000_000, 100_000, 50_000, 0, 0) > 0.0 + ); + } + // ========================================================================= // BASELINE / LEGACY TESTS // ========================================================================= @@ -2994,6 +3510,40 @@ mod tests { assert!(lookup.lookup("gemini-3-pro(weird)").is_none()); } + #[test] + fn test_exact_raw_built_in_parenthesized_key_precedes_tier_validation() { + let mut litellm = HashMap::new(); + litellm.insert( + "gpt-5.5(high)".into(), + ModelPricing { + input_cost_per_token: Some(0.00001), + output_cost_per_token: Some(0.00003), + ..Default::default() + }, + ); + let lookup = PricingLookup::new(litellm, HashMap::new(), HashMap::new()); + + let result = lookup.lookup("GPT-5.5(HIGH)").unwrap(); + assert_eq!(result.source, "LiteLLM"); + assert_eq!(result.matched_key, "gpt-5.5(high)"); + } + + #[test] + fn test_missing_exact_raw_invalid_parenthesized_built_in_key_fails_closed() { + let mut litellm = HashMap::new(); + litellm.insert( + "gpt-5.5".into(), + ModelPricing { + input_cost_per_token: Some(0.00001), + output_cost_per_token: Some(0.00003), + ..Default::default() + }, + ); + let lookup = PricingLookup::new(litellm, HashMap::new(), HashMap::new()); + + assert!(lookup.lookup("gpt-5.5(invalid)").is_none()); + } + #[test] fn test_parenthesized_reasoning_tier_cost_matches_base_model() { let lookup = create_lookup(); @@ -3635,21 +4185,25 @@ mod tests { #[test] fn test_force_source_litellm() { let lookup = create_lookup(); - let result = lookup - .lookup_with_source("gpt-4o", Some("litellm")) - .unwrap(); - assert_eq!(result.source, "LiteLLM"); - assert_eq!(result.matched_key, "gpt-4o"); + for force_source in ["litellm", "LiteLLM", "LITELLM"] { + let result = lookup + .lookup_with_source("gpt-4o", Some(force_source)) + .unwrap(); + assert_eq!(result.source, "LiteLLM"); + assert_eq!(result.matched_key, "gpt-4o"); + } } #[test] fn test_force_source_openrouter() { let lookup = create_lookup(); - let result = lookup - .lookup_with_source("gpt-4o", Some("openrouter")) - .unwrap(); - assert_eq!(result.source, "OpenRouter"); - assert_eq!(result.matched_key, "openai/gpt-4o"); + for force_source in ["openrouter", "OpenRouter", "OPENROUTER"] { + let result = lookup + .lookup_with_source("gpt-4o", Some(force_source)) + .unwrap(); + assert_eq!(result.source, "OpenRouter"); + assert_eq!(result.matched_key, "openai/gpt-4o"); + } } #[test] @@ -3680,6 +4234,202 @@ mod tests { assert!(lookup.lookup("nonexistent-model-xyz").is_none()); } + /// Regression (#831): router/proxy-assigned ids like `cx/gpt-5.5` (seen + /// from OpenCode's `omniroute` provider) carry a prefix outside the + /// curated `PROVIDER_PREFIXES` list, so the pricing lookup used to return + /// `None` instead of pricing the underlying `gpt-5.5` model. + #[test] + fn test_unknown_prefixed_model_id_strips_to_underlying_model() { + let lookup = create_lookup(); + let direct = lookup.lookup("gpt-5.5").unwrap(); + let prefixed = lookup.lookup("cx/gpt-5.5").unwrap(); + assert_eq!(prefixed.matched_key, direct.matched_key); + assert_eq!(prefixed.source, direct.source); + assert_eq!( + prefixed.pricing.input_cost_per_token, + direct.pricing.input_cost_per_token + ); + assert_eq!( + prefixed.pricing.output_cost_per_token, + direct.pricing.output_cost_per_token + ); + } + + /// Regression (#831): a dataset key that legitimately keeps its own + /// provider prefix (e.g. `anthropic/claude-fable-5`) must still resolve + /// through the exact/direct lookup before the generic fallback runs. + #[test] + fn test_known_prefixed_dataset_key_still_resolves_exactly() { + let lookup = claude_family_fixture(); + let result = lookup.lookup("anthropic/claude-fable-5").unwrap(); + assert_eq!(result.matched_key, "anthropic/claude-fable-5"); + } + + /// Regression (#831): an unrecognized prefix and an unrecognized + /// underlying model must still return `None` rather than fuzzy-matching an + /// unrelated pricing key. + #[test] + fn test_unknown_prefixed_unknown_model_stays_none() { + let lookup = create_lookup(); + assert!(lookup.lookup("unknown/nonexistent").is_none()); + } + + /// The generic fallback must reuse the local guarded pipeline, including + /// provider-hint source selection and cache-rate backfill, rather than + /// bypassing those local adaptations when composing a terminal suffix + /// fallback. + #[test] + fn test_generic_prefix_fallback_keeps_provider_selection_and_cache_backfill() { + let pricing = |input: f64, output: f64| ModelPricing { + input_cost_per_token: Some(input), + output_cost_per_token: Some(output), + ..Default::default() + }; + + let mut litellm = HashMap::new(); + litellm.insert("alpha-model".into(), pricing(0.000001, 0.000002)); + litellm.insert("anthropic/alpha-model".into(), pricing(0.000003, 0.000004)); + + let mut openrouter = HashMap::new(); + openrouter.insert( + "anthropic/alpha-model".into(), + ModelPricing { + input_cost_per_token: Some(0.000005), + output_cost_per_token: Some(0.000006), + cache_read_input_token_cost: Some(0.0000007), + cache_creation_input_token_cost: Some(0.0000008), + ..Default::default() + }, + ); + + let lookup = PricingLookup::new(litellm, openrouter, HashMap::new()); + let direct = lookup + .lookup_with_provider("alpha-model-high", Some("anthropic")) + .unwrap(); + let prefixed = lookup + .lookup_with_provider("cx/alpha-model-high", Some("anthropic")) + .unwrap(); + + for result in [&direct, &prefixed] { + assert_eq!(result.matched_key, "anthropic/alpha-model"); + assert_eq!(result.source, "LiteLLM"); + assert_eq!(result.pricing.input_cost_per_token, Some(0.000003)); + assert_eq!(result.pricing.output_cost_per_token, Some(0.000004)); + assert_eq!(result.pricing.cache_read_input_token_cost, Some(0.0000007)); + assert_eq!( + result.pricing.cache_creation_input_token_cost, + Some(0.0000008) + ); + } + } + + /// Regression: an existing suffixed pricing fixture must resolve through + /// the terminal segment before the older bounded prefix fallback runs. + #[test] + fn test_generic_prefix_fallback_composes_suffixed_pricing_fixture() { + let lookup = create_lookup(); + let direct = lookup.lookup("gpt-5-codex-max").unwrap(); + let routed = lookup.lookup("cx-router-edge/gpt-5-codex-max").unwrap(); + + assert_eq!(routed.matched_key, direct.matched_key); + assert_eq!(routed.source, direct.source); + assert_eq!( + routed.pricing.input_cost_per_token, + direct.pricing.input_cost_per_token + ); + assert_eq!( + routed.pricing.output_cost_per_token, + direct.pricing.output_cost_per_token + ); + assert_eq!( + routed.pricing.cache_read_input_token_cost, + direct.pricing.cache_read_input_token_cost + ); + assert_eq!( + routed.pricing.cache_creation_input_token_cost, + direct.pricing.cache_creation_input_token_cost + ); + } + + #[test] + fn test_generic_prefix_terminal_upstream_exact_precedes_cursor_model_part() { + let pricing = |input: f64| ModelPricing { + input_cost_per_token: Some(input), + ..Default::default() + }; + let mut litellm = HashMap::new(); + litellm.insert("gpt-5.3-codex".into(), pricing(0.000001)); + let mut cursor = HashMap::new(); + cursor.insert("gpt-5.3-codex".into(), pricing(0.000009)); + + let lookup = PricingLookup::new(litellm, HashMap::new(), cursor); + for model_id in ["gpt-5.3-codex", "cx/gpt-5.3-codex", "cx/gpt-5.3-codex-high"] { + let result = lookup.lookup(model_id).unwrap(); + assert_eq!(result.source, "LiteLLM"); + assert_eq!(result.matched_key, "gpt-5.3-codex"); + assert_eq!(result.pricing.input_cost_per_token, Some(0.000001)); + } + } + + #[test] + fn test_generic_prefix_fallback_unknown_terminal_suffix_stays_none() { + let lookup = create_lookup(); + assert!(lookup.lookup("cx/nonexistent-model-high").is_none()); + } + + #[test] + fn test_generic_prefix_fallback_preserves_full_path_suffix_precedence() { + let pricing = |input: f64| ModelPricing { + input_cost_per_token: Some(input), + ..Default::default() + }; + let mut litellm = HashMap::new(); + litellm.insert("azure_ai/grok-code-fast-1".into(), pricing(0.0035)); + litellm.insert("xai/grok-code-fast-1-0825".into(), pricing(0.0000002)); + + let lookup = PricingLookup::new(litellm, HashMap::new(), HashMap::new()); + let result = lookup.lookup("azure_ai/grok-code-fast-1-high").unwrap(); + + assert_eq!(result.matched_key, "azure_ai/grok-code-fast-1"); + assert_eq!(result.pricing.input_cost_per_token, Some(0.0035)); + } + + #[test] + fn test_generic_prefix_fallback_resolves_terminal_alias_like_direct_lookup() { + let lookup = create_lookup(); + let direct = lookup.lookup("k2p6").unwrap(); + let routed = lookup.lookup("cx/k2p6").unwrap(); + + assert_eq!(routed.matched_key, direct.matched_key); + assert_eq!(routed.source, direct.source); + assert_eq!( + routed.pricing.input_cost_per_token, + direct.pricing.input_cost_per_token + ); + assert_eq!( + routed.pricing.output_cost_per_token, + direct.pricing.output_cost_per_token + ); + } + + #[test] + fn test_generic_prefix_fallback_applies_terminal_reasoning_tier_guard() { + let lookup = create_lookup(); + let direct = lookup.lookup("cx/k2p6").unwrap(); + let tiered = lookup.lookup("cx/k2p6(high)").unwrap(); + + assert_eq!(tiered.matched_key, direct.matched_key); + assert!(lookup.lookup("cx/k2p6(invalid)").is_none()); + } + + #[test] + fn test_generic_prefix_helper_uses_terminal_segment() { + assert_eq!(strip_generic_provider_prefix("cx/gpt-5.5"), Some("gpt-5.5")); + assert_eq!(strip_generic_provider_prefix("a/b/c"), Some("c")); + assert_eq!(strip_generic_provider_prefix("gpt-5.5"), None); + assert_eq!(strip_generic_provider_prefix("cx/"), None); + } + #[test] fn test_fallback_suffix_lookup() { // Create a lookup with only the base model (no -codex variant) @@ -4989,4 +5739,167 @@ mod tests { "unknown hint via source_and_provider should behave like None" ); } + + #[test] + fn test_gpt_5_5_uses_full_request_long_context_rates() { + let lookup = create_lookup(); + let cost = lookup.calculate_cost("gpt-5.5", 270_540, 630, 7_936, 0, 0); + let expected = 270_540.0 * 0.000010 + 630.0 * 0.000045 + 7_936.0 * 0.000001; + assert!((cost - expected).abs() < 1e-12); + } + + #[test] + fn test_provider_prefixed_gpt_5_5_uses_full_request_long_context_rates() { + let mut litellm = HashMap::new(); + litellm.insert( + "openai/gpt-5.5".into(), + ModelPricing { + input_cost_per_token: Some(0.000005), + input_cost_per_token_above_272k_tokens: Some(0.000010), + output_cost_per_token: Some(0.000030), + output_cost_per_token_above_272k_tokens: Some(0.000045), + cache_read_input_token_cost: Some(0.0000005), + cache_read_input_token_cost_above_272k_tokens: Some(0.000001), + ..Default::default() + }, + ); + let lookup = PricingLookup::new(litellm, HashMap::new(), HashMap::new()); + let usage = TokenBreakdown { + input: 271_999, + output: 3, + cache_read: 2, + cache_write: 0, + reasoning: 4, + }; + + let cost = lookup.calculate_cost_with_provider("openai/gpt-5.5", Some("openai"), &usage); + let expected = 271_999.0 * 0.000010 + 7.0 * 0.000045 + 2.0 * 0.000001; + assert!((cost - expected).abs() < 1e-12); + } + + #[test] + fn test_gpt_5_5_threshold_includes_cache_read_and_reasoning() { + let lookup = create_lookup(); + let cost = lookup.calculate_cost("gpt-5.5", 271_999, 3, 2, 0, 4); + let expected = 271_999.0 * 0.000010 + 7.0 * 0.000045 + 2.0 * 0.000001; + assert!((cost - expected).abs() < 1e-12); + } + + #[test] + fn test_gpt_5_5_exact_threshold_stays_regular() { + let lookup = create_lookup(); + let cost = lookup.calculate_cost("gpt-5.5", 272_000, 3, 0, 0, 4); + let expected = 272_000.0 * 0.000005 + 7.0 * 0.000030; + assert!((cost - expected).abs() < 1e-12); + } + + #[test] + fn test_gpt_5_5_cache_write_does_not_select_long_context() { + let lookup = create_lookup(); + let without_cache_write = lookup.calculate_cost("gpt-5.5", 271_999, 10, 0, 0, 0); + let with_cache_write = lookup.calculate_cost("gpt-5.5", 271_999, 10, 0, 10_000, 0); + assert!((with_cache_write - without_cache_write).abs() < 1e-12); + } + + #[test] + fn test_full_request_long_context_policy_is_identity_scoped() { + let pricing = ModelPricing { + input_cost_per_token: Some(0.000005), + input_cost_per_token_above_272k_tokens: Some(0.000010), + ..Default::default() + }; + assert!(uses_full_session_long_context_tier(&lookup_result( + "gpt-5.5", + "LiteLLM", + pricing.clone() + ))); + assert!(uses_full_session_long_context_tier(&lookup_result( + "openai/gpt-5.5", + "LiteLLM", + pricing.clone() + ))); + assert!(uses_full_session_long_context_tier(&lookup_result( + "fugu-ultra", + "Sakana", + pricing.clone() + ))); + for (source, key) in [ + ("LiteLLM", "gpt-5.5-pro"), + ("OpenRouter", "gpt-5.5"), + ("Custom", "gpt-5.5"), + ("Sakana", "fugu"), + ] { + assert!(!uses_full_session_long_context_tier(&lookup_result( + key, + source, + pricing.clone() + ))); + } + } + + #[test] + fn test_non_verified_identity_keeps_marginal_tier_behavior() { + let result = lookup_result( + "gpt-5.5-pro", + "LiteLLM", + ModelPricing { + input_cost_per_token: Some(0.000005), + input_cost_per_token_above_272k_tokens: Some(0.000010), + output_cost_per_token: Some(0.000030), + output_cost_per_token_above_272k_tokens: Some(0.000045), + ..Default::default() + }, + ); + let cost = compute_cost_for_lookup_result( + &result, + &TokenBreakdown { + input: 272_001, + output: 1, + cache_read: 0, + cache_write: 0, + reasoning: 0, + }, + ); + let expected = 272_000.0 * 0.000005 + 1.0 * 0.000010 + 0.000030; + assert!((cost - expected).abs() < 1e-12); + } + + #[test] + fn test_every_terminal_custom_result_uses_claude_never_degrade_guard() { + let lookup = create_lookup(); + let mismatched = || { + Some(lookup_result( + "claude-opus-4", + "Custom", + ModelPricing { + input_cost_per_token: Some(0.000015), + output_cost_per_token: Some(0.000075), + ..Default::default() + }, + )) + }; + + for id in [ + "cx/claude-opus-4-7-thinking", + "cx/claude-sonnet-4-6-thinking", + "cx/claude-opus-4-60-thinking", + ] { + assert!( + lookup + .lookup_with_source_and_provider_and_terminal_custom(id, None, None, |_| { + mismatched() + }) + .is_none(), + "unsafe custom fallback resolved: {id}" + ); + } + assert!(lookup + .lookup_with_source_and_provider_and_terminal_custom( + "cx/claude-opus-4-7(invalid)", + None, + None, + |_| mismatched(), + ) + .is_none()); + } } diff --git a/vendor/tokscale-core/src/pricing/mod.rs b/vendor/tokscale-core/src/pricing/mod.rs index 55b6964..bb8016d 100644 --- a/vendor/tokscale-core/src/pricing/mod.rs +++ b/vendor/tokscale-core/src/pricing/mod.rs @@ -6,7 +6,7 @@ pub mod lookup; pub mod openrouter; use custom::CustomPricing; -use lookup::{compute_cost, LookupResult, PricingLookup}; +use lookup::{compute_cost_for_lookup_result, LookupResult, PricingLookup}; use std::collections::HashMap; use std::sync::Arc; use std::time::{Duration, Instant}; @@ -65,10 +65,11 @@ impl PricingService { ) -> Self { Self { custom, - lookup: PricingLookup::new( + lookup: PricingLookup::new_with_sakana( litellm_data, openrouter_data, Self::build_cursor_overrides(), + Self::build_sakana_overrides(), ), } } @@ -138,6 +139,27 @@ impl PricingService { overrides } + // Sakana's Fugu Ultra is not present in the upstream catalogs. Bare `fugu` + // remains intentionally unpriced because its router rate depends on the + // underlying model selected for each request. + fn build_sakana_overrides() -> HashMap { + let mut overrides = HashMap::with_capacity(1); + overrides.insert( + "fugu-ultra".to_string(), + ModelPricing { + input_cost_per_token: Some(5e-6), + output_cost_per_token: Some(3e-5), + cache_read_input_token_cost: Some(5e-7), + cache_creation_input_token_cost: None, + input_cost_per_token_above_272k_tokens: Some(1e-5), + output_cost_per_token_above_272k_tokens: Some(4.5e-5), + cache_read_input_token_cost_above_272k_tokens: Some(1e-6), + ..Default::default() + }, + ); + overrides + } + async fn fetch_inner() -> Result { let (litellm_result, openrouter_data) = tokio::join!(litellm::fetch(), openrouter::fetch_all_mapped()); @@ -219,19 +241,7 @@ impl PricingService { model_id: &str, force_source: Option<&str>, ) -> Option { - match force_source { - Some(source) if source.eq_ignore_ascii_case("custom") => { - return self.lookup_custom(model_id); - } - None => { - if let Some(result) = self.lookup_custom(model_id) { - return Some(result); - } - } - Some(_) => {} - } - - self.lookup.lookup_with_source(model_id, force_source) + self.lookup_with_source_and_provider(model_id, force_source, None) } pub fn lookup_with_source_and_provider( @@ -240,18 +250,38 @@ impl PricingService { force_source: Option<&str>, provider_id: Option<&str>, ) -> Option { - match force_source { - Some(source) if source.eq_ignore_ascii_case("custom") => { - return self.lookup_custom(model_id); + let force_custom = force_source.is_some_and(|source| source.eq_ignore_ascii_case("custom")); + + if force_custom { + // `PricingLookup` has no custom dataset of its own. Keep forced + // custom lookups fail-closed, while still allowing the same full-id + // and terminal custom seam used by automatic lookup. + if self.custom.is_empty() { + return None; } - None => { - if let Some(result) = self.lookup_custom(model_id) { - return Some(result); - } - } - Some(_) => {} + return self + .lookup + .lookup_with_source_and_provider_and_terminal_custom( + model_id, + Some("custom"), + provider_id, + |candidate| self.lookup_custom(candidate), + ); } + if force_source.is_none() && !self.custom.is_empty() { + return self + .lookup + .lookup_with_source_and_provider_and_terminal_custom( + model_id, + None, + provider_id, + |candidate| self.lookup_custom(candidate), + ); + } + + // With no custom data, keep the normal cached PricingLookup path. For + // forced upstream sources, never allow a custom callback to participate. self.lookup .lookup_with_source_and_provider(model_id, force_source, provider_id) } @@ -281,24 +311,22 @@ impl PricingService { provider_id: Option<&str>, usage: &TokenBreakdown, ) -> f64 { - if let Some(result) = self.custom.lookup_with_key(model_id) { - return compute_cost( - result.pricing, - usage.input, - usage.output, - usage.cache_read, - usage.cache_write, - usage.reasoning, - ); + if self.custom.is_empty() { + return self + .lookup + .calculate_cost_with_provider(model_id, provider_id, usage); } - self.lookup - .calculate_cost_with_provider(model_id, provider_id, usage) + let Some(result) = self.lookup_with_source_and_provider(model_id, None, provider_id) else { + return 0.0; + }; + + compute_cost_for_lookup_result(&result, usage) } fn lookup_custom(&self, model_id: &str) -> Option { self.custom - .lookup_with_key(model_id) + .lookup_exact_with_key(model_id) .map(|result| LookupResult { pricing: result.pricing.clone(), source: "Custom".into(), @@ -733,6 +761,215 @@ mod tests { assert_eq!(result.pricing.output_cost_per_token, Some(0.000008)); } + #[test] + fn custom_override_wins_over_built_in_terminal_for_routed_id() { + let mut custom = HashMap::new(); + custom.insert("gpt-5.5".into(), model_pricing(0.000002, 0.000008)); + let mut litellm = HashMap::new(); + litellm.insert("gpt-5.5".into(), model_pricing(0.00001, 0.00003)); + + let service = custom_service(custom, litellm, HashMap::new()); + let result = service.lookup_with_source("cx/gpt-5.5", None).unwrap(); + + assert_eq!(result.source, "Custom"); + assert_eq!(result.matched_key, "gpt-5.5"); + assert_eq!(result.pricing.input_cost_per_token, Some(0.000002)); + } + + #[test] + fn custom_terminal_raw_alias_key_matches_direct_and_forced_custom() { + let mut custom = HashMap::new(); + custom.insert("k2p6".into(), model_pricing(0.000002, 0.000008)); + let service = custom_service(custom, HashMap::new(), HashMap::new()); + + for force_source in [None, Some("custom")] { + let result = service.lookup_with_source("cx/k2p6", force_source).unwrap(); + assert_eq!(result.source, "Custom"); + assert_eq!(result.matched_key, "k2p6"); + assert_eq!(result.pricing.input_cost_per_token, Some(0.000002)); + } + } + + #[test] + fn custom_terminal_canonical_alias_key_matches_direct_and_forced_custom() { + let mut custom = HashMap::new(); + custom.insert("kimi-k2.6".into(), model_pricing(0.000003, 0.000012)); + let service = custom_service(custom, HashMap::new(), HashMap::new()); + + for force_source in [None, Some("custom")] { + let result = service.lookup_with_source("cx/k2p6", force_source).unwrap(); + assert_eq!(result.source, "Custom"); + assert_eq!(result.matched_key, "kimi-k2.6"); + assert_eq!(result.pricing.input_cost_per_token, Some(0.000003)); + } + } + + #[test] + fn custom_exact_parenthesized_key_wins_before_tier_stripping() { + let mut custom = HashMap::new(); + custom.insert("cx/gpt-5.5(high)".into(), model_pricing(0.000002, 0.000008)); + let service = custom_service(custom, HashMap::new(), HashMap::new()); + + for force_source in [None, Some("custom")] { + let result = service + .lookup_with_source("CX/GPT-5.5(HIGH)", force_source) + .unwrap(); + assert_eq!(result.source, "Custom"); + assert_eq!(result.matched_key, "cx/gpt-5.5(high)"); + } + } + + #[test] + fn custom_exact_invalid_parenthesized_key_wins_but_missing_key_fails_closed() { + let mut custom = HashMap::new(); + custom.insert( + "cx/not-a-real-model(invalid)".into(), + model_pricing(0.000002, 0.000008), + ); + let mut litellm = HashMap::new(); + litellm.insert("not-a-real-model".into(), model_pricing(0.00001, 0.00003)); + let service = custom_service(custom, litellm.clone(), HashMap::new()); + + for force_source in [None, Some("custom")] { + let result = service + .lookup_with_source("cx/not-a-real-model(invalid)", force_source) + .unwrap(); + assert_eq!(result.source, "Custom"); + assert_eq!(result.matched_key, "cx/not-a-real-model(invalid)"); + } + + let without_exact = custom_service(HashMap::new(), litellm, HashMap::new()); + for force_source in [None, Some("custom")] { + assert!(without_exact + .lookup_with_source("cx/not-a-real-model(invalid)", force_source) + .is_none()); + } + } + + #[test] + fn forced_upstream_sources_do_not_use_raw_custom_parenthesized_keys() { + let mut custom = HashMap::new(); + custom.insert("cx/gpt-5.5(high)".into(), model_pricing(0.000002, 0.000008)); + let mut litellm = HashMap::new(); + litellm.insert("gpt-5.5".into(), model_pricing(0.00001, 0.00003)); + let mut openrouter = HashMap::new(); + openrouter.insert("gpt-5.5".into(), model_pricing(0.000003, 0.000012)); + let service = custom_service(custom, litellm, openrouter); + + let litellm_result = service + .lookup_with_source("cx/gpt-5.5(high)", Some("litellm")) + .unwrap(); + assert_eq!(litellm_result.source, "LiteLLM"); + assert_eq!(litellm_result.matched_key, "gpt-5.5"); + + let openrouter_result = service + .lookup_with_source("cx/gpt-5.5(high)", Some("openrouter")) + .unwrap(); + assert_eq!(openrouter_result.source, "OpenRouter"); + assert_eq!(openrouter_result.matched_key, "gpt-5.5"); + } + + #[test] + fn custom_terminal_override_does_not_steal_full_id_exact() { + let mut custom = HashMap::new(); + custom.insert("gpt-5.5".into(), model_pricing(0.000002, 0.000008)); + let mut litellm = HashMap::new(); + litellm.insert("cx/gpt-5.5".into(), model_pricing(0.00001, 0.00003)); + + let service = custom_service(custom, litellm, HashMap::new()); + let result = service.lookup_with_source("cx/gpt-5.5", None).unwrap(); + + assert_eq!(result.source, "LiteLLM"); + assert_eq!(result.matched_key, "cx/gpt-5.5"); + assert_eq!(result.pricing.input_cost_per_token, Some(0.00001)); + } + + #[test] + fn custom_terminal_override_does_not_steal_full_path_suffix_exact() { + let mut custom = HashMap::new(); + custom.insert("grok-code-fast-1".into(), model_pricing(0.000002, 0.000008)); + let mut litellm = HashMap::new(); + litellm.insert( + "azure_ai/grok-code-fast-1".into(), + model_pricing(0.00001, 0.00003), + ); + + let service = custom_service(custom, litellm, HashMap::new()); + let result = service + .lookup_with_source("azure_ai/grok-code-fast-1-high", None) + .unwrap(); + + assert_eq!(result.source, "LiteLLM"); + assert_eq!(result.matched_key, "azure_ai/grok-code-fast-1"); + assert_eq!(result.pricing.input_cost_per_token, Some(0.00001)); + } + + #[test] + fn custom_terminal_override_respects_forced_sources() { + let mut custom = HashMap::new(); + custom.insert("gpt-5.5".into(), model_pricing(0.000002, 0.000008)); + let mut litellm = HashMap::new(); + litellm.insert("gpt-5.5".into(), model_pricing(0.00001, 0.00003)); + let mut openrouter = HashMap::new(); + openrouter.insert("gpt-5.5".into(), model_pricing(0.000003, 0.000012)); + + let service = custom_service(custom, litellm, openrouter); + + let litellm_result = service + .lookup_with_source("cx/gpt-5.5", Some("litellm")) + .unwrap(); + assert_eq!(litellm_result.source, "LiteLLM"); + assert_eq!(litellm_result.pricing.input_cost_per_token, Some(0.00001)); + + let openrouter_result = service + .lookup_with_source("cx/gpt-5.5", Some("openrouter")) + .unwrap(); + assert_eq!(openrouter_result.source, "OpenRouter"); + assert_eq!( + openrouter_result.pricing.input_cost_per_token, + Some(0.000003) + ); + + let custom_result = service + .lookup_with_source("cx/gpt-5.5", Some("custom")) + .unwrap(); + assert_eq!(custom_result.source, "Custom"); + assert_eq!(custom_result.pricing.input_cost_per_token, Some(0.000002)); + } + + #[test] + fn custom_terminal_override_does_not_bypass_provider_scoped_fail_closed() { + let mut terminal_custom = HashMap::new(); + terminal_custom.insert("deepseek-v4-pro".into(), model_pricing(0.000002, 0.000008)); + let service = custom_service(terminal_custom, HashMap::new(), HashMap::new()); + + for scope in ["models", "routers"] { + assert!(service + .lookup_with_source(&format!("accounts/fireworks/{scope}/deepseek-v4-pro"), None,) + .is_none()); + } + + let full_key = "accounts/fireworks/models/deepseek-v4-pro"; + let mut exact_custom = HashMap::new(); + exact_custom.insert(full_key.into(), model_pricing(0.000003, 0.000012)); + let exact_service = custom_service(exact_custom, HashMap::new(), HashMap::new()); + let result = exact_service.lookup_with_source(full_key, None).unwrap(); + + assert_eq!(result.source, "Custom"); + assert_eq!(result.matched_key, full_key); + } + + #[test] + fn custom_terminal_override_unknown_stays_none() { + let mut custom = HashMap::new(); + custom.insert("gpt-5.5".into(), model_pricing(0.000002, 0.000008)); + + let service = custom_service(custom, HashMap::new(), HashMap::new()); + assert!(service + .lookup_with_source("cx/not-a-real-model-high", None) + .is_none()); + } + #[test] fn custom_override_respects_force_source() { let mut custom = HashMap::new(); @@ -745,13 +982,13 @@ mod tests { let service = custom_service(custom, litellm, openrouter); let litellm_result = service - .lookup_with_source("gpt-4o", Some("litellm")) + .lookup_with_source("gpt-4o", Some("LiteLLM")) .unwrap(); assert_eq!(litellm_result.source, "LiteLLM"); assert_eq!(litellm_result.pricing.input_cost_per_token, Some(0.00001)); let openrouter_result = service - .lookup_with_source("gpt-4o", Some("openrouter")) + .lookup_with_source("gpt-4o", Some("OpenRouter")) .unwrap(); assert_eq!(openrouter_result.source, "OpenRouter"); assert_eq!( @@ -760,7 +997,7 @@ mod tests { ); let custom_result = service - .lookup_with_source("gpt-4o", Some("custom")) + .lookup_with_source("gpt-4o", Some("Custom")) .unwrap(); assert_eq!(custom_result.source, "Custom"); assert_eq!(custom_result.pricing.input_cost_per_token, Some(0.000002)); @@ -802,7 +1039,7 @@ mod tests { } #[test] - fn custom_override_normalized_match_wins() { + fn custom_override_non_scoped_terminal_match_wins() { let mut custom = HashMap::new(); custom.insert("kimi-k2p6".into(), model_pricing(0.00000095, 0.000004)); let mut litellm = HashMap::new(); @@ -810,7 +1047,7 @@ mod tests { let service = custom_service(custom, litellm, HashMap::new()); let result = service - .lookup_with_source("accounts/fireworks/models/kimi-k2p6", None) + .lookup_with_source("hf:moonshotai/kimi-k2p6", None) .unwrap(); assert_eq!(result.source, "Custom"); @@ -889,4 +1126,183 @@ mod tests { let expected = 1_000_000.0 * 0.000002 + 100_000.0 * 0.000008; assert!((cost - expected).abs() < 1e-10); } + + #[test] + fn custom_calculate_cost_uses_terminal_override() { + let mut custom = HashMap::new(); + custom.insert("gpt-5.5".into(), model_pricing(0.000002, 0.000008)); + let mut litellm = HashMap::new(); + litellm.insert("gpt-5.5".into(), model_pricing(0.00001, 0.00003)); + + let service = custom_service(custom, litellm, HashMap::new()); + let cost = service.calculate_cost("cx/gpt-5.5", 1_000_000, 100_000, 0, 0, 0); + + let expected = 1_000_000.0 * 0.000002 + 100_000.0 * 0.000008; + assert!((cost - expected).abs() < 1e-10); + } + + #[test] + fn test_sakana_returns_pricing_for_fugu_ultra() { + let service = PricingService::new(HashMap::new(), HashMap::new()); + let result = service.lookup_with_source("fugu-ultra", None).unwrap(); + assert_eq!(result.source, "Sakana"); + assert_eq!(result.matched_key, "fugu-ultra"); + assert_eq!(result.pricing.input_cost_per_token, Some(5e-6)); + assert_eq!(result.pricing.output_cost_per_token, Some(3e-5)); + assert_eq!(result.pricing.cache_read_input_token_cost, Some(5e-7)); + assert_eq!( + result.pricing.input_cost_per_token_above_272k_tokens, + Some(1e-5) + ); + assert_eq!( + result.pricing.output_cost_per_token_above_272k_tokens, + Some(4.5e-5) + ); + assert_eq!( + result.pricing.cache_read_input_token_cost_above_272k_tokens, + Some(1e-6) + ); + } + + #[test] + fn test_sakana_bare_fugu_is_unpriced() { + let service = PricingService::new(HashMap::new(), HashMap::new()); + assert!(service.lookup_with_source("fugu", None).is_none()); + } + + #[test] + fn test_sakana_yields_to_litellm_exact() { + let mut litellm = HashMap::new(); + litellm.insert( + "fugu-ultra".into(), + ModelPricing { + input_cost_per_token: Some(0.001), + output_cost_per_token: Some(0.002), + ..Default::default() + }, + ); + + let service = PricingService::new(litellm, HashMap::new()); + let result = service.lookup_with_source("fugu-ultra", None).unwrap(); + assert_eq!(result.source, "LiteLLM"); + assert_eq!(result.pricing.input_cost_per_token, Some(0.001)); + } + + #[test] + fn test_sakana_override_requires_sakana_provider_identity() { + let service = PricingService::new(HashMap::new(), HashMap::new()); + + for (model_id, provider_id) in [ + ("fugu-ultra", "openrouter"), + ("openrouter/fugu-ultra", "openrouter"), + ("openrouter/fugu-ultra", "sakana"), + ("sakana/fugu-ultra", "openrouter"), + ("fugu-ultra", "sakana/openrouter"), + ("fugu-ultra", "openrouter/sakana"), + ("fugu-ultra", "unknown"), + ("fugu-ultra", "unknown/sakana"), + ("fugu-ultra", "sakana/unknown"), + ("fugu-ultra", "sakana/"), + ("fugu-ultra", "/sakana"), + ("fugu-ultra", "sakana//sakana"), + ("fugu-ultra", "sakana.openrouter"), + ("fugu-ultra", ""), + ("fugu-ultra", " "), + ] { + assert!( + service + .lookup_with_source_and_provider(model_id, None, Some(provider_id)) + .is_none(), + "non-Sakana identity resolved: model={model_id} provider={provider_id}" + ); + } + + assert!(service + .lookup_with_source_and_provider("sakana/openrouter/fugu-ultra", None, None) + .is_none()); + + let result = service + .lookup_with_source_and_provider("sakana/fugu-ultra", None, Some("sakana")) + .unwrap(); + assert_eq!(result.source, "Sakana"); + assert_eq!(result.matched_key, "fugu-ultra"); + + let repeated = service + .lookup_with_source_and_provider( + "sakana/sakana/fugu-ultra", + None, + Some("sakana/sakana"), + ) + .unwrap(); + assert_eq!(repeated.source, "Sakana"); + assert_eq!(repeated.matched_key, "fugu-ultra"); + + let usage = TokenBreakdown { + input: 1, + ..Default::default() + }; + assert_eq!( + service.calculate_cost_with_provider("fugu-ultra", None, &usage), + 5e-6 + ); + assert_eq!( + service.calculate_cost_with_provider("fugu-ultra", Some("unknown"), &usage), + 0.0 + ); + assert_eq!( + service.calculate_cost_with_provider("fugu-ultra", None, &usage), + 5e-6 + ); + } + + #[test] + fn test_sakana_fugu_ultra_request_tier_and_alias() { + let service = PricingService::new(HashMap::new(), HashMap::new()); + let regular = service.calculate_cost("fugu-ultra", 272_000, 3, 0, 0, 4); + let regular_expected = 272_000.0 * 5e-6 + 7.0 * 3e-5; + assert!((regular - regular_expected).abs() < 1e-12); + + let long = service.calculate_cost("fugu-ultra-20260615", 271_999, 3, 2, 0, 4); + let long_expected = 271_999.0 * 1e-5 + 7.0 * 4.5e-5 + 2.0 * 1e-6; + assert!((long - long_expected).abs() < 1e-12); + } + + #[test] + fn test_sakana_cache_write_does_not_select_long_context() { + let service = PricingService::new(HashMap::new(), HashMap::new()); + let without_cache_write = service.calculate_cost("fugu-ultra", 271_999, 10, 0, 0, 0); + let with_cache_write = service.calculate_cost("fugu-ultra", 271_999, 10, 0, 10_000, 0); + assert!((with_cache_write - without_cache_write).abs() < 1e-12); + } + + #[test] + fn test_sakana_cold_and_warm_costs_are_deterministic() { + let service = PricingService::new(HashMap::new(), HashMap::new()); + let cold = service.calculate_cost("fugu-ultra", 270_540, 630, 7_936, 0, 0); + let warm = service.calculate_cost("fugu-ultra", 270_540, 630, 7_936, 0, 0); + assert_eq!(cold.to_bits(), warm.to_bits()); + } + + #[test] + fn test_verified_litellm_long_context_survives_custom_pipeline() { + let mut custom = HashMap::new(); + custom.insert("unrelated-model".into(), model_pricing(0.000002, 0.000008)); + let mut litellm = HashMap::new(); + litellm.insert( + "gpt-5.5".into(), + ModelPricing { + input_cost_per_token: Some(5e-6), + input_cost_per_token_above_272k_tokens: Some(1e-5), + output_cost_per_token: Some(3e-5), + output_cost_per_token_above_272k_tokens: Some(4.5e-5), + cache_read_input_token_cost: Some(5e-7), + cache_read_input_token_cost_above_272k_tokens: Some(1e-6), + ..Default::default() + }, + ); + let service = custom_service(custom, litellm, HashMap::new()); + let cost = service.calculate_cost("gpt-5.5", 271_999, 3, 2, 0, 4); + let expected = 271_999.0 * 1e-5 + 7.0 * 4.5e-5 + 2.0 * 1e-6; + assert!((cost - expected).abs() < 1e-12); + } } diff --git a/vendor/tokscale-core/src/provider_identity.rs b/vendor/tokscale-core/src/provider_identity.rs index 7a008a9..802e6be 100644 --- a/vendor/tokscale-core/src/provider_identity.rs +++ b/vendor/tokscale-core/src/provider_identity.rs @@ -171,6 +171,25 @@ pub fn inferred_provider_from_model(model: &str) -> Option<&'static str> { return Some("qwen"); } + // Kimi (Moonshot AI) — `kimi`, `kimi-k2.5`, `kimi-code` variants + if contains_delimited(&lower, "kimi") { + return Some("moonshotai"); + } + // MiMo (Xiaomi) — `mimo-v2.5` etc. + if contains_delimited(&lower, "mimo") { + return Some("xiaomi"); + } + // GLM (Zhipu AI / Zai) — `glm-4.6`, `glm-5.2` etc. + if contains_delimited(&lower, "glm") { + return Some("zai"); + } + + // Sakana's Fugu model line. Provider identity is independent of whether + // the bare router model has a recoverable fixed price. + if lower.contains("fugu") { + return Some("sakana"); + } + None } @@ -289,6 +308,40 @@ mod tests { assert_eq!(inferred_provider_from_model("unknown-model"), None); } + #[test] + fn test_inferred_provider_normalizes_kimi_mimo_and_glm() { + assert_eq!( + inferred_provider_from_model("kimi-k2.5"), + Some("moonshotai") + ); + assert_eq!( + inferred_provider_from_model("moonshotai/kimi-code"), + Some("moonshotai") + ); + assert_eq!(inferred_provider_from_model("mimo-v2.5"), Some("xiaomi")); + assert_eq!(inferred_provider_from_model("glm-4.6"), Some("zai")); + } + + #[test] + fn test_inferred_provider_delimiters_avoid_kimi_mimo_glm_false_positives() { + assert_eq!(inferred_provider_from_model("kimiko"), None); + assert_eq!(inferred_provider_from_model("mimosa"), None); + assert_eq!(inferred_provider_from_model("aglm"), None); + } + + #[test] + fn test_inferred_provider_fugu_maps_to_sakana() { + assert_eq!(inferred_provider_from_model("fugu"), Some("sakana")); + assert_eq!(inferred_provider_from_model("fugu-ultra"), Some("sakana")); + assert_eq!(inferred_provider_from_model("Fugu"), Some("sakana")); + assert_eq!(inferred_provider_from_model("FUGU-ULTRA"), Some("sakana")); + } + + #[test] + fn test_provider_tags_preserves_sakana() { + assert_eq!(provider_tags("sakana"), vec!["sakana"]); + } + #[test] fn test_inferred_provider_no_false_positives() { assert_eq!(inferred_provider_from_model("protocol1-fast"), None); diff --git a/vendor/tokscale-core/src/scanner.rs b/vendor/tokscale-core/src/scanner.rs index a0811b5..a51bc3e 100644 --- a/vendor/tokscale-core/src/scanner.rs +++ b/vendor/tokscale-core/src/scanner.rs @@ -7,7 +7,7 @@ use std::collections::{BTreeMap, HashSet}; use std::path::{Path, PathBuf}; use walkdir::WalkDir; -use crate::clients::ClientId; +use crate::clients::{ClientId, PathRoot}; use crate::sessions::{normalize_workspace_key, workspace_label_from_key}; use serde::{Deserialize, Serialize}; use serde_json::Value; @@ -79,6 +79,7 @@ pub struct ScanResult { /// `opencode-nightly.db`, etc. See upstream logic in opencode's /// `packages/opencode/src/storage/db.ts` (`getChannelPath`). pub opencode_dbs: Vec, + pub copilot_desktop_db: Option, pub synthetic_db: Option, pub kilo_db: Option, pub hermes_db: Option, @@ -95,6 +96,7 @@ impl Default for ScanResult { Self { files: std::array::from_fn(|_| Vec::new()), opencode_dbs: Vec::new(), + copilot_desktop_db: None, synthetic_db: None, kilo_db: None, hermes_db: None, @@ -231,9 +233,51 @@ pub fn copilot_exporter_path() -> Option { copilot_exporter_path_with_env_strategy(true) } +// Kiro keeps legacy snapshots at `/*.chat`, execution records at +// `//*`, and workspace-session JSON at +// `workspace-sessions//*.json`. Do not recurse into mirrored project +// trees or treat root-level JSON/extensionless files as usage sources. +fn is_kiro_globalstorage_artifact(root: &Path, path: &Path) -> bool { + let Ok(relative) = path.strip_prefix(root) else { + return false; + }; + let components: Vec<_> = relative + .components() + .map(|component| component.as_os_str().to_string_lossy()) + .collect(); + + match components.as_slice() { + [workspace, file] if workspace != "workspace-sessions" => file.ends_with(".chat"), + [bucket, _, file] if bucket == "workspace-sessions" => file.ends_with(".json"), + [workspace, _, file] if workspace != "workspace-sessions" => { + file.ends_with(".json") || Path::new(file.as_ref()).extension().is_none() + } + _ => false, + } +} + +// Kiro structured sessions have one exact topology below the sessions root: +// `/sess_/session.json`. Reject deeper mirrored trees so an +// unrelated nested `sess_*` directory cannot become an automatic source. +fn is_kiro_ide_session_artifact(root: &Path, path: &Path) -> bool { + let Ok(relative) = path.strip_prefix(root) else { + return false; + }; + let components: Vec<_> = relative + .components() + .map(|component| component.as_os_str().to_string_lossy()) + .collect(); + + matches!( + components.as_slice(), + [_, session, file] if session.starts_with("sess_") && file == "session.json" + ) +} + /// Scan a single directory for session files pub fn scan_directory(root: &str, pattern: &str) -> Vec { - if !std::path::Path::new(root).exists() { + let root_path = std::path::Path::new(root); + if !root_path.exists() { return Vec::new(); } @@ -243,7 +287,13 @@ pub fn scan_directory(root: &str, pattern: &str) -> Vec { .filter_map(|e| e.ok()) .filter(|e| { let path = e.path(); - if !path.is_file() { + // WalkDir already knows the entry type from the directory read, so + // trust it for the common regular-file case and avoid a redundant + // stat() per file. Symlinks still follow metadata to preserve the + // existing file-symlink behavior. + let file_type = e.file_type(); + let is_file = file_type.is_file() || (file_type.is_symlink() && path.is_file()); + if !is_file { return false; } @@ -299,11 +349,14 @@ pub fn scan_directory(root: &str, pattern: &str) -> Vec { } "T-*.json" => file_name.starts_with("T-") && file_name.ends_with(".json"), "*.settings.json" => file_name.ends_with(".settings.json"), + "kiro-globalstorage" => is_kiro_globalstorage_artifact(root_path, path), + "kiro-ide-session" => is_kiro_ide_session_artifact(root_path, path), "sessions.json" => file_name == "sessions.json", "wire.jsonl" => file_name == "wire.jsonl", - // Grok Build ACP session updates under - // ~/.grok/sessions///updates.jsonl + "events.jsonl" => file_name == "events.jsonl", + // Grok Build ACP session updates and unified inference log. "updates.jsonl" => file_name == "updates.jsonl", + "unified.jsonl" => file_name == "unified.jsonl", "ui_messages.json" => file_name == "ui_messages.json", "session-usage.json" => file_name == "session-usage.json", "chat-messages.json" => file_name == "chat-messages.json", @@ -444,6 +497,38 @@ pub(crate) fn discover_hermes_profile_state_dbs(hermes_home: &Path) -> Vec Vec { + let mut homes = vec![PathBuf::from( + ClientId::Hermes + .data() + .root + .resolve_with_env_strategy(home_dir, use_env_roots), + )]; + + let hermes_home_set = use_env_roots + && std::env::var("HERMES_HOME") + .map(|value| !value.trim().is_empty()) + .unwrap_or(false); + if !hermes_home_set { + if cfg!(target_os = "windows") && use_env_roots { + if let Some(local_app_data) = + std::env::var_os("LOCALAPPDATA").filter(|value| !value.is_empty()) + { + homes.push(PathBuf::from(local_app_data).join("hermes")); + } + } + homes.push(PathBuf::from(home_dir).join("AppData/Local/hermes")); + } + + homes +} + /// Claude desktop "Cowork" (local-agent-mode) writes standard Claude Code /// transcripts under a fresh per-session directory: /// @@ -667,11 +752,38 @@ fn supports_extra_dir_scanning(client_id: ClientId) -> bool { ) } +fn push_grok_unified_log_candidates(candidates: &mut Vec, scan_root: &Path) { + candidates.push(scan_root.join("logs/unified.jsonl")); + if scan_root.file_name().and_then(|name| name.to_str()) == Some("sessions") { + if let Some(grok_home) = scan_root.parent() { + candidates.push(grok_home.join("logs/unified.jsonl")); + } + } +} + +fn grok_unified_log_path_from_updates(updates_path: &Path) -> Option { + updates_path + .ancestors() + .find(|ancestor| ancestor.file_name().and_then(|name| name.to_str()) == Some("sessions")) + .and_then(Path::parent) + .map(|grok_home| grok_home.join("logs/unified.jsonl")) +} + fn push_unique_scan_task( tasks: &mut Vec<(ClientId, String, &'static str)>, seen: &mut HashSet<(ClientId, PathBuf)>, client_id: ClientId, raw_path: impl Into, +) { + push_unique_scan_task_with_pattern(tasks, seen, client_id, raw_path, client_id.data().pattern); +} + +fn push_unique_scan_task_with_pattern( + tasks: &mut Vec<(ClientId, String, &'static str)>, + seen: &mut HashSet<(ClientId, PathBuf)>, + client_id: ClientId, + raw_path: impl Into, + pattern: &'static str, ) { let raw_path = raw_path.into(); if raw_path.as_os_str().is_empty() { @@ -680,11 +792,24 @@ fn push_unique_scan_task( let key = std::fs::canonicalize(&raw_path).unwrap_or_else(|_| raw_path.clone()); if seen.insert((client_id, key)) { - let pattern = client_id.data().pattern; tasks.push((client_id, raw_path.to_string_lossy().to_string(), pattern)); } } +#[cfg(target_os = "macos")] +fn kiro_global_storage_roots(home_dir: &str) -> [PathBuf; 2] { + [ + PathBuf::from(format!( + "{}/Library/Application Support/Kiro/User/globalStorage/kiro.kiroagent", + home_dir + )), + PathBuf::from(format!( + "{}/Library/Application Support/kiro/User/globalStorage/kiro.kiroagent", + home_dir + )), + ] +} + /// Merge user-configured OpenCode db paths from [`ScannerSettings`] into the /// auto-discovered list, in-place. /// @@ -814,12 +939,71 @@ fn scan_all_clients_with_env_strategy_inner( push_unique_scan_task(&mut tasks, &mut seen_scan_roots, *client_id, path); } + if enabled.contains(&ClientId::Kimi) { + let kimi_code_home = if use_env_roots { + std::env::var("KIMI_CODE_HOME") + .ok() + .filter(|value| !value.trim().is_empty()) + .unwrap_or_else(|| format!("{}/.kimi-code", home_dir)) + } else { + format!("{}/.kimi-code", home_dir) + }; + push_unique_scan_task( + &mut tasks, + &mut seen_scan_roots, + ClientId::Kimi, + format!("{}/sessions", kimi_code_home), + ); + } + + let mut grok_unified_paths = Vec::new(); + if enabled.contains(&ClientId::Grok) { + let grok_sessions = PathBuf::from( + ClientId::Grok + .data() + .resolve_path_with_env_strategy(home_dir, use_env_roots), + ); + push_grok_unified_log_candidates(&mut grok_unified_paths, &grok_sessions); + } + + // Register built-in Kiro IDE roots before user-configured or environment + // extras. Otherwise a default `*.json` task for an overlapping root can + // reserve it first and silently miss `.chat` and extensionless records. + #[cfg(target_os = "macos")] + if enabled.contains(&ClientId::Kiro) { + for root in kiro_global_storage_roots(home_dir) { + push_unique_scan_task_with_pattern( + &mut tasks, + &mut seen_scan_roots, + ClientId::Kiro, + root, + "kiro-globalstorage", + ); + } + } + + if enabled.contains(&ClientId::Kiro) { + push_unique_scan_task_with_pattern( + &mut tasks, + &mut seen_scan_roots, + ClientId::Kiro, + PathBuf::from(format!("{}/.kiro/sessions", home_dir)), + "kiro-ide-session", + ); + } + for (client_id, path) in extra_scan_paths_for(scanner_settings, &enabled) { warn_if_escapes_home(Path::new(home_dir), client_id, &path); + if client_id == ClientId::Grok { + push_grok_unified_log_candidates(&mut grok_unified_paths, &path); + } push_unique_scan_task(&mut tasks, &mut seen_scan_roots, client_id, path); } for (client_id, path) in built_in_extra_scan_paths_for(home_dir, &enabled) { + if client_id == ClientId::Grok { + push_grok_unified_log_candidates(&mut grok_unified_paths, &path); + } push_unique_scan_task(&mut tasks, &mut seen_scan_roots, client_id, path); } @@ -828,17 +1012,17 @@ fn scan_all_clients_with_env_strategy_inner( if use_env_roots { let extra_dirs_val = std::env::var("TOKSCALE_EXTRA_DIRS").unwrap_or_default(); for (client_id, path) in parse_extra_dirs(&extra_dirs_val, &enabled) { - warn_if_escapes_home(Path::new(home_dir), client_id, &PathBuf::from(&path)); + let path = PathBuf::from(path); + warn_if_escapes_home(Path::new(home_dir), client_id, &path); + if client_id == ClientId::Grok { + push_grok_unified_log_candidates(&mut grok_unified_paths, &path); + } push_unique_scan_task(&mut tasks, &mut seen_scan_roots, client_id, path); } } if enabled.contains(&ClientId::OpenCode) { - let xdg_data = if use_env_roots { - std::env::var("XDG_DATA_HOME").unwrap_or_else(|_| format!("{}/.local/share", home_dir)) - } else { - format!("{}/.local/share", home_dir) - }; + let xdg_data = PathRoot::XdgData.resolve_with_env_strategy(home_dir, use_env_roots); // OpenCode 1.2+: SQLite database(s) at ~/.local/share/opencode/opencode*.db // @@ -960,11 +1144,7 @@ fn scan_all_clients_with_env_strategy_inner( } if include_synthetic { - let xdg_data = if use_env_roots { - std::env::var("XDG_DATA_HOME").unwrap_or_else(|_| format!("{}/.local/share", home_dir)) - } else { - format!("{}/.local/share", home_dir) - }; + let xdg_data = PathRoot::XdgData.resolve_with_env_strategy(home_dir, use_env_roots); let octofriend_db_path = PathBuf::from(format!("{}/octofriend/sqlite.db", xdg_data)); if octofriend_db_path.exists() { result.synthetic_db = Some(octofriend_db_path); @@ -1038,20 +1218,21 @@ fn scan_all_clients_with_env_strategy_inner( } if enabled.contains(&ClientId::Hermes) { - let hermes_db_path = PathBuf::from( - ClientId::Hermes - .data() - .resolve_path_with_env_strategy(home_dir, use_env_roots), - ); - let hermes_home = hermes_db_path.parent().map(Path::to_path_buf); - if hermes_db_path.is_file() { - result.hermes_db = Some(hermes_db_path); - } - if let Some(hermes_home) = hermes_home { - result - .get_mut(ClientId::Hermes) - .extend(discover_hermes_profile_state_dbs(&hermes_home)); + let mut extra_dbs = Vec::new(); + for hermes_home in hermes_home_candidates(home_dir, use_env_roots) { + let default_db = hermes_home.join("state.db"); + if default_db.is_file() { + if result.hermes_db.is_none() { + result.hermes_db = Some(default_db); + } else if result.hermes_db.as_ref() != Some(&default_db) { + extra_dbs.push(default_db); + } + } + extra_dbs.extend(discover_hermes_profile_state_dbs(&hermes_home)); } + extra_dbs.sort_unstable(); + extra_dbs.dedup(); + result.get_mut(ClientId::Hermes).extend(extra_dbs); } if enabled.contains(&ClientId::Goose) { @@ -1122,6 +1303,12 @@ fn scan_all_clients_with_env_strategy_inner( result.zed_db = Some(macos_path); } } + if !use_env_roots && result.zed_db.is_none() { + let windows_path = PathBuf::from(home_dir).join("AppData/Local/Zed/threads/threads.db"); + if windows_path.is_file() { + result.zed_db = Some(windows_path); + } + } #[cfg(target_os = "windows")] if use_env_roots && result.zed_db.is_none() { if let Some(local_app_data) = dirs::data_local_dir() { @@ -1196,19 +1383,41 @@ fn scan_all_clients_with_env_strategy_inner( }) .collect(); - // Aggregate results, deduplicating file paths across overlapping directories + // Aggregate results, deduplicating physical files across overlapping roots. let mut seen: HashSet = HashSet::new(); for (client_id, files) in scan_results { for file in files { - if seen.insert(file.clone()) { + let key = std::fs::canonicalize(&file).unwrap_or_else(|_| file.clone()); + if seen.insert(key) { result.get_mut(client_id).push(file); } } } + grok_unified_paths.extend( + result + .get(ClientId::Grok) + .iter() + .filter_map(|path| grok_unified_log_path_from_updates(path)), + ); + + for path in grok_unified_paths { + let key = std::fs::canonicalize(&path).unwrap_or_else(|_| path.clone()); + if path.is_file() && seen.insert(key) { + result.get_mut(ClientId::Grok).push(path); + } + } + result.get_mut(ClientId::Grok).sort_unstable(); + if enabled.contains(&ClientId::Copilot) { + let desktop_db = PathBuf::from(home_dir).join(".copilot/data.db"); + if desktop_db.is_file() { + result.copilot_desktop_db = Some(desktop_db); + } + if let Some(path) = copilot_exporter_path_with_env_strategy(use_env_roots) { - if path.is_file() && seen.insert(path.clone()) { + let key = std::fs::canonicalize(&path).unwrap_or_else(|_| path.clone()); + if path.is_file() && seen.insert(key) { let copilot_files = result.get_mut(ClientId::Copilot); copilot_files.push(path); copilot_files.sort_unstable(); @@ -1400,6 +1609,19 @@ mod tests { assert!(json_files.iter().all(|p| p.extension().unwrap() == "json")); } + #[cfg(unix)] + #[test] + fn test_scan_directory_keeps_file_symlinks() { + let dir = TempDir::new().unwrap(); + let target = dir.path().join("target.json"); + let alias = dir.path().join("alias.json"); + File::create(&target).unwrap(); + std::os::unix::fs::symlink(&target, &alias).unwrap(); + + let files = scan_directory(dir.path().to_str().unwrap(), "*.json"); + assert_eq!(files, vec![alias, target]); + } + #[test] fn test_scan_directory_json_or_jsonl_pattern() { let dir = TempDir::new().unwrap(); @@ -1633,12 +1855,42 @@ mod tests { .unwrap(); } - fn setup_mock_grok_dir(base: &std::path::Path) { - let grok_session = base.join(".grok/sessions/%2Ftmp%2Fproject/session-uuid-1"); + fn setup_mock_kimi_code_dir(base: &std::path::Path) -> PathBuf { + let wire = base.join(".kimi-code/sessions/workspace/session-code/agents/main/wire.jsonl"); + fs::create_dir_all(wire.parent().unwrap()).unwrap(); + File::create(&wire).unwrap(); + wire + } + + fn setup_mock_junie_dir(base: &std::path::Path) -> PathBuf { + let events = base.join(".junie/sessions/session-1/events.jsonl"); + fs::create_dir_all(events.parent().unwrap()).unwrap(); + File::create(&events).unwrap(); + events + } + + fn setup_mock_opencodereview_dir(base: &std::path::Path) -> PathBuf { + let session = base.join(".opencodereview/sessions/repo/session-1.jsonl"); + fs::create_dir_all(session.parent().unwrap()).unwrap(); + File::create(&session).unwrap(); + session + } + + fn setup_mock_grok_home(grok_home: &std::path::Path) { + let grok_session = grok_home.join("sessions/%2Ftmp%2Fproject/session-uuid-1"); fs::create_dir_all(&grok_session).unwrap(); let mut file = File::create(grok_session.join("updates.jsonl")).unwrap(); file.write_all(b"{\"method\":\"session/update\"}\n") .unwrap(); + + let grok_logs = grok_home.join("logs"); + fs::create_dir_all(grok_logs.join("archive")).unwrap(); + File::create(grok_logs.join("unified.jsonl")).unwrap(); + File::create(grok_logs.join("archive/unified.jsonl")).unwrap(); + } + + fn setup_mock_grok_dir(base: &std::path::Path) { + setup_mock_grok_home(&base.join(".grok")); } fn setup_mock_openclaw_dir(base: &std::path::Path) { @@ -1776,6 +2028,29 @@ mod tests { assert!(result.get(ClientId::Gemini).is_empty()); } + #[test] + #[serial] + fn test_scan_all_clients_opencode_empty_xdg_uses_home_fallback() { + let mut _xdg = EnvGuard::capture(&["XDG_DATA_HOME"]); + _xdg.set("XDG_DATA_HOME", ""); + + let dir = TempDir::new().unwrap(); + let home = dir.path(); + setup_mock_opencode_dir(home); + let data_dir = home.join(".local/share/opencode"); + let db = data_dir.join("opencode.db"); + File::create(&db).unwrap(); + + let result = scan_without_extra_dirs(home.to_str().unwrap(), &["opencode".to_string()]); + let json_dir = home.join(".local/share/opencode/storage/message"); + assert_eq!(result.opencode_json_dir, Some(json_dir.clone())); + assert_eq!( + result.get(ClientId::OpenCode), + &[json_dir.join("proj1/msg_001.json")] + ); + assert_eq!(result.opencode_dbs, vec![db]); + } + #[test] #[serial] fn test_scan_all_clients_opencode_home_override_ignores_xdg_env() { @@ -2256,7 +2531,90 @@ mod tests { } #[test] - fn test_scan_all_clients_with_scanner_settings_merges_zed_extra_threads_db() { + fn test_scan_discovers_hermes_windows_local_appdata_home() { + let dir = TempDir::new().unwrap(); + let home = dir.path(); + + let windows_home = home.join("AppData/Local/hermes"); + fs::create_dir_all(&windows_home).unwrap(); + let default_db = windows_home.join("state.db"); + File::create(&default_db).unwrap(); + + let profile_dir = windows_home.join("profiles/research"); + fs::create_dir_all(&profile_dir).unwrap(); + let profile_db = profile_dir.join("state.db"); + File::create(&profile_db).unwrap(); + + let result = scan_all_clients_with_scanner_settings( + home.to_str().unwrap(), + &["hermes".to_string()], + false, + &ScannerSettings::default(), + ); + + assert_eq!(result.hermes_db.as_ref(), Some(&default_db)); + assert_eq!(result.hermes_db_paths(), vec![default_db, profile_db]); + } + + #[test] + #[serial] + fn test_explicit_hermes_home_does_not_widen_to_windows_fallback() { + let mut _hermes = EnvGuard::capture(&["HERMES_HOME", "TOKSCALE_EXTRA_DIRS"]); + _hermes.remove("TOKSCALE_EXTRA_DIRS"); + let dir = TempDir::new().unwrap(); + let home = dir.path(); + + let explicit_home = home.join("custom-hermes-home"); + fs::create_dir_all(&explicit_home).unwrap(); + let explicit_db = explicit_home.join("state.db"); + File::create(&explicit_db).unwrap(); + + let windows_home = home.join("AppData/Local/hermes"); + fs::create_dir_all(&windows_home).unwrap(); + File::create(windows_home.join("state.db")).unwrap(); + + _hermes.set("HERMES_HOME", &explicit_home); + let result = scan_all_clients_with_scanner_settings( + home.to_str().unwrap(), + &["hermes".to_string()], + true, + &ScannerSettings::default(), + ); + + assert_eq!(result.hermes_db.as_ref(), Some(&explicit_db)); + assert_eq!(result.hermes_db_paths(), vec![explicit_db]); + } + + #[cfg(unix)] + #[test] + fn test_hermes_windows_fallback_dedups_physical_database_overlaps() { + let dir = TempDir::new().unwrap(); + let home = dir.path(); + + let windows_home = home.join("AppData/Local/hermes"); + fs::create_dir_all(windows_home.join("profiles/research")).unwrap(); + File::create(windows_home.join("state.db")).unwrap(); + File::create(windows_home.join("profiles/research/state.db")).unwrap(); + std::os::unix::fs::symlink(&windows_home, home.join(".hermes")).unwrap(); + + let result = scan_all_clients_with_scanner_settings( + home.to_str().unwrap(), + &["hermes".to_string()], + false, + &ScannerSettings::default(), + ); + + let paths = result.hermes_db_paths(); + assert_eq!(paths.len(), 2, "physical dbs must be parsed exactly once"); + let canonical: HashSet = paths + .iter() + .map(|path| std::fs::canonicalize(path).unwrap()) + .collect(); + assert_eq!(canonical.len(), 2); + } + + #[test] + fn test_scan_all_clients_with_scanner_settings_discovers_zed_windows_local_appdata_home() { let dir = TempDir::new().unwrap(); let home = dir.path(); @@ -2265,6 +2623,26 @@ mod tests { let threads_db = windows_threads_dir.join("threads.db"); File::create(&threads_db).unwrap(); + let result = scan_all_clients_with_scanner_settings( + home.to_str().unwrap(), + &["zed".to_string()], + false, + &ScannerSettings::default(), + ); + + assert_eq!(result.zed_db.as_ref(), Some(&threads_db)); + } + + #[test] + fn test_scan_all_clients_with_scanner_settings_merges_zed_extra_threads_db() { + let dir = TempDir::new().unwrap(); + let home = dir.path(); + + let windows_threads_dir = home.join("custom-zed/threads"); + fs::create_dir_all(&windows_threads_dir).unwrap(); + let threads_db = windows_threads_dir.join("threads.db"); + File::create(&threads_db).unwrap(); + let settings: ScannerSettings = serde_json::from_value(serde_json::json!({ "extraScanPaths": { "zed": [windows_threads_dir] @@ -2774,6 +3152,24 @@ mod tests { assert!(result.get(ClientId::Copilot)[0].ends_with("copilot.jsonl")); } + #[test] + fn test_scan_all_clients_copilot_discovers_desktop_database() { + let dir = TempDir::new().unwrap(); + let home = dir.path(); + fs::create_dir_all(home.join(".copilot")).unwrap(); + let desktop_db = home.join(".copilot/data.db"); + File::create(&desktop_db).unwrap(); + + let result = scan_all_clients_with_env_strategy( + home.to_str().unwrap(), + &["copilot".to_string()], + false, + ); + + assert_eq!(result.copilot_desktop_db.as_ref(), Some(&desktop_db)); + assert!(result.get(ClientId::Copilot).is_empty()); + } + #[test] #[serial] fn test_scan_all_clients_copilot_includes_explicit_exporter_file() { @@ -3110,22 +3506,95 @@ mod tests { } #[test] - fn test_scan_all_clients_kimi() { + fn test_scan_all_clients_kimi_and_kimi_code() { let dir = TempDir::new().unwrap(); let home = dir.path(); setup_mock_kimi_dir(home); + let kimi_code = setup_mock_kimi_code_dir(home); let result = scan_all_clients_with_env_strategy( home.to_str().unwrap(), &["kimi".to_string()], false, ); - assert_eq!(result.get(ClientId::Kimi).len(), 1); - assert!(result.get(ClientId::Kimi)[0].ends_with("wire.jsonl")); + assert_eq!(result.get(ClientId::Kimi).len(), 2); + assert!(result + .get(ClientId::Kimi) + .iter() + .any(|path| path == &kimi_code)); + assert!(result + .get(ClientId::Kimi) + .iter() + .all(|path| path.ends_with("wire.jsonl"))); assert!(result.get(ClientId::OpenCode).is_empty()); assert!(result.get(ClientId::Claude).is_empty()); } + #[test] + #[serial] + fn test_scan_all_clients_kimi_code_home_obeys_env_strategy() { + let dir = TempDir::new().unwrap(); + let home = dir.path().join("home"); + let override_home = dir.path().join("override"); + fs::create_dir_all(&home).unwrap(); + let default_wire = setup_mock_kimi_code_dir(&home); + let override_wire = + override_home.join("sessions/workspace/session-code/agents/main/wire.jsonl"); + fs::create_dir_all(override_wire.parent().unwrap()).unwrap(); + File::create(&override_wire).unwrap(); + + let mut env = EnvGuard::capture(&["KIMI_CODE_HOME"]); + env.set("KIMI_CODE_HOME", &override_home); + let with_env = + scan_all_clients_with_env_strategy(home.to_str().unwrap(), &["kimi".to_string()], true); + assert_eq!( + with_env.get(ClientId::Kimi), + std::slice::from_ref(&override_wire) + ); + + for empty_override in ["", " "] { + env.set("KIMI_CODE_HOME", empty_override); + let fallback = scan_all_clients_with_env_strategy( + home.to_str().unwrap(), + &["kimi".to_string()], + true, + ); + assert_eq!( + fallback.get(ClientId::Kimi), + std::slice::from_ref(&default_wire) + ); + } + + let without_env = scan_all_clients_with_env_strategy( + home.to_str().unwrap(), + &["kimi".to_string()], + false, + ); + assert_eq!( + without_env.get(ClientId::Kimi), + std::slice::from_ref(&default_wire) + ); + } + + #[test] + fn test_scan_all_clients_junie_and_opencodereview() { + let dir = TempDir::new().unwrap(); + let home = dir.path(); + let junie = setup_mock_junie_dir(home); + let review = setup_mock_opencodereview_dir(home); + + let result = scan_all_clients_with_env_strategy( + home.to_str().unwrap(), + &["junie".to_string(), "opencodereview".to_string()], + false, + ); + assert_eq!(result.get(ClientId::Junie), std::slice::from_ref(&junie)); + assert_eq!( + result.get(ClientId::OpenCodeReview), + std::slice::from_ref(&review) + ); + } + #[test] fn test_scan_directory_updates_jsonl_pattern() { let dir = TempDir::new().unwrap(); @@ -3154,12 +3623,84 @@ mod tests { &["grok".to_string()], false, ); - assert_eq!(result.get(ClientId::Grok).len(), 1); - assert!(result.get(ClientId::Grok)[0].ends_with("updates.jsonl")); + assert_eq!(result.get(ClientId::Grok).len(), 2); + assert!(result + .get(ClientId::Grok) + .iter() + .any(|path| path.ends_with("updates.jsonl"))); + assert!(result + .get(ClientId::Grok) + .iter() + .any(|path| path == &home.join(".grok/logs/unified.jsonl"))); + assert!(!result + .get(ClientId::Grok) + .iter() + .any(|path| path == &home.join(".grok/logs/archive/unified.jsonl"))); assert!(result.get(ClientId::OpenCode).is_empty()); assert!(result.get(ClientId::Claude).is_empty()); } + #[test] + #[serial] + fn test_scan_all_clients_grok_discovers_configured_unified_logs_once() { + let mut env = EnvGuard::capture(&["GROK_HOME", "TOKSCALE_EXTRA_DIRS"]); + env.remove("GROK_HOME"); + + let dir = TempDir::new().unwrap(); + let home = dir.path(); + let primary_home = home.join(".grok"); + let settings_home = home.join("settings-grok"); + let env_root = home.join("env-root"); + let env_home = env_root.join("nested/.grok"); + setup_mock_grok_home(&primary_home); + setup_mock_grok_home(&settings_home); + setup_mock_grok_home(&env_home); + + let primary_sessions = primary_home.join("sessions"); + let settings_sessions = settings_home.join("sessions"); + let settings = ScannerSettings { + extra_scan_paths: BTreeMap::from([( + "grok".to_string(), + vec![primary_sessions.clone(), settings_home.clone()], + )]), + ..Default::default() + }; + env.set( + "TOKSCALE_EXTRA_DIRS", + format!( + "grok:{},grok:{}", + settings_sessions.display(), + env_root.display() + ), + ); + + let result = scan_all_clients_with_scanner_settings( + home.to_str().unwrap(), + &["grok".to_string()], + true, + &settings, + ); + let files = result.get(ClientId::Grok); + assert_eq!(files.len(), 6); + + for grok_home in [&primary_home, &settings_home, &env_home] { + for expected in [ + grok_home.join("sessions/%2Ftmp%2Fproject/session-uuid-1/updates.jsonl"), + grok_home.join("logs/unified.jsonl"), + ] { + assert_eq!( + files.iter().filter(|path| *path == &expected).count(), + 1, + "expected {} exactly once in {files:?}", + expected.display() + ); + } + assert!(!files + .iter() + .any(|path| path == &grok_home.join("logs/archive/unified.jsonl"))); + } + } + #[test] fn test_scan_all_clients_roocode() { let dir = TempDir::new().unwrap(); @@ -3505,4 +4046,156 @@ mod tests { "audit.jsonl must never be scanned: {claude_files:?}" ); } + + #[test] + #[cfg(target_os = "macos")] + fn m15a_globalstorage_scanner_discovers_file() { + let home = TempDir::new().unwrap(); + let root = home + .path() + .join("Library/Application Support/Kiro/User/globalStorage/kiro.kiroagent"); + let file = root.join("workspace-a/conversation.chat"); + fs::create_dir_all(file.parent().unwrap()).unwrap(); + File::create(&file).unwrap(); + + let result = scan_all_clients_with_env_strategy( + home.path().to_str().unwrap(), + &["kiro".to_string()], + false, + ); + + assert!( + result.get(ClientId::Kiro).contains(&file), + "globalStorage fixture must be discovered; old scanner returned no file" + ); + } + + #[test] + fn test_scan_directory_kiro_globalstorage_pattern_filters_sidecars() { + let dir = TempDir::new().unwrap(); + let root = dir.path().join("kiro.kiroagent"); + let workspace = root.join("workspace-a"); + let execution_store = workspace.join("execution-store"); + let workspace_session = root.join("workspace-sessions/workspace-a"); + fs::create_dir_all(&execution_store).unwrap(); + fs::create_dir_all(&workspace_session).unwrap(); + File::create(workspace.join("snapshot.chat")).unwrap(); + File::create(workspace.join("project.json")).unwrap(); + File::create(workspace.join("notes")).unwrap(); + File::create(execution_store.join("execution.json")).unwrap(); + File::create(execution_store.join("execution-record")).unwrap(); + fs::create_dir_all(execution_store.join("project")).unwrap(); + File::create(execution_store.join("project/mirror.json")).unwrap(); + File::create(workspace_session.join("session.json")).unwrap(); + fs::create_dir_all(workspace_session.join("project")).unwrap(); + File::create(workspace_session.join("project/mirror.json")).unwrap(); + File::create(workspace.join("index.sqlite")).unwrap(); + File::create(workspace.join("notes.txt")).unwrap(); + + let files = scan_directory(root.to_str().unwrap(), "kiro-globalstorage"); + let relative: Vec<_> = files + .iter() + .map(|path| path.strip_prefix(&root).unwrap().to_path_buf()) + .collect(); + + assert_eq!( + relative, + vec![ + PathBuf::from("workspace-a") + .join("execution-store") + .join("execution-record"), + PathBuf::from("workspace-a") + .join("execution-store") + .join("execution.json"), + PathBuf::from("workspace-a").join("snapshot.chat"), + PathBuf::from("workspace-sessions") + .join("workspace-a") + .join("session.json"), + ] + ); + } + + #[test] + fn m15b_kiro_ide_scanner_discovers_only_session_anchors() { + let home = TempDir::new().unwrap(); + let sessions_root = home.path().join(".kiro/sessions"); + let sess_dir = sessions_root.join("workspace-a/sess_02f1c107"); + fs::create_dir_all(&sess_dir).unwrap(); + let session = sess_dir.join("session.json"); + let messages = sess_dir.join("messages.jsonl"); + File::create(&session).unwrap(); + File::create(&messages).unwrap(); + + let cli_dir = sessions_root.join("cli"); + fs::create_dir_all(&cli_dir).unwrap(); + File::create(cli_dir.join("session.json")).unwrap(); + File::create(sessions_root.join("workspace-a/session.json")).unwrap(); + let nested_mirror = sessions_root.join("workspace-a/mirror/sess_nested"); + fs::create_dir_all(&nested_mirror).unwrap(); + File::create(nested_mirror.join("session.json")).unwrap(); + + assert_eq!( + scan_directory(sessions_root.to_str().unwrap(), "kiro-ide-session"), + vec![session.clone()] + ); + + let result = scan_all_clients_with_env_strategy( + home.path().to_str().unwrap(), + &["kiro".to_string()], + false, + ); + assert!(result.get(ClientId::Kiro).contains(&session)); + assert!(!result.get(ClientId::Kiro).contains(&messages)); + } + + #[cfg(target_os = "macos")] + #[test] + #[serial] + fn test_kiro_globalstorage_roots_precede_extra_paths_and_dedup_canonical_files() { + let mut extra = EnvGuard::capture(&["TOKSCALE_EXTRA_DIRS"]); + let home = TempDir::new().unwrap(); + let uppercase_root = home + .path() + .join("Library/Application Support/Kiro/User/globalStorage/kiro.kiroagent"); + let lowercase_root = home + .path() + .join("Library/Application Support/kiro/User/globalStorage/kiro.kiroagent"); + let file = uppercase_root.join("workspace-a/transcript.chat"); + fs::create_dir_all(file.parent().unwrap()).unwrap(); + File::create(&file).unwrap(); + let execution = uppercase_root.join("workspace-a/execution-store/execution"); + fs::create_dir_all(execution.parent().unwrap()).unwrap(); + File::create(&execution).unwrap(); + + // The two casing roots resolve to the same physical root. The scanner + // must register the first root once and not duplicate its files. + if !lowercase_root.exists() { + fs::create_dir_all(lowercase_root.parent().unwrap()).unwrap(); + std::os::unix::fs::symlink(&uppercase_root, &lowercase_root).unwrap(); + } + let settings: ScannerSettings = serde_json::from_value(serde_json::json!({ + "extraScanPaths": {"kiro": [lowercase_root]} + })) + .unwrap(); + extra.set( + "TOKSCALE_EXTRA_DIRS", + format!("kiro:{}", uppercase_root.display()), + ); + + let result = scan_all_clients_with_scanner_settings( + home.path().to_str().unwrap(), + &["kiro".to_string()], + true, + &settings, + ); + assert_eq!(result.get(ClientId::Kiro), &vec![execution, file.clone()]); + + let disabled = scan_all_clients_with_scanner_settings( + home.path().to_str().unwrap(), + &["claude".to_string()], + true, + &settings, + ); + assert!(disabled.get(ClientId::Kiro).is_empty()); + } } diff --git a/vendor/tokscale-core/src/sessionize.rs b/vendor/tokscale-core/src/sessionize.rs index 5565580..91632d0 100644 --- a/vendor/tokscale-core/src/sessionize.rs +++ b/vendor/tokscale-core/src/sessionize.rs @@ -98,7 +98,7 @@ pub fn sessionize(messages: &[UnifiedMessage], idle_gap_ms: i64) -> Vec`. @@ -455,6 +455,7 @@ mod tests { message_count: 1, agent: None, dedup_key: None, + dedup_aliases: Vec::new(), is_turn_start: false, duration_ms: None, } @@ -759,6 +760,7 @@ mod tests { message_count, agent: None, dedup_key: None, + dedup_aliases: Vec::new(), is_turn_start: false, duration_ms: None, } @@ -873,6 +875,28 @@ mod tests { assert_eq!(reference[0].message_count, result[0].message_count, "message_count"); } + #[test] + fn test_sessionize_preserves_explicit_zero_message_count() { + let msgs = vec![ + make_msg_full("grok", "s1", 1_000_000, 100, 10, 0.01, 1), + make_msg_full("grok", "s1", 1_001_000, 200, 20, 0.02, 0), + ]; + + let result = sessionize(&msgs, DEFAULT_IDLE_GAP_MS); + assert_eq!(result[0].message_count, 1); + assert_eq!(result[0].tokens.total(), 330); + assert_eq!(result[0].wall_duration_ms, 1_000); + + let mut acc = SessionizeAccumulator::new(); + for message in &msgs { + acc.feed(message); + } + let acc_result = acc.finalize(DEFAULT_IDLE_GAP_MS); + assert_eq!(acc_result[0].message_count, 1); + assert_eq!(acc_result[0].tokens.total(), 330); + assert_eq!(acc_result[0].wall_duration_ms, 1_000); + } + #[test] fn test_sessionize_saturates_overflowing_token_fold() { // A parser can clamp an untrusted token count to i64::MAX (e.g. diff --git a/vendor/tokscale-core/src/sessions/antigravity.rs b/vendor/tokscale-core/src/sessions/antigravity.rs index ed5fe65..2b33c6d 100644 --- a/vendor/tokscale-core/src/sessions/antigravity.rs +++ b/vendor/tokscale-core/src/sessions/antigravity.rs @@ -155,4 +155,21 @@ mod tests { assert_eq!(messages[0].model_id, "claude-opus-4-6"); assert_eq!(messages[0].provider_id, "anthropic"); } + + #[test] + fn parse_usage_row_resolves_current_placeholder_models() { + let input = r#"{"type":"usage","sessionId":"abc","modelId":"model_placeholder_m84","timestamp":1711200000000,"input":12,"output":4,"cacheRead":2,"cacheWrite":0,"reasoning":1} +{"type":"usage","sessionId":"abc","modelId":"model_placeholder_m16","timestamp":1711200000001,"input":8,"output":3,"cacheRead":0,"cacheWrite":0,"reasoning":0} +"#; + + let path = tempfile::NamedTempFile::new().unwrap(); + std::fs::write(path.path(), input).unwrap(); + + let messages = parse_antigravity_file(path.path()); + assert_eq!(messages.len(), 2); + assert_eq!(messages[0].model_id, "gemini-3-flash-preview"); + assert_eq!(messages[0].provider_id, "google"); + assert_eq!(messages[1].model_id, "gemini-3.1-pro"); + assert_eq!(messages[1].provider_id, "google"); + } } diff --git a/vendor/tokscale-core/src/sessions/antigravity_cli.rs b/vendor/tokscale-core/src/sessions/antigravity_cli.rs index d5d00b0..5044866 100644 --- a/vendor/tokscale-core/src/sessions/antigravity_cli.rs +++ b/vendor/tokscale-core/src/sessions/antigravity_cli.rs @@ -410,6 +410,10 @@ mod tests { } fn build_gen_metadata() -> Vec { + build_gen_metadata_with_model("gemini-3-flash-a") + } + + fn build_gen_metadata_with_model(model: &str) -> Vec { // usage message (#4 of chatModel) let mut usage = Vec::new(); usage.extend(enc_varint(1, 1132)); // fixed system prompt @@ -422,7 +426,7 @@ mod tests { // chatModel message (#1 of gen_metadata) let mut chat_model = Vec::new(); chat_model.extend(enc_len(4, &usage)); - chat_model.extend(enc_len(19, b"gemini-3-flash-a")); + chat_model.extend(enc_len(19, model.as_bytes())); enc_len(1, &chat_model) } @@ -492,7 +496,10 @@ mod tests { assert_eq!(message.client, "antigravity-cli"); // `gemini-3-flash-a` (raw #19 responseModel) is alias-resolved to the // priced canonical model so cost lookups don't fall through to 0. - assert_eq!(message.model_id, "gemini-3-flash-preview"); + // Per upstream (models.ts@603e3ea), this legacy M132 response model + // belongs to the Gemini 3.5 Flash High tier, not the retired preview + // family. + assert_eq!(message.model_id, "gemini-3.5-flash-high"); assert_eq!(message.provider_id, "google"); assert_eq!(message.session_id, "session-test"); assert_eq!(message.tokens.input, 1632); // 1132 + 500 @@ -508,6 +515,17 @@ mod tests { assert_eq!(message.workspace_label.as_deref(), Some("obsidian-vault")); } + #[test] + fn resolves_current_antigravity_cli_response_model() { + let blob = build_gen_metadata_with_model("gemini-3-flash-agent"); + let mut seen = HashSet::new(); + + let message = parse_gen_metadata(&blob, "session", 1_000, &mut seen).unwrap(); + + assert_eq!(message.model_id, "gemini-3.5-flash-high"); + assert_eq!(message.provider_id, "google"); + } + #[test] fn per_generation_timestamp_overrides_session_fallback() { // chatModel.#9.#4 = {#1: seconds, #2: nanos} is the per-turn wall-clock @@ -600,7 +618,7 @@ mod tests { // dataset, which is unavailable in unit tests). assert_eq!( pricing::aliases::resolve_alias("gemini-3-flash-a"), - Some("gemini-3-flash-preview") + Some("gemini-3.5-flash-high") ); } diff --git a/vendor/tokscale-core/src/sessions/claudecode.rs b/vendor/tokscale-core/src/sessions/claudecode.rs index bf82d25..956cfc2 100644 --- a/vendor/tokscale-core/src/sessions/claudecode.rs +++ b/vendor/tokscale-core/src/sessions/claudecode.rs @@ -146,6 +146,13 @@ fn is_workflow_journal(path: &Path) -> bool { .any(|ancestor| ancestor.file_name().and_then(|n| n.to_str()) == Some("subagents")) } +fn is_in_transcripts_dir(path: &Path) -> bool { + path.parent() + .and_then(|p| p.file_name()) + .and_then(|n| n.to_str()) + == Some("transcripts") +} + /// Locate the parent main-session JSONL for a sidechain transcript. /// /// Nested layout: `.../projects///subagents/agent-X.jsonl` @@ -448,6 +455,15 @@ pub fn parse_claude_file_with_cache_and_home( .unwrap_or("unknown") .to_string(); + // Bare transcripts (files under ~/.claude/transcripts/ with no workspace/project + // context) must not use char-based token estimation. These files may be written by + // third-party tools (e.g. OpenCode) that log tool outputs without Claude API usage + // metadata. Estimating tokens from their content would double-count usage already + // tracked by the originating client's own parser. Explicit tool-result token counts + // are still honored — only the char-based fallback estimate is suppressed. + let is_bare_transcript = + is_in_transcripts_dir(path) && cc_mirror_metadata.is_none() && workspace_key.is_none(); + let fallback_timestamp = file_modified_timestamp_ms(path); if path.extension().and_then(|s| s.to_str()) == Some("json") { @@ -538,6 +554,7 @@ pub fn parse_claude_file_with_cache_and_home( workspace_key: workspace_key.clone(), workspace_label: workspace_label.clone(), sidechain_agent: sidechain_agent.clone(), + allow_char_estimate: !is_bare_transcript, }, ); @@ -611,7 +628,6 @@ pub fn parse_claude_file_with_cache_and_home( &mut messages[existing_idx], &usage, parse_claude_entry_timestamp(entry.timestamp.as_deref()), - pending_request_start_timestamp_ms, ); if let Some(choice) = duplicate_provider_choice { update_claude_provider_id( @@ -631,7 +647,6 @@ pub fn parse_claude_file_with_cache_and_home( &mut messages[existing_idx], &usage, parse_claude_entry_timestamp(entry.timestamp.as_deref()), - pending_request_start_timestamp_ms, ); if let Some(choice) = duplicate_provider_choice { update_claude_provider_id( @@ -666,7 +681,8 @@ pub fn parse_claude_file_with_cache_and_home( let model = canonicalize_claude_model(&raw_model); let parsed_timestamp = parse_claude_entry_timestamp(entry.timestamp.as_deref()); - let timestamp = parsed_timestamp.unwrap_or(fallback_timestamp); + let timestamp = pending_request_start_timestamp_ms + .unwrap_or_else(|| parsed_timestamp.unwrap_or(fallback_timestamp)); let duration_ms = duration_between_ms(pending_request_start_timestamp_ms, parsed_timestamp); @@ -868,7 +884,6 @@ fn merge_claude_duplicate( existing: &mut UnifiedMessage, usage: &ClaudeUsage, parsed_timestamp: Option, - request_start_timestamp_ms: Option, ) { // Per-field max merge: each token field is updated independently. let t = &mut existing.tokens; @@ -883,21 +898,13 @@ fn merge_claude_duplicate( if let Some(timestamp_ms) = parsed_timestamp { if timestamp_ms >= existing.timestamp { - // Recover the original request-start timestamp from the existing - // message's recorded duration. The parent loop clears - // `pending_request_start_timestamp_ms` after the first chunk of a - // message commits (so a NEW message with no preceding user doesn't - // inflate by reusing a stale start), which would otherwise blank - // out streaming duplicates' duration. Recovering from - // `existing.timestamp - existing.duration_ms` keeps the duration - // honest for late chunks of the same logical message. - let recovered_start = existing - .duration_ms - .map(|d| existing.timestamp - d) - .or(request_start_timestamp_ms); - existing.set_timestamp(timestamp_ms); - if let Some(new_duration) = duration_between_ms(recovered_start, Some(timestamp_ms)) { - existing.duration_ms = Some(new_duration); + let new_duration = timestamp_ms.saturating_sub(existing.timestamp); + if new_duration > 0 { + // Duplicates can arrive out of order (e.g. late-processed + // streaming chunks), so never let a later-processed duplicate + // with an earlier completion timestamp shrink a duration + // already established by another duplicate. + existing.duration_ms = Some(existing.duration_ms.unwrap_or(0).max(new_duration)); } } } @@ -930,6 +937,12 @@ struct ClaudeToolResultContext<'a> { workspace_key: Option, workspace_label: Option, sidechain_agent: Option, + /// Whether char-based token estimation may be used as a fallback when no + /// explicit tool-result token count is present. Bare transcripts (see + /// `is_bare_transcript`) set this to `false` to avoid double-counting + /// usage already tracked by the originating client's own parser, while + /// still honoring any explicit tool-result token counts. + allow_char_estimate: bool, } fn extract_claude_tool_result_message( @@ -937,7 +950,7 @@ fn extract_claude_tool_result_message( context: ClaudeToolResultContext<'_>, ) -> Option { let value: Value = serde_json::from_str(line).ok()?; - let usage = extract_claude_tool_result_usage(&value)?; + let usage = extract_claude_tool_result_usage(&value, context.allow_char_estimate)?; let raw_model = extract_claude_model(&value) .or_else(|| { @@ -997,7 +1010,10 @@ fn extract_claude_tool_result_message( Some(message) } -fn extract_claude_tool_result_usage(value: &Value) -> Option { +fn extract_claude_tool_result_usage( + value: &Value, + allow_char_estimate: bool, +) -> Option { let mut total_tokens = 0; let mut first_dedup_id: Option = None; let mut seen_ids = HashSet::new(); @@ -1012,7 +1028,8 @@ fn extract_claude_tool_result_usage(value: &Value) -> Option Option { .or_else(|| extract_string(tool_result.get("tool_result_id"))) } -fn extract_tool_result_input_tokens(tool_result: &Value) -> Option { +fn extract_tool_result_input_tokens(tool_result: &Value, allow_char_estimate: bool) -> Option { explicit_tool_result_input_tokens(tool_result).or_else(|| { + if !allow_char_estimate { + return None; + } let chars = tool_result_output_char_count(tool_result); (chars > 0).then(|| estimate_tokens_from_chars(chars)) }) @@ -1910,11 +1930,51 @@ mod tests { assert_eq!(messages.len(), 1); assert_eq!(messages[0].tokens.output, 250); - assert_eq!(messages[0].timestamp, 1_733_047_203_500); + assert_eq!(messages[0].timestamp, 1_733_047_200_000); assert_eq!(messages[0].duration_ms, Some(3500)); assert_eq!(messages[0].dedup_key.as_deref(), Some("message:msg_stream")); } + #[test] + fn test_dedup_merge_duration_is_monotonic_across_out_of_order_duplicates() { + // Regression: several streaming duplicates of one message can be + // processed out of order (e.g. a late-arriving chunk carrying an + // earlier completion timestamp than one already merged). The start + // anchor (existing.timestamp) must survive every merge, and + // duration_ms must never shrink below a value already established by + // an earlier-processed duplicate — it must track the latest + // (largest) end timestamp seen so far. + let content = r#"{"type":"user","timestamp":"2024-12-01T10:00:00.000Z","message":{"content":"Hello"}} +{"type":"assistant","timestamp":"2024-12-01T10:00:01.000Z","requestId":"req_multi","message":{"id":"msg_multi","model":"claude-3-5-sonnet","usage":{"input_tokens":10,"output_tokens":30}}} +{"type":"assistant","timestamp":"2024-12-01T10:00:05.000Z","requestId":"req_multi","message":{"id":"msg_multi","model":"claude-3-5-sonnet","usage":{"input_tokens":10,"output_tokens":100}}} +{"type":"assistant","timestamp":"2024-12-01T10:00:02.000Z","requestId":"req_multi","message":{"id":"msg_multi","model":"claude-3-5-sonnet","usage":{"input_tokens":10,"output_tokens":50}}} +{"type":"assistant","timestamp":"2024-12-01T10:00:07.000Z","requestId":"req_multi","message":{"id":"msg_multi","model":"claude-3-5-sonnet","usage":{"input_tokens":10,"output_tokens":200}}}"#; + + let file = create_test_file(content); + let messages = parse_claude_file(file.path()); + + assert_eq!( + messages.len(), + 1, + "all streaming duplicates should collapse to one message" + ); + assert_eq!( + messages[0].timestamp, 1_733_047_200_000, + "the start anchor must survive every merge (the user entry's timestamp)" + ); + assert_eq!( + messages[0].duration_ms, + Some(7_000), + "duration_ms must equal the latest end timestamp minus the start \ + anchor (7s), not shrink when an out-of-order duplicate with an \ + earlier timestamp is merged" + ); + assert_eq!( + messages[0].tokens.output, 200, + "token fields keep the per-field max across all duplicates" + ); + } + #[test] fn test_pending_request_start_is_cleared_between_assistant_messages() { // Regression: previously, the user-entry timestamp was set into @@ -2371,6 +2431,85 @@ mod tests { ); } + #[test] + fn test_bare_transcript_with_tool_outputs_is_not_estimated() { + let content = r#"{"type":"tool_use","timestamp":"2026-04-01T10:00:00.000Z","tool_name":"read","tool_input":{"filePath":"/src/main.rs"}} +{"type":"tool_result","timestamp":"2026-04-01T10:00:01.000Z","tool_name":"read","tool_input":{"filePath":"/src/main.rs"},"tool_output":{"output":"fn main() {\n println!(\"Hello, world!\");\n}\n"}} +{"type":"tool_use","timestamp":"2026-04-01T10:00:02.000Z","tool_name":"bash","tool_input":{"command":"cargo build"}} +{"type":"tool_result","timestamp":"2026-04-01T10:00:03.000Z","tool_name":"bash","tool_input":{"command":"cargo build"},"tool_output":{"output":" Compiling myproject v0.1.0\n Finished dev [unoptimized + debuginfo] target(s) in 2.34s\n"}}"#; + let (_dir, path) = create_transcript_file(content, "ses_aabbccdd11223344556677889.jsonl"); + + let messages = parse_claude_file(&path); + + assert!( + messages.is_empty(), + "bare transcripts with only tool outputs must not produce estimated token messages" + ); + } + + #[test] + fn test_project_transcript_with_tool_outputs_is_estimated() { + let content = r#"{"type":"tool_result","timestamp":"2026-04-01T10:00:01.000Z","message":{"role":"user","content":[{"type":"tool_result","tool_use_id":"toolu_001","content":[{"type":"text","text":"fn main() { println!(\"hello\"); }"}]}]}}"#; + let temp_dir = tempfile::tempdir().unwrap(); + let path = temp_dir + .path() + .join(".claude") + .join("projects") + .join("myproject") + .join("ses_project123.jsonl"); + std::fs::create_dir_all(path.parent().unwrap()).unwrap(); + std::fs::write(&path, content).unwrap(); + + let messages = parse_claude_file(&path); + + assert!( + !messages.is_empty(), + "project transcripts with tool results should still estimate tokens" + ); + } + + #[test] + fn test_bare_transcript_with_explicit_tool_result_tokens_is_counted() { + // Bare transcripts must not char-estimate tokens, but explicit tool-result + // token counts (e.g. reported by the originating client) should still be honored. + let content = r#"{"type":"tool_result","timestamp":"2026-04-01T10:00:01.000Z","tool_name":"read","input_tokens":42,"tool_output":{"output":"fn main() {\n println!(\"Hello, world!\");\n}\n"}}"#; + let (_dir, path) = create_transcript_file(content, "ses_explicit112233445566778899.jsonl"); + + let messages = parse_claude_file(&path); + + assert_eq!( + messages.len(), + 1, + "bare transcripts must still count explicit tool-result token usage" + ); + assert_eq!(messages[0].tokens.input, 42); + } + + #[test] + fn test_transcripts_dir_under_project_is_not_treated_as_bare() { + // A `transcripts/` directory nested under a resolvable `projects//` path + // must not be treated as a bare transcript, since its workspace can still be + // attributed. Char-based estimation should proceed normally. + let content = r#"{"type":"tool_result","timestamp":"2026-04-01T10:00:01.000Z","tool_name":"read","tool_output":{"output":"fn main() {\n println!(\"Hello, world!\");\n}\n"}}"#; + let temp_dir = tempfile::tempdir().unwrap(); + let path = temp_dir + .path() + .join("projects") + .join("myproject") + .join("transcripts") + .join("ses_scoped112233445566778899.jsonl"); + std::fs::create_dir_all(path.parent().unwrap()).unwrap(); + std::fs::write(&path, content).unwrap(); + + let messages = parse_claude_file(&path); + + assert!( + !messages.is_empty(), + "transcripts nested under a resolvable projects// path should still be estimated" + ); + assert_eq!(messages[0].workspace_key, Some("myproject".to_string())); + } + // --- Sidechain / Agent tracking tests --- /// Helper: create a sidechain JSONL file and optional meta sidecar in a nested layout. diff --git a/vendor/tokscale-core/src/sessions/codex.rs b/vendor/tokscale-core/src/sessions/codex.rs index e48f8f5..9f85005 100644 --- a/vendor/tokscale-core/src/sessions/codex.rs +++ b/vendor/tokscale-core/src/sessions/codex.rs @@ -41,6 +41,17 @@ pub struct CodexPayload { pub model_info: Option, pub info: Option, pub turn_id: Option, + /// Unix timestamp (seconds) from `task_started` events. Legacy Codex turns + /// may use UUID v4 ids, so this is their only causal ordering signal. + /// Confirmed against codex-rs (`TurnStartedEvent::started_at`, serialized + /// under `task_started`): documented as "Unix timestamp (in seconds)", + /// `Option`. Deserialized leniently anyway: int/float values coerce + /// to `i64`, and any other JSON type (string, object, ...) decodes as + /// `None` rather than failing the whole `task_started` entry. A strict + /// `Option` would make a wrong-typed value reject deserialization of + /// the entire JSONL line, silently dropping the rest of that payload too. + #[serde(default, deserialize_with = "deserialize_lenient_i64")] + pub started_at: Option, pub source: Option, /// Thread origin from session_meta. `"user"` marks a human-initiated fork /// (e.g. a VS Code "fork conversation"), which replays parent history but @@ -59,6 +70,21 @@ pub struct CodexPayload { pub message: Option, } +/// Lenient `Option` deserializer for `CodexPayload::started_at`. Coerces +/// JSON integers and floats to `i64`; any other type (string, bool, object, +/// array) or `null`/absent decodes as `None` instead of failing the entry. +fn deserialize_lenient_i64<'de, D>(deserializer: D) -> Result, D::Error> +where + D: serde::Deserializer<'de>, +{ + let value = Option::::deserialize(deserializer)?; + Ok(value.and_then(|v| { + v.as_i64() + .or_else(|| v.as_u64().map(|u| u as i64)) + .or_else(|| v.as_f64().map(|f| f as i64)) + })) +} + #[derive(Debug, Deserialize)] pub struct CodexModelInfo { pub slug: Option, @@ -337,12 +363,23 @@ fn parse_codex_reader( && payload.payload_type.as_deref() == Some("task_started") { // The child's own turn is introduced by a - // `task_started`; remember its turn_id so the gate can - // recognize the child's own same-millisecond turn. - if let Some(turn_id) = payload.turn_id.as_deref() { - state - .forked_child_task_started_turn_ids - .insert(turn_id.to_string()); + // `task_started`; remember it only when its id or + // timestamp places it at/after the child session. + // Nested child logs can replay ancestor task_started + // events before the child's live turn. + if forked_child_task_starts_own_session( + &state, + payload.turn_id.as_deref(), + payload.started_at, + ) { + // Safety of this branch is coupled to + // `forked_child_task_starts_own_session` + // returning false when `turn_id` is `None`. + if let Some(turn_id) = payload.turn_id.as_deref() { + state + .forked_child_task_started_turn_ids + .insert(turn_id.to_string()); + } } } if entry.entry_type == "session_meta" { @@ -457,6 +494,15 @@ fn parse_codex_reader( { if codex_message_is_human_turn(payload.message.as_deref()) { state.pending_turn_start = true; + // Defensively reset the start-anchor cursor here too. + // Normally `turn_context` resets it every turn (see + // above), but a resumed/compacted session can emit a + // `token_count` after this `user_message` with no + // intervening `turn_context`. Without this reset the + // cursor would still hold the previous turn's last + // token time and bridge backward across the idle gap. + state.last_accepted_token_timestamp_ms = + parse_codex_entry_timestamp(entry.timestamp.as_deref()); } handled = true; } @@ -557,7 +603,9 @@ fn parse_codex_reader( state.previous_totals = next_totals; let parsed_timestamp = parse_codex_entry_timestamp(entry.timestamp.as_deref()); - let timestamp = parsed_timestamp.unwrap_or(fallback_timestamp); + let timestamp = state + .last_accepted_token_timestamp_ms + .unwrap_or_else(|| parsed_timestamp.unwrap_or(fallback_timestamp)); let duration_ms = duration_between_ms( state.last_accepted_token_timestamp_ms, parsed_timestamp, @@ -743,7 +791,12 @@ fn forked_child_turn_starts_own_session(state: &CodexParseState, turn_id: Option match (turn_id, codex_uuid_v7_order_key(child_session_id)) { (Some(turn_id), Some(child_key)) => { let Some(turn_key) = codex_uuid_v7_order_key(turn_id) else { - return true; + // Nested child logs can replay legacy UUID v4 turns from an + // ancestor. Only a child-local task_started event may end the + // replay gate for a non-v7 subagent turn. Human forks do not + // emit task_started, so retain their existing fallback. + return state.forked_child_is_user_fork + || state.forked_child_task_started_turn_ids.contains(turn_id); }; // Compare only the UUID v7 48-bit millisecond timestamp (the first // 12 hex of the order key), not the full id. The child's own turn is @@ -785,6 +838,37 @@ fn forked_child_turn_starts_own_session(state: &CodexParseState, turn_id: Option } } +fn forked_child_task_starts_own_session( + state: &CodexParseState, + turn_id: Option<&str>, + started_at: Option, +) -> bool { + let (Some(turn_id), Some(child_session_id)) = + (turn_id, state.forked_child_session_id.as_deref()) + else { + return false; + }; + let Some(child_key) = codex_uuid_v7_order_key(child_session_id) else { + return true; + }; + + if let Some(turn_key) = codex_uuid_v7_order_key(turn_id) { + return turn_key[..12] >= child_key[..12]; + } + + let Some(started_at) = started_at else { + return false; + }; + let Ok(child_started_at_ms) = i64::from_str_radix(&child_key[..12], 16) else { + return false; + }; + + // `child_started_at_ms / 1000` floors to the child's fork second, so a + // legacy replay whose `started_at` lands in that same integer second + // (but strictly before the child's sub-second fork instant) is admitted. + started_at >= child_started_at_ms / 1000 +} + fn codex_uuid_v7_order_key(id: &str) -> Option { let mut parts = id.split('-'); let first = parts.next()?; @@ -2278,6 +2362,130 @@ mod tests { assert_eq!(child_messages[0].tokens.output, 2); } + #[test] + fn test_nested_child_skips_replayed_legacy_uuid_v4_turn() { + // Nested Codex child logs can replay an ancestor turn whose legacy UUID + // v4 id cannot be ordered against the child's UUID v7 session id. Its + // task_started timestamp still predates the child, so it must not open + // the gate or count the inherited token snapshot. + let child = create_test_file(concat!( + r#"{"timestamp":"2026-05-05T21:52:10.197Z","type":"session_meta","payload":{"id":"019e5c03-1f5d-7000-8000-000000000001","forked_from_id":"019e5c03-0000-7000-8000-000000000001","source":{"subagent":{"thread_spawn":{"parent_thread_id":"019e5c03-0000-7000-8000-000000000001","depth":2}}},"thread_source":"subagent","model_provider":"openai","cwd":"/repo"}}"#, + "\n", + r#"{"timestamp":"2026-05-05T21:52:10.197Z","type":"session_meta","payload":{"id":"019e5c03-0000-7000-8000-000000000001","forked_from_id":"019e5b00-0000-7000-8000-000000000001","source":{"subagent":{"thread_spawn":{"parent_thread_id":"019e5b00-0000-7000-8000-000000000001","depth":1}}},"thread_source":"subagent","model_provider":"openai","cwd":"/repo"}}"#, + "\n", + r#"{"timestamp":"2026-05-05T21:52:10.197Z","type":"session_meta","payload":{"id":"019e5b00-0000-7000-8000-000000000001","source":"cli","thread_source":"user","model_provider":"openai","cwd":"/repo"}}"#, + "\n", + r#"{"timestamp":"2026-05-05T21:52:10.197Z","type":"event_msg","payload":{"type":"task_started","turn_id":"81d2f55b-894b-4d67-b75b-436ead477f65","started_at":1778017800}}"#, + "\n", + r#"{"timestamp":"2026-05-05T21:52:10.197Z","type":"turn_context","payload":{"turn_id":"81d2f55b-894b-4d67-b75b-436ead477f65","model":"gpt-5.5","cwd":"/repo"}}"#, + "\n", + r#"{"timestamp":"2026-05-05T21:52:10.198Z","type":"event_msg","payload":{"type":"token_count","info":{"total_token_usage":{"input_tokens":300,"output_tokens":30,"total_tokens":330},"last_token_usage":{"input_tokens":300,"output_tokens":30,"total_tokens":330}}}}"#, + "\n", + r#"{"timestamp":"2026-05-05T21:52:10.610Z","type":"event_msg","payload":{"type":"task_started","turn_id":"019e5c03-2100-7000-8000-000000000001","started_at":1779660169}}"#, + "\n", + r#"{"timestamp":"2026-05-05T21:52:10.611Z","type":"turn_context","payload":{"turn_id":"019e5c03-2100-7000-8000-000000000001","model":"gpt-5.5","cwd":"/repo"}}"#, + "\n", + r#"{"timestamp":"2026-05-05T21:52:10.612Z","type":"event_msg","payload":{"type":"token_count","info":{"total_token_usage":{"input_tokens":320,"output_tokens":32,"total_tokens":352},"last_token_usage":{"input_tokens":20,"output_tokens":2,"total_tokens":22}}}}"#, + "\n" + )); + + let child_messages = parse_codex_file(child.path()); + + assert_eq!(child_messages.len(), 1); + assert_eq!(child_messages[0].tokens.input, 20); + assert_eq!(child_messages[0].tokens.output, 2); + } + + #[test] + fn test_forked_child_legacy_turn_pins_seconds_unit_contract() { + // Pins the `started_at` unit contract used by + // `forked_child_task_starts_own_session`: it compares against the + // child's fork second (`started_at >= child_started_at_ms / 1000`), + // so a legacy replayed turn timestamped exactly one second before + // the child's fork second must stay rejected, while one landing on + // that same second must be admitted. The child's UUID v7 id here + // (`018bcfe5-6800-...`) encodes ms=1700000000000, i.e. + // floor(ms/1000) == 1700000000, a round multiple of 1000 so there is + // no ambiguity from the ms->s truncation. + let child = create_test_file(concat!( + r#"{"timestamp":"2023-11-14T22:13:20.000Z","type":"session_meta","payload":{"id":"018bcfe5-6800-7000-8000-000000000001","forked_from_id":"018bcfe5-0000-7000-8000-000000000001","source":{"subagent":{"thread_spawn":{"parent_thread_id":"018bcfe5-0000-7000-8000-000000000001","depth":2}}},"thread_source":"subagent","model_provider":"openai","cwd":"/repo"}}"#, + "\n", + r#"{"timestamp":"2023-11-14T22:13:20.000Z","type":"session_meta","payload":{"id":"018bcfe5-0000-7000-8000-000000000001","forked_from_id":"018bcfe4-0000-7000-8000-000000000001","source":{"subagent":{"thread_spawn":{"parent_thread_id":"018bcfe4-0000-7000-8000-000000000001","depth":1}}},"thread_source":"subagent","model_provider":"openai","cwd":"/repo"}}"#, + "\n", + r#"{"timestamp":"2023-11-14T22:13:20.000Z","type":"session_meta","payload":{"id":"018bcfe4-0000-7000-8000-000000000001","source":"cli","thread_source":"user","model_provider":"openai","cwd":"/repo"}}"#, + "\n", + // legacy turn one second BEFORE the child's fork second + // (1700000000) -- must NOT open the gate or count its token + // snapshot. + r#"{"timestamp":"2023-11-14T22:13:19.000Z","type":"event_msg","payload":{"type":"task_started","turn_id":"71d2f55b-894b-4d67-b75b-436ead477f65","started_at":1699999999}}"#, + "\n", + r#"{"timestamp":"2023-11-14T22:13:19.000Z","type":"turn_context","payload":{"turn_id":"71d2f55b-894b-4d67-b75b-436ead477f65","model":"gpt-5.5","cwd":"/repo"}}"#, + "\n", + r#"{"timestamp":"2023-11-14T22:13:19.100Z","type":"event_msg","payload":{"type":"token_count","info":{"total_token_usage":{"input_tokens":300,"output_tokens":30,"total_tokens":330},"last_token_usage":{"input_tokens":300,"output_tokens":30,"total_tokens":330}}}}"#, + "\n", + // legacy turn exactly AT the child's fork second -- must be + // admitted. + r#"{"timestamp":"2023-11-14T22:13:20.100Z","type":"event_msg","payload":{"type":"task_started","turn_id":"82d2f55b-894b-4d67-b75b-436ead477f66","started_at":1700000000}}"#, + "\n", + r#"{"timestamp":"2023-11-14T22:13:20.200Z","type":"turn_context","payload":{"turn_id":"82d2f55b-894b-4d67-b75b-436ead477f66","model":"gpt-5.5","cwd":"/repo"}}"#, + "\n", + r#"{"timestamp":"2023-11-14T22:13:20.300Z","type":"event_msg","payload":{"type":"token_count","info":{"total_token_usage":{"input_tokens":320,"output_tokens":32,"total_tokens":352},"last_token_usage":{"input_tokens":20,"output_tokens":2,"total_tokens":22}}}}"#, + "\n" + )); + + let child_messages = parse_codex_file(child.path()); + + assert_eq!(child_messages.len(), 1); + assert_eq!(child_messages[0].tokens.input, 20); + assert_eq!(child_messages[0].tokens.output, 2); + } + + #[test] + fn test_forked_child_task_started_non_numeric_started_at_does_not_fail_parsing() { + // A `task_started` event with a non-integer `started_at` (e.g. a + // string, from a malformed or unexpected log) must not fail + // deserialization of the whole JSONL line -- it should decode with + // `started_at: None`, which keeps the replay gate closed (same as a + // missing timestamp) and still allows the rest of the file, + // including a valid subsequent `task_started`, to parse normally. + let child = create_test_file(concat!( + r#"{"timestamp":"2026-05-05T21:52:10.000Z","type":"session_meta","payload":{"id":"019e5c03-1e99-7000-8000-000000000001","forked_from_id":"019e5b00-0000-7000-8000-000000000001","source":{"subagent":{"thread_spawn":{"parent_thread_id":"019e5b00-0000-7000-8000-000000000001","depth":1}}},"model_provider":"openai","agent_nickname":"worker","cwd":"/repo"}}"#, + "\n", + r#"{"timestamp":"2026-05-05T21:52:10.000Z","type":"session_meta","payload":{"id":"019e5b00-0000-7000-8000-000000000001","source":"vscode","model_provider":"openai","cwd":"/repo"}}"#, + "\n", + r#"{"timestamp":"2026-05-05T21:52:10.100Z","type":"turn_context","payload":{"turn_id":"019e5b00-0001-7000-8000-000000000001","model":"gpt-5.5","cwd":"/repo"}}"#, + "\n", + r#"{"timestamp":"2026-05-05T21:52:10.200Z","type":"event_msg","payload":{"type":"token_count","info":{"total_token_usage":{"input_tokens":300,"output_tokens":30,"total_tokens":330},"last_token_usage":{"input_tokens":300,"output_tokens":30,"total_tokens":330}}}}"#, + "\n", + // legacy task_started with a malformed (non-numeric) started_at + // -- must decode with started_at: None (not fail the whole + // entry), so the gate stays closed rather than opening on a + // wrong-typed value. + r#"{"timestamp":"2026-05-05T21:52:15.000Z","type":"event_msg","payload":{"type":"task_started","turn_id":"71d2f55b-894b-4d67-b75b-436ead477f65","started_at":"not-a-number"}}"#, + "\n", + // the child's real own turn, announced by a well-formed + // task_started -- proves the malformed line above didn't corrupt + // parser state or halt parsing of the rest of the file. + r#"{"timestamp":"2026-05-05T21:52:20.000Z","type":"event_msg","payload":{"type":"task_started","turn_id":"019e5c03-6425-7000-8000-000000000001"}}"#, + "\n", + r#"{"timestamp":"2026-05-05T21:52:20.100Z","type":"turn_context","payload":{"turn_id":"019e5c03-6425-7000-8000-000000000001","model":"gpt-5.5","cwd":"/repo"}}"#, + "\n", + r#"{"timestamp":"2026-05-05T21:52:20.200Z","type":"event_msg","payload":{"type":"token_count","info":{"total_token_usage":{"input_tokens":320,"output_tokens":32,"total_tokens":352},"last_token_usage":{"input_tokens":20,"output_tokens":2,"total_tokens":22}}}}"#, + "\n" + )); + + let parsed = parse_codex_file_incremental(child.path(), 0, CodexParseState::default()); + + // The malformed line did not abort file-level parsing. + assert!(parsed.parse_succeeded); + // The malformed task_started did not open the gate; only the + // well-formed one that follows (admitted via the UUID v7 ordering + // path) did. + assert_eq!(parsed.messages.len(), 1); + assert_eq!(parsed.messages[0].tokens.input, 20); + assert_eq!(parsed.messages[0].tokens.output, 2); + } + #[test] fn test_forked_child_incremental_state_skips_inherited_prefix() { let file = create_test_file(concat!( @@ -2641,4 +2849,64 @@ mod tests { "the pending flag is consumed once applied" ); } + + #[test] + fn test_token_count_timestamp_is_start_anchored() { + let line1 = r#"{"timestamp":"1970-01-01T00:00:01Z","type":"turn_context","payload":{"model":"gpt-5.2"}}"#; + let line2 = r#"{"timestamp":"1970-01-01T00:00:01.005Z","type":"event_msg","payload":{"type":"token_count","info":{"total_token_usage":{"input_tokens":10,"cached_input_tokens":2,"output_tokens":3},"last_token_usage":{"input_tokens":10,"cached_input_tokens":2,"output_tokens":3}}}}"#; + let content = format!("{}\n{}", line1, line2); + let file = create_test_file(&content); + + let messages = parse_codex_file(file.path()); + + assert_eq!(messages.len(), 1); + assert_eq!( + messages[0].timestamp, 1_000, + "timestamp must be the turn_context start (1000ms epoch)" + ); + assert_eq!( + messages[0].duration_ms, + Some(5), + "duration_ms must span from turn start to token_count event (5ms)" + ); + } + + #[test] + fn test_user_message_without_turn_context_anchors_at_user_message() { + // Regression: a resumed/compacted session can emit a human + // `user_message` followed directly by a `token_count` with no + // intervening `turn_context` (which normally resets the start-anchor + // cursor every turn). Before this fix, the token_count would anchor + // at the previous turn's last accepted token timestamp instead of + // this user message, bridging backward across the idle gap between + // turns. + let line1 = r#"{"timestamp":"1970-01-01T00:00:01Z","type":"turn_context","payload":{"model":"gpt-5.2"}}"#; + let line2 = r#"{"timestamp":"1970-01-01T00:00:01.100Z","type":"event_msg","payload":{"type":"token_count","info":{"total_token_usage":{"input_tokens":10,"cached_input_tokens":2,"output_tokens":3},"last_token_usage":{"input_tokens":10,"cached_input_tokens":2,"output_tokens":3}}}}"#; + // A long idle gap follows: the session resumes with a human + // user_message but no fresh turn_context before the next token_count. + let line3 = r#"{"timestamp":"1970-01-01T01:00:00Z","type":"event_msg","payload":{"type":"user_message","message":"still there?"}}"#; + let line4 = r#"{"timestamp":"1970-01-01T01:00:00.500Z","type":"event_msg","payload":{"type":"token_count","info":{"total_token_usage":{"input_tokens":20,"cached_input_tokens":4,"output_tokens":6},"last_token_usage":{"input_tokens":10,"cached_input_tokens":2,"output_tokens":3}}}}"#; + let content = [line1, line2, line3, line4].join("\n"); + let file = create_test_file(&content); + + let messages = parse_codex_file(file.path()); + + assert_eq!(messages.len(), 2); + assert_eq!( + messages[1].timestamp, + parse_codex_entry_timestamp(Some("1970-01-01T01:00:00Z")).unwrap(), + "the second token_count must anchor at the user_message, not the \ + previous turn's last accepted token timestamp" + ); + assert_eq!( + messages[1].duration_ms, + Some(500), + "duration_ms must span from the user_message to its token_count \ + (500ms), not bridge backward across the idle gap" + ); + assert!( + messages[1].is_turn_start, + "the deferred turn-start marker must still apply" + ); + } } diff --git a/vendor/tokscale-core/src/sessions/copilot.rs b/vendor/tokscale-core/src/sessions/copilot.rs index 0d4e9d9..2424da6 100644 --- a/vendor/tokscale-core/src/sessions/copilot.rs +++ b/vendor/tokscale-core/src/sessions/copilot.rs @@ -55,21 +55,35 @@ pub fn parse_copilot_file(path: &Path) -> Vec { let agent_turn_response_ids = candidate_response_ids(&candidates, CopilotUsageSource::AgentTurnLog); - candidates - .into_iter() - .filter(|candidate| { - should_emit_candidate( - candidate, - &chat_traces, - &inference_traces, - &agent_turn_traces, - &chat_response_ids, - &inference_response_ids, - &agent_turn_response_ids, - ) - }) - .map(CopilotUsageCandidate::into_message) - .collect() + let mut messages = Vec::with_capacity(candidates.len()); + let mut index_by_dedup_key = HashMap::new(); + + for candidate in candidates.into_iter().filter(|candidate| { + should_emit_candidate( + candidate, + &chat_traces, + &inference_traces, + &agent_turn_traces, + &chat_response_ids, + &inference_response_ids, + &agent_turn_response_ids, + ) + }) { + let message = candidate.into_message(); + let Some(dedup_key) = message.dedup_key.clone() else { + messages.push(message); + continue; + }; + + if let Some(&existing_index) = index_by_dedup_key.get(&dedup_key) { + merge_copilot_duplicate(&mut messages[existing_index], message); + } else { + index_by_dedup_key.insert(dedup_key, messages.len()); + messages.push(message); + } + } + + messages } #[derive(Clone, Copy, Eq, PartialEq)] @@ -127,6 +141,43 @@ impl CopilotUsageCandidate { } } +fn merge_copilot_duplicate(existing: &mut UnifiedMessage, duplicate: UnifiedMessage) { + // Copilot exports can repeat one span as its OTEL record is updated. Keep + // the most complete token buckets, but never add replayed usage together. + if existing.agent.is_none() { + existing.agent = duplicate.agent.clone(); + } + existing.tokens.input = existing.tokens.input.max(duplicate.tokens.input); + existing.tokens.output = existing.tokens.output.max(duplicate.tokens.output); + existing.tokens.cache_read = existing.tokens.cache_read.max(duplicate.tokens.cache_read); + existing.tokens.cache_write = existing + .tokens + .cache_write + .max(duplicate.tokens.cache_write); + existing.tokens.reasoning = existing.tokens.reasoning.max(duplicate.tokens.reasoning); + + // Sessionization needs the activity start, not the response completion. A + // repeated span can carry a different start and duration, so merge the + // observed time envelope without adding replayed duration. + let existing_start = existing.timestamp; + let duplicate_start = duplicate.timestamp; + let merged_start = existing_start.min(duplicate_start); + let merged_end = existing + .duration_ms + .map(|duration| existing_start.saturating_add(duration)) + .into_iter() + .chain( + duplicate + .duration_ms + .map(|duration| duplicate_start.saturating_add(duration)), + ) + .max(); + if merged_start != existing_start { + existing.set_timestamp(merged_start); + } + existing.duration_ms = merged_end.map(|end| end.saturating_sub(merged_start)); +} + fn collect_trace_contexts(records: &[Value]) -> HashMap { let mut contexts = HashMap::new(); @@ -691,7 +742,7 @@ fn attr_i64_first(attributes: &Map, keys: &[&str]) -> i64 { .unwrap_or(0) } -fn normalize_input_tokens( +pub(crate) fn normalize_input_tokens( input: i64, output: i64, cache_read: i64, @@ -754,9 +805,14 @@ fn value_as_i64(value: &Value) -> Option { fn timestamp_ms_from_record(value: &Value) -> Option { value - .get("endTime") + .get("startTime") .and_then(timestamp_ms_from_value) - .or_else(|| value.get("startTime").and_then(timestamp_ms_from_value)) + .or_else(|| { + // When only endTime is available, back-calculate the start if duration is known. + let end_ms = value.get("endTime").and_then(timestamp_ms_from_value)?; + let duration = duration_ms_from_record(value).unwrap_or(0); + Some(end_ms.saturating_sub(duration)) + }) .or_else(|| value.get("hrTime").and_then(timestamp_ms_from_value)) .or_else(|| value.get("_hrTime").and_then(timestamp_ms_from_value)) .or_else(|| value.get("time").and_then(timestamp_ms_from_value)) @@ -870,11 +926,46 @@ mod tests { assert_eq!(message.tokens.output, 281); assert_eq!(message.tokens.cache_read, 123); assert_eq!(message.tokens.reasoning, 128); - assert_eq!(message.timestamp, 1_775_934_264_967); + assert_eq!(message.timestamp, 1_775_934_260_133); assert_eq!(message.duration_ms, Some(4834)); assert_eq!(message.dedup_key.as_deref(), Some("trace-1:span-1")); } + #[test] + fn test_parse_copilot_back_calculates_start_from_end_and_duration() { + let content = r#"{"type":"span","traceId":"trace-back","spanId":"span-back","name":"chat gpt-5.4-mini","endTime":[1775934266,0],"duration":[7,0],"attributes":{"gen_ai.operation.name":"chat","gen_ai.response.model":"gpt-5.4-mini","gen_ai.usage.input_tokens":10,"gen_ai.usage.output_tokens":5}}"#; + let file = create_test_file(content); + + let messages = parse_copilot_file(file.path()); + + assert_eq!(messages.len(), 1); + assert_eq!(messages[0].timestamp, 1_775_934_259_000); + assert_eq!(messages[0].duration_ms, Some(7_000)); + } + + #[test] + fn test_parse_copilot_merges_duplicate_spans_without_additive_duration() { + let content = r#"{"type":"span","traceId":"trace-duplicate","spanId":"span-duplicate","name":"chat gpt-5.4-mini","startTime":[1775934260,0],"endTime":[1775934270,0],"attributes":{"gen_ai.operation.name":"chat","gen_ai.response.model":"gpt-5.4-mini","gen_ai.usage.input_tokens":10,"gen_ai.usage.output_tokens":20}} +{"type":"span","traceId":"trace-duplicate","spanId":"span-duplicate","name":"chat gpt-5.4-mini","startTime":[1775934255,0],"endTime":[1775934261,0],"attributes":{"gen_ai.operation.name":"chat","gen_ai.response.model":"gpt-5.4-mini","gen_ai.usage.input_tokens":30,"gen_ai.usage.output_tokens":50}}"#; + let file = create_test_file(content); + + let messages = parse_copilot_file(file.path()); + + assert_eq!( + messages.len(), + 1, + "duplicate spans must collapse before folds" + ); + assert_eq!(messages[0].timestamp, 1_775_934_255_000); + assert_eq!( + messages[0].duration_ms, + Some(15_000), + "the merged span must preserve the latest observed endpoint" + ); + assert_eq!(messages[0].tokens.input, 30); + assert_eq!(messages[0].tokens.output, 50); + } + #[test] fn test_parse_copilot_ignores_non_chat_spans() { let content = r#"{"type":"span","traceId":"trace-1","spanId":"tool-1","name":"execute_tool rg","attributes":{"gen_ai.operation.name":"execute_tool","gen_ai.tool.name":"rg"}} diff --git a/vendor/tokscale-core/src/sessions/copilot_desktop.rs b/vendor/tokscale-core/src/sessions/copilot_desktop.rs new file mode 100644 index 0000000..0f3cf96 --- /dev/null +++ b/vendor/tokscale-core/src/sessions/copilot_desktop.rs @@ -0,0 +1,504 @@ +//! GitHub Copilot Desktop SQLite parser. +//! +//! The macOS desktop app stores aggregate token totals in `~/.copilot/data.db` +//! and per-session event metadata in `~/.copilot/session-state/{session_id}`. + +use super::{normalize_workspace_key, workspace_label_from_key, UnifiedMessage}; +use crate::provider_identity::inferred_provider_from_model; +use chrono::{DateTime, NaiveDateTime}; +use rusqlite::{Connection, OpenFlags}; +use serde_json::Value; +use std::io::{BufRead, BufReader}; +use std::path::{Path, PathBuf}; +use tracing::warn; + +#[derive(Debug)] +struct CopilotDesktopSessionRow { + id: String, + model: Option, + total_input_tokens: i64, + total_output_tokens: i64, + total_cached_tokens: i64, + total_reasoning_tokens: i64, + created_at: Option, + agent: Option, +} + +#[derive(Debug, Default)] +struct SessionStateMetadata { + model: Option, + cwd: Option, +} + +pub fn parse_copilot_desktop_db(db_path: &Path) -> Vec { + let conn = match Connection::open_with_flags( + db_path, + OpenFlags::SQLITE_OPEN_READ_ONLY | OpenFlags::SQLITE_OPEN_NO_MUTEX, + ) { + Ok(conn) => conn, + Err(err) => { + warn!( + db_path = %db_path.display(), + error = %err, + "Failed to open Copilot Desktop database" + ); + return Vec::new(); + } + }; + + let mut stmt = match conn.prepare( + r#" + SELECT + id, + model, + total_input_tokens, + total_output_tokens, + total_cached_tokens, + total_reasoning_tokens, + created_at, + agent + FROM sessions + WHERE total_input_tokens > 0 + OR total_output_tokens > 0 + OR total_cached_tokens > 0 + OR total_reasoning_tokens > 0 + "#, + ) { + Ok(stmt) => stmt, + Err(err) => { + warn!( + db_path = %db_path.display(), + error = %err, + "Failed to prepare Copilot Desktop sessions query" + ); + return Vec::new(); + } + }; + + let rows = match stmt.query_map([], |row| { + Ok(CopilotDesktopSessionRow { + id: row.get(0)?, + model: row.get(1)?, + total_input_tokens: row.get::<_, Option>(2)?.unwrap_or(0), + total_output_tokens: row.get::<_, Option>(3)?.unwrap_or(0), + total_cached_tokens: row.get::<_, Option>(4)?.unwrap_or(0), + total_reasoning_tokens: row.get::<_, Option>(5)?.unwrap_or(0), + created_at: row.get(6)?, + agent: row.get(7)?, + }) + }) { + Ok(rows) => rows, + Err(err) => { + warn!( + db_path = %db_path.display(), + error = %err, + "Failed to execute Copilot Desktop sessions query" + ); + return Vec::new(); + } + }; + + rows.filter_map(|row| match row { + Ok(row) => Some(session_row_to_message(db_path, row)), + Err(err) => { + warn!( + db_path = %db_path.display(), + error = %err, + "Failed to decode Copilot Desktop session row" + ); + None + } + }) + .collect() +} + +pub(crate) fn session_state_event_paths(db_path: &Path) -> std::io::Result> { + let Some(copilot_root) = db_path.parent() else { + return Ok(Vec::new()); + }; + let entries = match std::fs::read_dir(copilot_root.join("session-state")) { + Ok(entries) => entries, + Err(error) if error.kind() == std::io::ErrorKind::NotFound => return Ok(Vec::new()), + Err(error) => return Err(error), + }; + + let mut paths = Vec::new(); + for entry in entries { + let entry = entry?; + if !entry.file_type()?.is_dir() { + return Err(std::io::Error::other( + "Copilot Desktop session-state entry is not a directory", + )); + } + let path = entry.path().join("events.jsonl"); + match std::fs::metadata(&path) { + Ok(metadata) if metadata.is_file() => { + std::fs::File::open(&path)?; + paths.push(path); + } + Ok(_) => {} + Err(error) if error.kind() == std::io::ErrorKind::NotFound => {} + Err(error) => return Err(error), + } + } + paths.sort_unstable(); + paths.dedup(); + Ok(paths) +} + +fn session_row_to_message(db_path: &Path, row: CopilotDesktopSessionRow) -> UnifiedMessage { + let metadata = read_session_state_metadata(db_path, &row.id); + let model_id = metadata + .model + .as_deref() + .or(row.model.as_deref()) + .map(str::trim) + .filter(|model| !model.is_empty()) + .unwrap_or("auto") + .to_string(); + let provider_id = inferred_provider_from_model(&model_id) + .unwrap_or("github-copilot") + .to_string(); + + let timestamp_ms = row + .created_at + .as_deref() + .and_then(parse_iso8601_timestamp_ms) + .unwrap_or_else(|| { + warn!( + session_id = %row.id, + created_at = ?row.created_at, + "Copilot Desktop session has unparseable created_at; defaulting to 0" + ); + 0 + }); + + let mut message = UnifiedMessage::new_with_dedup( + "copilot", + model_id, + provider_id, + row.id.clone(), + timestamp_ms, + super::copilot::normalize_input_tokens( + row.total_input_tokens, + row.total_output_tokens, + row.total_cached_tokens, + 0, + row.total_reasoning_tokens, + ), + 0.0, + Some(format!("copilot-desktop:{}", row.id)), + ); + message.agent = row + .agent + .map(|agent| agent.trim().to_string()) + .filter(|agent| !agent.is_empty()); + + if let Some(workspace_key) = metadata.cwd.as_deref().and_then(normalize_workspace_key) { + let workspace_label = workspace_label_from_key(&workspace_key); + message.set_workspace(Some(workspace_key), workspace_label); + } + + message +} + +fn read_session_state_metadata(db_path: &Path, session_id: &str) -> SessionStateMetadata { + let Some(copilot_root) = db_path.parent() else { + return SessionStateMetadata::default(); + }; + let events_path = copilot_root + .join("session-state") + .join(session_id) + .join("events.jsonl"); + + read_events_metadata(&events_path) +} + +fn read_events_metadata(events_path: &Path) -> SessionStateMetadata { + let file = match std::fs::File::open(events_path) { + Ok(file) => file, + Err(_) => return SessionStateMetadata::default(), + }; + + let mut metadata = SessionStateMetadata::default(); + for line in BufReader::new(file).lines().map_while(Result::ok) { + let trimmed = line.trim(); + if trimmed.is_empty() { + continue; + } + + let Ok(event) = serde_json::from_str::(trimmed) else { + continue; + }; + let Some(event_type) = event.get("type").and_then(Value::as_str) else { + continue; + }; + + match event_type { + "session.start" if metadata.cwd.is_none() => { + metadata.cwd = event + .pointer("/data/context/cwd") + .and_then(Value::as_str) + .map(str::trim) + .filter(|cwd| !cwd.is_empty()) + .map(str::to_string); + } + "session.model_change" => { + if let Some(model) = event + .pointer("/data/newModel") + .and_then(Value::as_str) + .map(str::trim) + .filter(|model| !model.is_empty() && model != &"auto") + { + metadata.model = Some(model.to_string()); + } + } + _ => {} + } + } + + metadata +} + +fn parse_iso8601_timestamp_ms(value: &str) -> Option { + DateTime::parse_from_rfc3339(value) + .map(|timestamp| timestamp.timestamp_millis()) + .ok() + .or_else(|| { + NaiveDateTime::parse_from_str(value, "%Y-%m-%d %H:%M:%S") + .ok() + .map(|timestamp| timestamp.and_utc().timestamp_millis()) + }) + .or_else(|| { + NaiveDateTime::parse_from_str(value, "%Y-%m-%d %H:%M:%S%.f") + .ok() + .map(|timestamp| timestamp.and_utc().timestamp_millis()) + }) + .or_else(|| { + NaiveDateTime::parse_from_str(value, "%Y-%m-%dT%H:%M:%S%.f") + .ok() + .map(|timestamp| timestamp.and_utc().timestamp_millis()) + }) + .or_else(|| { + let numeric = value.parse::().ok()?; + if numeric > 10_000_000_000 { + Some(numeric) + } else { + Some(numeric.saturating_mul(1000)) + } + }) +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::sessions::CostSource; + use rusqlite::{params, Connection}; + use std::fs::{self, File}; + use std::io::Write; + + fn create_copilot_desktop_db(path: &Path) -> Connection { + let conn = Connection::open(path).unwrap(); + conn.execute_batch( + r#" + CREATE TABLE sessions ( + id TEXT, + model TEXT, + total_input_tokens INTEGER, + total_output_tokens INTEGER, + total_cached_tokens INTEGER, + total_reasoning_tokens INTEGER, + total_nano_aiu INTEGER, + created_at TEXT, + agent TEXT + ); + "#, + ) + .unwrap(); + conn + } + + #[allow(clippy::too_many_arguments)] + fn insert_session( + conn: &Connection, + id: &str, + model: &str, + input: i64, + output: i64, + cached: i64, + reasoning: i64, + nano_aiu: i64, + agent: Option<&str>, + ) { + conn.execute( + r#" + INSERT INTO sessions ( + id, model, total_input_tokens, total_output_tokens, + total_cached_tokens, total_reasoning_tokens, total_nano_aiu, + created_at, agent + ) VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9) + "#, + params![ + id, + model, + input, + output, + cached, + reasoning, + nano_aiu, + "2026-07-01T12:34:56Z", + agent + ], + ) + .unwrap(); + } + + fn write_events(root: &Path, session_id: &str, lines: &[&str]) -> PathBuf { + let events_dir = root.join("session-state").join(session_id); + fs::create_dir_all(&events_dir).unwrap(); + let path = events_dir.join("events.jsonl"); + let mut file = File::create(&path).unwrap(); + for line in lines { + writeln!(file, "{line}").unwrap(); + } + path + } + + #[test] + fn parse_copilot_desktop_db_reads_token_sessions_and_agent() { + let dir = tempfile::tempdir().unwrap(); + let db_path = dir.path().join("data.db"); + let conn = create_copilot_desktop_db(&db_path); + insert_session( + &conn, + "session-1", + "gpt-5.1-codex", + 100, + 50, + 25, + 10, + 0, + Some(" github.copilot.default "), + ); + drop(conn); + + let messages = parse_copilot_desktop_db(&db_path); + + assert_eq!(messages.len(), 1); + let message = &messages[0]; + assert_eq!(message.client, "copilot"); + assert_eq!(message.model_id, "gpt-5.1-codex"); + assert_eq!(message.provider_id, "openai"); + assert_eq!(message.session_id, "session-1"); + assert_eq!(message.timestamp, 1_782_909_296_000); + assert_eq!(message.tokens.input, 75); + assert_eq!(message.tokens.output, 50); + assert_eq!(message.tokens.cache_read, 25); + assert_eq!(message.tokens.cache_write, 0); + assert_eq!(message.tokens.reasoning, 10); + assert_eq!(message.agent.as_deref(), Some("github.copilot.default")); + assert_eq!(message.cost, 0.0); + assert_eq!(message.cost_source, CostSource::Unknown); + assert_eq!( + message.dedup_key.as_deref(), + Some("copilot-desktop:session-1") + ); + } + + #[test] + fn parse_copilot_desktop_db_skips_aiu_only_sessions() { + let dir = tempfile::tempdir().unwrap(); + let db_path = dir.path().join("data.db"); + let conn = create_copilot_desktop_db(&db_path); + insert_session(&conn, "session-1", "auto", 0, 0, 0, 0, 42, None); + drop(conn); + + assert!(parse_copilot_desktop_db(&db_path).is_empty()); + } + + #[test] + fn parse_copilot_desktop_db_enriches_model_and_workspace_from_events() { + let dir = tempfile::tempdir().unwrap(); + let db_path = dir.path().join("data.db"); + let conn = create_copilot_desktop_db(&db_path); + insert_session(&conn, "session-1", "auto", 100, 50, 0, 0, 0, None); + drop(conn); + write_events( + dir.path(), + "session-1", + &[ + r#"{"type":"session.start","data":{"context":{"cwd":"/Users/alice/project"}}}"#, + r#"{"type":"session.model_change","data":{"newModel":"claude-sonnet-4-5"}}"#, + ], + ); + + let messages = parse_copilot_desktop_db(&db_path); + + assert_eq!(messages.len(), 1); + let message = &messages[0]; + assert_eq!(message.model_id, "claude-sonnet-4-5"); + assert_eq!(message.provider_id, "anthropic"); + assert_eq!(message.workspace_label.as_deref(), Some("project")); + } + + #[test] + fn session_state_dependencies_are_sorted_and_token_independent() { + let dir = tempfile::tempdir().unwrap(); + let db_path = dir.path().join("data.db"); + File::create(&db_path).unwrap(); + let z = write_events(dir.path(), "z-session", &["{}"]); + let a = write_events(dir.path(), "a-session", &["{}"]); + fs::create_dir_all(dir.path().join("session-state/missing-events")).unwrap(); + + assert_eq!(session_state_event_paths(&db_path).unwrap(), vec![a, z]); + } + + #[test] + fn session_state_dependency_probe_reports_non_missing_errors() { + let dir = tempfile::tempdir().unwrap(); + let db_path = dir.path().join("data.db"); + File::create(&db_path).unwrap(); + fs::create_dir_all(dir.path().join("session-state")).unwrap(); + File::create(dir.path().join("session-state/not-a-directory")).unwrap(); + + assert!(session_state_event_paths(&db_path).is_err()); + } + + #[cfg(unix)] + #[test] + fn session_state_dependency_probe_reports_unreadable_events() { + use std::os::unix::fs::PermissionsExt; + + let dir = tempfile::tempdir().unwrap(); + let db_path = dir.path().join("data.db"); + File::create(&db_path).unwrap(); + let events = write_events(dir.path(), "session-1", &["{}"]); + fs::set_permissions(&events, fs::Permissions::from_mode(0o000)).unwrap(); + + let result = session_state_event_paths(&db_path); + fs::set_permissions(&events, fs::Permissions::from_mode(0o600)).unwrap(); + + assert!(result.is_err()); + } + + #[test] + fn parse_timestamp_handles_sqlite_and_numeric_forms() { + assert_eq!( + parse_iso8601_timestamp_ms("2026-07-01 12:34:56.789"), + Some(1_782_909_296_789) + ); + assert_eq!( + parse_iso8601_timestamp_ms("2026-07-01T12:34:56.789"), + Some(1_782_909_296_789) + ); + assert_eq!( + parse_iso8601_timestamp_ms("1782909296"), + Some(1_782_909_296_000) + ); + assert_eq!( + parse_iso8601_timestamp_ms("1782909296789"), + Some(1_782_909_296_789) + ); + assert_eq!(parse_iso8601_timestamp_ms("not-a-timestamp"), None); + } +} diff --git a/vendor/tokscale-core/src/sessions/gjc.rs b/vendor/tokscale-core/src/sessions/gjc.rs index ee42e6e..813309b 100644 --- a/vendor/tokscale-core/src/sessions/gjc.rs +++ b/vendor/tokscale-core/src/sessions/gjc.rs @@ -179,12 +179,12 @@ pub fn parse_gjc_file(path: &Path) -> Vec { None => continue, }; - // A missing provider field is recoverable: infer it from the model name - // (and fall back to "gjc") rather than dropping a message that carries - // valid tokens. + // A missing/blank provider field is recoverable: infer it from the + // model name (and fall back to "gjc") rather than dropping a message + // that carries valid tokens. let provider = match message.provider { - Some(p) => p, - None => inferred_provider_from_model(&model) + Some(p) if !p.is_empty() => p, + _ => inferred_provider_from_model(&model) .unwrap_or("gjc") .to_string(), }; @@ -563,6 +563,19 @@ not valid json at all ); } + #[test] + fn test_parse_gjc_infers_provider_from_model_when_blank() { + // A blank provider must follow the same inference path as an absent + // provider, rather than leaking an empty provider identity downstream. + let content = r#"{"type":"session","id":"gjc_ses_blank_provider","cwd":"/tmp"} +{"type":"message","id":"msg_blank_provider","message":{"role":"assistant","model":"gpt-5","provider":"","timestamp":1700000001000,"usage":{"input":10,"output":5,"cost":{"total":0.02}}}}"#; + let file = create_test_file(content); + let messages = parse_gjc_file(file.path()); + + assert_eq!(messages.len(), 1); + assert_eq!(messages[0].provider_id, "openai"); + } + /// (e) Negative token values are clamped to >= 0. #[test] fn test_adv_negative_token_values_clamped_to_zero() { diff --git a/vendor/tokscale-core/src/sessions/grok.rs b/vendor/tokscale-core/src/sessions/grok.rs index 0499518..d1e3f9d 100644 --- a/vendor/tokscale-core/src/sessions/grok.rs +++ b/vendor/tokscale-core/src/sessions/grok.rs @@ -2,27 +2,46 @@ //! //! Grok Build writes JSON-RPC session updates under //! `~/.grok/sessions///updates.jsonl`. -//! Session rollups also land in sibling `signals.json` (including -//! `totalTokensBeforeCompaction` and `contextTokensUsed`). Current update -//! logs expose cumulative `totalTokens` counters without a stable -//! input/output split, so this parser records per-turn positive total-token -//! deltas as input tokens and reconciles any remaining `signals.json` total -//! so compacted sessions are not under-counted. +//! +//! **Primary path (modern logs):** each `sessionUpdate: turn_completed` carries +//! authoritative per-turn `params.update.usage` (`inputTokens`, `outputTokens`, +//! `reasoningTokens`, `cachedReadTokens`, optional `modelUsage`, +//! `costUsdTicks`). Those are the real API totals; multi-call turns re-bill +//! context, so they dwarf the context-window counter. +//! +//! **Legacy / live path:** older logs (and the open turn before +//! `turn_completed`) only expose cumulative `params._meta.totalTokens`, which +//! tracks **context occupancy**, not cumulative spend. Positive deltas of that +//! counter are recorded as input tokens; local compaction epochs bank large +//! rewinds; sibling `signals.json` can reconcile remaining context undercount +//! when the usage path is unavailable. +//! +//! **Unified log:** recent releases also write per-inference token buckets to +//! the global `~/.grok/logs/unified.jsonl`, which replaces legacy/update rows +//! for covered sessions via `prefer_unified_log_messages`. + use super::utils::{ extract_i64, extract_string, file_modified_timestamp_ms, parse_timestamp_value, read_file_or_none, }; -use super::{normalize_workspace_key, workspace_label_from_key, UnifiedMessage}; +use super::{normalize_workspace_key, workspace_label_from_key, CostSource, UnifiedMessage}; use crate::TokenBreakdown; use serde_json::Value; +use std::collections::{HashMap, HashSet}; use std::io::{BufRead, BufReader}; use std::path::{Path, PathBuf}; + const CLIENT_ID: &str = "grok"; const PROVIDER_ID: &str = "xai"; const UNKNOWN_MODEL: &str = "grok-unknown"; const COMPACTION_MIN_DROP_TOKENS: i64 = 32_000; +const UNIFIED_LOG_DEDUP_PREFIX: &str = "grok-unified:"; +/// xAI stamps `costUsdTicks` as nano-USD (`1e9` ticks = $1). Observed magnitudes +/// match that scale on real sessions; if the unit is ever proven different, only +/// cost mapping changes. +const COST_USD_TICKS_PER_DOLLAR: f64 = 1_000_000_000.0; #[derive(Debug, Clone)] struct GrokMetadata { @@ -110,6 +129,47 @@ impl ActiveTurn { } } +#[derive(Debug, Clone)] +struct ParsedUsage { + input_tokens: i64, + output_tokens: i64, + reasoning_tokens: i64, + cached_read_tokens: i64, + cost_usd_ticks: i64, + api_duration_ms: Option, + model_id: Option, +} + +impl ParsedUsage { + fn has_positive_tokens(&self) -> bool { + self.input_tokens > 0 + || self.output_tokens > 0 + || self.reasoning_tokens > 0 + || self.cached_read_tokens > 0 + } + + /// Grok's `inputTokens` includes cache reads. TokenBar totals sum every + /// bucket, so net uncached input into `input` and put cache in `cache_read`. + fn into_token_breakdown(&self) -> TokenBreakdown { + let cache_read = self.cached_read_tokens.max(0); + let input = self.input_tokens.max(0).saturating_sub(cache_read); + TokenBreakdown { + input, + output: self.output_tokens.max(0), + cache_read, + cache_write: 0, + reasoning: self.reasoning_tokens.max(0), + } + } + + fn cost_usd(&self) -> Option { + if self.cost_usd_ticks <= 0 { + return None; + } + Some(self.cost_usd_ticks as f64 / COST_USD_TICKS_PER_DOLLAR) + } +} + pub fn parse_grok_updates_file(path: &Path) -> Vec { if path.file_name().and_then(|name| name.to_str()) != Some("updates.jsonl") { return Vec::new(); @@ -121,7 +181,8 @@ pub fn parse_grok_updates_file(path: &Path) -> Vec { Err(_) => return Vec::new(), }; - let mut messages = Vec::new(); + let mut usage_messages = Vec::new(); + let mut context_messages = Vec::new(); let mut current_model = metadata .model_id .clone() @@ -130,6 +191,15 @@ pub fn parse_grok_updates_file(path: &Path) -> Vec { let mut last_total_timestamp = metadata.timestamp; let mut active_turn: Option = None; let mut turn_index = 0usize; + let mut usage_turn_index = 0usize; + let mut saw_usage = false; + // Context baseline after the latest completed usage turn so a live open turn + // only counts post-completion context growth (no double-count). + let mut context_baseline_after_usage: Option = None; + // After saw_usage, a user_message_chunk with no known context counter must + // not open ActiveTurn at 0 (full occupancy would double-count completed + // usage). Defer until the first post-usage totalTokens establishes baseline. + let mut pending_post_usage_live_partial = false; for line in BufReader::new(file).lines().map_while(Result::ok) { if line.trim().is_empty() { @@ -150,20 +220,78 @@ pub fn parse_grok_updates_file(path: &Path) -> Vec { } let timestamp = extract_timestamp_ms(&value).unwrap_or(metadata.timestamp); - if is_user_message_chunk(&value) { - if let Some(turn) = active_turn.take() { - if let Some(message) = turn.into_message(&metadata) { - messages.push(message); + + // Primary path: authoritative per-turn API usage. + if is_turn_completed(&value) { + if let Some(usage_value) = get_path(&value, &["params", "update", "usage"]) { + let prompt_id = get_path(&value, &["params", "update", "prompt_id"]) + .and_then(|v| extract_string(Some(v))); + let emitted = emit_usage_messages( + usage_value, + &metadata, + ¤t_model, + timestamp, + usage_turn_index, + prompt_id.as_deref(), + ); + if !emitted.is_empty() { + saw_usage = true; + usage_turn_index = usage_turn_index.saturating_add(1); + usage_messages.extend(emitted); + // Close any open context turn without emitting — usage owns + // the completed turn's totals. + active_turn = None; + pending_post_usage_live_partial = false; + // If this line also carries a context counter, advance the + // baseline from it but do not open a new ActiveTurn from + // the completed turn's context growth (would double-count). + if let Some(total) = extract_total_tokens(&value).filter(|t| *t >= 0) { + context_baseline_after_usage = Some(total); + last_total = Some(total); + last_total_timestamp = timestamp; + } else if let Some(total) = last_total { + context_baseline_after_usage = Some(total); + } + continue; } } + } - active_turn = Some(ActiveTurn::new( - last_total.unwrap_or(0), - timestamp, - current_model.clone(), - turn_index, - )); - turn_index = turn_index.saturating_add(1); + if is_user_message_chunk(&value) { + if !saw_usage { + if let Some(turn) = active_turn.take() { + if let Some(message) = turn.into_message(&metadata) { + context_messages.push(message); + } + } + pending_post_usage_live_partial = false; + let baseline = last_total.unwrap_or(0); + active_turn = Some(ActiveTurn::new( + baseline, + timestamp, + current_model.clone(), + turn_index, + )); + turn_index = turn_index.saturating_add(1); + } else { + // Usage path: discard any open context turn (completed turns are + // owned by usage). Never zero-base a post-usage live partial — + // without a known context counter, wait for the first post-usage + // totalTokens instead of treating full occupancy as delta. + if let Some(baseline) = context_baseline_after_usage.or(last_total) { + pending_post_usage_live_partial = false; + active_turn = Some(ActiveTurn::new( + baseline, + timestamp, + current_model.clone(), + turn_index, + )); + turn_index = turn_index.saturating_add(1); + } else { + pending_post_usage_live_partial = true; + active_turn = None; + } + } } let Some(total_tokens) = extract_total_tokens(&value) else { @@ -177,6 +305,27 @@ pub fn parse_grok_updates_file(path: &Path) -> Vec { Some(previous) if total_tokens < previous => { if is_compaction_reset(previous, total_tokens) { if active_turn.is_none() { + // After usage, only user_message_chunk may open a live + // partial. Compaction without an open user turn only + // advances the counter baseline. + if saw_usage { + last_total_timestamp = timestamp; + last_total = Some(total_tokens); + context_baseline_after_usage = Some(total_tokens); + // First known counter after deferred user turn: + // open at the post-compaction total (baseline only). + if pending_post_usage_live_partial { + pending_post_usage_live_partial = false; + active_turn = Some(ActiveTurn::new( + total_tokens, + timestamp, + current_model.clone(), + turn_index, + )); + turn_index = turn_index.saturating_add(1); + } + continue; + } let mut turn = ActiveTurn::new( 0, last_total_timestamp, @@ -203,8 +352,32 @@ pub fn parse_grok_updates_file(path: &Path) -> Vec { } Some(previous) => { if active_turn.is_none() { + // After saw_usage / completed usage, do not open a new + // ActiveTurn from non-user updates that only bump + // `_meta.totalTokens`. Only user_message_chunk may start a + // post-usage live partial; context-only growth advances the + // baseline so the next user turn does not inherit it. + if saw_usage { + last_total_timestamp = timestamp; + last_total = Some(total_tokens); + context_baseline_after_usage = Some(total_tokens); + if pending_post_usage_live_partial { + // Deferred user turn: first/next counter is baseline + // only (no zero-base full-occupancy emission). + pending_post_usage_live_partial = false; + active_turn = Some(ActiveTurn::new( + total_tokens, + timestamp, + current_model.clone(), + turn_index, + )); + turn_index = turn_index.saturating_add(1); + } + continue; + } + let baseline = previous; active_turn = Some(ActiveTurn::new( - previous, + baseline, timestamp, current_model.clone(), turn_index, @@ -218,22 +391,59 @@ pub fn parse_grok_updates_file(path: &Path) -> Vec { last_total = Some(total_tokens); } None => { - if let Some(turn) = active_turn.as_mut() { - turn.observe_total(total_tokens, timestamp); - } + // First context counter observed in this file. last_total_timestamp = timestamp; last_total = Some(total_tokens); + if saw_usage { + context_baseline_after_usage = Some(total_tokens); + if pending_post_usage_live_partial { + // User already opened a post-usage turn with no known + // baseline: use this first counter as baseline only + // (delta 0), not as growth from zero. + pending_post_usage_live_partial = false; + active_turn = Some(ActiveTurn::new( + total_tokens, + timestamp, + current_model.clone(), + turn_index, + )); + turn_index = turn_index.saturating_add(1); + } else if let Some(turn) = active_turn.as_mut() { + // Should not zero-base after usage; observe only if a + // turn already exists with a real baseline. + turn.observe_total(total_tokens, timestamp); + } + } else if let Some(turn) = active_turn.as_mut() { + turn.observe_total(total_tokens, timestamp); + } + } + } + } + + if saw_usage { + // Live partial: open turn after the last completed usage only. + if let Some(turn) = active_turn { + if let Some(message) = turn.into_message(&metadata) { + usage_messages.push(message); } } + // signals.json is context occupancy — do not reconcile against usage. + // Keep any legacy context turns that completed before the first usage + // record (pre-upgrade prefix of a mixed session). + if context_messages.is_empty() { + return usage_messages; + } + context_messages.extend(usage_messages); + return context_messages; } if let Some(turn) = active_turn { if let Some(message) = turn.into_message(&metadata) { - messages.push(message); + context_messages.push(message); } } - if messages.is_empty() { + if context_messages.is_empty() { if let Some(total_tokens) = last_total.filter(|tokens| *tokens > 0) { let aggregate_turn = ActiveTurn { baseline_total: 0, @@ -244,13 +454,454 @@ pub fn parse_grok_updates_file(path: &Path) -> Vec { turn_index: 0, }; if let Some(message) = aggregate_turn.into_message(&metadata) { - messages.push(message); + context_messages.push(message); + } + } + } + + append_signals_reconciliation(path, &metadata, &mut context_messages, ¤t_model); + context_messages +} + +fn emit_usage_messages( + usage_value: &Value, + metadata: &GrokMetadata, + fallback_model: &str, + timestamp: i64, + turn_index: usize, + prompt_id: Option<&str>, +) -> Vec { + let top = parse_usage_object(usage_value, None); + let mut rows: Vec = Vec::new(); + + if let Some(model_usage) = usage_value.get("modelUsage").and_then(|v| v.as_object()) { + for (model_id, entry) in model_usage { + let parsed = parse_usage_object(entry, Some(model_id.as_str())); + if parsed.has_positive_tokens() { + rows.push(parsed); + } + } + } + + // True only when parent costUsdTicks was copied onto a modelUsage row. + // Multi-model zero-cost siblings may then be ProviderReported $0; otherwise + // an omitting sibling stays Unknown for apply_pricing estimation. + let mut parent_cost_inherited = false; + + if rows.is_empty() { + if top.has_positive_tokens() { + rows.push(top); + } else { + return Vec::new(); + } + } else { + // Model entries sometimes only carry token buckets while the parent + // usage object holds costUsdTicks / apiDurationMs. Inherit so we do not + // drop provider-reported cost. Single-model: full inherit. Multi-model: + // put parent totals on the first row only when no model entry has them + // (avoids double-counting the parent cost across models). + parent_cost_inherited = inherit_top_level_cost_and_duration(&mut rows, &top); + } + + let turn_key = prompt_id + .filter(|id| !id.trim().is_empty()) + .map(|id| format!("turn:{id}")) + .unwrap_or_else(|| format!("turn:{turn_index}")); + + let multi_model = rows.len() > 1; + // Only when multi-model inherits parent cost onto the first row must + // remaining zero-cost siblings be ProviderReported at $0 so pricing cannot + // stack estimates on top of that once-inherited total. If a model entry has + // its own costUsdTicks and a sibling simply omits cost (no parent inherit), + // the omitting row stays Unknown so apply_pricing can estimate it. + let mut messages = Vec::with_capacity(rows.len()); + for (model_i, row) in rows.into_iter().enumerate() { + let model_id = row + .model_id + .clone() + .filter(|m| !m.trim().is_empty()) + .or_else(|| metadata.model_id.clone()) + .filter(|m| !m.trim().is_empty()) + .unwrap_or_else(|| fallback_model.to_string()); + + let tokens = row.into_token_breakdown(); + if tokens.input == 0 + && tokens.output == 0 + && tokens.cache_read == 0 + && tokens.reasoning == 0 + { + continue; + } + + let dedup_key = if multi_model { + format!("grok:{}:{turn_key}:{model_id}", metadata.session_id) + } else { + format!("grok:{}:{turn_key}", metadata.session_id) + }; + + let (cost, cost_source) = match row.cost_usd() { + Some(cost) => (cost, CostSource::ProviderReported), + None if multi_model && parent_cost_inherited => { + (0.0, CostSource::ProviderReported) + } + None => (0.0, CostSource::Unknown), + }; + + let mut message = UnifiedMessage::new_with_dedup( + CLIENT_ID, + model_id, + PROVIDER_ID, + metadata.session_id.clone(), + timestamp, + tokens, + cost, + Some(dedup_key), + ); + if cost_source == CostSource::ProviderReported { + message.mark_provider_reported_cost(); + } + if let Some(duration_ms) = row.api_duration_ms.filter(|ms| *ms > 0) { + message.duration_ms = Some(duration_ms); + } + message.set_workspace( + metadata.workspace_key.clone(), + metadata.workspace_label.clone(), + ); + // Match unified-log semantics: only the first model row of a turn + // counts as a user message so report sums do not overcount. + message.is_turn_start = model_i == 0; + message.message_count = i32::from(message.is_turn_start); + messages.push(message); + } + + messages +} + +/// When `modelUsage` rows omit cost/duration, copy them from the parent usage +/// object. Single-model inherits fully; multi-model places parent totals on +/// the first row only if every model entry is missing them. +/// +/// Returns `true` when parent `costUsdTicks` was copied onto a row (so multi- +/// model zero-cost siblings may be sealed as ProviderReported $0). +fn inherit_top_level_cost_and_duration(rows: &mut [ParsedUsage], top: &ParsedUsage) -> bool { + if rows.is_empty() { + return false; + } + let mut cost_inherited = false; + if rows.len() == 1 { + if rows[0].cost_usd_ticks <= 0 && top.cost_usd_ticks > 0 { + rows[0].cost_usd_ticks = top.cost_usd_ticks; + cost_inherited = true; + } + if rows[0].api_duration_ms.is_none() { + rows[0].api_duration_ms = top.api_duration_ms; + } + return cost_inherited; + } + if top.cost_usd_ticks > 0 && rows.iter().all(|r| r.cost_usd_ticks <= 0) { + rows[0].cost_usd_ticks = top.cost_usd_ticks; + cost_inherited = true; + } + if top.api_duration_ms.is_some() && rows.iter().all(|r| r.api_duration_ms.is_none()) { + rows[0].api_duration_ms = top.api_duration_ms; + } + cost_inherited +} + +/// Parses Grok Build's append-only unified log. Each `shell.turn.inference_done` +/// record reports a prompt total that includes cached prompt tokens and a +/// completion total that includes reasoning tokens. Tokscale stores the +/// non-overlapping component buckets so their sum remains the source total. +pub fn parse_grok_unified_log_file(path: &Path) -> Vec { + if path.file_name().and_then(|name| name.to_str()) != Some("unified.jsonl") { + return Vec::new(); + } + + let file = match std::fs::File::open(path) { + Ok(file) => file, + Err(_) => return Vec::new(), + }; + + let fallback_timestamp = file_modified_timestamp_ms(path); + let mut fallback_model_by_pid = HashMap::new(); + let mut model_by_pid_and_session = HashMap::new(); + let mut model_by_session = HashMap::new(); + let mut seen = HashSet::new(); + let mut messages = Vec::new(); + + for line in BufReader::new(file).lines().map_while(Result::ok) { + if line.trim().is_empty() { + continue; + } + + let Ok(value) = serde_json::from_str::(&line) else { + continue; + }; + + if let Some(pid) = unified_log_process_start_pid(&value) { + // The unified log survives process restarts, so an OS-reused PID + // must not inherit model authority from the previous process. + fallback_model_by_pid.remove(&pid); + model_by_pid_and_session.retain(|(model_pid, _), _| *model_pid != pid); + continue; + } + + if let Some((pid, model_session_id, model_id)) = unified_log_model_change(&value) { + match (pid, model_session_id) { + (Some(pid), Some(session_id)) => { + model_by_pid_and_session.insert((pid, session_id), model_id); + } + (None, Some(session_id)) => { + model_by_pid_and_session + .retain(|(_, existing_session), _| existing_session != &session_id); + model_by_session.insert(session_id, model_id); + } + (Some(pid), None) => { + fallback_model_by_pid.insert(pid, model_id); + } + (None, None) => {} + } + continue; + } + + if value.get("msg").and_then(Value::as_str) != Some("shell.turn.inference_done") { + continue; + } + + let Some(session_id) = + extract_string(value.get("sid")).filter(|session_id| !session_id.trim().is_empty()) + else { + continue; + }; + let Some(context) = value.get("ctx") else { + continue; + }; + let Some(prompt_tokens) = required_non_negative_i64(context.get("prompt_tokens")) else { + continue; + }; + let Some(mut cached_prompt_tokens) = + optional_non_negative_i64(context.get("cached_prompt_tokens")) + else { + continue; + }; + let Some(completion_tokens) = required_non_negative_i64(context.get("completion_tokens")) + else { + continue; + }; + let Some(reasoning_tokens) = optional_non_negative_i64(context.get("reasoning_tokens")) + else { + continue; + }; + cached_prompt_tokens = cached_prompt_tokens.min(prompt_tokens); + + let loop_index = match context.get("loop_index") { + Some(value) => { + let Some(loop_index) = required_non_negative_i64(Some(value)) else { + continue; + }; + loop_index + } + None => 1, + }; + let Some(pid) = optional_non_negative_i64(value.get("pid")) else { + continue; + }; + let timestamp = value + .get("ts") + .and_then(parse_timestamp_value) + .unwrap_or(fallback_timestamp); + let reasoning = reasoning_tokens.min(completion_tokens); + let dedup_key = format!( + "{UNIFIED_LOG_DEDUP_PREFIX}{session_id}:{timestamp}:{pid}:{loop_index}:{prompt_tokens}:{cached_prompt_tokens}:{completion_tokens}:{reasoning_tokens}" + ); + if !seen.insert(dedup_key.clone()) { + continue; + } + + let model_id = model_by_pid_and_session + .get(&(pid, session_id.clone())) + .or_else(|| model_by_session.get(&session_id)) + .or_else(|| fallback_model_by_pid.get(&pid)) + .cloned() + .unwrap_or_else(|| UNKNOWN_MODEL.to_string()); + let mut message = UnifiedMessage::new_with_dedup( + CLIENT_ID, + model_id, + PROVIDER_ID, + session_id, + timestamp, + TokenBreakdown { + input: prompt_tokens.saturating_sub(cached_prompt_tokens), + output: completion_tokens.saturating_sub(reasoning), + cache_read: cached_prompt_tokens, + cache_write: 0, + reasoning, + }, + 0.0, + Some(dedup_key), + ); + // The unified log records one inference for each tool-loop iteration. + // In observed Grok logs, loop one starts the user turn; later loops do + // not represent additional user interactions or messages. + message.is_turn_start = loop_index == 1; + message.message_count = i32::from(message.is_turn_start); + messages.push(message); + } + + messages +} + +/// Dispatches between Grok's legacy per-session updates and its newer unified +/// log without accepting unrelated JSONL files under the Grok home directory. +pub fn parse_grok_file(path: &Path) -> Vec { + match path.file_name().and_then(|name| name.to_str()) { + Some("updates.jsonl") => parse_grok_updates_file(path), + Some("unified.jsonl") => parse_grok_unified_log_file(path), + _ => Vec::new(), + } +} + +/// Uses the richer, per-inference unified log for sessions it covers. Legacy +/// updates remain a fallback for sessions absent from that log, avoiding an +/// additive merge of two representations of the same activity. +pub fn prefer_unified_log_messages(mut messages: Vec) -> Vec { + let unified_sessions: HashSet = messages + .iter() + .filter(|message| is_unified_log_message(message)) + .map(|message| message.session_id.clone()) + .collect(); + + if unified_sessions.is_empty() { + return messages; + } + + let mut legacy_models = HashMap::new(); + let mut legacy_workspaces = HashMap::new(); + for message in messages + .iter() + .filter(|message| !is_unified_log_message(message)) + { + if message.model_id != UNKNOWN_MODEL { + match legacy_models.entry(message.session_id.clone()) { + std::collections::hash_map::Entry::Vacant(entry) => { + entry.insert(Some(message.model_id.clone())); + } + std::collections::hash_map::Entry::Occupied(mut entry) => { + if entry.get().as_ref() != Some(&message.model_id) { + entry.insert(None); + } + } + } + } + + let workspace = ( + message.workspace_key.clone(), + message.workspace_label.clone(), + ); + if workspace == (None, None) { + continue; + } + + match legacy_workspaces.entry(message.session_id.clone()) { + std::collections::hash_map::Entry::Vacant(entry) => { + entry.insert(Some(workspace)); + } + std::collections::hash_map::Entry::Occupied(mut entry) => { + if entry.get().as_ref() != Some(&workspace) { + entry.insert(None); + } + } + } + } + + for message in messages + .iter_mut() + .filter(|message| is_unified_log_message(message)) + { + if message.model_id == UNKNOWN_MODEL { + if let Some(Some(model_id)) = legacy_models.get(&message.session_id) { + message.model_id = model_id.clone(); + } + } + if message.workspace_key.is_none() && message.workspace_label.is_none() { + if let Some(Some((workspace_key, workspace_label))) = + legacy_workspaces.get(&message.session_id) + { + message.set_workspace(workspace_key.clone(), workspace_label.clone()); } } } - append_signals_reconciliation(path, &metadata, &mut messages, ¤t_model); messages + .into_iter() + .filter(|message| { + is_unified_log_message(message) || !unified_sessions.contains(&message.session_id) + }) + .collect() +} + +fn is_unified_log_message(message: &UnifiedMessage) -> bool { + message + .dedup_key + .as_deref() + .is_some_and(|key| key.starts_with(UNIFIED_LOG_DEDUP_PREFIX)) +} + +fn unified_log_process_start_pid(value: &Value) -> Option { + if value.get("msg").and_then(Value::as_str) != Some("AuthManager::new") { + return None; + } + required_non_negative_i64(value.get("pid")) +} + +fn unified_log_model_change(value: &Value) -> Option<(Option, Option, String)> { + let pid = match value.get("pid") { + Some(value) => Some(required_non_negative_i64(Some(value))?), + None => None, + }; + let context = value.get("ctx")?; + let model_id = match value.get("msg").and_then(Value::as_str)? { + "model changed" => extract_string(context.get("model")), + "model catalog: notifying clients" => extract_string(context.get("current_model_id")), + "backend_search: model switch" => extract_string(context.get("new_model")) + .or_else(|| extract_string(context.get("model"))) + .or_else(|| extract_string(context.get("current_model_id"))), + "subagent model resolved" => { + extract_string(context.get("model_id")).or_else(|| extract_string(context.get("model"))) + } + _ => None, + }?; + + let session_id = + extract_string(value.get("sid")).filter(|session_id| !session_id.trim().is_empty()); + (!model_id.trim().is_empty() && (pid.is_some() || session_id.is_some())) + .then_some((pid, session_id, model_id)) +} + +fn required_non_negative_i64(value: Option<&Value>) -> Option { + extract_i64(value).filter(|value| *value >= 0) +} + +fn optional_non_negative_i64(value: Option<&Value>) -> Option { + match value { + Some(value) => required_non_negative_i64(Some(value)), + None => Some(0), + } +} + +fn parse_usage_object(value: &Value, model_id: Option<&str>) -> ParsedUsage { + ParsedUsage { + input_tokens: non_negative_i64(value.get("inputTokens")), + output_tokens: non_negative_i64(value.get("outputTokens")), + reasoning_tokens: non_negative_i64(value.get("reasoningTokens")), + cached_read_tokens: non_negative_i64(value.get("cachedReadTokens")), + cost_usd_ticks: non_negative_i64(value.get("costUsdTicks")), + api_duration_ms: extract_i64(value.get("apiDurationMs")).filter(|ms| *ms > 0), + model_id: model_id + .map(str::to_string) + .or_else(|| extract_string(value.get("modelId"))) + .or_else(|| extract_string(value.get("model"))), + } } fn is_compaction_reset(previous: i64, current: i64) -> bool { @@ -477,12 +1128,14 @@ fn extract_model_id(value: &Value) -> Option { } fn extract_total_tokens(value: &Value) -> Option { + // Context occupancy counter only. Do NOT read params.update.usage.totalTokens + // here — that is per-turn API spend handled by the usage path, not a + // cumulative context counter. for path in [ &["params", "_meta", "totalTokens"][..], &["params", "update", "_meta", "totalTokens"][..], &["params", "update", "totalTokens"][..], &["params", "totalTokens"][..], - &["usage", "totalTokens"][..], &["totalTokens"][..], ] { if let Some(total) = get_path(value, path).and_then(|value| extract_i64(Some(value))) { @@ -512,6 +1165,11 @@ fn is_user_message_chunk(value: &Value) -> bool { == Some("user_message_chunk") } +fn is_turn_completed(value: &Value) -> bool { + get_path(value, &["params", "update", "sessionUpdate"]).and_then(|value| value.as_str()) + == Some("turn_completed") +} + fn get_path<'a>(value: &'a Value, path: &[&str]) -> Option<&'a Value> { path.iter() .try_fold(value, |current, key| current.get(*key)) @@ -575,6 +1233,289 @@ mod tests { (temp, updates_path) } + fn write_unified_fixture(unified_jsonl: &str) -> (tempfile::TempDir, PathBuf) { + let temp = tempfile::TempDir::new().unwrap(); + let logs_dir = temp.path().join(".grok/logs"); + std::fs::create_dir_all(&logs_dir).unwrap(); + let path = logs_dir.join("unified.jsonl"); + std::fs::write(&path, unified_jsonl).unwrap(); + (temp, path) + } + + fn test_message(session_id: &str, dedup_key: &str) -> UnifiedMessage { + UnifiedMessage::new_with_dedup( + CLIENT_ID, + "grok-build", + PROVIDER_ID, + session_id, + 1_700_000_000_000, + TokenBreakdown::default(), + 0.0, + Some(dedup_key.to_string()), + ) + } + + #[test] + fn parses_unified_log_token_breakdown_without_double_counting_reasoning() { + let (_temp, path) = write_unified_fixture( + r#"{"ts":"2023-11-14T22:13:19Z","pid":17,"sid":"session-1","msg":"model changed","ctx":{"model":"grok-composer-2.5-fast"}} +{"ts":"2023-11-14T22:13:19Z","pid":17,"msg":"model catalog: notifying clients","ctx":{"current_model_id":"grok-4.5"}} +{"ts":"2023-11-14T22:13:20Z","pid":17,"sid":"session-1","msg":"shell.turn.inference_done","ctx":{"loop_index":1,"prompt_tokens":100,"cached_prompt_tokens":60,"completion_tokens":25,"reasoning_tokens":5}} +{"ts":"2023-11-14T22:13:21Z","pid":17,"sid":"session-1","msg":"shell.turn.inference_done","ctx":{"loop_index":2,"prompt_tokens":80,"cached_prompt_tokens":0,"completion_tokens":12,"reasoning_tokens":0}} +{"ts":"2023-11-14T22:13:20Z","pid":17,"sid":"session-1","msg":"shell.turn.inference_done","ctx":{"loop_index":1,"prompt_tokens":100,"cached_prompt_tokens":60,"completion_tokens":25,"reasoning_tokens":5}} +{"ts":"2023-11-14T22:13:22Z","pid":17,"sid":"session-1","msg":"shell.turn.inference_done","ctx":{"loop_index":3,"prompt_tokens":10,"cached_prompt_tokens":11,"completion_tokens":1,"reasoning_tokens":0}}"#, + ); + + let messages = parse_grok_unified_log_file(&path); + + assert_eq!(messages.len(), 3); + assert_eq!(messages[0].client, CLIENT_ID); + assert_eq!(messages[0].model_id, "grok-composer-2.5-fast"); + assert_eq!(messages[0].session_id, "session-1"); + assert_eq!(messages[0].tokens.input, 40); + assert_eq!(messages[0].tokens.cache_read, 60); + assert_eq!(messages[0].tokens.output, 20); + assert_eq!(messages[0].tokens.reasoning, 5); + assert_eq!(messages[0].tokens.total(), 125); + assert_eq!(messages[0].message_count, 1); + assert!(messages[0].is_turn_start); + assert_eq!(messages[1].tokens.input, 80); + assert_eq!(messages[1].tokens.output, 12); + assert_eq!(messages[1].message_count, 0); + assert!(!messages[1].is_turn_start); + assert_eq!(messages[2].tokens.input, 0); + assert_eq!(messages[2].tokens.cache_read, 10); + assert_eq!(messages[2].tokens.output, 1); + assert_eq!(messages[2].tokens.total(), 11); + assert_eq!(messages[2].message_count, 0); + assert!(!messages[2].is_turn_start); + } + + #[test] + fn unified_log_counts_missing_loop_index_as_first_loop() { + let (_temp, path) = write_unified_fixture( + r#"{"ts":"2023-11-14T22:13:20Z","pid":17,"sid":"session-1","msg":"shell.turn.inference_done","ctx":{"prompt_tokens":100,"completion_tokens":25}} +{"ts":"2023-11-14T22:13:21Z","pid":17,"sid":"session-1","msg":"shell.turn.inference_done","ctx":{"loop_index":2,"prompt_tokens":80,"completion_tokens":12}}"#, + ); + + let messages = parse_grok_unified_log_file(&path); + + assert_eq!(messages.len(), 2); + assert_eq!(messages[0].message_count, 1); + assert!(messages[0].is_turn_start); + assert_eq!(messages[1].message_count, 0); + assert!(!messages[1].is_turn_start); + } + + #[test] + fn unified_log_keeps_distinct_inferences_that_share_base_identity() { + let (_temp, path) = write_unified_fixture( + r#"{"ts":"2023-11-14T22:13:20Z","pid":17,"sid":"session-1","msg":"shell.turn.inference_done","ctx":{"loop_index":1,"prompt_tokens":100,"cached_prompt_tokens":60,"completion_tokens":25,"reasoning_tokens":5}} +{"ts":"2023-11-14T22:13:20Z","pid":17,"sid":"session-1","msg":"shell.turn.inference_done","ctx":{"loop_index":1,"prompt_tokens":120,"cached_prompt_tokens":70,"completion_tokens":30,"reasoning_tokens":6}} +{"ts":"2023-11-14T22:13:20Z","pid":17,"sid":"session-1","msg":"shell.turn.inference_done","ctx":{"loop_index":1,"prompt_tokens":100,"cached_prompt_tokens":60,"completion_tokens":25,"reasoning_tokens":5}}"#, + ); + + let messages = parse_grok_unified_log_file(&path); + + assert_eq!(messages.len(), 2); + assert_ne!(messages[0].dedup_key, messages[1].dedup_key); + assert_eq!( + messages + .iter() + .map(|message| message.tokens.total()) + .sum::(), + 275 + ); + } + + #[test] + fn unified_log_applies_pidless_session_model_switch() { + let (_temp, path) = write_unified_fixture( + r#"{"ts":"2023-11-14T22:13:18Z","pid":17,"msg":"model catalog: notifying clients","ctx":{"current_model_id":"grok-4.5"}} +{"ts":"2023-11-14T22:13:19Z","pid":17,"sid":"session-with-model-event","msg":"model changed","ctx":{"model":"grok-composer-2.5-fast"}} +{"ts":"2023-11-14T22:13:20Z","pid":17,"sid":"session-with-model-event","msg":"shell.turn.inference_done","ctx":{"loop_index":1,"prompt_tokens":10,"completion_tokens":1}} +{"ts":"2023-11-14T22:13:21Z","sid":"session-with-model-event","msg":"model changed","ctx":{"model":"grok-4.1-fast"}} +{"ts":"2023-11-14T22:13:22Z","pid":17,"sid":"session-with-model-event","msg":"shell.turn.inference_done","ctx":{"loop_index":2,"prompt_tokens":15,"completion_tokens":2}} +{"ts":"2023-11-14T22:13:23Z","pid":17,"sid":"session-without-model-event","msg":"shell.turn.inference_done","ctx":{"loop_index":1,"prompt_tokens":20,"completion_tokens":2}}"#, + ); + + let messages = parse_grok_unified_log_file(&path); + + assert_eq!(messages.len(), 3); + assert_eq!(messages[0].model_id, "grok-composer-2.5-fast"); + assert_eq!(messages[1].model_id, "grok-4.1-fast"); + assert_eq!(messages[2].model_id, "grok-4.5"); + } + + #[test] + fn unified_log_expires_pid_scoped_models_on_process_restart() { + let (_temp, path) = write_unified_fixture( + r#"{"ts":"2023-11-14T22:13:17Z","sid":"session-stable","msg":"model changed","ctx":{"model":"grok-session"}} +{"ts":"2023-11-14T22:13:18Z","pid":17,"msg":"model catalog: notifying clients","ctx":{"current_model_id":"grok-old"}} +{"ts":"2023-11-14T22:13:19Z","pid":17,"sid":"session-old","msg":"shell.turn.inference_done","ctx":{"loop_index":1,"prompt_tokens":10,"completion_tokens":1}} +{"ts":"2023-11-14T22:13:20Z","pid":17,"msg":"AuthManager::new","src":"shell","ctx":{}} +{"ts":"2023-11-14T22:13:21Z","pid":17,"sid":"session-stable","msg":"shell.turn.inference_done","ctx":{"loop_index":1,"prompt_tokens":15,"completion_tokens":1}} +{"ts":"2023-11-14T22:13:22Z","pid":17,"sid":"session-new","msg":"shell.turn.inference_done","ctx":{"loop_index":1,"prompt_tokens":20,"completion_tokens":2}} +{"ts":"2023-11-14T22:13:23Z","pid":17,"msg":"model catalog: notifying clients","ctx":{"current_model_id":"grok-new"}} +{"ts":"2023-11-14T22:13:24Z","pid":17,"sid":"session-new","msg":"shell.turn.inference_done","ctx":{"loop_index":2,"prompt_tokens":30,"completion_tokens":3}}"#, + ); + + let messages = parse_grok_unified_log_file(&path); + + assert_eq!(messages.len(), 4); + assert_eq!(messages[0].model_id, "grok-old"); + assert_eq!(messages[1].model_id, "grok-session"); + assert_eq!(messages[2].model_id, UNKNOWN_MODEL); + assert_eq!(messages[3].model_id, "grok-new"); + } + + #[test] + fn selector_suppresses_covered_legacy_without_dropping_partial_fallback() { + let mut covered_legacy = test_message("covered", "grok:covered:0"); + covered_legacy.tokens = TokenBreakdown { + input: 900, + output: 80, + cache_read: 70, + cache_write: 60, + reasoning: 50, + }; + covered_legacy.message_count = 7; + covered_legacy.set_workspace( + Some("/tmp/project".to_string()), + Some("project".to_string()), + ); + + let mut legacy_only = test_message("legacy-only", "grok:legacy-only:0"); + legacy_only.tokens.input = 17; + legacy_only.message_count = 3; + + let mut covered_unified = test_message("covered", "grok-unified:covered:1:17:1"); + covered_unified.model_id = UNKNOWN_MODEL.to_string(); + covered_unified.tokens = TokenBreakdown { + input: 40, + output: 20, + cache_read: 60, + cache_write: 0, + reasoning: 5, + }; + covered_unified.message_count = 1; + + let raw = vec![covered_legacy, legacy_only, covered_unified]; + let selected = prefer_unified_log_messages(raw.clone()); + + assert_eq!(selected.len(), 2); + let covered = selected + .iter() + .find(|message| message.session_id == "covered" && is_unified_log_message(message)) + .unwrap(); + assert_eq!(covered.model_id, "grok-build"); + assert_eq!(covered.workspace_key.as_deref(), Some("/tmp/project")); + assert_eq!(covered.workspace_label.as_deref(), Some("project")); + assert!(selected + .iter() + .any(|message| message.session_id == "legacy-only")); + let token_buckets = + selected + .iter() + .fold(TokenBreakdown::default(), |mut total, message| { + total.input += message.tokens.input; + total.output += message.tokens.output; + total.cache_read += message.tokens.cache_read; + total.cache_write += message.tokens.cache_write; + total.reasoning += message.tokens.reasoning; + total + }); + assert_eq!( + token_buckets, + TokenBreakdown { + input: 57, + output: 20, + cache_read: 60, + cache_write: 0, + reasoning: 5, + } + ); + assert_eq!(token_buckets.total(), 142); + assert_eq!( + selected + .iter() + .map(|message| message.message_count) + .sum::(), + 4 + ); + assert_ne!( + raw.iter() + .map(|message| message.tokens.total()) + .sum::(), + 142, + "additive legacy + unified handling would double-count the covered session" + ); + } + + #[test] + fn selector_result_set_is_input_order_independent() { + let mut legacy = test_message("covered", "grok:covered:0"); + legacy.tokens.input = 999; + legacy.message_count = 9; + let mut unified = test_message("covered", "grok-unified:covered:1:17:1"); + unified.tokens.cache_read = 12; + unified.tokens.reasoning = 3; + let fallback = test_message("legacy-only", "grok:legacy-only:0"); + + let forward = + prefer_unified_log_messages(vec![legacy.clone(), unified.clone(), fallback.clone()]); + let reverse = prefer_unified_log_messages(vec![fallback, unified, legacy]); + + let signature = |messages: Vec| { + let mut signature: Vec<_> = messages + .into_iter() + .map(|message| { + ( + message.dedup_key.unwrap(), + message.tokens, + message.message_count, + ) + }) + .collect(); + signature.sort_by(|left, right| left.0.cmp(&right.0)); + signature + }; + + assert_eq!(signature(forward), signature(reverse)); + } + + #[test] + fn selector_keeps_unknown_model_when_legacy_models_conflict() { + let mut legacy_a = test_message("covered", "grok:covered:0"); + legacy_a.model_id = "grok-model-a".to_string(); + let mut legacy_b = test_message("covered", "grok:covered:1"); + legacy_b.model_id = "grok-model-b".to_string(); + let mut unified = test_message("covered", "grok-unified:covered:1:17:1"); + unified.model_id = UNKNOWN_MODEL.to_string(); + + for raw in [ + vec![legacy_a.clone(), legacy_b.clone(), unified.clone()], + vec![legacy_b.clone(), unified.clone(), legacy_a.clone()], + ] { + let selected = prefer_unified_log_messages(raw); + assert_eq!(selected.len(), 1); + assert_eq!(selected[0].model_id, UNKNOWN_MODEL); + } + } + + fn token_all(messages: &[UnifiedMessage]) -> i64 { + messages + .iter() + .map(|m| { + m.tokens.input + + m.tokens.output + + m.tokens.cache_read + + m.tokens.cache_write + + m.tokens.reasoning + }) + .sum() + } + #[test] fn parses_grok_total_token_deltas_by_turn() { let (_temp, path) = write_fixture( @@ -823,4 +1764,453 @@ mod tests { assert_eq!(messages[1].tokens.input, 200); assert_eq!(messages[1].model_id, "grok-composer-2.5-fast"); } + + #[test] + fn prefers_turn_completed_usage_over_context_counter() { + // Context peaks at 236879 (old undercount); two completed turns report + // far larger API usage. Primary path must sum usage, not context growth. + let (_temp, path) = write_fixture( + r#"{"method":"session/update","params":{"sessionId":"session-1","update":{"sessionUpdate":"user_message_chunk"},"_meta":{"agentTimestampMs":1700000000000}}} +{"method":"session/update","params":{"sessionId":"session-1","update":{"sessionUpdate":"agent_thought_chunk"},"_meta":{"totalTokens":100000,"agentTimestampMs":1700000001000}}} +{"method":"session/update","params":{"sessionId":"session-1","update":{"sessionUpdate":"turn_completed","prompt_id":"p1","usage":{"inputTokens":372794,"outputTokens":8021,"totalTokens":380815,"cachedReadTokens":314112,"reasoningTokens":4144,"modelCalls":8,"apiDurationMs":136943,"costUsdTicks":2597236000,"modelUsage":{"grok-4.5-build":{"inputTokens":372794,"outputTokens":8021,"totalTokens":380815,"cachedReadTokens":314112,"reasoningTokens":4144,"modelCalls":8,"apiDurationMs":136943,"costUsdTicks":2597236000}}}},"_meta":{"agentTimestampMs":1700000002000}}} +{"method":"session/update","params":{"sessionId":"session-1","update":{"sessionUpdate":"user_message_chunk"},"_meta":{"agentTimestampMs":1700000003000}}} +{"method":"session/update","params":{"sessionId":"session-1","update":{"sessionUpdate":"agent_message_chunk"},"_meta":{"totalTokens":236879,"agentTimestampMs":1700000004000}}} +{"method":"session/update","params":{"sessionId":"session-1","update":{"sessionUpdate":"turn_completed","prompt_id":"p2","usage":{"inputTokens":1000,"outputTokens":200,"totalTokens":1200,"cachedReadTokens":400,"reasoningTokens":50,"costUsdTicks":1000000000,"modelUsage":{"grok-4.5-build":{"inputTokens":1000,"outputTokens":200,"totalTokens":1200,"cachedReadTokens":400,"reasoningTokens":50,"costUsdTicks":1000000000}}}},"_meta":{"agentTimestampMs":1700000005000}}}"#, + Some(r#"{"current_model_id":"grok-4.5"}"#), + Some( + r#"{"primaryModelId":"grok-4.5","contextTokensUsed":236879,"totalTokensBeforeCompaction":0}"#, + ), + ); + + let messages = parse_grok_updates_file(&path); + assert_eq!(messages.len(), 2); + assert_eq!(messages[0].model_id, "grok-4.5-build"); + // input nets out cache: 372794 - 314112 = 58682 + assert_eq!(messages[0].tokens.input, 58682); + assert_eq!(messages[0].tokens.output, 8021); + assert_eq!(messages[0].tokens.cache_read, 314112); + assert_eq!(messages[0].tokens.reasoning, 4144); + assert_eq!(messages[0].duration_ms, Some(136943)); + assert_eq!(messages[0].cost_source, CostSource::ProviderReported); + assert!((messages[0].cost - 2.597236).abs() < 1e-9); + assert_eq!( + messages[0].dedup_key.as_deref(), + Some("grok:session-1:turn:p1") + ); + assert!(messages[0].is_turn_start); + + assert_eq!(messages[1].tokens.input, 600); // 1000 - 400 + assert_eq!(messages[1].tokens.output, 200); + assert_eq!(messages[1].tokens.cache_read, 400); + assert_eq!(messages[1].tokens.reasoning, 50); + assert!((messages[1].cost - 1.0).abs() < 1e-12); + + // in+out+cache+reason across both turns == raw Grok input+output+reason + // (cache is netted from input then re-added in the total). + assert_eq!(token_all(&messages), 372794 + 8021 + 4144 + 1000 + 200 + 50); + // Must not collapse to context peak (~236k) or signals context. + assert!(token_all(&messages) > 380_000); + assert!(messages.iter().all(|m| m.dedup_key.as_deref() != Some("grok:session-1:signals"))); + } + + #[test] + fn usage_path_emits_live_partial_after_last_completed_turn() { + let (_temp, path) = write_fixture( + r#"{"method":"session/update","params":{"sessionId":"session-1","update":{"sessionUpdate":"user_message_chunk"},"_meta":{"agentTimestampMs":1700000000000}}} +{"method":"session/update","params":{"sessionId":"session-1","update":{"sessionUpdate":"agent_message_chunk"},"_meta":{"totalTokens":50000,"agentTimestampMs":1700000001000}}} +{"method":"session/update","params":{"sessionId":"session-1","update":{"sessionUpdate":"turn_completed","prompt_id":"p1","usage":{"inputTokens":10000,"outputTokens":100,"totalTokens":10100,"cachedReadTokens":0,"reasoningTokens":0}},"_meta":{"agentTimestampMs":1700000002000}}} +{"method":"session/update","params":{"sessionId":"session-1","update":{"sessionUpdate":"user_message_chunk"},"_meta":{"agentTimestampMs":1700000003000}}} +{"method":"session/update","params":{"sessionId":"session-1","update":{"sessionUpdate":"agent_message_chunk"},"_meta":{"totalTokens":55000,"agentTimestampMs":1700000004000}}}"#, + None, + None, + ); + + let messages = parse_grok_updates_file(&path); + assert_eq!(messages.len(), 2); + // Completed usage turn. + assert_eq!(messages[0].tokens.input, 10000); + assert_eq!(messages[0].tokens.output, 100); + // Live open turn: context grew 50000 -> 55000 after completion. + assert_eq!(messages[1].tokens.input, 5000); + assert_eq!(messages[1].tokens.output, 0); + assert!(messages[1].is_turn_start); + } + + #[test] + fn usage_path_does_not_double_count_completed_context_growth() { + // Context grows a lot during the turn; only usage should count for that + // completed turn (not both usage and the context delta). + let (_temp, path) = write_fixture( + r#"{"method":"session/update","params":{"sessionId":"session-1","update":{"sessionUpdate":"user_message_chunk"},"_meta":{"agentTimestampMs":1700000000000}}} +{"method":"session/update","params":{"sessionId":"session-1","update":{"sessionUpdate":"agent_message_chunk"},"_meta":{"totalTokens":200000,"agentTimestampMs":1700000001000}}} +{"method":"session/update","params":{"sessionId":"session-1","update":{"sessionUpdate":"turn_completed","usage":{"inputTokens":500,"outputTokens":50,"totalTokens":550,"cachedReadTokens":0,"reasoningTokens":0}},"_meta":{"agentTimestampMs":1700000002000}}}"#, + None, + None, + ); + + let messages = parse_grok_updates_file(&path); + assert_eq!(messages.len(), 1); + assert_eq!(messages[0].tokens.input, 500); + assert_eq!(messages[0].tokens.output, 50); + assert_eq!(token_all(&messages), 550); + } + + #[test] + fn splits_model_usage_rows() { + let (_temp, path) = write_fixture( + r#"{"method":"session/update","params":{"sessionId":"session-1","update":{"sessionUpdate":"turn_completed","prompt_id":"p-split","usage":{"inputTokens":300,"outputTokens":30,"totalTokens":330,"cachedReadTokens":0,"reasoningTokens":0,"modelUsage":{"grok-a":{"inputTokens":200,"outputTokens":20,"totalTokens":220,"cachedReadTokens":0,"reasoningTokens":0},"grok-b":{"inputTokens":100,"outputTokens":10,"totalTokens":110,"cachedReadTokens":0,"reasoningTokens":0}}}},"_meta":{"agentTimestampMs":1700000000000}}}"#, + None, + None, + ); + + let mut messages = parse_grok_updates_file(&path); + messages.sort_by(|a, b| a.model_id.cmp(&b.model_id)); + assert_eq!(messages.len(), 2); + assert_eq!(messages[0].model_id, "grok-a"); + assert_eq!(messages[0].tokens.input, 200); + assert_eq!(messages[1].model_id, "grok-b"); + assert_eq!(messages[1].tokens.input, 100); + assert!(messages + .iter() + .all(|m| m.dedup_key.as_deref().unwrap().contains("turn:p-split:"))); + // Exactly one turn_start among the split rows. + assert_eq!(messages.iter().filter(|m| m.is_turn_start).count(), 1); + // message_count must not overcount a single turn split across models. + assert_eq!( + messages.iter().map(|m| m.message_count).sum::(), + 1 + ); + assert_eq!( + messages + .iter() + .find(|m| m.is_turn_start) + .map(|m| m.message_count), + Some(1) + ); + assert!(messages + .iter() + .filter(|m| !m.is_turn_start) + .all(|m| m.message_count == 0)); + } + + #[test] + fn skips_context_counter_on_usage_line_with_total_tokens() { + // turn_completed carries both usage and _meta.totalTokens. After the + // usage path clears active_turn, the context arm must not reopen a + // turn from that same line (would double-count completed growth). + let (_temp, path) = write_fixture( + r#"{"method":"session/update","params":{"sessionId":"session-1","update":{"sessionUpdate":"user_message_chunk"},"_meta":{"agentTimestampMs":1700000000000}}} +{"method":"session/update","params":{"sessionId":"session-1","update":{"sessionUpdate":"agent_message_chunk"},"_meta":{"totalTokens":100000,"agentTimestampMs":1700000001000}}} +{"method":"session/update","params":{"sessionId":"session-1","update":{"sessionUpdate":"turn_completed","prompt_id":"p1","usage":{"inputTokens":500,"outputTokens":50,"totalTokens":550,"cachedReadTokens":0,"reasoningTokens":0}},"_meta":{"totalTokens":200000,"agentTimestampMs":1700000002000}}}"#, + None, + None, + ); + + let messages = parse_grok_updates_file(&path); + assert_eq!(messages.len(), 1); + assert_eq!(messages[0].tokens.input, 500); + assert_eq!(messages[0].tokens.output, 50); + assert_eq!(token_all(&messages), 550); + } + + #[test] + fn preserves_legacy_context_turns_before_first_usage() { + // Pre-upgrade turns only have context counters; a later turn has + // turn_completed.usage. Keep the early legacy deltas, then usage. + let (_temp, path) = write_fixture( + r#"{"method":"session/update","params":{"sessionId":"session-1","update":{"sessionUpdate":"user_message_chunk"},"_meta":{"agentTimestampMs":1700000000000}}} +{"method":"session/update","params":{"sessionId":"session-1","update":{"sessionUpdate":"agent_message_chunk"},"_meta":{"totalTokens":100,"agentTimestampMs":1700000001000}}} +{"method":"session/update","params":{"sessionId":"session-1","update":{"sessionUpdate":"user_message_chunk"},"_meta":{"agentTimestampMs":1700000002000}}} +{"method":"session/update","params":{"sessionId":"session-1","update":{"sessionUpdate":"agent_message_chunk"},"_meta":{"totalTokens":250,"agentTimestampMs":1700000003000}}} +{"method":"session/update","params":{"sessionId":"session-1","update":{"sessionUpdate":"user_message_chunk"},"_meta":{"agentTimestampMs":1700000004000}}} +{"method":"session/update","params":{"sessionId":"session-1","update":{"sessionUpdate":"agent_message_chunk"},"_meta":{"totalTokens":300,"agentTimestampMs":1700000005000}}} +{"method":"session/update","params":{"sessionId":"session-1","update":{"sessionUpdate":"turn_completed","prompt_id":"p-new","usage":{"inputTokens":1000,"outputTokens":20,"totalTokens":1020,"cachedReadTokens":0,"reasoningTokens":0}},"_meta":{"agentTimestampMs":1700000006000}}}"#, + None, + None, + ); + + let messages = parse_grok_updates_file(&path); + // First legacy turn: 0→100; second: 100→250; then usage 1000+20. + // (Third user_message_chunk opens a turn that is closed without emit + // when usage arrives for that turn — only completed legacy turns + // before saw_usage are kept via the user_message_chunk flush.) + assert!( + messages.len() >= 3, + "expected legacy turns + usage, got {}", + messages.len() + ); + assert_eq!(messages[0].tokens.input, 100); + assert_eq!(messages[0].tokens.output, 0); + assert_eq!(messages[1].tokens.input, 150); + assert_eq!(messages[1].tokens.output, 0); + let usage = messages + .iter() + .find(|m| m.dedup_key.as_deref() == Some("grok:session-1:turn:p-new")) + .expect("usage turn"); + assert_eq!(usage.tokens.input, 1000); + assert_eq!(usage.tokens.output, 20); + } + + #[test] + fn inherits_parent_cost_and_duration_onto_single_model_usage() { + // Parent usage has cost/duration; model entry only has token buckets. + let (_temp, path) = write_fixture( + r#"{"method":"session/update","params":{"sessionId":"session-1","update":{"sessionUpdate":"turn_completed","prompt_id":"p-cost","usage":{"inputTokens":100,"outputTokens":10,"totalTokens":110,"cachedReadTokens":0,"reasoningTokens":0,"apiDurationMs":42000,"costUsdTicks":2500000000,"modelUsage":{"grok-4.5-build":{"inputTokens":100,"outputTokens":10,"totalTokens":110,"cachedReadTokens":0,"reasoningTokens":0}}}},"_meta":{"agentTimestampMs":1700000000000}}}"#, + None, + None, + ); + + let messages = parse_grok_updates_file(&path); + assert_eq!(messages.len(), 1); + assert_eq!(messages[0].model_id, "grok-4.5-build"); + assert_eq!(messages[0].tokens.input, 100); + assert_eq!(messages[0].tokens.output, 10); + assert_eq!(messages[0].duration_ms, Some(42000)); + assert_eq!(messages[0].cost_source, CostSource::ProviderReported); + assert!((messages[0].cost - 2.5).abs() < 1e-12); + assert_eq!(messages[0].message_count, 1); + } + + #[test] + fn post_usage_context_partials_require_user_message_chunk() { + // After turn_completed.usage, non-user updates that only bump + // `_meta.totalTokens` must not open a live partial ActiveTurn. + let (_temp, path) = write_fixture( + r#"{"method":"session/update","params":{"sessionId":"session-1","update":{"sessionUpdate":"user_message_chunk"},"_meta":{"agentTimestampMs":1700000000000}}} +{"method":"session/update","params":{"sessionId":"session-1","update":{"sessionUpdate":"agent_message_chunk"},"_meta":{"totalTokens":50000,"agentTimestampMs":1700000001000}}} +{"method":"session/update","params":{"sessionId":"session-1","update":{"sessionUpdate":"turn_completed","prompt_id":"p1","usage":{"inputTokens":10000,"outputTokens":100,"totalTokens":10100,"cachedReadTokens":0,"reasoningTokens":0}},"_meta":{"agentTimestampMs":1700000002000}}} +{"method":"session/update","params":{"sessionId":"session-1","update":{"sessionUpdate":"agent_message_chunk"},"_meta":{"totalTokens":60000,"agentTimestampMs":1700000003000}}} +{"method":"session/update","params":{"sessionId":"session-1","update":{"sessionUpdate":"agent_thought_chunk"},"_meta":{"totalTokens":65000,"agentTimestampMs":1700000004000}}}"#, + None, + None, + ); + + let messages = parse_grok_updates_file(&path); + assert_eq!( + messages.len(), + 1, + "context-only growth after usage must not emit a partial, got {messages:?}" + ); + assert_eq!(messages[0].tokens.input, 10000); + assert_eq!(messages[0].tokens.output, 100); + assert_eq!(token_all(&messages), 10100); + } + + #[test] + fn post_usage_user_turn_baseline_skips_inter_turn_context_growth() { + // Context-only growth between completed usage and the next user turn + // advances the baseline only; the live partial starts from the latest + // counter, not the post-usage snapshot. + let (_temp, path) = write_fixture( + r#"{"method":"session/update","params":{"sessionId":"session-1","update":{"sessionUpdate":"user_message_chunk"},"_meta":{"agentTimestampMs":1700000000000}}} +{"method":"session/update","params":{"sessionId":"session-1","update":{"sessionUpdate":"agent_message_chunk"},"_meta":{"totalTokens":50000,"agentTimestampMs":1700000001000}}} +{"method":"session/update","params":{"sessionId":"session-1","update":{"sessionUpdate":"turn_completed","prompt_id":"p1","usage":{"inputTokens":10000,"outputTokens":100,"totalTokens":10100,"cachedReadTokens":0,"reasoningTokens":0}},"_meta":{"agentTimestampMs":1700000002000}}} +{"method":"session/update","params":{"sessionId":"session-1","update":{"sessionUpdate":"agent_message_chunk"},"_meta":{"totalTokens":60000,"agentTimestampMs":1700000003000}}} +{"method":"session/update","params":{"sessionId":"session-1","update":{"sessionUpdate":"user_message_chunk"},"_meta":{"agentTimestampMs":1700000004000}}} +{"method":"session/update","params":{"sessionId":"session-1","update":{"sessionUpdate":"agent_message_chunk"},"_meta":{"totalTokens":65000,"agentTimestampMs":1700000005000}}}"#, + None, + None, + ); + + let messages = parse_grok_updates_file(&path); + assert_eq!(messages.len(), 2); + assert_eq!(messages[0].tokens.input, 10000); + assert_eq!(messages[0].tokens.output, 100); + // Live partial: 65000 - 60000 (inter-turn growth was baselined away). + assert_eq!(messages[1].tokens.input, 5000); + assert_eq!(messages[1].tokens.output, 0); + assert!(messages[1].is_turn_start); + } + + #[test] + fn multi_model_parent_cost_blocks_sibling_repricing() { + // Parent cost lands on the first modelUsage row only; siblings must be + // ProviderReported at $0 so a later pricing pass cannot estimate them + // on top of the parent total. + let (_temp, path) = write_fixture( + r#"{"method":"session/update","params":{"sessionId":"session-1","update":{"sessionUpdate":"turn_completed","prompt_id":"p-multi-cost","usage":{"inputTokens":300,"outputTokens":30,"totalTokens":330,"cachedReadTokens":0,"reasoningTokens":0,"costUsdTicks":2500000000,"modelUsage":{"grok-a":{"inputTokens":200,"outputTokens":20,"totalTokens":220,"cachedReadTokens":0,"reasoningTokens":0},"grok-b":{"inputTokens":100,"outputTokens":10,"totalTokens":110,"cachedReadTokens":0,"reasoningTokens":0}}}},"_meta":{"agentTimestampMs":1700000000000}}}"#, + None, + None, + ); + + let mut messages = parse_grok_updates_file(&path); + messages.sort_by(|a, b| a.model_id.cmp(&b.model_id)); + assert_eq!(messages.len(), 2); + assert_eq!(messages[0].model_id, "grok-a"); + assert_eq!(messages[1].model_id, "grok-b"); + + let parent_cost = 2.5; + let total_before: f64 = messages.iter().map(|m| m.cost).sum(); + assert!( + (total_before - parent_cost).abs() < 1e-12, + "parent cost must appear once before pricing, got {total_before}" + ); + assert!(messages + .iter() + .all(|m| m.cost_source == CostSource::ProviderReported)); + assert!(messages.iter().all(|m| m.has_authoritative_cost())); + // Exactly one row carries the parent dollars; the sibling is $0. + assert_eq!( + messages.iter().filter(|m| m.cost > 0.0).count(), + 1, + "only the first inherited row should carry parent cost" + ); + assert!(messages.iter().any(|m| m.cost == 0.0)); + + // Hermetic pricing pass mirroring apply_pricing_if_available: models + // are resolvable at non-zero rates, but authoritative rows must not + // be overwritten. + let mut litellm = std::collections::HashMap::new(); + litellm.insert( + "grok-a".to_string(), + crate::pricing::ModelPricing { + input_cost_per_token: Some(0.01), + output_cost_per_token: Some(0.02), + ..Default::default() + }, + ); + litellm.insert( + "grok-b".to_string(), + crate::pricing::ModelPricing { + input_cost_per_token: Some(0.01), + output_cost_per_token: Some(0.02), + ..Default::default() + }, + ); + let pricing = crate::pricing::PricingService::new(litellm, std::collections::HashMap::new()); + for message in &mut messages { + if message.has_authoritative_cost() { + continue; + } + let calculated = pricing.calculate_cost_with_provider( + &message.model_id, + Some(&message.provider_id), + &message.tokens, + ); + if calculated > 0.0 { + message.cost = calculated; + message.mark_estimated_cost(); + } + } + + let total_after: f64 = messages.iter().map(|m| m.cost).sum(); + assert!( + (total_after - parent_cost).abs() < 1e-12, + "total cost must equal parent only once after pricing, got {total_after}" + ); + assert!(messages + .iter() + .all(|m| m.cost_source == CostSource::ProviderReported)); + } + + #[test] + fn post_usage_live_partial_does_not_zero_base_without_counter() { + // Completed usage with no _meta.totalTokens ever seen, then a new user + // turn whose first context counter is full occupancy. Must not emit + // that occupancy as delta-from-0 (would double-count completed usage). + let (_temp, path) = write_fixture( + r#"{"method":"session/update","params":{"sessionId":"session-1","update":{"sessionUpdate":"turn_completed","prompt_id":"p1","usage":{"inputTokens":10000,"outputTokens":100,"totalTokens":10100,"cachedReadTokens":0,"reasoningTokens":0}},"_meta":{"agentTimestampMs":1700000000000}}} +{"method":"session/update","params":{"sessionId":"session-1","update":{"sessionUpdate":"user_message_chunk"},"_meta":{"agentTimestampMs":1700000001000}}} +{"method":"session/update","params":{"sessionId":"session-1","update":{"sessionUpdate":"agent_message_chunk"},"_meta":{"totalTokens":50000,"agentTimestampMs":1700000002000}}} +{"method":"session/update","params":{"sessionId":"session-1","update":{"sessionUpdate":"agent_message_chunk"},"_meta":{"totalTokens":55000,"agentTimestampMs":1700000003000}}}"#, + None, + None, + ); + + let messages = parse_grok_updates_file(&path); + assert_eq!(messages.len(), 2, "usage + live partial only, got {messages:?}"); + assert_eq!(messages[0].tokens.input, 10000); + assert_eq!(messages[0].tokens.output, 100); + // First counter (50000) is baseline only; live partial is 55000-50000. + assert_eq!( + messages[1].tokens.input, 5000, + "must not emit full occupancy 50000 as zero-based input" + ); + assert_eq!(messages[1].tokens.output, 0); + assert!(messages[1].is_turn_start); + assert_eq!(token_all(&messages), 10000 + 100 + 5000); + } + + #[test] + fn post_usage_live_partial_waits_for_baseline_when_only_first_counter() { + // Same shape but open turn ends at the first post-usage counter: that + // counter establishes baseline only, so no live partial is emitted. + let (_temp, path) = write_fixture( + r#"{"method":"session/update","params":{"sessionId":"session-1","update":{"sessionUpdate":"turn_completed","prompt_id":"p1","usage":{"inputTokens":8000,"outputTokens":80,"totalTokens":8080,"cachedReadTokens":0,"reasoningTokens":0}},"_meta":{"agentTimestampMs":1700000000000}}} +{"method":"session/update","params":{"sessionId":"session-1","update":{"sessionUpdate":"user_message_chunk"},"_meta":{"agentTimestampMs":1700000001000}}} +{"method":"session/update","params":{"sessionId":"session-1","update":{"sessionUpdate":"agent_message_chunk"},"_meta":{"totalTokens":42000,"agentTimestampMs":1700000002000}}}"#, + None, + None, + ); + + let messages = parse_grok_updates_file(&path); + assert_eq!( + messages.len(), + 1, + "first post-usage counter is baseline only, got {messages:?}" + ); + assert_eq!(messages[0].tokens.input, 8000); + assert_eq!(messages[0].tokens.output, 80); + assert_eq!(token_all(&messages), 8080); + } + + #[test] + fn multi_model_own_cost_leaves_omitting_sibling_unknown() { + // One modelUsage entry has its own costUsdTicks; the other omits cost + // and parent does not cover both (parent has no cost). Omitting sibling + // must stay Unknown so apply_pricing can estimate — not ProviderReported $0. + let (_temp, path) = write_fixture( + r#"{"method":"session/update","params":{"sessionId":"session-1","update":{"sessionUpdate":"turn_completed","prompt_id":"p-partial-cost","usage":{"inputTokens":300,"outputTokens":30,"totalTokens":330,"cachedReadTokens":0,"reasoningTokens":0,"modelUsage":{"grok-a":{"inputTokens":200,"outputTokens":20,"totalTokens":220,"cachedReadTokens":0,"reasoningTokens":0,"costUsdTicks":1500000000},"grok-b":{"inputTokens":100,"outputTokens":10,"totalTokens":110,"cachedReadTokens":0,"reasoningTokens":0}}}},"_meta":{"agentTimestampMs":1700000000000}}}"#, + None, + None, + ); + + let mut messages = parse_grok_updates_file(&path); + messages.sort_by(|a, b| a.model_id.cmp(&b.model_id)); + assert_eq!(messages.len(), 2); + let a = messages.iter().find(|m| m.model_id == "grok-a").unwrap(); + let b = messages.iter().find(|m| m.model_id == "grok-b").unwrap(); + assert_eq!(a.cost_source, CostSource::ProviderReported); + assert!((a.cost - 1.5).abs() < 1e-12); + assert_eq!( + b.cost_source, + CostSource::Unknown, + "omitting sibling must stay Unknown when parent cost was not inherited" + ); + assert_eq!(b.cost, 0.0); + assert!(!b.has_authoritative_cost()); + + // Pricing may estimate the Unknown sibling; authoritative row stays put. + let mut litellm = std::collections::HashMap::new(); + litellm.insert( + "grok-b".to_string(), + crate::pricing::ModelPricing { + input_cost_per_token: Some(0.01), + output_cost_per_token: Some(0.02), + ..Default::default() + }, + ); + let pricing = crate::pricing::PricingService::new(litellm, std::collections::HashMap::new()); + for message in &mut messages { + if message.has_authoritative_cost() { + continue; + } + let calculated = pricing.calculate_cost_with_provider( + &message.model_id, + Some(&message.provider_id), + &message.tokens, + ); + if calculated > 0.0 { + message.cost = calculated; + message.mark_estimated_cost(); + } + } + let a = messages.iter().find(|m| m.model_id == "grok-a").unwrap(); + let b = messages.iter().find(|m| m.model_id == "grok-b").unwrap(); + assert_eq!(a.cost_source, CostSource::ProviderReported); + assert!((a.cost - 1.5).abs() < 1e-12); + assert_eq!(b.cost_source, CostSource::Estimated); + // 100 input * 0.01 + 10 output * 0.02 = 1.2 + assert!((b.cost - 1.2).abs() < 1e-12); + } } diff --git a/vendor/tokscale-core/src/sessions/jcode.rs b/vendor/tokscale-core/src/sessions/jcode.rs index 7c9d3c9..b890c55 100644 --- a/vendor/tokscale-core/src/sessions/jcode.rs +++ b/vendor/tokscale-core/src/sessions/jcode.rs @@ -4,7 +4,7 @@ //! Jcode stores authoritative assistant token usage on messages under //! `token_usage`; user/tool messages without usage are skipped. -use super::utils::{file_modified_timestamp_ms, parse_timestamp_str}; +use super::utils::{back_anchor_timestamp, file_modified_timestamp_ms, parse_timestamp_str}; use super::{normalize_workspace_key, workspace_label_from_key, UnifiedMessage}; use crate::{provider_identity, TokenBreakdown}; use serde::Deserialize; @@ -181,11 +181,29 @@ fn parse_jcode_messages( if tokens.total() <= 0 { return None; } - let timestamp = message - .timestamp - .as_deref() - .and_then(parse_timestamp_str) - .unwrap_or(fallback_timestamp); + // `explicit_timestamp` is the message's own recorded `timestamp` + // field, as opposed to `fallback_timestamp` (a session/file-level + // fallback used when it's absent or unparseable). + let explicit_timestamp = message.timestamp.as_deref().and_then(parse_timestamp_str); + let recorded_timestamp = explicit_timestamp.unwrap_or(fallback_timestamp); + // The assistant message's `timestamp` is written once the message + // (including `token_usage`) is finalized, i.e. the turn's *end*, + // not its start. `tool_duration_ms` is that turn's elapsed time, + // so `sessionize()`'s `[timestamp, timestamp + duration_ms]` span + // would otherwise project forward past completion into phantom + // idle time. Back-calculate the start anchor the same way #890 + // did for Copilot's `endTime`-only records. + // + // Only do this when `explicit_timestamp` is a real recorded end + // timestamp: when it's absent, `recorded_timestamp` is the + // session/file-level fallback, not this message's own completion + // time, and subtracting `tool_duration_ms` from it would shift + // the message into the wrong day rather than anchor it correctly. + let duration_ms = message.tool_duration_ms.filter(|duration| *duration > 0); + let timestamp = match (explicit_timestamp, duration_ms) { + (Some(end), Some(duration)) => back_anchor_timestamp(end, duration), + _ => recorded_timestamp, + }; let mut unified = UnifiedMessage::new_with_dedup( "jcode", context.model.clone(), @@ -196,7 +214,7 @@ fn parse_jcode_messages( 0.0, Some(dedup_key), ); - unified.duration_ms = message.tool_duration_ms.filter(|duration| *duration > 0); + unified.duration_ms = duration_ms; if !is_replacement && message.role.as_deref() == Some("assistant") && context.pending_turn_start @@ -660,4 +678,40 @@ mod tests { let turn_count = messages.iter().filter(|m| m.is_turn_start).count(); assert_eq!(turn_count, 2); } + + #[test] + fn test_tool_duration_timestamp_is_start_anchored() { + // Regression (follow-up to #890): an assistant message's `timestamp` + // is written once the message (including `token_usage`) is + // finalized, i.e. the turn's *end*, not its start. `tool_duration_ms` + // is that turn's elapsed time, so sessionize()'s + // `[timestamp, timestamp + duration_ms]` span would otherwise project + // forward past the actual completion into phantom idle time. The + // parser must back-calculate the start anchor instead. + let file = tempfile::NamedTempFile::new().unwrap(); + std::fs::write( + file.path(), + r#"{ + "id":"session_test", + "model":"snapshot-model", + "messages":[ + {"id":"assistant_1","role":"assistant","timestamp":"2026-06-16T12:00:05Z","token_usage":{"input_tokens":100,"output_tokens":10},"tool_duration_ms":2000} + ] +}"#, + ) + .unwrap(); + + let messages = parse_jcode_file(file.path()); + assert_eq!(messages.len(), 1); + assert_eq!( + messages[0].timestamp, + parse_timestamp_str("2026-06-16T12:00:03Z").unwrap(), + "timestamp must be back-calculated to the turn start (end - duration)" + ); + assert_eq!( + messages[0].duration_ms, + Some(2000), + "duration_ms must still span from start to the recorded end timestamp" + ); + } } diff --git a/vendor/tokscale-core/src/sessions/junie.rs b/vendor/tokscale-core/src/sessions/junie.rs new file mode 100644 index 0000000..9765a83 --- /dev/null +++ b/vendor/tokscale-core/src/sessions/junie.rs @@ -0,0 +1,397 @@ +//! Junie session parser +//! +//! Junie stores local sessions under `~/.junie/sessions//events.jsonl`. + +use super::utils::{back_anchor_timestamp, file_modified_timestamp_ms}; +use super::UnifiedMessage; +use crate::{pricing, provider_identity, TokenBreakdown}; +use chrono::{Local, LocalResult, NaiveDateTime, TimeZone}; +use serde_json::Value; +use std::collections::HashSet; +use std::io::{BufRead, BufReader}; +use std::path::Path; + +const USAGE_EVENT_KIND: &str = "LlmResponseMetadataEvent"; +const USER_PROMPT_KIND: &str = "UserPromptEvent"; +const SKIP_EVENT_KINDS: &[&str] = &[ + "AgentStateUpdatedEvent", + "AgentCurrentStatusUpdatedEvent", + "AgentPatchCreatedEvent", +]; + +pub fn parse_junie_file(path: &Path) -> Vec { + let file = match std::fs::File::open(path) { + Ok(file) => file, + Err(_) => return Vec::new(), + }; + + let session_id = session_id_from_path(path); + let default_timestamp = + session_timestamp_from_id(&session_id).unwrap_or_else(|| file_modified_timestamp_ms(path)); + let mut pending_turn_start = false; + let mut messages = Vec::new(); + let mut seen = HashSet::new(); + + for line in BufReader::new(file).lines() { + let Ok(line) = line else { + continue; + }; + // State snapshots can be large and do not carry usage rows. This is a + // cheap pre-filter; the parsed event kind below makes the final choice. + if !line.contains(USAGE_EVENT_KIND) && !line.contains(USER_PROMPT_KIND) { + continue; + } + + let Ok(value) = serde_json::from_str::(&line) else { + continue; + }; + if let Some(kind) = parsed_event_kind(&value) { + if SKIP_EVENT_KINDS.contains(&kind) { + continue; + } + } + if event_kind(&value) == Some(USER_PROMPT_KIND) { + pending_turn_start = true; + continue; + } + + let Some(agent_event) = value + .pointer("/event/agentEvent") + .filter(|event| string_field(event, "kind") == Some(USAGE_EVENT_KIND)) + else { + continue; + }; + + let explicit_timestamp = + number_field(&value, "timestampMs").filter(|timestamp| *timestamp > 0); + let timestamp = explicit_timestamp.unwrap_or(default_timestamp); + let agent = agent_name(agent_event); + let Some(usages) = agent_event.get("modelUsage").and_then(Value::as_array) else { + pending_turn_start = false; + continue; + }; + + let mut turn_start_assigned = false; + for (usage_index, usage) in usages.iter().enumerate() { + let Some(model_raw) = string_field(usage, "model") else { + continue; + }; + let model_id = pricing::aliases::resolve_alias(model_raw) + .unwrap_or(model_raw) + .to_string(); + let provider_id = provider_from_usage(usage, &model_id); + let tokens = tokens_from_usage(usage); + let provider_cost = + float_field(usage, "cost").filter(|cost| cost.is_finite() && *cost >= 0.0); + let cost = provider_cost.unwrap_or(0.0); + if tokens.total() == 0 && cost == 0.0 { + continue; + } + + let dedup_key = format!( + "junie:{session_id}:{timestamp}:{model_id}:{}:{}:{}:{}:{}:{:.12}:{usage_index}", + tokens.input, + tokens.output, + tokens.cache_read, + tokens.cache_write, + tokens.reasoning, + cost + ); + if !seen.insert(dedup_key.clone()) { + continue; + } + + let duration_ms = number_field(usage, "time").filter(|duration| *duration > 0); + let start_timestamp = match (explicit_timestamp, duration_ms) { + (Some(end), Some(duration)) => back_anchor_timestamp(end, duration), + _ => timestamp, + }; + + let mut message = UnifiedMessage::new_with_agent( + "junie", + model_id, + provider_id, + &session_id, + start_timestamp, + tokens, + cost, + agent.clone(), + ); + message.dedup_key = Some(dedup_key); + message.duration_ms = duration_ms; + if provider_cost.is_some() { + message.mark_provider_reported_cost(); + } + if pending_turn_start && !turn_start_assigned { + message.is_turn_start = true; + turn_start_assigned = true; + } + messages.push(message); + } + // A prompt belongs to the next usage event, not to a later event if this + // response had no countable usage. + pending_turn_start = false; + } + + messages +} + +fn session_id_from_path(path: &Path) -> String { + path.parent() + .and_then(Path::file_name) + .and_then(|name| name.to_str()) + .filter(|name| !name.trim().is_empty()) + .unwrap_or("unknown") + .to_string() +} + +fn session_timestamp_from_id(session_id: &str) -> Option { + let mut parts = session_id.split('-'); + if parts.next()? != "session" { + return None; + } + let date = parts.next()?; + let time = parts.next()?; + if date.len() != 6 + || time.len() != 6 + || !date.bytes().all(|byte| byte.is_ascii_digit()) + || !time.bytes().all(|byte| byte.is_ascii_digit()) + { + return None; + } + + let naive = NaiveDateTime::parse_from_str(&format!("{date}{time}"), "%y%m%d%H%M%S").ok()?; + match Local.from_local_datetime(&naive) { + LocalResult::Single(datetime) => Some(datetime.timestamp_millis()), + LocalResult::Ambiguous(earliest, _) => Some(earliest.timestamp_millis()), + LocalResult::None => None, + } +} + +fn event_kind(value: &Value) -> Option<&str> { + string_field(value, "kind") +} + +fn parsed_event_kind(value: &Value) -> Option<&str> { + event_kind(value).or_else(|| { + value + .pointer("/event/agentEvent") + .and_then(|event| string_field(event, "kind")) + }) +} + +fn agent_name(agent_event: &Value) -> Option { + let agent = agent_event.get("agent")?; + string_field(agent, "name") + .or_else(|| string_field(agent, "id")) + .map(str::to_string) +} + +fn provider_from_usage(usage: &Value, model_id: &str) -> String { + string_field(usage, "provider") + .and_then(provider_identity::canonical_provider) + .or_else(|| provider_identity::inferred_provider_from_model(model_id).map(str::to_string)) + .unwrap_or_else(|| "junie".to_string()) +} + +fn tokens_from_usage(usage: &Value) -> TokenBreakdown { + TokenBreakdown { + input: first_number_field(usage, &["inputTokens", "input"]), + output: first_number_field(usage, &["outputTokens", "output"]), + cache_read: first_number_field( + usage, + &["cacheInputTokens", "cacheReadInputTokens", "cacheRead"], + ), + cache_write: first_number_field( + usage, + &[ + "cacheCreateTokens", + "cacheCreationInputTokens", + "cacheWrite", + ], + ), + reasoning: first_number_field( + usage, + &["reasoningTokens", "reasoningOutputTokens", "thinkingTokens"], + ), + } +} + +fn string_field<'a>(value: &'a Value, field: &str) -> Option<&'a str> { + value + .get(field) + .and_then(Value::as_str) + .map(str::trim) + .filter(|value| !value.is_empty()) +} + +fn first_number_field(value: &Value, fields: &[&str]) -> i64 { + fields + .iter() + .find_map(|field| number_field(value, field)) + .unwrap_or(0) +} + +fn number_field(value: &Value, field: &str) -> Option { + number_value(value.get(field)?) +} + +fn number_value(value: &Value) -> Option { + if let Some(value) = value.as_i64() { + return Some(value.max(0)); + } + if let Some(value) = value.as_u64() { + return Some(value.min(i64::MAX as u64) as i64); + } + if let Some(value) = value.as_f64() { + return value.is_finite().then_some(value.max(0.0) as i64); + } + value + .as_str() + .and_then(|value| value.trim().parse::().ok()) + .and_then(|value| value.is_finite().then_some(value.max(0.0) as i64)) +} + +fn float_field(value: &Value, field: &str) -> Option { + let value = value.get(field)?; + if let Some(number) = value.as_f64() { + return Some(number); + } + value.as_str()?.trim().parse().ok() +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::sessions::CostSource; + use std::io::Write; + use tempfile::TempDir; + + fn parse_events(content: &str) -> Vec { + let dir = TempDir::new().unwrap(); + let session_dir = dir.path().join("session-250622-101010"); + std::fs::create_dir_all(&session_dir).unwrap(); + let path = session_dir.join("events.jsonl"); + let mut file = std::fs::File::create(&path).unwrap(); + file.write_all(content.as_bytes()).unwrap(); + file.flush().unwrap(); + parse_junie_file(&path) + } + + fn usage_event(timestamp_ms: i64, model: &str, input: i64, output: i64) -> String { + format!( + r#"{{"timestampMs":{timestamp_ms},"event":{{"agentEvent":{{"kind":"LlmResponseMetadataEvent","modelUsage":[{{"model":"{model}","inputTokens":{input},"outputTokens":{output}}}]}}}}}}"# + ) + } + + #[test] + fn provider_reported_cost_is_authoritative() { + let content = r#"{"timestampMs":1750000005000,"event":{"agentEvent":{"kind":"LlmResponseMetadataEvent","modelUsage":[{"model":"gpt-5","inputTokens":100,"outputTokens":50,"cost":0.125}]}}}"#; + let messages = parse_events(content); + + assert_eq!(messages.len(), 1); + assert_eq!(messages[0].cost, 0.125); + assert_eq!(messages[0].cost_source, CostSource::ProviderReported); + } + + #[test] + fn missing_cost_remains_estimable() { + let content = usage_event(1_750_000_000_000, "gpt-5", 100, 50); + let messages = parse_events(&content); + + assert_eq!(messages.len(), 1); + assert_eq!(messages[0].cost, 0.0); + assert_eq!(messages[0].cost_source, CostSource::Unknown); + } + + #[test] + fn prompt_ownership_marks_only_first_usage_event() { + let content = format!( + "{}\n{}\n{}\n", + r#"{"kind":"UserPromptEvent"}"#, + usage_event(1_750_000_000_000, "gpt-5", 100, 50), + usage_event(1_750_000_100_000, "gpt-5", 200, 60), + ); + let messages = parse_events(&content); + + assert_eq!(messages.len(), 2); + assert!(messages[0].is_turn_start); + assert!(!messages[1].is_turn_start); + } + + #[test] + fn prompt_does_not_leak_past_empty_response() { + let empty_usage = r#"{"timestampMs":1750000000000,"event":{"agentEvent":{"kind":"LlmResponseMetadataEvent","modelUsage":[{"model":"gpt-5","inputTokens":0,"outputTokens":0}]}}}"#; + let content = format!( + "{}\n{}\n{}\n{}\n", + r#"{"kind":"UserPromptEvent"}"#, + empty_usage, + r#"{"kind":"UserPromptEvent"}"#, + usage_event(1_750_000_100_000, "gpt-5", 100, 50), + ); + let messages = parse_events(&content); + + assert_eq!(messages.len(), 1); + assert!(messages[0].is_turn_start); + } + + #[test] + fn prompt_does_not_leak_past_response_without_usage_array() { + let content = format!( + "{}\n{}\n{}\n", + r#"{"kind":"UserPromptEvent"}"#, + r#"{"timestampMs":1750000000000,"event":{"agentEvent":{"kind":"LlmResponseMetadataEvent"}}}"#, + usage_event(1_750_000_100_000, "gpt-5", 100, 50), + ); + let messages = parse_events(&content); + + assert_eq!(messages.len(), 1); + assert!(!messages[0].is_turn_start); + } + + #[test] + fn start_anchor_uses_recorded_end_timestamp() { + let content = r#"{"timestampMs":1750000005000,"event":{"agentEvent":{"kind":"LlmResponseMetadataEvent","modelUsage":[{"model":"gpt-5","inputTokens":100,"outputTokens":50,"time":2000}]}}}"#; + let messages = parse_events(content); + + assert_eq!(messages.len(), 1); + assert_eq!(messages[0].timestamp, 1_750_000_003_000); + assert_eq!(messages[0].duration_ms, Some(2000)); + } + + #[test] + fn missing_timestamp_does_not_anchor_from_session_fallback() { + let content = r#"{"event":{"agentEvent":{"kind":"LlmResponseMetadataEvent","modelUsage":[{"model":"gpt-5","inputTokens":100,"outputTokens":50,"time":2000}]}}}"#; + let messages = parse_events(content); + + assert_eq!(messages.len(), 1); + assert_eq!( + messages[0].timestamp, + session_timestamp_from_id("session-250622-101010").unwrap() + ); + assert_eq!(messages[0].duration_ms, Some(2000)); + } + + #[test] + fn usage_line_mentioning_skipped_kind_keeps_prompt_ownership() { + let content = format!( + "{}\n{}\n", + r#"{"kind":"UserPromptEvent","prompt":"please review AgentStateUpdatedEvent handling"}"#, + usage_event(1_750_000_000_000, "gpt-5", 100, 50), + ); + let messages = parse_events(&content); + assert_eq!(messages.len(), 1); + assert!(messages[0].is_turn_start); + } + + #[test] + fn skipped_kind_is_checked_after_parsing() { + let content = format!( + "{}\n{}\n", + r#"{"kind":"AgentStateUpdatedEvent","event":{"agentEvent":{"kind":"LlmResponseMetadataEvent","modelUsage":[{"model":"gpt-5","inputTokens":100,"outputTokens":50}]}}}"#, + usage_event(1_750_000_000_000, "gpt-5", 100, 50), + ); + let messages = parse_events(&content); + assert_eq!(messages.len(), 1); + } +} diff --git a/vendor/tokscale-core/src/sessions/kimi.rs b/vendor/tokscale-core/src/sessions/kimi.rs index 1fda959..bf1a113 100644 --- a/vendor/tokscale-core/src/sessions/kimi.rs +++ b/vendor/tokscale-core/src/sessions/kimi.rs @@ -1,13 +1,19 @@ -//! Kimi CLI session parser +//! Kimi CLI / Kimi Code session parser //! -//! Parses wire.jsonl files from ~/.kimi/sessions/[GROUP_ID]/[SESSION_UUID]/wire.jsonl -//! Token data comes from StatusUpdate messages in the wire protocol. +//! Parses wire.jsonl from both `kimi-cli` and `kimi-code`. +//! +//! ~/.kimi/sessions/[GROUP_ID]/[SESSION_UUID]/wire.jsonl +//! Token data comes from StatusUpdate messages. +//! +//! ~/.kimi-code/sessions/[WORKSPACE]/[SESSION]/agents/[AGENT]/wire.jsonl +//! Token data comes from usage.record lines. use super::utils::file_modified_timestamp_ms; use super::UnifiedMessage; use crate::TokenBreakdown; use serde::Deserialize; -use std::collections::HashMap; +use serde_json::Value; +use std::collections::{HashMap, HashSet}; use std::io::{BufRead, BufReader}; use std::path::{Path, PathBuf}; @@ -34,20 +40,54 @@ struct StatusPayload { message_id: Option, } +/// Token usage counts shared by both wire formats. +/// +/// Legacy kimi-cli StatusUpdate payloads use snake_case field names; +/// kimi-code usage.record lines use the camelCase aliases. #[derive(Debug, Deserialize)] struct TokenUsage { + #[serde(alias = "inputOther")] input_other: Option, output: Option, + #[serde(alias = "inputCacheRead")] input_cache_read: Option, + #[serde(alias = "inputCacheCreation")] input_cache_creation: Option, } +impl TokenUsage { + /// Clamp negative counts to zero and build a breakdown. + /// Returns `None` when every count is zero so callers can skip the entry. + fn to_breakdown(&self) -> Option { + let input = self.input_other.unwrap_or(0).max(0); + let output = self.output.unwrap_or(0).max(0); + let cache_read = self.input_cache_read.unwrap_or(0).max(0); + let cache_write = self.input_cache_creation.unwrap_or(0).max(0); + + if input == 0 && output == 0 && cache_read == 0 && cache_write == 0 { + return None; + } + + Some(TokenBreakdown { + input, + output, + cache_read, + cache_write, + // Kimi wire protocols do not expose reasoning tokens; all reasoning included in output. + reasoning: 0, + }) + } +} + /// Default model name when config.json is not available const DEFAULT_MODEL: &str = "kimi-for-coding"; const DEFAULT_PROVIDER: &str = "moonshot"; /// Locate the legacy Kimi CLI config consumed by `parse_kimi_file`. pub(crate) fn kimi_config_path(wire_path: &Path) -> Option { + if is_kimi_code_path(wire_path) { + return None; + } let sessions_dir = wire_path.parent()?.parent()?.parent()?; Some(sessions_dir.parent()?.join("config.json")) } @@ -78,6 +118,166 @@ fn extract_session_id(path: &Path) -> String { .to_string() } +/// Check whether a wire.jsonl path belongs to kimi-code. +/// +/// kimi-code writes `/sessions/WORKSPACE/SESSION/agents/AGENT/wire.jsonl` +/// while legacy kimi-cli writes `/sessions/GROUP/UUID/wire.jsonl`, so the +/// grandparent directory component (`agents`) distinguishes the formats. The +/// layout under the root is created by kimi-code itself, so this holds for the +/// default `~/.kimi-code` root and custom `KIMI_CODE_HOME` roots alike. +pub fn is_kimi_code_path(path: &Path) -> bool { + path.parent() + .and_then(|agent_dir| agent_dir.parent()) + .and_then(|agents_dir| agents_dir.file_name()) + .is_some_and(|name| name == "agents") +} + +/// Extract session ID from a kimi-code wire.jsonl path. +/// Path format: ~/.kimi-code/sessions/WORKSPACE/SESSION_UUID/agents/AGENT/wire.jsonl +fn extract_session_id_from_kimi_code_path(path: &Path) -> String { + path.parent() + .and_then(|p| p.parent()) + .and_then(|p| p.parent()) + .and_then(|p| p.file_name()) + .and_then(|n| n.to_str()) + .unwrap_or("unknown") + .to_string() +} + +/// Strip the `kimi-code/` prefix from model IDs emitted by kimi-code. +fn normalize_kimi_code_model(model: &str) -> String { + model + .strip_prefix("kimi-code/") + .unwrap_or(model) + .to_string() +} + +/// Kimi Code wire.jsonl line structure. +#[derive(Debug, Deserialize)] +struct KimiCodeWireLine { + #[serde(rename = "type")] + line_type: String, + model: Option, + usage: Option, + #[serde(rename = "usageScope")] + usage_scope: Option, + time: Option, + #[serde(default)] + id: Option, + #[serde(default, rename = "eventId")] + event_id: Option, + #[serde(default, rename = "turnId")] + turn_id: Option, + #[serde(default, rename = "requestId")] + request_id: Option, + #[serde(default, rename = "messageId")] + message_id: Option, +} + +/// Build a stable, explainable identity for a Kimi Code usage record. +/// +/// The session, event timestamp/identity, model, scope, and complete token +/// payload are retained so an exact replay collapses without merging two calls +/// that happen to have the same token counts but belong to different turns. +fn identity_value(value: &Option) -> String { + value.as_ref().map(ToString::to_string).unwrap_or_default() +} + +fn kimi_code_dedup_key( + session_id: &str, + wire_line: &KimiCodeWireLine, + model: &str, + tokens: &TokenBreakdown, +) -> String { + format!( + "kimi:{session_id}:code:{line_type}:{scope}:{time}:{id}:{event_id}:{turn_id}:{request_id}:{message_id}:{model}:{input}:{output}:{cache_read}:{cache_write}:{reasoning}", + line_type = wire_line.line_type, + scope = wire_line.usage_scope.as_deref().unwrap_or(""), + time = wire_line.time.unwrap_or(0), + id = identity_value(&wire_line.id), + event_id = identity_value(&wire_line.event_id), + turn_id = identity_value(&wire_line.turn_id), + request_id = identity_value(&wire_line.request_id), + message_id = identity_value(&wire_line.message_id), + input = tokens.input, + output = tokens.output, + cache_read = tokens.cache_read, + cache_write = tokens.cache_write, + reasoning = tokens.reasoning, + ) +} + +/// Parse a Kimi Code wire.jsonl file. +pub fn parse_kimi_code_file(path: &Path) -> Vec { + let file = match std::fs::File::open(path) { + Ok(f) => f, + Err(_) => return Vec::new(), + }; + + let session_id = extract_session_id_from_kimi_code_path(path); + let fallback_timestamp = file_modified_timestamp_ms(path); + let reader = BufReader::new(file); + let mut messages: Vec = Vec::new(); + let mut seen = HashSet::new(); + + for line in reader.lines() { + let line = match line { + Ok(l) => l, + Err(_) => continue, + }; + let trimmed = line.trim(); + if trimmed.is_empty() { + continue; + } + + let mut bytes = trimmed.as_bytes().to_vec(); + let wire_line = match simd_json::from_slice::(&mut bytes) { + Ok(wl) => wl, + Err(_) => continue, + }; + + // `step.end` also carries usage, but duplicates the same usage.record + // emitted in the same turn, so only usage.record is counted. + if wire_line.line_type != "usage.record" { + continue; + } + + // Missing and session-scoped usage is bookkeeping such as compaction; + // only explicit turn-scoped records represent billable calls. + if wire_line.usage_scope.as_deref() != Some("turn") { + continue; + } + + let Some(tokens) = wire_line.usage.as_ref().and_then(TokenUsage::to_breakdown) else { + continue; + }; + + let model = wire_line + .model + .as_deref() + .map(normalize_kimi_code_model) + .unwrap_or_else(|| DEFAULT_MODEL.to_string()); + let timestamp_ms = wire_line.time.unwrap_or(fallback_timestamp); + let dedup_key = kimi_code_dedup_key(&session_id, &wire_line, &model, &tokens); + if !seen.insert(dedup_key.clone()) { + continue; + } + + messages.push(UnifiedMessage::new_with_dedup( + "kimi", + model, + DEFAULT_PROVIDER, + session_id.clone(), + timestamp_ms, + tokens, + 0.0, + Some(dedup_key), + )); + } + + messages +} + /// Parse a Kimi CLI wire.jsonl file pub fn parse_kimi_file(path: &Path) -> Vec { let file = match std::fs::File::open(path) { @@ -140,15 +340,9 @@ pub fn parse_kimi_file(path: &Path) -> Vec { .map(|ts| (ts * 1000.0) as i64) .unwrap_or_else(|| file_modified_timestamp_ms(path)); - let input = token_usage.input_other.unwrap_or(0).max(0); - let output = token_usage.output.unwrap_or(0).max(0); - let cache_read = token_usage.input_cache_read.unwrap_or(0).max(0); - let cache_write = token_usage.input_cache_creation.unwrap_or(0).max(0); - - // Skip entries with zero tokens - if input + output + cache_read + cache_write == 0 { + let Some(tokens) = token_usage.to_breakdown() else { continue; - } + }; let dedup_key = payload.message_id; @@ -158,14 +352,7 @@ pub fn parse_kimi_file(path: &Path) -> Vec { DEFAULT_PROVIDER, session_id.clone(), timestamp_ms, - TokenBreakdown { - input, - output, - cache_read, - cache_write, - // Kimi wire protocol does not expose reasoning tokens; all reasoning included in output - reasoning: 0, - }, + tokens, 0.0, dedup_key, ); @@ -224,6 +411,12 @@ mod tests { kimi_config_path(wire), Some(PathBuf::from("root/.kimi/config.json")) ); + assert_eq!( + kimi_config_path(Path::new( + "root/.kimi-code/sessions/workspace/session/agents/main/wire.jsonl" + )), + None + ); assert_eq!(kimi_config_path(Path::new("wire.jsonl")), None); } @@ -395,4 +588,145 @@ not valid json at all assert_eq!(messages.len(), 1); assert_eq!(messages[0].tokens.input, 100); } + + fn create_kimi_code_test_file(content: &str) -> (tempfile::TempDir, PathBuf) { + let dir = tempfile::tempdir().unwrap(); + let path = dir + .path() + .join(".kimi-code") + .join("sessions") + .join("workspace") + .join("session-code") + .join("agents") + .join("main") + .join("wire.jsonl"); + std::fs::create_dir_all(path.parent().unwrap()).unwrap(); + std::fs::write(&path, content).unwrap(); + (dir, path) + } + + fn kimi_code_usage(time: i64, turn_id: &str) -> String { + format!( + r#"{{"type":"usage.record","model":"kimi-code/kimi-for-coding","usage":{{"inputOther":100,"output":50,"inputCacheRead":10,"inputCacheCreation":0}},"usageScope":"turn","time":{time},"turnId":"{turn_id}"}}"# + ) + } + + #[test] + fn kimi_code_is_old_parser_zero_new_parser_nonzero() { + let content = kimi_code_usage(1_780_319_377_014, "turn-a"); + let (_dir, path) = create_kimi_code_test_file(&content); + + assert!(parse_kimi_file(&path).is_empty()); + let messages = parse_kimi_code_file(&path); + assert_eq!(messages.len(), 1); + assert_eq!(messages[0].client, "kimi"); + assert_eq!(messages[0].model_id, "kimi-for-coding"); + assert_eq!(messages[0].provider_id, "moonshot"); + assert_eq!(messages[0].session_id, "session-code"); + assert_eq!(messages[0].tokens.input, 100); + assert_eq!(messages[0].tokens.output, 50); + assert_eq!(messages[0].tokens.cache_read, 10); + } + + #[test] + fn kimi_code_extreme_tokens_do_not_overflow_zero_check() { + let content = format!( + r#"{{"type":"usage.record","model":"kimi-code/kimi-for-coding","usage":{{"inputOther":{},"output":1,"inputCacheRead":{},"inputCacheCreation":1}},"usageScope":"turn","time":1,"turnId":"turn-extreme"}}"#, + i64::MAX, + i64::MAX, + ); + let (_dir, path) = create_kimi_code_test_file(&content); + + let messages = parse_kimi_code_file(&path); + assert_eq!(messages.len(), 1); + assert_eq!(messages[0].tokens.input, i64::MAX); + assert_eq!(messages[0].tokens.output, 1); + assert_eq!(messages[0].tokens.cache_read, i64::MAX); + assert_eq!(messages[0].tokens.cache_write, 1); + assert_eq!(messages[0].tokens.total(), i64::MAX); + } + + #[test] + fn kimi_code_counts_only_explicit_turn_scope() { + let content = format!( + "{}\n{}\n{}\n", + r#"{"type":"usage.record","usage":{"inputOther":999,"output":999},"usageScope":"session","time":1}"#, + r#"{"type":"usage.record","usage":{"inputOther":888,"output":888},"time":2}"#, + kimi_code_usage(3, "turn-a"), + ); + let (_dir, path) = create_kimi_code_test_file(&content); + + let messages = parse_kimi_code_file(&path); + assert_eq!(messages.len(), 1); + assert_eq!(messages[0].tokens.input, 100); + assert_eq!(messages[0].timestamp, 3); + } + + #[test] + fn kimi_code_replay_dedup_keeps_distinct_turns() { + let replay = kimi_code_usage(10, "turn-a"); + let distinct_turn = kimi_code_usage(10, "turn-b"); + let later_turn = kimi_code_usage(11, "turn-c"); + let content = format!("{replay}\n{replay}\n{distinct_turn}\n{later_turn}\n"); + let (_dir, path) = create_kimi_code_test_file(&content); + + let messages = parse_kimi_code_file(&path); + assert_eq!(messages.len(), 3); + assert_eq!(messages[0].timestamp, 10); + assert_eq!(messages[1].timestamp, 10); + assert_eq!(messages[2].timestamp, 11); + assert_ne!(messages[0].dedup_key, messages[1].dedup_key); + assert_ne!(messages[1].dedup_key, messages[2].dedup_key); + } + + #[test] + fn kimi_legacy_and_kimi_code_paths_coexist_under_one_client() { + let dir = tempfile::tempdir().unwrap(); + let legacy_path = dir + .path() + .join(".kimi") + .join("sessions") + .join("group") + .join("legacy-session") + .join("wire.jsonl"); + std::fs::create_dir_all(legacy_path.parent().unwrap()).unwrap(); + std::fs::write( + &legacy_path, + r#"{"timestamp":1770983410.0,"message":{"type":"StatusUpdate","payload":{"token_usage":{"input_other":7,"output":3},"message_id":"legacy-1"}}}"#, + ) + .unwrap(); + let code_path = dir + .path() + .join(".kimi-code") + .join("sessions") + .join("workspace") + .join("code-session") + .join("agents") + .join("main") + .join("wire.jsonl"); + std::fs::create_dir_all(code_path.parent().unwrap()).unwrap(); + std::fs::write(&code_path, kimi_code_usage(20, "turn-code")).unwrap(); + + let legacy = parse_kimi_file(&legacy_path); + let code = parse_kimi_code_file(&code_path); + assert_eq!(legacy.len(), 1); + assert_eq!(legacy[0].session_id, "legacy-session"); + assert_eq!(legacy[0].tokens.input, 7); + assert_eq!(code.len(), 1); + assert_eq!(code[0].session_id, "code-session"); + assert_eq!(code[0].tokens.input, 100); + } + + #[test] + fn kimi_code_path_detection_uses_agents_topology() { + assert!(is_kimi_code_path(Path::new( + "/home/user/.kimi-code/sessions/workspace/session/agents/main/wire.jsonl" + ))); + assert!(!is_kimi_code_path(Path::new( + "/home/user/.kimi-code/sessions/workspace/session/wire.jsonl" + ))); + assert!(!is_kimi_code_path(Path::new( + "/home/user/.kimi/sessions/group/session/wire.jsonl" + ))); + } } diff --git a/vendor/tokscale-core/src/sessions/kiro.rs b/vendor/tokscale-core/src/sessions/kiro.rs index 639c04f..d0b586a 100644 --- a/vendor/tokscale-core/src/sessions/kiro.rs +++ b/vendor/tokscale-core/src/sessions/kiro.rs @@ -1,19 +1,23 @@ //! Kiro session parser //! -//! Parses session data from two sources: -//! 1. File-based: ~/.kiro/sessions/cli/*.json + *.jsonl -//! 2. SQLite-based: ~/Library/Application Support/kiro-cli/data.sqlite3 +//! Parses session data from four Kiro sources: +//! 1. CLI JSON + same-stem JSONL: ~/.kiro/sessions/cli/*.json + *.jsonl +//! 2. macOS IDE globalStorage: snapshot, execution, and workspace-session files +//! 3. CLI SQLite: ~/Library/Application Support/kiro-cli/data.sqlite3 //! (conversations_v2 table with history[*].request_metadata) +//! 4. Kiro IDE structured sessions: ~/.kiro/sessions//sess_*/ +//! session.json plus its sibling messages.jsonl. //! -//! Turn-level token counts are currently zero in both sources, so usage is -//! estimated from context_usage_percentage * context_window (input) and -//! response_size / 4 (output). +//! Token counts from all sources are surfaced as estimated usage. The estimates +//! use the source's available context, character, or metadata hints; downstream +//! must not treat them as provider-reported token counts. -use super::utils::file_modified_timestamp_ms; +use super::utils::{back_anchor_timestamp, file_modified_timestamp_ms}; use super::{normalize_workspace_key, workspace_label_from_key, UnifiedMessage}; use crate::TokenBreakdown; use rusqlite::Connection; use serde::Deserialize; +use serde_json::Value; use std::collections::HashMap; use std::io::{BufRead, BufReader}; use std::path::{Path, PathBuf}; @@ -22,6 +26,7 @@ use tracing::warn; const CLIENT_ID: &str = "kiro"; const PROVIDER_ID: &str = "amazon-bedrock"; const UNKNOWN_MODEL: &str = "unknown"; +const IDE_UNKNOWN_MODEL: &str = "auto"; #[derive(Debug, Deserialize)] struct KiroSessionHeader { @@ -93,12 +98,41 @@ struct KiroMessageContent { prompt_timestamp_ms: Option, } -/// Return the same-stem message sidecar consumed by a Kiro CLI session header. +#[derive(Debug, Deserialize)] +struct KiroIdeSession { + id: Option, + #[serde(rename = "modelId")] + model_id: Option, + #[serde(rename = "workspacePaths")] + workspace_paths: Option>, + #[serde(rename = "createdAt")] + created_at: Option, + #[serde(rename = "lastModifiedAt")] + last_modified_at: Option, +} + +/// Return the message sidecar consumed by a Kiro file source: same-stem JSONL +/// for CLI headers and sibling `messages.jsonl` for IDE sessions. GlobalStorage +/// and `.chat` artifacts are self-contained. pub(crate) fn kiro_related_messages_path(session_path: &Path) -> Option { + if is_kiro_ide_session_path(session_path) { + return Some(session_path.with_file_name("messages.jsonl")); + } + if is_kiro_global_storage_source(session_path) { + return None; + } Some(session_path.with_extension("jsonl")) } pub fn parse_kiro_file(path: &Path) -> Vec { + if is_kiro_ide_session_path(path) { + return parse_kiro_ide_session_file(path); + } + + if is_kiro_global_storage_source(path) { + return parse_kiro_global_storage_file(path); + } + let fallback_timestamp = file_modified_timestamp_ms(path); let mut json_bytes = match std::fs::read(path) { @@ -308,15 +342,17 @@ fn duration_between_ms(start_ms: Option, end_ms: Option) -> Option 0).then_some(duration) } +fn numeric_epoch_to_millis(timestamp: f64) -> i64 { + if timestamp.abs() < 1_000_000_000_000.0 { + seconds_to_millis(timestamp) + } else { + timestamp as i64 + } +} + fn parse_timestamp_value(value: Option<&serde_json::Value>) -> Option { match value? { - serde_json::Value::Number(number) => number.as_f64().map(|timestamp| { - if timestamp.abs() < 1_000_000_000_000.0 { - seconds_to_millis(timestamp) - } else { - timestamp as i64 - } - }), + serde_json::Value::Number(number) => number.as_f64().map(numeric_epoch_to_millis), serde_json::Value::String(timestamp) => chrono::DateTime::parse_from_rfc3339(timestamp) .ok() .map(|dt| dt.timestamp_millis()) @@ -325,6 +361,16 @@ fn parse_timestamp_value(value: Option<&serde_json::Value>) -> Option { } } +fn parse_execution_timestamp_value(value: Option<&serde_json::Value>) -> Option { + match value? { + serde_json::Value::String(timestamp) => chrono::DateTime::parse_from_rfc3339(timestamp) + .ok() + .map(|dt| dt.timestamp_millis()) + .or_else(|| timestamp.parse::().ok().map(numeric_epoch_to_millis)), + value => parse_timestamp_value(Some(value)), + } +} + fn session_id_from_path(path: &Path) -> String { path.file_stem() .and_then(|name| name.to_str()) @@ -332,6 +378,890 @@ fn session_id_from_path(path: &Path) -> String { .to_string() } +fn is_kiro_global_storage_path(path: &Path) -> bool { + let mut saw_global_storage = false; + let mut saw_extension_root = false; + for component in path.components() { + let component = component.as_os_str().to_string_lossy(); + saw_global_storage |= component == "globalStorage"; + saw_extension_root |= component == "kiro.kiroagent"; + } + saw_global_storage && saw_extension_root +} + +fn is_kiro_chat_path(path: &Path) -> bool { + path.extension() + .and_then(|extension| extension.to_str()) + .is_some_and(|extension| extension.eq_ignore_ascii_case("chat")) +} + +pub(crate) fn is_kiro_global_storage_source(path: &Path) -> bool { + is_kiro_global_storage_path(path) || is_kiro_chat_path(path) +} + +/// A Kiro IDE session file is `session.json` sitting inside a `sess_` +/// directory (`~/.kiro/sessions//sess_/session.json`). The +/// `sess_` parent requirement keeps this from matching the CLI layout, whose +/// arbitrary `~/.kiro/sessions/cli/*.json` files share the same tree. +pub(crate) fn is_kiro_ide_session_path(path: &Path) -> bool { + let is_session_json = path + .file_name() + .and_then(|name| name.to_str()) + .is_some_and(|name| name == "session.json"); + if !is_session_json { + return false; + } + path.parent() + .and_then(|parent| parent.file_name()) + .and_then(|name| name.to_str()) + .is_some_and(|name| name.starts_with("sess_")) +} + +/// Parse the Kiro IDE session layout: `session.json` (metadata) plus the +/// sibling `messages.jsonl` (conversation). +/// +/// The IDE records no per-turn token usage in these files, so structured turns +/// estimate input from context usage and output from generated text/tool args. +/// Flat role-based JSONL remains a single aggregated fallback for older files. +fn parse_kiro_ide_session_file(path: &Path) -> Vec { + let fallback_timestamp = file_modified_timestamp_ms(path); + + let session_json = match std::fs::read_to_string(path) { + Ok(contents) => contents, + Err(_) => return Vec::new(), + }; + let session: KiroIdeSession = match serde_json::from_str(&session_json) { + Ok(session) => session, + Err(_) => return Vec::new(), + }; + + let sess_dir = path.parent(); + let sess_dir_name = sess_dir + .and_then(|dir| dir.file_name()) + .and_then(|name| name.to_str()); + let session_id = session + .id + .filter(|id| !id.trim().is_empty()) + .or_else(|| sess_dir_name.map(ToOwned::to_owned)) + .unwrap_or_else(|| "unknown".to_string()); + + let session_model_id = session.model_id.filter(|model| !model.trim().is_empty()); + let workspace_path = session + .workspace_paths + .as_ref() + .and_then(|paths| paths.first()) + .map(String::as_str); + let workspace_from_dir = sess_dir + .and_then(|dir| dir.parent()) + .and_then(|workspace| workspace.file_name()) + .and_then(|name| name.to_str()) + .map(ToOwned::to_owned); + let workspace = workspace_path.map(ToOwned::to_owned).or(workspace_from_dir); + let workspace_key = workspace.as_deref().and_then(normalize_workspace_key); + let workspace_label = workspace_key.as_deref().and_then(workspace_label_from_key); + + let messages_path = path.with_file_name("messages.jsonl"); + if !messages_path.is_file() { + return Vec::new(); + } + let jsonl_file = match std::fs::File::open(&messages_path) { + Ok(file) => file, + Err(_) => return Vec::new(), + }; + let reader = BufReader::new(jsonl_file); + + const DEFAULT_CONTEXT_WINDOW: i64 = 200_000; + + #[derive(Default)] + struct IdeTurn { + prompt_chars: usize, + assistant_chars: usize, + prompt_timestamp_ms: Option, + end_timestamp_ms: Option, + context_usage_percentage: f64, + elapsed_ms: Option, + } + + let mut turns: Vec = Vec::new(); + let mut current_turn: Option = None; + let mut has_structured_format = false; + + let mut flat_counts = KiroSnapshotTextCounts::default(); + let mut flat_model_id: Option = None; + let mut flat_assistant_turns = 0i32; + + for line in reader.lines() { + let line = match line { + Ok(line) => line, + Err(_) => continue, + }; + let trimmed = line.trim(); + if trimmed.is_empty() { + continue; + } + + let entry: Value = match serde_json::from_str(trimmed) { + Ok(value) => value, + Err(_) => continue, + }; + + if let Some(payload) = entry.get("payload") { + if let Some(message_type) = payload.get("type").and_then(Value::as_str) { + has_structured_format = true; + let timestamp_ms = + entry + .get("timestamp") + .and_then(Value::as_str) + .and_then(|timestamp| { + chrono::DateTime::parse_from_rfc3339(timestamp) + .ok() + .map(|date| date.timestamp_millis()) + }); + + match message_type { + "user" => { + let chars = payload + .get("content") + .and_then(Value::as_str) + .map(str::chars) + .map(Iterator::count) + .unwrap_or(0); + let turn = current_turn.get_or_insert_with(IdeTurn::default); + turn.prompt_chars += chars; + if turn.prompt_timestamp_ms.is_none() { + turn.prompt_timestamp_ms = timestamp_ms; + } + } + "assistant" => { + let chars = payload + .get("content") + .and_then(Value::as_str) + .map(str::chars) + .map(Iterator::count) + .unwrap_or(0); + if let Some(turn) = current_turn.as_mut() { + turn.assistant_chars += chars; + } + } + "tool_call" => { + let args_chars = payload + .get("args") + .map(|args| match args { + Value::String(value) => value.chars().count(), + value => value.to_string().chars().count(), + }) + .unwrap_or(0); + if let Some(turn) = current_turn.as_mut() { + turn.assistant_chars += args_chars; + } + } + "session_metadata" + if payload.get("key").and_then(Value::as_str) == Some("contextUsage") => + { + if let Some(percentage) = payload + .get("value") + .and_then(|value| value.get("usagePercentage")) + .and_then(Value::as_f64) + { + if let Some(turn) = current_turn.as_mut() { + turn.context_usage_percentage = percentage; + } + } + } + "usage_summary" => { + if let Some(elapsed) = payload.get("elapsedTime").and_then(Value::as_i64) { + if let Some(turn) = current_turn.as_mut() { + turn.elapsed_ms = Some(elapsed); + } + } + } + "turn_end" => { + if let Some(turn) = current_turn.as_mut() { + turn.end_timestamp_ms = timestamp_ms; + } + if let Some(turn) = current_turn.take() { + if turn.prompt_chars > 0 || turn.assistant_chars > 0 { + turns.push(turn); + } + } + } + _ => {} + } + continue; + } + } + + if flat_model_id.is_none() { + flat_model_id = find_kiro_snapshot_model_id(&entry); + } + let assistant_before = flat_counts.assistant_chars; + collect_kiro_snapshot_text(&entry, &mut flat_counts, None); + if flat_counts.assistant_chars > assistant_before { + flat_assistant_turns += 1; + } + } + + if let Some(turn) = current_turn.take() { + if turn.prompt_chars > 0 || turn.assistant_chars > 0 { + turns.push(turn); + } + } + + if has_structured_format && !turns.is_empty() { + let model_id = session_model_id + .clone() + .unwrap_or_else(|| IDE_UNKNOWN_MODEL.to_string()); + return turns + .into_iter() + .enumerate() + .filter_map(|(index, turn)| { + let input = if turn.context_usage_percentage > 0.0 { + ((DEFAULT_CONTEXT_WINDOW as f64) * turn.context_usage_percentage / 100.0) as i64 + } else { + estimate_tokens(turn.prompt_chars) + }; + let output = estimate_tokens(turn.assistant_chars); + if input + output == 0 { + return None; + } + + let duration_ms = turn.elapsed_ms.or_else(|| { + duration_between_ms(turn.prompt_timestamp_ms, turn.end_timestamp_ms) + }); + let timestamp = turn + .prompt_timestamp_ms + .or_else(|| match (turn.end_timestamp_ms, duration_ms) { + (Some(end), Some(elapsed)) => Some(back_anchor_timestamp(end, elapsed)), + _ => None, + }) + .or(turn.end_timestamp_ms) + .unwrap_or(fallback_timestamp); + + let mut message = UnifiedMessage::new_with_dedup( + CLIENT_ID, + model_id.clone(), + PROVIDER_ID, + session_id.clone(), + timestamp, + TokenBreakdown { + input, + output, + cache_read: 0, + cache_write: 0, + reasoning: 0, + }, + 0.0, + Some(format!("{}:ide:{}", session_id, index)), + ); + message.message_count = 1; + message.is_turn_start = true; + message.duration_ms = duration_ms; + message.set_workspace(workspace_key.clone(), workspace_label.clone()); + Some(message) + }) + .collect(); + } + + let input = estimate_tokens(flat_counts.prompt_chars); + let output = estimate_tokens(flat_counts.assistant_chars); + if input + output == 0 { + return Vec::new(); + } + + let created_value = session + .created_at + .as_deref() + .map(|timestamp| Value::String(timestamp.to_string())); + let created_ms = parse_timestamp_value(created_value.as_ref()); + let modified_value = session + .last_modified_at + .as_deref() + .map(|timestamp| Value::String(timestamp.to_string())); + let modified_ms = parse_timestamp_value(modified_value.as_ref()); + let timestamp = created_ms.or(modified_ms).unwrap_or(fallback_timestamp); + let duration_ms = duration_between_ms(created_ms, modified_ms); + let model_id = session_model_id + .or(flat_model_id) + .filter(|model| !model.trim().is_empty()) + .unwrap_or_else(|| IDE_UNKNOWN_MODEL.to_string()); + + let mut message = UnifiedMessage::new_with_dedup( + CLIENT_ID, + model_id, + PROVIDER_ID, + session_id.clone(), + timestamp, + TokenBreakdown { + input, + output, + cache_read: 0, + cache_write: 0, + reasoning: 0, + }, + 0.0, + Some(format!("{}:ide-session", session_id)), + ); + message.message_count = flat_assistant_turns.max(1); + message.duration_ms = duration_ms; + message.is_turn_start = true; + message.set_workspace(workspace_key, workspace_label); + vec![message] +} + +#[derive(Debug, Default)] +struct KiroSnapshotTextCounts { + prompt_chars: usize, + assistant_chars: usize, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +enum KiroSnapshotRole { + Prompt, + Assistant, +} + +fn collect_kiro_snapshot_text( + value: &Value, + counts: &mut KiroSnapshotTextCounts, + mut role: Option, +) { + match value { + Value::Object(map) => { + // `role` is authoritative, so unknown roles clear inheritance. `type` + // also labels neutral payload nodes such as `text`; only explicit + // tool/unknown types clear the inherited conversation role. + if let Some(kind) = map.get("role").and_then(Value::as_str) { + role = match kind { + "user" | "prompt" | "human" => Some(KiroSnapshotRole::Prompt), + "assistant" | "response" | "bot" => Some(KiroSnapshotRole::Assistant), + _ => None, + }; + } + if let Some(kind) = map.get("type").and_then(Value::as_str) { + role = match kind { + "user" | "prompt" | "human" => Some(KiroSnapshotRole::Prompt), + "assistant" | "response" | "bot" => Some(KiroSnapshotRole::Assistant), + "tool" | "unknown" => None, + _ => role, + }; + } + + // These keys are aliases within each group. Equal subtrees are one + // payload; distinct subtrees are all meaningful conversation data. + for group in [ + &["prompt", "response", "content", "text", "message"][..], + &[ + "messages", + "conversation", + "chat", + "transcript", + "entries", + "events", + "history", + ][..], + &["parts", "items", "nodes"][..], + ] { + let mut visited: Vec<&Value> = Vec::new(); + for key in group { + if let Some(item) = map.get(*key) { + if visited.contains(&item) { + continue; + } + visited.push(item); + collect_kiro_snapshot_text(item, counts, role); + } + } + } + } + Value::Array(items) => { + for item in items { + collect_kiro_snapshot_text(item, counts, role); + } + } + Value::String(text) => match role { + Some(KiroSnapshotRole::Prompt) => counts.prompt_chars += text.chars().count(), + Some(KiroSnapshotRole::Assistant) => counts.assistant_chars += text.chars().count(), + None => {} + }, + _ => {} + } +} + +fn find_kiro_snapshot_model_id(value: &Value) -> Option { + fn is_pseudo_model(model: &str) -> bool { + matches!( + model.to_ascii_lowercase().as_str(), + "agent" | "auto" | "qdev" + ) + } + + match value { + Value::Object(map) => { + for key in ["model_id", "modelId", "model"] { + if let Some(model) = map.get(key).and_then(Value::as_str) { + let model = model.trim(); + if !model.is_empty() && !is_pseudo_model(model) { + return Some(model.to_string()); + } + } + } + for key in [ + "messages", + "conversation", + "chat", + "transcript", + "entries", + "events", + "history", + "prompt", + "response", + "content", + "text", + "message", + "parts", + "items", + "nodes", + "promptLogs", + "completionOptions", + ] { + if let Some(item) = map.get(key) { + if let Some(model) = find_kiro_snapshot_model_id(item) { + return Some(model); + } + } + } + None + } + Value::Array(items) => items.iter().find_map(find_kiro_snapshot_model_id), + _ => None, + } +} + +fn kiro_workspace_session_workspace(path: &Path) -> Option { + if path.extension().and_then(|extension| extension.to_str()) != Some("json") { + return None; + } + + let mut components = path + .components() + .map(|component| component.as_os_str().to_string_lossy().into_owned()); + while let Some(component) = components.next() { + if component == "kiro.kiroagent" { + if components.next()?.as_str() != "workspace-sessions" { + return None; + } + let workspace = components.next()?; + components.next()?; + return components.next().is_none().then_some(workspace); + } + } + None +} + +fn kiro_global_storage_workspace(path: &Path) -> Option { + if let Some(workspace) = kiro_workspace_session_workspace(path) { + return Some(workspace); + } + + let mut components = path + .components() + .map(|component| component.as_os_str().to_string_lossy().into_owned()); + while let Some(component) = components.next() { + if component == "kiro.kiroagent" { + return components + .next() + .filter(|workspace| workspace != "workspace-sessions"); + } + } + None +} + +fn parse_kiro_global_storage_file(path: &Path) -> Vec { + let fallback_timestamp = file_modified_timestamp_ms(path); + let json = match std::fs::read_to_string(path) { + Ok(contents) => contents, + Err(_) => return Vec::new(), + }; + let value: Value = match serde_json::from_str(&json) { + Ok(value) => value, + Err(_) => return Vec::new(), + }; + + if let Some(messages) = try_parse_kiro_execution_file(&value, path) { + return messages; + } + if value.get("executions").is_some() && value.get("version").is_some() { + return Vec::new(); + } + if kiro_workspace_session_workspace(path).is_some() { + if let Some(messages) = try_parse_kiro_workspace_session(&value, path, fallback_timestamp) { + return messages; + } + } + // Generic role/content traversal is valid only for legacy `.chat` snapshots. + // JSON and extensionless sources must match an execution or workspace-session + // shape above, otherwise mirrored project data could be counted as usage. + if path.extension().and_then(|extension| extension.to_str()) != Some("chat") { + return Vec::new(); + } + + let file_stem = path + .file_stem() + .and_then(|name| name.to_str()) + .unwrap_or("unknown"); + let workspace = kiro_global_storage_workspace(path); + let workspace_key = workspace.as_deref().and_then(normalize_workspace_key); + let workspace_label = workspace_key.as_deref().and_then(workspace_label_from_key); + let session_id = match workspace.as_deref() { + Some(workspace) => format!("{workspace}/{file_stem}"), + None => file_stem.to_string(), + }; + let model_id = find_kiro_snapshot_model_id(&value).unwrap_or_else(|| "auto".to_string()); + + let mut counts = KiroSnapshotTextCounts::default(); + collect_kiro_snapshot_text(&value, &mut counts, None); + let input = estimate_tokens(counts.prompt_chars); + let output = estimate_tokens(counts.assistant_chars); + if input + output == 0 { + return Vec::new(); + } + + let dedup_key = match value.get("executionId").and_then(Value::as_str) { + Some(execution_id) => format!("{session_id}:globalstorage:exec:{execution_id}"), + None => format!("{session_id}:globalstorage"), + }; + let mut message = UnifiedMessage::new_with_dedup( + CLIENT_ID, + model_id, + PROVIDER_ID, + session_id, + fallback_timestamp, + TokenBreakdown { + input, + output, + cache_read: 0, + cache_write: 0, + reasoning: 0, + }, + 0.0, + Some(dedup_key), + ); + message.message_count = 1; + message.is_turn_start = true; + message.set_workspace(workspace_key, workspace_label); + vec![message] +} + +fn try_parse_kiro_execution_file(value: &Value, path: &Path) -> Option> { + let obj = value.as_object()?; + let execution_id = obj.get("executionId")?.as_str()?; + let actions: &[Value] = match obj.get("actions") { + Some(actions) => actions.as_array()?.as_slice(), + None if path.extension().and_then(|extension| extension.to_str()) != Some("chat") + && obj.get("status").is_some() + && obj + .get("context") + .and_then(|context| context.get("messages")) + .and_then(Value::as_array) + .is_some() => + { + &[] + } + None => return None, + }; + if obj.get("status").and_then(Value::as_str) != Some("succeed") { + return Some(Vec::new()); + } + + let session_id = obj + .get("chatSessionId") + .and_then(Value::as_str) + .unwrap_or(execution_id) + .to_string(); + let start_time = parse_execution_timestamp_value(obj.get("startTime")); + let timestamp = start_time.unwrap_or_else(|| file_modified_timestamp_ms(path)); + let end_time = parse_execution_timestamp_value(obj.get("endTime")); + let duration_ms = duration_between_ms(start_time.or(Some(timestamp)), end_time); + + let output_chars: usize = actions + .iter() + .filter(|action| { + matches!( + action.get("actionType").and_then(Value::as_str), + Some("say") | Some("reasoning") + ) + }) + .filter_map(|action| action.get("output")) + .map(|output| { + output + .as_str() + .map(str::chars) + .map(Iterator::count) + .or_else(|| { + output + .get("message") + .and_then(Value::as_str) + .map(|text| text.chars().count()) + }) + .unwrap_or(0) + }) + .sum(); + + let context_input_chars = obj + .get("context") + .and_then(|context| context.get("messages")) + .and_then(Value::as_array) + .map(|messages| { + messages + .iter() + .filter_map(|message| message.get("entries").and_then(Value::as_array)) + .flatten() + .filter(|entry| entry.get("type").and_then(Value::as_str) == Some("text")) + .filter_map(|entry| entry.get("text").and_then(Value::as_str)) + .map(|text| text.chars().count()) + .sum::() + }) + .unwrap_or(0); + let input_data_chars = obj + .get("input") + .and_then(|input| input.get("data")) + .and_then(|data| data.get("messages")) + .and_then(Value::as_array) + .map(|messages| { + messages + .iter() + .map(|message| { + if let Some(parts) = message.get("content").and_then(Value::as_array) { + parts + .iter() + .filter(|part| part.get("type").and_then(Value::as_str) == Some("text")) + .filter_map(|part| part.get("text").and_then(Value::as_str)) + .map(|text| text.chars().count()) + .sum() + } else { + message + .get("content") + .and_then(Value::as_str) + .map(|text| text.chars().count()) + .unwrap_or(0) + } + }) + .sum::() + }) + .unwrap_or(0); + + let input = estimate_tokens(context_input_chars + input_data_chars); + let output = estimate_tokens(output_chars); + if input + output == 0 { + return Some(Vec::new()); + } + + let model_id = find_kiro_snapshot_model_id(value).unwrap_or_else(|| "auto".to_string()); + let workspace = kiro_global_storage_workspace(path); + let workspace_key = workspace.as_deref().and_then(normalize_workspace_key); + let workspace_label = workspace_key.as_deref().and_then(workspace_label_from_key); + let mut message = UnifiedMessage::new_with_dedup( + CLIENT_ID, + model_id, + PROVIDER_ID, + session_id, + timestamp, + TokenBreakdown { + input, + output, + cache_read: 0, + cache_write: 0, + reasoning: 0, + }, + 0.0, + Some(format!("execution:{execution_id}")), + ); + message.message_count = 1; + message.is_turn_start = true; + message.duration_ms = duration_ms; + message.set_workspace(workspace_key, workspace_label); + Some(vec![message]) +} + +fn try_parse_kiro_workspace_session( + value: &Value, + path: &Path, + fallback_timestamp: i64, +) -> Option> { + let history = value.get("history")?.as_array()?; + if value.get("sessionId").is_none() && value.get("selectedModel").is_none() { + return None; + } + + let file_stem = path + .file_stem() + .and_then(|name| name.to_str()) + .unwrap_or("unknown"); + let workspace = kiro_global_storage_workspace(path); + let workspace_key = workspace.as_deref().and_then(normalize_workspace_key); + let workspace_label = workspace_key.as_deref().and_then(workspace_label_from_key); + let session_id = value + .get("sessionId") + .and_then(Value::as_str) + .map(ToOwned::to_owned) + .unwrap_or_else(|| match workspace.as_deref() { + Some(workspace) => format!("{workspace}/{file_stem}"), + None => file_stem.to_string(), + }); + let model_id = value + .get("selectedModel") + .and_then(Value::as_str) + .filter(|model| !model.trim().is_empty()) + .unwrap_or("auto") + .to_string(); + + let mut prompt_chars = 0usize; + let mut prompt_log_count = 0i32; + let mut assistant_chars = 0usize; + for entry in history { + if let Some(prompt_logs) = entry.get("promptLogs").and_then(Value::as_array) { + for prompt_log in prompt_logs { + if let Some(prompt) = prompt_log.get("prompt").and_then(Value::as_str) { + prompt_chars += prompt.chars().count(); + prompt_log_count += 1; + } + } + } + if entry + .get("message") + .and_then(|message| message.get("role")) + .and_then(Value::as_str) + == Some("assistant") + { + assistant_chars += entry + .get("message") + .and_then(|message| message.get("content")) + .and_then(Value::as_str) + .map(|text| text.chars().count()) + .unwrap_or(0); + } + } + if prompt_chars == 0 { + return None; + } + + let input = estimate_tokens(prompt_chars); + let output = estimate_tokens(assistant_chars); + if input + output == 0 { + return Some(Vec::new()); + } + let mut message = UnifiedMessage::new_with_dedup( + CLIENT_ID, + model_id, + PROVIDER_ID, + session_id.clone(), + fallback_timestamp, + TokenBreakdown { + input, + output, + cache_read: 0, + cache_write: 0, + reasoning: 0, + }, + 0.0, + Some(format!("{session_id}:workspace-session")), + ); + message.message_count = prompt_log_count.max(1); + message.is_turn_start = true; + message.set_workspace(workspace_key, workspace_label); + Some(vec![message]) +} + +/// Merge Kiro file sources with exact globalStorage execution precedence. +/// Source paths stay attached through suppression and deduplication: only IDE +/// sources seed execution suppression, and identical keys deduplicate only +/// within the IDE or CLI cohort rather than colliding across them. +/// The match is exact for `(workspace, executionId)`; legacy snapshots use the +/// execution's `(workspace, chatSessionId)` and workspace-session artifacts use +/// the session id globally because they live under a separate storage subtree. +pub(crate) fn merge_kiro_source_messages( + sources: Vec<(PathBuf, Vec)>, +) -> Vec { + let mut executed_sessions: std::collections::HashSet<(Option, String)> = + std::collections::HashSet::new(); + let mut executed_ids: std::collections::HashSet<(Option, String)> = + std::collections::HashSet::new(); + let mut executed_session_ids: std::collections::HashSet = + std::collections::HashSet::new(); + let mut tagged_messages = Vec::new(); + + for (path, source_messages) in sources { + let is_global_storage_source = is_kiro_global_storage_source(&path); + if is_global_storage_source { + for message in &source_messages { + let Some(execution_id) = message + .dedup_key + .as_deref() + .and_then(|key| key.strip_prefix("execution:")) + else { + continue; + }; + executed_sessions + .insert((message.workspace_key.clone(), message.session_id.clone())); + executed_ids.insert((message.workspace_key.clone(), execution_id.to_string())); + executed_session_ids.insert(message.session_id.clone()); + } + } + tagged_messages.extend( + source_messages + .into_iter() + .map(|message| (is_global_storage_source, message)), + ); + } + + let mut seen_keys: std::collections::HashSet<(bool, Option, String)> = + std::collections::HashSet::new(); + tagged_messages + .into_iter() + .filter(|(is_global_storage_source, message)| { + if !*is_global_storage_source { + return true; + } + let Some(key) = message.dedup_key.as_deref() else { + return true; + }; + if let Some((_, execution_id)) = key.split_once(":globalstorage:exec:") { + return !executed_ids + .contains(&(message.workspace_key.clone(), execution_id.to_string())); + } + if key.ends_with(":workspace-session") { + return !executed_session_ids.contains(&message.session_id); + } + if !key.ends_with(":globalstorage") { + return true; + } + let stem = message + .session_id + .rsplit('/') + .next() + .unwrap_or(&message.session_id); + !executed_sessions.contains(&(message.workspace_key.clone(), stem.to_string())) + }) + .filter(|(is_global_storage_source, message)| { + message.dedup_key.as_ref().is_none_or(|key| { + let execution_workspace = + if *is_global_storage_source && key.starts_with("execution:") { + message.workspace_key.clone() + } else { + None + }; + key.is_empty() + || seen_keys.insert(( + *is_global_storage_source, + execution_workspace, + key.clone(), + )) + }) + }) + .map(|(_, message)| message) + .collect() +} + pub fn parse_kiro_sqlite(db_path: &Path) -> Vec { let conn = match Connection::open_with_flags( db_path, @@ -486,6 +1416,8 @@ struct KiroDbRequestMetadata { #[cfg(test)] mod tests { use super::*; + use std::collections::HashSet; + use std::fs; use std::io::Write; use tempfile::TempDir; @@ -537,6 +1469,22 @@ mod tests { assert_eq!(messages[0].workspace_label, Some("project".to_string())); } + #[test] + fn m15b_keeps_existing_cli_unknown_model_default() { + let dir = TempDir::new().unwrap(); + let json = r#"{"session_id":"session-unknown","session_state":{"conversation_metadata":{"user_turn_metadatas":[{"message_ids":["assistant-unknown"]}]}}}"#; + let jsonl = concat!( + "{\"kind\":\"Prompt\",\"data\":{\"message_id\":\"prompt-unknown\",\"content\":[{\"kind\":\"text\",\"data\":\"hello\"}]}}\n", + "{\"kind\":\"AssistantMessage\",\"data\":{\"message_id\":\"assistant-unknown\",\"content\":[{\"kind\":\"text\",\"data\":\"answer\"}]}}\n", + ); + let path = create_session_files(&dir, "session-unknown", json, jsonl); + + let messages = parse_kiro_file(&path); + + assert_eq!(messages.len(), 1); + assert_eq!(messages[0].model_id, "unknown"); + } + #[test] fn test_parse_kiro_skips_zero_content_turns() { let dir = TempDir::new().unwrap(); @@ -603,4 +1551,727 @@ not valid json at all assert_eq!(messages[0].tokens.input, 100); assert_eq!(messages[0].tokens.output, 10); } + + #[test] + fn m15a_globalstorage_snapshot_parser_emits_usage() { + let dir = TempDir::new().unwrap(); + let path = dir.path().join( + "Library/Application Support/Kiro/User/globalStorage/kiro.kiroagent/workspace-a/conversation.chat", + ); + std::fs::create_dir_all(path.parent().unwrap()).unwrap(); + std::fs::write( + &path, + r#"{ + "model": "claude-sonnet-4-5", + "messages": [ + {"role": "user", "content": "hello from Kiro"}, + {"role": "assistant", "content": "response from Kiro"} + ] + }"#, + ) + .unwrap(); + + let messages = parse_kiro_file(&path); + + assert!( + !messages.is_empty(), + "globalStorage fixture must produce usage; old parser returned empty" + ); + } + + fn globalstorage_path(dir: &TempDir, relative: &str) -> PathBuf { + let path = dir + .path() + .join("Library/Application Support/Kiro/User/globalStorage/kiro.kiroagent") + .join(relative); + fs::create_dir_all(path.parent().unwrap()).unwrap(); + path + } + + fn create_ide_session_files( + dir: &TempDir, + workspace: &str, + sess_dir: &str, + session_json: &str, + messages_jsonl: &str, + ) -> PathBuf { + let sess_path = dir + .path() + .join(".kiro/sessions") + .join(workspace) + .join(sess_dir); + fs::create_dir_all(&sess_path).unwrap(); + let session_path = sess_path.join("session.json"); + fs::write(&session_path, session_json).unwrap(); + fs::write(sess_path.join("messages.jsonl"), messages_jsonl).unwrap(); + session_path + } + + #[test] + fn m15b_structured_ide_session_emits_turns_with_context_and_tool_output() { + let session_json = r#"{ + "schemaVersion": "1.0.0", + "id": "sess_structured", + "modelId": "claude-opus-4.6", + "workspacePaths": ["/Users/alice/project"] + }"#; + let messages_jsonl = concat!( + "{\"timestamp\":\"2026-06-20T10:00:00Z\",\"payload\":{\"type\":\"user\",\"content\":\"hello\"}}\n", + "{\"payload\":{\"type\":\"session_metadata\",\"key\":\"contextUsage\",\"value\":{\"usagePercentage\":12.5}}}\n", + "{\"payload\":{\"type\":\"assistant\",\"content\":\"answer\"}}\n", + "{\"payload\":{\"type\":\"tool_call\",\"args\":\"toolargs\"}}\n", + "{\"payload\":{\"type\":\"usage_summary\",\"elapsedTime\":2500}}\n", + "{\"timestamp\":\"2026-06-20T10:00:02.500Z\",\"payload\":{\"type\":\"turn_end\"}}\n", + "{\"timestamp\":\"2026-06-20T10:01:00Z\",\"payload\":{\"type\":\"user\",\"content\":\"next\"}}\n", + "{\"payload\":{\"type\":\"assistant\",\"content\":\"done\"}}\n", + "{\"timestamp\":\"2026-06-20T10:01:01Z\",\"payload\":{\"type\":\"turn_end\"}}\n", + ); + let dir = TempDir::new().unwrap(); + let path = create_ide_session_files( + &dir, + "session-workspace", + "sess_structured", + session_json, + messages_jsonl, + ); + + let messages = parse_kiro_file(&path); + + assert_eq!(messages.len(), 2); + assert_eq!(messages[0].model_id, "claude-opus-4.6"); + assert_eq!( + messages[0].workspace_key.as_deref(), + Some("/Users/alice/project") + ); + assert_eq!(messages[0].workspace_label.as_deref(), Some("project")); + assert_eq!(messages[0].tokens.input, 25_000); + assert_eq!(messages[0].tokens.output, 4); + assert_eq!(messages[0].duration_ms, Some(2_500)); + assert_eq!(messages[0].message_count, 1); + assert_eq!( + messages[0].dedup_key.as_deref(), + Some("sess_structured:ide:0") + ); + assert_eq!(messages[1].tokens.input, 1); + assert_eq!(messages[1].tokens.output, 1); + assert_eq!(messages[1].message_count, 1); + assert_eq!( + messages[1].dedup_key.as_deref(), + Some("sess_structured:ide:1") + ); + } + + #[test] + fn m15b_flat_ide_session_fallback_preserves_model_and_workspace() { + let session_json = r#"{ + "schemaVersion": "1.0.0", + "id": "sess_flat", + "modelId": "session-model", + "workspacePaths": ["C:\\work\\repo"], + "createdAt": "2026-06-20T10:00:00Z", + "lastModifiedAt": "2026-06-20T10:00:01Z" + }"#; + let messages_jsonl = concat!( + "{\"role\":\"user\",\"content\":\"hello\"}\n", + "{\"role\":\"assistant\",\"model\":\"line-model\",\"content\":\"world\"}\n", + ); + let dir = TempDir::new().unwrap(); + let path = create_ide_session_files( + &dir, + "dir-workspace", + "sess_flat", + session_json, + messages_jsonl, + ); + + let messages = parse_kiro_file(&path); + + assert_eq!(messages.len(), 1); + assert_eq!(messages[0].model_id, "session-model"); + assert_eq!(messages[0].workspace_key.as_deref(), Some("C:/work/repo")); + assert_eq!(messages[0].workspace_label.as_deref(), Some("repo")); + assert_eq!(messages[0].tokens.input, 2); + assert_eq!(messages[0].tokens.output, 2); + assert_eq!(messages[0].message_count, 1); + assert_eq!(messages[0].duration_ms, Some(1_000)); + } + + #[test] + fn m15b_ide_session_handles_absent_and_malformed_siblings() { + let dir = TempDir::new().unwrap(); + let missing = create_ide_session_files( + &dir, + "workspace", + "sess_missing", + r#"{"id":"sess_missing"}"#, + "{\"role\":\"user\",\"content\":\"ignored\"}\n", + ); + fs::remove_file(missing.with_file_name("messages.jsonl")).unwrap(); + assert!(parse_kiro_file(&missing).is_empty()); + + let malformed = create_ide_session_files( + &dir, + "workspace", + "sess_malformed", + r#"{"id":"sess_malformed"}"#, + "not json\n{\"role\":\"user\",\"content\":\"hello\"}\n{\"role\":\"assistant\",\"content\":\"world\"}\n", + ); + assert!(is_kiro_ide_session_path(&malformed)); + assert!(!is_kiro_ide_session_path(Path::new( + ".kiro/sessions/cli/session.json" + ))); + let messages = parse_kiro_file(&malformed); + assert_eq!(messages.len(), 1); + assert_eq!(messages[0].tokens.input, 2); + assert_eq!(messages[0].tokens.output, 2); + } + + #[test] + fn m15b_structured_turn_missing_prompt_timestamp_back_calculates_anchor() { + let session_json = r#"{"id":"sess_anchor"}"#; + let messages_jsonl = concat!( + "{\"payload\":{\"type\":\"user\",\"content\":\"hello world\"}}\n", + "{\"payload\":{\"type\":\"assistant\",\"content\":\"response text\"}}\n", + "{\"payload\":{\"type\":\"usage_summary\",\"elapsedTime\":5000}}\n", + "{\"payload\":{\"type\":\"turn_end\"},\"timestamp\":\"2026-06-20T10:00:05Z\"}\n", + ); + let dir = TempDir::new().unwrap(); + let path = create_ide_session_files( + &dir, + "workspace", + "sess_anchor", + session_json, + messages_jsonl, + ); + + let messages = parse_kiro_file(&path); + + let end = chrono::DateTime::parse_from_rfc3339("2026-06-20T10:00:05Z") + .unwrap() + .timestamp_millis(); + assert_eq!(messages.len(), 1); + assert_eq!(messages[0].model_id, "auto"); + assert_eq!(messages[0].timestamp, end - 5_000); + assert_eq!(messages[0].duration_ms, Some(5_000)); + } + + fn make_test_message( + client: &str, + session_id: &str, + dedup_key: &str, + workspace: Option<&str>, + ) -> UnifiedMessage { + let mut message = UnifiedMessage::new_with_dedup( + client, + "auto".to_string(), + PROVIDER_ID, + session_id.to_string(), + 1_770_983_426_000, + TokenBreakdown { + input: 10, + output: 2, + cache_read: 0, + cache_write: 0, + reasoning: 0, + }, + 0.0, + Some(dedup_key.to_string()), + ); + message.set_workspace(workspace.map(str::to_string), workspace.map(str::to_string)); + message + } + + #[test] + fn test_parse_kiro_globalstorage_snapshot_roles_aliases_and_model() { + let dir = TempDir::new().unwrap(); + let path = globalstorage_path(&dir, "workspace-a/snapshot.chat"); + fs::write( + &path, + r#"{ + "model": "auto", + "completionOptions": {"modelId": "claude-sonnet-4-5"}, + "messages": [ + {"role": "user", "content": "abcd", "text": "abcd"}, + {"role": "assistant", "content": "1234", "text": "1234"}, + {"role": "tool", "content": "tool context must not count"}, + {"role": "unknown", "content": "unknown must not count"} + ], + "history": [{"role": "human", "content": "efghij"}] + }"#, + ) + .unwrap(); + + let messages = parse_kiro_file(&path); + + assert_eq!(messages.len(), 1); + assert_eq!(messages[0].model_id, "claude-sonnet-4-5"); + assert_eq!(messages[0].tokens.input, 3); + assert_eq!(messages[0].tokens.output, 1); + assert_eq!(messages[0].workspace_key.as_deref(), Some("workspace-a")); + assert_eq!(messages[0].workspace_label.as_deref(), Some("workspace-a")); + assert_eq!( + messages[0].dedup_key.as_deref(), + Some("workspace-a/snapshot:globalstorage") + ); + } + + #[test] + fn test_collect_kiro_snapshot_text_value_dedup_and_distinct_subtrees() { + let value: Value = serde_json::from_str( + r#"{ + "prompt": {"role": "user", "text": "hello"}, + "response": {"role": "assistant", "text": "world"}, + "messages": [{"role": "user", "content": "alpha"}], + "entries": [{"role": "user", "content": "alpha"}], + "history": [{"role": "user", "content": "bravo"}], + "parts": [{"role": "assistant", "text": "cd"}], + "items": [{"role": "assistant", "text": "cd"}] + }"#, + ) + .unwrap(); + let mut counts = KiroSnapshotTextCounts::default(); + collect_kiro_snapshot_text(&value, &mut counts, None); + + assert_eq!(counts.prompt_chars, 5 + 5 + 5); + assert_eq!(counts.assistant_chars, 5 + 2); + } + + #[test] + fn test_collect_kiro_snapshot_text_nested_tool_and_unknown_override_parent_role() { + let value: Value = serde_json::from_str( + r#"{ + "messages": [ + {"role": "user", "content": {"parts": [ + {"type": "tool", "text": "tool input must not count"}, + {"type": "text", "text": "ABCD"} + ]}}, + {"role": "assistant", "content": {"parts": [ + {"type": "tool", "text": "tool output must not count"}, + {"type": "unknown", "text": "unknown type must not count"}, + {"role": "tool", "text": "tool role must not count"}, + {"role": "unknown", "text": "unknown role must not count"}, + {"type": "text", "text": "WXYZ"} + ]}} + ] + }"#, + ) + .unwrap(); + let mut counts = KiroSnapshotTextCounts::default(); + collect_kiro_snapshot_text(&value, &mut counts, None); + + assert_eq!(counts.prompt_chars, 4); + assert_eq!(counts.assistant_chars, 4); + } + + #[test] + fn test_find_kiro_snapshot_model_skips_pseudo_models_recursively() { + let value: Value = serde_json::from_str( + r#"{ + "model_id": "agent", + "promptLogs": [{"model": "auto"}], + "conversation": [{"completionOptions": {"modelId": "qdev"}}], + "history": [{"model": "claude-sonnet-4-5"}] + }"#, + ) + .unwrap(); + assert_eq!( + find_kiro_snapshot_model_id(&value).as_deref(), + Some("claude-sonnet-4-5") + ); + } + + #[test] + fn test_parse_kiro_execution_supports_input_output_shapes_and_model_duration() { + let dir = TempDir::new().unwrap(); + let path = globalstorage_path(&dir, "workspace-a/execution-store/execution-one"); + fs::write( + &path, + r#"{ + "executionId": "exec-one", + "chatSessionId": "chat-one", + "status": "succeed", + "startTime": "2026-02-13T12:00:00Z", + "endTime": 1770984001500.0, + "completionOptions": {"model": "claude-sonnet-4-5"}, + "context": {"messages": [{"entries": [ + {"type": "text", "text": "context text"}, + {"type": "image", "text": "ignored"} + ]}]}, + "input": {"data": {"messages": [ + {"content": "string input"}, + {"content": [ + {"type": "text", "text": "part input"}, + {"type": "image", "text": "ignored"} + ]} + ]}}, + "actions": [ + {"actionType": "say", "output": "answer"}, + {"actionType": "reasoning", "output": {"message": "thinking"}}, + {"actionType": "tool", "output": "ignored"} + ] + }"#, + ) + .unwrap(); + + let messages = parse_kiro_file(&path); + + assert_eq!(messages.len(), 1); + assert_eq!(messages[0].session_id, "chat-one"); + assert_eq!(messages[0].model_id, "claude-sonnet-4-5"); + assert_eq!(messages[0].tokens.input, 9); + assert_eq!(messages[0].tokens.output, 4); + assert_eq!(messages[0].timestamp, 1770984000000); + assert_eq!(messages[0].duration_ms, Some(1500)); + assert_eq!(messages[0].dedup_key.as_deref(), Some("execution:exec-one")); + assert_eq!(messages[0].workspace_key.as_deref(), Some("workspace-a")); + } + + #[test] + fn test_parse_kiro_actionless_execution_uses_context_messages() { + let dir = TempDir::new().unwrap(); + let path = globalstorage_path(&dir, "workspace-a/execution-store/actionless"); + fs::write( + &path, + r#"{ + "executionId": "exec-actionless", + "chatSessionId": "chat-actionless", + "status": "succeed", + "startTime": "1770983426000", + "endTime": "1770983427500", + "context": {"messages": [{"entries": [ + {"type": "text", "text": "actionless input"}, + {"type": "image", "text": "ignored"} + ]}]} + }"#, + ) + .unwrap(); + + let messages = parse_kiro_file(&path); + + assert_eq!(messages.len(), 1); + assert_eq!(messages[0].session_id, "chat-actionless"); + assert_eq!(messages[0].tokens.input, 4); + assert_eq!(messages[0].tokens.output, 0); + assert_eq!(messages[0].timestamp, 1770983426000); + assert_eq!(messages[0].duration_ms, Some(1500)); + assert_eq!( + messages[0].dedup_key.as_deref(), + Some("execution:exec-actionless") + ); + } + + #[test] + fn test_parse_kiro_status_bearing_chat_without_actions_stays_snapshot() { + let dir = TempDir::new().unwrap(); + let path = globalstorage_path(&dir, "workspace-a/status.chat"); + fs::write( + &path, + r#"{ + "executionId": "snapshot-execution", + "status": "succeed", + "messages": [ + {"role": "user", "content": "0123456789abcdef"}, + {"role": "assistant", "content": "response"} + ] + }"#, + ) + .unwrap(); + + let messages = parse_kiro_file(&path); + + assert_eq!(messages.len(), 1); + assert_eq!(messages[0].tokens.input, 4); + assert_eq!(messages[0].tokens.output, 2); + assert_eq!( + messages[0].dedup_key.as_deref(), + Some("workspace-a/status:globalstorage:exec:snapshot-execution") + ); + } + + #[test] + fn test_parse_kiro_execution_success_failed_and_container_artifacts() { + let dir = TempDir::new().unwrap(); + let success = globalstorage_path(&dir, "workspace-a/execution-store/success.json"); + fs::write( + &success, + r#"{"executionId":"success","status":"succeed","startTime":1770983426,"endTime":1770983427500,"actions":[{"actionType":"say","output":{"message":"answer"}}],"input":{"data":{"messages":[{"content":"question"}]}}}"#, + ) + .unwrap(); + let failed = globalstorage_path(&dir, "workspace-a/execution-store/failed.json"); + fs::write( + &failed, + r#"{"executionId":"failed","status":"failed","actions":[{"actionType":"say","output":"answer"}],"input":{"data":{"messages":[{"content":"question"}]}}}"#, + ) + .unwrap(); + let container = globalstorage_path(&dir, "workspace-a/execution-store/executions.json"); + fs::write(&container, r#"{"version":1,"executions":[]}"#).unwrap(); + + let success_messages = parse_kiro_file(&success); + assert_eq!(success_messages.len(), 1); + assert_eq!(success_messages[0].timestamp, 1770983426000); + assert!(parse_kiro_file(&failed).is_empty()); + assert!(parse_kiro_file(&container).is_empty()); + } + + #[test] + fn test_parse_kiro_globalstorage_skips_non_session_json_and_extensionless_files() { + let dir = TempDir::new().unwrap(); + let project_json = globalstorage_path(&dir, "workspace-a/project.json"); + let mirrored_file = globalstorage_path(&dir, "workspace-a/project-store/mirror"); + let mirrored_session = globalstorage_path(&dir, "workspace-a/project-store/session.json"); + let body = r#"{"messages":[{"role":"user","content":"must not count"}]}"#; + fs::write(&project_json, body).unwrap(); + fs::write(&mirrored_file, body).unwrap(); + fs::write( + &mirrored_session, + r#"{"sessionId":"project-session","history":[{"promptLogs":[{"prompt":"must not count"}]}]}"#, + ) + .unwrap(); + + assert!(parse_kiro_file(&project_json).is_empty()); + assert!(parse_kiro_file(&mirrored_file).is_empty()); + assert!(parse_kiro_file(&mirrored_session).is_empty()); + } + + #[test] + fn test_parse_kiro_workspace_session_promptlogs() { + let dir = TempDir::new().unwrap(); + let path = globalstorage_path(&dir, "workspace-sessions/workspace-a/session.json"); + fs::write( + &path, + r#"{ + "sessionId": "session-1", + "selectedModel": "claude-sonnet-4", + "history": [ + {"promptLogs": [{"prompt": "0123456789"}]}, + {"promptLogs": [{"prompt": "abcdefghijklmnopqrst"}], "message": {"role": "assistant", "content": "answer"}} + ] + }"#, + ) + .unwrap(); + + let messages = parse_kiro_file(&path); + + assert_eq!(messages.len(), 1); + assert_eq!(messages[0].model_id, "claude-sonnet-4"); + assert_eq!(messages[0].tokens.input, 8); + assert_eq!(messages[0].tokens.output, 2); + assert_eq!(messages[0].message_count, 2); + assert_eq!(messages[0].workspace_key.as_deref(), Some("workspace-a")); + assert_eq!(messages[0].workspace_label.as_deref(), Some("workspace-a")); + assert_eq!( + messages[0].dedup_key.as_deref(), + Some("session-1:workspace-session") + ); + } + + #[test] + fn test_kiro_related_messages_path_is_none_only_for_globalstorage() { + let dir = TempDir::new().unwrap(); + let cli = dir.path().join("session.json"); + let chat = globalstorage_path(&dir, "workspace-a/session.chat"); + let extensionless = globalstorage_path(&dir, "workspace-a/execution-store/execution"); + assert_eq!( + kiro_related_messages_path(&cli), + Some(dir.path().join("session.jsonl")) + ); + assert_eq!(kiro_related_messages_path(&chat), None); + assert_eq!(kiro_related_messages_path(&extensionless), None); + } + + #[test] + fn test_kiro_suppression_ignores_cli_execution_prefix_collision() { + let dir = TempDir::new().unwrap(); + let cli = create_session_files( + &dir, + "cli", + r#"{"session_id":"execution","cwd":"workspace-a","session_state":{"rts_model_state":{"model_info":{"model_id":"cli-model"}},"conversation_metadata":{"user_turn_metadatas":[{"input_token_count":1}]}}}"#, + "", + ); + let snapshot = globalstorage_path(&dir, "workspace-a/snapshot.chat"); + fs::write( + &snapshot, + r#"{"executionId":"0","messages":[{"role":"user","content":"ABCD"}]}"#, + ) + .unwrap(); + + let cli_messages = parse_kiro_file(&cli); + let snapshot_messages = parse_kiro_file(&snapshot); + let kept = + merge_kiro_source_messages(vec![(cli, cli_messages), (snapshot, snapshot_messages)]); + let keys: HashSet<_> = kept + .iter() + .filter_map(|message| message.dedup_key.as_deref()) + .collect(); + + assert_eq!(kept.len(), 2); + assert!(keys.contains("execution:0")); + assert!(keys.contains("workspace-a/snapshot:globalstorage:exec:0")); + } + + #[test] + fn m15a_globalstorage_and_m15b_ide_sources_remain_distinct() { + let dir = TempDir::new().unwrap(); + let ide = create_ide_session_files( + &dir, + "workspace-a", + "sess_overlap", + r#"{"id":"sess_overlap","modelId":"claude-opus-4.6"}"#, + concat!( + "{\"payload\":{\"type\":\"user\",\"content\":\"question\"}}\n", + "{\"payload\":{\"type\":\"assistant\",\"content\":\"answer\"}}\n", + "{\"payload\":{\"type\":\"turn_end\"}}\n", + ), + ); + let execution = globalstorage_path(&dir, "workspace-a/execution-store/execution"); + fs::write( + &execution, + r#"{ + "executionId":"exec-overlap", + "chatSessionId":"chat-overlap", + "status":"succeed", + "startTime":1770983426000, + "actions":[{"actionType":"say","output":"answer"}], + "context":{"messages":[{"entries":[{"type":"text","text":"question"}]}]} + }"#, + ) + .unwrap(); + + let execution_messages = parse_kiro_file(&execution); + let ide_messages = parse_kiro_file(&ide); + assert_eq!(execution_messages.len(), 1); + assert_eq!(ide_messages.len(), 1); + let kept = merge_kiro_source_messages(vec![ + (execution.clone(), execution_messages), + (ide.clone(), ide_messages), + ]); + let keys: HashSet<_> = kept + .iter() + .filter_map(|message| message.dedup_key.as_deref()) + .collect(); + + assert_eq!(kept.len(), 2); + assert!(keys.contains("execution:exec-overlap")); + assert!(keys.contains("sess_overlap:ide:0")); + } + + #[test] + fn test_kiro_execution_dedup_is_workspace_scoped() { + let dir = TempDir::new().unwrap(); + let kept = merge_kiro_source_messages(vec![ + ( + globalstorage_path(&dir, "ws-a/execution-store/execution"), + vec![make_test_message( + CLIENT_ID, + "chat-a", + "execution:shared-id", + Some("ws-a"), + )], + ), + ( + globalstorage_path(&dir, "ws-b/execution-store/execution"), + vec![make_test_message( + CLIENT_ID, + "chat-b", + "execution:shared-id", + Some("ws-b"), + )], + ), + ]); + let workspaces: HashSet<_> = kept + .iter() + .filter_map(|message| message.workspace_key.as_deref()) + .collect(); + + assert_eq!(kept.len(), 2); + assert_eq!(workspaces, HashSet::from(["ws-a", "ws-b"])); + } + + #[test] + fn test_kiro_suppression_is_workspace_scoped_and_preserves_other_lanes() { + let dir = TempDir::new().unwrap(); + let global_messages = vec![ + make_test_message( + CLIENT_ID, + "ws-a/chat-a", + "ws-a/chat-a:globalstorage:exec:exec-1", + Some("ws-a"), + ), + make_test_message( + CLIENT_ID, + "ws-a/chat-a", + "ws-a/chat-a:globalstorage", + Some("ws-a"), + ), + make_test_message(CLIENT_ID, "chat-a", "execution:exec-1", Some("ws-a")), + make_test_message( + CLIENT_ID, + "ws-b/chat-a", + "ws-b/chat-a:globalstorage:exec:exec-1", + Some("ws-b"), + ), + make_test_message( + CLIENT_ID, + "ws-a/other", + "ws-a/other:globalstorage", + Some("ws-a"), + ), + make_test_message( + CLIENT_ID, + "ws-a/failed", + "ws-a/failed:globalstorage:exec:exec-failed", + Some("ws-a"), + ), + make_test_message( + CLIENT_ID, + "ws-a/unmatched", + "ws-a/unmatched:globalstorage:exec:exec-missing", + Some("ws-a"), + ), + make_test_message( + CLIENT_ID, + "session-1", + "session-1:workspace-session", + Some("ws-a"), + ), + make_test_message( + CLIENT_ID, + "session-1", + "execution:exec-session", + Some("ws-b"), + ), + make_test_message( + CLIENT_ID, + "session-2", + "session-2:workspace-session", + Some("ws-a"), + ), + ]; + let other_messages = vec![ + make_test_message(CLIENT_ID, "chat-a", "cli-session:0", Some("ws-a")), + make_test_message(CLIENT_ID, "chat-a", "sqlite-session:0", Some("ws-a")), + ]; + + let kept = merge_kiro_source_messages(vec![ + (globalstorage_path(&dir, "ws-a/execution"), global_messages), + (dir.path().join("cli.json"), other_messages), + ]); + let keys: HashSet<_> = kept + .iter() + .filter_map(|message| message.dedup_key.as_deref()) + .collect(); + + assert!(!keys.contains("ws-a/chat-a:globalstorage:exec:exec-1")); + assert!(!keys.contains("ws-a/chat-a:globalstorage")); + assert!(keys.contains("execution:exec-1")); + assert!(keys.contains("ws-b/chat-a:globalstorage:exec:exec-1")); + assert!(keys.contains("ws-a/other:globalstorage")); + assert!(keys.contains("ws-a/failed:globalstorage:exec:exec-failed")); + assert!(keys.contains("ws-a/unmatched:globalstorage:exec:exec-missing")); + assert!(!keys.contains("session-1:workspace-session")); + assert!(keys.contains("execution:exec-session")); + assert!(keys.contains("session-2:workspace-session")); + assert!(keys.contains("cli-session:0")); + assert!(keys.contains("sqlite-session:0")); + } } diff --git a/vendor/tokscale-core/src/sessions/mod.rs b/vendor/tokscale-core/src/sessions/mod.rs index 3049bdb..1a07321 100644 --- a/vendor/tokscale-core/src/sessions/mod.rs +++ b/vendor/tokscale-core/src/sessions/mod.rs @@ -10,6 +10,7 @@ pub mod cline; pub mod codebuff; pub mod codex; pub mod copilot; +pub mod copilot_desktop; pub mod crush; pub mod cursor; pub mod droid; @@ -19,6 +20,7 @@ pub mod goose; pub mod grok; pub mod hermes; pub mod jcode; +pub mod junie; pub mod kilo; pub mod kilocode; pub mod kimi; @@ -27,6 +29,7 @@ pub mod micode; pub mod mux; pub mod openclaw; pub mod opencode; +pub mod opencodereview; pub mod pi; pub mod qwen; pub mod roocode; @@ -67,6 +70,11 @@ pub struct UnifiedMessage { pub message_count: i32, pub agent: Option, pub dedup_key: Option, + /// Alternate source keys for the same logical message. OpenCode uses these + /// when one migrated copy has an embedded id and another only has a row/file + /// fallback; ordinary clients leave this empty. + #[serde(default)] + pub dedup_aliases: Vec, /// True if this message is the first assistant response after a user turn. /// Used to count user interaction turns (as opposed to API message count). #[serde(default)] @@ -346,6 +354,7 @@ impl UnifiedMessage { message_count: default_message_count(), agent, dedup_key, + dedup_aliases: Vec::new(), is_turn_start: false, } } diff --git a/vendor/tokscale-core/src/sessions/opencode.rs b/vendor/tokscale-core/src/sessions/opencode.rs index 5e5e684..1f7f8c9 100644 --- a/vendor/tokscale-core/src/sessions/opencode.rs +++ b/vendor/tokscale-core/src/sessions/opencode.rs @@ -6,7 +6,7 @@ use super::utils::{open_readonly_sqlite, read_file_or_none}; use super::{ - normalize_opencode_agent_name, normalize_workspace_key, workspace_label_from_key, + normalize_opencode_agent_name, normalize_workspace_key, workspace_label_from_key, CostSource, UnifiedMessage, }; use crate::{provider_identity, TokenBreakdown}; @@ -16,7 +16,14 @@ use serde::{Deserialize, Serialize}; use std::collections::HashMap; use std::path::Path; -/// OpenCode message structure (from JSON files and SQLite data column) +/// OpenCode message structure (from JSON files and SQLite data column). +/// +/// Handles two on-disk shapes: +/// - **v1** (`opencode.db` `message` table, legacy JSON files): a `role` +/// field, and top-level `modelID` / `providerID` strings. +/// - **v2** (`opencode-next.db` `session_message` table): no `role` field +/// (the row's `type` column carries it), and the model identifiers nested +/// under a `model` object (`model.id` / `model.providerID`). #[derive(Debug, Deserialize)] #[allow(dead_code)] pub struct OpenCodeMessage { @@ -24,11 +31,17 @@ pub struct OpenCodeMessage { pub id: Option, #[serde(rename = "sessionID", default)] pub session_id: Option, - pub role: String, - #[serde(rename = "modelID")] + /// Absent in v2 `session_message` rows (the `type` column is the role + /// there and the SQL query already filters to `assistant`). + #[serde(default)] + pub role: Option, + #[serde(rename = "modelID", default)] pub model_id: Option, - #[serde(rename = "providerID")] + #[serde(rename = "providerID", default)] pub provider_id: Option, + /// v2 nests model + provider under a `model` object. + #[serde(default)] + pub model: Option, pub cost: Option, pub tokens: Option, pub time: OpenCodeTime, @@ -38,6 +51,40 @@ pub struct OpenCodeMessage { pub path: Option, } +impl OpenCodeMessage { + /// Resolve the model id from the top-level v1 field or the nested v2 + /// `model.id`, preferring the explicit top-level value when both exist. + fn resolve_model_id(&self) -> Option { + self.model_id + .clone() + .or_else(|| self.model.as_ref().and_then(|m| m.id.clone())) + } + + /// Resolve the provider id from the top-level v1 field or the nested v2 + /// `model.providerID`, preferring the explicit top-level value. + fn resolve_provider_id(&self) -> Option { + self.provider_id + .clone() + .or_else(|| self.model.as_ref().and_then(|m| m.provider_id.clone())) + } + + /// True when this row is an assistant turn. v1 rows carry an explicit + /// `role`; v2 rows omit it and are pre-filtered by the SQL `type` column, + /// so a missing role is treated as assistant. + fn is_assistant(&self) -> bool { + self.role.as_deref().is_none_or(|role| role == "assistant") + } +} + +/// v2 nested model descriptor: `{"id": "...", "providerID": "...", ...}`. +#[derive(Debug, Deserialize)] +pub struct OpenCodeModel { + #[serde(default)] + pub id: Option, + #[serde(rename = "providerID", default)] + pub provider_id: Option, +} + #[derive(Debug, Deserialize)] pub struct OpenCodePath { pub root: Option, @@ -88,13 +135,21 @@ struct OpenCodeSqliteFingerprint { reasoning: i64, cache_read: i64, cache_write: i64, - cost_bits: u64, + // Cost is deliberately absent: it changes authority, not payload identity. agent: Option, } #[derive(Debug, Clone)] struct OpenCodeSqliteDedupState { - has_embedded_message_id: bool, + /// The entry's embedded (`$.id`) message id, if any. Two rows that share + /// every fingerprint field but carry *different* embedded ids are distinct + /// messages, not fork copies, and must not be merged. A fork copies the id, + /// so equal ids (or an id absent on either side) still merge. + message_id: Option, + /// Whether an id-less v1 row has supplied its `message.id` fallback. It is + /// promoted when no embedded id exists, or retained as an alias when v2 did + /// carry one, so both JSON and sibling-database identities remain reachable. + has_v1_fallback_key: bool, has_workspace_conflict: bool, } @@ -160,7 +215,13 @@ pub fn parse_opencode_file(path: &Path) -> Option { let msg: OpenCodeMessage = simd_json::from_slice(&mut bytes).ok()?; - if msg.role != "assistant" { + // OpenCode JSON files (v1) always carry an explicit role, so require it to + // be "assistant" here. Missing-role acceptance (is_assistant) is reserved + // for the v2 `session_message` SQLite path, whose SQL already filters + // `type = 'assistant'`; applying it to files would count a role-less or + // malformed file as assistant usage (previously it was skipped when the + // required `role` field failed to deserialize). + if msg.role.as_deref() != Some("assistant") { return None; } @@ -169,8 +230,15 @@ pub fn parse_opencode_file(path: &Path) -> Option { .as_ref() .and_then(|path| path.root.as_deref()) .map(str::to_string); + // Resolve model + provider before moving any fields out of `msg`, since + // both borrow the whole struct to fall back onto the nested `model` object. + let model_id = msg.resolve_model_id()?; + let provider_id = msg + .resolve_provider_id() + .unwrap_or_else(|| "unknown".to_string()); + let provider_id = provider_identity::canonical_provider(&provider_id).unwrap_or(provider_id); + let tokens = msg.tokens?; - let model_id = msg.model_id?; let agent_or_mode = msg.mode.or(msg.agent); let agent = agent_or_mode.map(|a| normalize_opencode_agent_name(&a)); @@ -182,9 +250,6 @@ pub fn parse_opencode_file(path: &Path) -> Option { .and_then(|s| s.to_str()) .map(|s| s.to_string()) }); - - let provider_id = msg.provider_id.unwrap_or_else(|| "unknown".to_string()); - let provider_id = provider_identity::canonical_provider(&provider_id).unwrap_or(provider_id); let cost = embedded_cost(msg.cost); let mut unified = UnifiedMessage::new_with_agent( @@ -210,65 +275,60 @@ pub fn parse_opencode_file(path: &Path) -> Option { Some(unified) } -pub fn parse_opencode_sqlite(db_path: &Path) -> Vec { - let Some(conn) = open_readonly_sqlite(db_path) else { - return Vec::new(); - }; - - let modern_query = r#" - SELECT m.id, m.session_id, m.data, NULLIF(s.directory, '') AS workspace_root - FROM message m - LEFT JOIN session s ON s.id = m.session_id - WHERE json_extract(m.data, '$.role') = 'assistant' - AND json_extract(m.data, '$.tokens') IS NOT NULL - ORDER BY m.id, m.session_id - "#; - - let legacy_query = r#" - SELECT m.id, m.session_id, m.data, NULL AS workspace_root - FROM message m - WHERE json_extract(m.data, '$.role') = 'assistant' - AND json_extract(m.data, '$.tokens') IS NOT NULL - ORDER BY m.id, m.session_id - "#; - - let mut stmt = match conn - .prepare(modern_query) - .or_else(|_| conn.prepare(legacy_query)) - { - Ok(s) => s, - Err(_) => return Vec::new(), - }; +/// Column layout shared by every OpenCode SQLite query variant: +/// `(row_id, session_id, data_json, workspace_root)`. +type OpenCodeSqliteRow = (String, String, String, Option); - let rows = match stmt.query_map([], |row| { - let id: String = row.get(0)?; - let session_id: String = row.get(1)?; - let data_json: String = row.get(2)?; - let workspace_root: Option = row.get(3)?; - Ok((id, session_id, data_json, workspace_root)) - }) { - Ok(r) => r, - Err(_) => return Vec::new(), +fn merge_opencode_dedup_key( + message: &mut UnifiedMessage, + incoming_key: Option, + prefer_incoming: bool, +) { + let Some(incoming_key) = incoming_key else { + return; }; + if message.dedup_key.as_ref() == Some(&incoming_key) { + return; + } + if prefer_incoming { + if let Some(previous) = message.dedup_key.replace(incoming_key.clone()) { + if !message.dedup_aliases.contains(&previous) { + message.dedup_aliases.push(previous); + } + } + message.dedup_aliases.retain(|alias| alias != &incoming_key); + } else if !message.dedup_aliases.contains(&incoming_key) { + message.dedup_aliases.push(incoming_key); + } +} - let mut messages: Vec = Vec::new(); - let mut fingerprint_indices: HashMap = HashMap::new(); - let mut dedup_states: Vec = Vec::new(); +/// Accumulates parsed assistant messages across OpenCode's v1 (`message`) and +/// v2 (`session_message`) tables, applying fingerprint-based deduplication so +/// forked-history copies — and any overlap between the two tables — collapse +/// into a single entry. A fingerprint maps to a *list* of entries, one per +/// distinct embedded message id, so two genuinely different messages that +/// happen to collide on every fingerprint field are kept apart. +#[derive(Default)] +struct OpenCodeSqliteAccumulator { + messages: Vec, + fingerprint_indices: HashMap>, + dedup_states: Vec, +} - for row_result in rows { - let (row_id, session_id, data_json, row_workspace_root) = match row_result { - Ok(r) => r, - Err(_) => continue, - }; +impl OpenCodeSqliteAccumulator { + /// Parse one SQLite row's JSON payload and merge it into the accumulator, + /// deduplicating against previously ingested rows. + fn ingest_row(&mut self, row: OpenCodeSqliteRow, is_v1: bool) { + let (row_id, session_id, data_json, row_workspace_root) = row; let mut bytes = data_json.into_bytes(); let msg: OpenCodeMessage = match simd_json::from_slice(&mut bytes) { Ok(m) => m, - Err(_) => continue, + Err(_) => return, }; - if msg.role != "assistant" { - continue; + if !msg.is_assistant() { + return; } let message_id = msg.id.clone(); @@ -279,19 +339,21 @@ pub fn parse_opencode_sqlite(db_path: &Path) -> Vec { .map(str::to_string); let tokens = match msg.tokens { - Some(t) => t, - None => continue, + Some(ref t) => t, + None => return, }; - let model_id = match msg.model_id { + let model_id = match msg.resolve_model_id() { Some(m) => m, - None => continue, + None => return, }; - let provider_id = msg.provider_id.unwrap_or_else(|| "unknown".to_string()); + let provider_id = msg + .resolve_provider_id() + .unwrap_or_else(|| "unknown".to_string()); let provider_id = provider_identity::canonical_provider(&provider_id).unwrap_or(provider_id); - let agent_or_mode = msg.mode.or(msg.agent); + let agent_or_mode = msg.mode.clone().or_else(|| msg.agent.clone()); let agent = agent_or_mode.map(|a| normalize_opencode_agent_name(&a)); let input = tokens.input.max(0); let output = tokens.output.max(0); @@ -310,7 +372,6 @@ pub fn parse_opencode_sqlite(db_path: &Path) -> Vec { reasoning, cache_read, cache_write, - cost_bits: cost.to_bits(), agent: agent.clone(), }; @@ -338,25 +399,148 @@ pub fn parse_opencode_sqlite(db_path: &Path) -> Vec { set_workspace_from_root(&mut unified, workspace_root); mark_opencode_cost_source(&mut unified); - if let Some(index) = fingerprint_indices.get(&fingerprint).copied() { - let dedup_state = &mut dedup_states[index]; - if message_id.is_some() && !dedup_state.has_embedded_message_id { - dedup_state.has_embedded_message_id = true; - messages[index].dedup_key = unified.dedup_key; + // Among entries sharing this fingerprint, merge into the first one that + // is NOT a definitively-different message -- i.e. skip any whose stored + // embedded id conflicts with this row's. The SQL row id remains a source + // fallback, not a hard identity: forked sessions can assign a new row id + // to the same id-less payload. (Cloning the small index list avoids + // holding a borrow of `fingerprint_indices` while we read `dedup_states`.) + let candidate = { + let slots = self + .fingerprint_indices + .get(&fingerprint) + .cloned() + .unwrap_or_default(); + slots.into_iter().find(|&index| { + !matches!( + (&self.dedup_states[index].message_id, &message_id), + (Some(existing), Some(incoming)) if existing != incoming + ) + }) + }; + + if let Some(index) = candidate { + let dedup_state = &mut self.dedup_states[index]; + // First copy carrying an embedded id promotes the entry's stable + // dedup key (and records the id so later rows can be told apart). + let promote_embedded_id = message_id.is_some() && dedup_state.message_id.is_none(); + if promote_embedded_id { + dedup_state.message_id = message_id.clone(); + } + let promote_v1_fallback = message_id.is_none() + && is_v1 + && dedup_state.message_id.is_none() + && !dedup_state.has_v1_fallback_key; + if message_id.is_none() && is_v1 { + dedup_state.has_v1_fallback_key = true; } - merge_duplicate_workspace(&mut messages[index], dedup_state, workspace_root); - continue; + merge_opencode_dedup_key( + &mut self.messages[index], + unified.dedup_key.clone(), + promote_embedded_id || promote_v1_fallback, + ); + if self.messages[index].cost_source != CostSource::ProviderReported + && unified.cost_source == CostSource::ProviderReported + { + self.messages[index].cost = unified.cost; + self.messages[index].cost_source = unified.cost_source; + } + merge_duplicate_workspace(&mut self.messages[index], dedup_state, workspace_root); + return; } - dedup_states.push(OpenCodeSqliteDedupState { - has_embedded_message_id: message_id.is_some(), + let new_index = self.messages.len(); + self.dedup_states.push(OpenCodeSqliteDedupState { + message_id: message_id.clone(), + has_v1_fallback_key: is_v1 && message_id.is_none(), has_workspace_conflict: false, }); - fingerprint_indices.insert(fingerprint, messages.len()); - messages.push(unified); + self.fingerprint_indices + .entry(fingerprint) + .or_default() + .push(new_index); + self.messages.push(unified); + } +} + +/// Run one query (whose columns are `id, session_id, data, workspace_root`) +/// against `conn` and feed every row into `acc`. A prepare/query failure — for +/// example a table that does not exist in this schema variant — is treated as +/// "no rows", so callers can attempt several schema variants against the same +/// database without an error aborting the scan. +fn collect_opencode_rows( + conn: &rusqlite::Connection, + query: &str, + acc: &mut OpenCodeSqliteAccumulator, + is_v1: bool, +) { + let mut stmt = match conn.prepare(query) { + Ok(s) => s, + Err(_) => return, + }; + + let rows = match stmt.query_map([], |row| { + let id: String = row.get(0)?; + let session_id: String = row.get(1)?; + let data_json: String = row.get(2)?; + let workspace_root: Option = row.get(3)?; + Ok((id, session_id, data_json, workspace_root)) + }) { + Ok(r) => r, + Err(_) => return, + }; + + for row_result in rows.flatten() { + acc.ingest_row(row_result, is_v1); + } +} + +pub fn parse_opencode_sqlite(db_path: &Path) -> Vec { + let Some(conn) = open_readonly_sqlite(db_path) else { + return Vec::new(); + }; + + let mut acc = OpenCodeSqliteAccumulator::default(); + + // OpenCode v2 (`opencode-next.db`): per-message rows live in + // `session_message`, keyed by a `type` column, with model + provider nested + // under `$.model`. Absent in v1 databases, where the prepare fails and this + // is a no-op. + let v2_query = r#" + SELECT sm.id, sm.session_id, sm.data, NULLIF(s.directory, '') AS workspace_root + FROM session_message sm + LEFT JOIN session s ON s.id = sm.session_id + WHERE sm.type = 'assistant' + AND json_extract(sm.data, '$.tokens') IS NOT NULL + ORDER BY sm.id, sm.session_id + "#; + collect_opencode_rows(&conn, v2_query, &mut acc, false); + + // OpenCode v1 (`opencode.db`, 1.2+): per-message rows in `message`, role in + // the JSON `$.role`. The `session` join supplies the workspace directory; + // the legacy variant drops it for databases without a `session` table. + let v1_modern_query = r#" + SELECT m.id, m.session_id, m.data, NULLIF(s.directory, '') AS workspace_root + FROM message m + LEFT JOIN session s ON s.id = m.session_id + WHERE json_extract(m.data, '$.role') = 'assistant' + AND json_extract(m.data, '$.tokens') IS NOT NULL + ORDER BY m.id, m.session_id + "#; + let v1_legacy_query = r#" + SELECT m.id, m.session_id, m.data, NULL AS workspace_root + FROM message m + WHERE json_extract(m.data, '$.role') = 'assistant' + AND json_extract(m.data, '$.tokens') IS NOT NULL + ORDER BY m.id, m.session_id + "#; + if conn.prepare(v1_modern_query).is_ok() { + collect_opencode_rows(&conn, v1_modern_query, &mut acc, true); + } else { + collect_opencode_rows(&conn, v1_legacy_query, &mut acc, true); } - messages + acc.messages } // ============================================================================= @@ -526,6 +710,318 @@ mod tests { conn } + /// Build a database shaped like OpenCode v2 (`opencode-next.db`): an empty + /// `message` table plus the `session_message` + `session` tables that hold + /// the real per-message data. Mirrors the columns tokscale actually reads. + fn create_opencode_v2_sqlite_db(db_path: &Path) -> Connection { + let conn = Connection::open(db_path).unwrap(); + conn.execute_batch( + "CREATE TABLE message ( + id TEXT PRIMARY KEY, + session_id TEXT NOT NULL, + data TEXT NOT NULL + ); + CREATE TABLE session ( + id TEXT PRIMARY KEY, + directory TEXT NOT NULL + ); + CREATE TABLE session_message ( + id TEXT PRIMARY KEY, + session_id TEXT NOT NULL, + type TEXT NOT NULL, + data TEXT NOT NULL + );", + ) + .unwrap(); + conn + } + + /// A representative v2 assistant payload: no `role` field, model + provider + /// nested under `$.model`, integer timestamps. + const V2_ASSISTANT_DATA: &str = r#"{ + "time": { "created": 1783882279705, "completed": 1783882279943 }, + "agent": "build", + "model": { "id": "claude-sonnet-4", "providerID": "anthropic", "variant": "default" }, + "content": [], + "finish": "stop", + "cost": 0.0123, + "tokens": { + "input": 5519, + "output": 20, + "reasoning": 23, + "cache": { "read": 100, "write": 50 } + } + }"#; + + #[test] + fn test_deserialize_v2_message_resolves_nested_model() { + let mut bytes = V2_ASSISTANT_DATA.as_bytes().to_vec(); + let msg: OpenCodeMessage = simd_json::from_slice(&mut bytes).unwrap(); + + assert_eq!(msg.role, None, "v2 payloads carry no role field"); + assert!(msg.is_assistant(), "missing role defaults to assistant"); + assert_eq!(msg.resolve_model_id().as_deref(), Some("claude-sonnet-4")); + assert_eq!(msg.resolve_provider_id().as_deref(), Some("anthropic")); + assert_eq!(msg.agent.as_deref(), Some("build")); + } + + #[test] + fn test_top_level_model_id_takes_precedence_over_nested() { + let json = r#"{ + "role": "assistant", + "modelID": "top-level-model", + "providerID": "top-level-provider", + "model": { "id": "nested-model", "providerID": "nested-provider" }, + "tokens": { "input": 1, "output": 1, "reasoning": 0, "cache": { "read": 0, "write": 0 } }, + "time": { "created": 1700000000000.0 } + }"#; + let mut bytes = json.as_bytes().to_vec(); + let msg: OpenCodeMessage = simd_json::from_slice(&mut bytes).unwrap(); + + assert_eq!(msg.resolve_model_id().as_deref(), Some("top-level-model")); + assert_eq!( + msg.resolve_provider_id().as_deref(), + Some("top-level-provider") + ); + } + + #[test] + fn test_parse_v2_session_message_reads_tokens_and_workspace() { + let dir = tempfile::tempdir().unwrap(); + let db_path = dir.path().join("opencode-next.db"); + + let conn = create_opencode_v2_sqlite_db(&db_path); + conn.execute( + "INSERT INTO session (id, directory) VALUES (?1, ?2)", + rusqlite::params!["ses_v2", "/Users/alice/opencode-v2-repo"], + ) + .unwrap(); + conn.execute( + "INSERT INTO session_message (id, session_id, type, data) VALUES (?1, ?2, ?3, ?4)", + rusqlite::params!["msg_v2_001", "ses_v2", "assistant", V2_ASSISTANT_DATA], + ) + .unwrap(); + drop(conn); + + let messages = parse_opencode_sqlite(&db_path); + assert_eq!(messages.len(), 1, "v2 assistant row should be parsed"); + let msg = &messages[0]; + assert_eq!(msg.model_id, "claude-sonnet-4"); + assert_eq!(msg.provider_id, "anthropic"); + assert_eq!(msg.tokens.input, 5519); + assert_eq!(msg.tokens.output, 20); + assert_eq!(msg.tokens.reasoning, 23); + assert_eq!(msg.tokens.cache_read, 100); + assert_eq!(msg.tokens.cache_write, 50); + assert_eq!(msg.duration_ms, Some(238)); + assert_eq!( + msg.workspace_key.as_deref(), + Some("/Users/alice/opencode-v2-repo"), + "workspace should come from session.directory" + ); + assert_eq!(msg.workspace_label.as_deref(), Some("opencode-v2-repo")); + assert_eq!( + msg.dedup_key.as_deref(), + Some("msg_v2_001"), + "v2 dedup_key falls back to the session_message row id" + ); + assert_eq!( + msg.cost_source, + crate::sessions::CostSource::ProviderReported + ); + } + + #[test] + fn test_parse_opencode_sqlite_deduplicates_v1_v2_overlap() { + let dir = tempfile::tempdir().unwrap(); + let db_path = dir.path().join("opencode-next.db"); + let conn = create_opencode_v2_sqlite_db(&db_path); + let v1_data = r#"{ + "role": "assistant", + "modelID": "claude-sonnet-4", + "providerID": "anthropic", + "agent": "build", + "cost": 0.0123, + "tokens": { "input": 5519, "output": 20, "reasoning": 23, "cache": { "read": 100, "write": 50 } }, + "time": { "created": 1783882279705, "completed": 1783882279943 } + }"#; + conn.execute( + "INSERT INTO message (id, session_id, data) VALUES (?1, ?2, ?3)", + rusqlite::params!["v1_row", "ses_overlap", v1_data], + ) + .unwrap(); + let v2_data = V2_ASSISTANT_DATA.replace("\"cost\": 0.0123", "\"cost\": 0.0"); + conn.execute( + "INSERT INTO session_message (id, session_id, type, data) VALUES (?1, ?2, ?3, ?4)", + rusqlite::params!["v2_row", "ses_overlap", "assistant", v2_data], + ) + .unwrap(); + drop(conn); + + let messages = parse_opencode_sqlite(&db_path); + assert_eq!(messages.len(), 1, "v1/v2 overlap should be counted once"); + assert_eq!(messages[0].tokens.input, 5519); + assert_eq!( + messages[0].dedup_key.as_deref(), + Some("v1_row"), + "an id-less overlap must retain the v1 row key used by legacy JSON filenames" + ); + assert_eq!(messages[0].dedup_aliases, vec!["v2_row"]); + assert_eq!(messages[0].cost, 0.0123); + assert_eq!(messages[0].cost_source, CostSource::ProviderReported); + } + + #[test] + fn test_parse_v2_skips_non_assistant_and_tokenless_rows() { + let dir = tempfile::tempdir().unwrap(); + let db_path = dir.path().join("opencode-next.db"); + + let conn = create_opencode_v2_sqlite_db(&db_path); + let user_data = r#"{ "time": { "created": 1783882279705 }, "content": [] }"#; + let tokenless = r#"{ "time": { "created": 1783882279705 }, "model": { "id": "m", "providerID": "p" } }"#; + conn.execute( + "INSERT INTO session_message (id, session_id, type, data) VALUES (?1, ?2, ?3, ?4)", + rusqlite::params!["msg_ok", "ses_v2", "assistant", V2_ASSISTANT_DATA], + ) + .unwrap(); + conn.execute( + "INSERT INTO session_message (id, session_id, type, data) VALUES (?1, ?2, ?3, ?4)", + rusqlite::params!["msg_user", "ses_v2", "user", user_data], + ) + .unwrap(); + conn.execute( + "INSERT INTO session_message (id, session_id, type, data) VALUES (?1, ?2, ?3, ?4)", + rusqlite::params!["msg_synthetic", "ses_v2", "synthetic", V2_ASSISTANT_DATA], + ) + .unwrap(); + conn.execute( + "INSERT INTO session_message (id, session_id, type, data) VALUES (?1, ?2, ?3, ?4)", + rusqlite::params!["msg_no_tokens", "ses_v2", "assistant", tokenless], + ) + .unwrap(); + drop(conn); + + let messages = parse_opencode_sqlite(&db_path); + assert_eq!( + messages.len(), + 1, + "only the assistant row with tokens should parse" + ); + assert_eq!(messages[0].dedup_key.as_deref(), Some("msg_ok")); + } + + #[test] + fn test_parse_v2_negative_tokens_clamped() { + let dir = tempfile::tempdir().unwrap(); + let db_path = dir.path().join("opencode-next.db"); + + let conn = create_opencode_v2_sqlite_db(&db_path); + let negative = r#"{ + "time": { "created": 1783882279705 }, + "model": { "id": "claude-sonnet-4", "providerID": "anthropic" }, + "cost": -1.0, + "tokens": { "input": -100, "output": -50, "reasoning": -25, "cache": { "read": -200, "write": -10 } } + }"#; + conn.execute( + "INSERT INTO session_message (id, session_id, type, data) VALUES (?1, ?2, ?3, ?4)", + rusqlite::params!["msg_neg", "ses_v2", "assistant", negative], + ) + .unwrap(); + drop(conn); + + let messages = parse_opencode_sqlite(&db_path); + assert_eq!(messages.len(), 1); + let msg = &messages[0]; + assert_eq!(msg.tokens.input, 0); + assert_eq!(msg.tokens.output, 0); + assert_eq!(msg.tokens.reasoning, 0); + assert_eq!(msg.tokens.cache_read, 0); + assert_eq!(msg.tokens.cache_write, 0); + assert!(msg.cost >= 0.0); + } + + #[test] + fn test_parse_v2_deduplicates_forked_session_message_history() { + let dir = tempfile::tempdir().unwrap(); + let db_path = dir.path().join("opencode-next.db"); + + let conn = create_opencode_v2_sqlite_db(&db_path); + // Forking can assign a new SQL row id while copying the same id-less + // payload, so row ids cannot be hard identities for this overlap. + conn.execute( + "INSERT INTO session_message (id, session_id, type, data) VALUES (?1, ?2, ?3, ?4)", + rusqlite::params!["root_row", "root_session", "assistant", V2_ASSISTANT_DATA], + ) + .unwrap(); + conn.execute( + "INSERT INTO session_message (id, session_id, type, data) VALUES (?1, ?2, ?3, ?4)", + rusqlite::params!["fork_row", "fork_session", "assistant", V2_ASSISTANT_DATA], + ) + .unwrap(); + drop(conn); + + let messages = parse_opencode_sqlite(&db_path); + assert_eq!( + messages.len(), + 1, + "forked copies of the same assistant turn collapse inside v2 parsing" + ); + } + + #[test] + fn test_distinct_embedded_ids_are_not_merged_despite_fingerprint_collision() { + // Two genuinely different assistant messages can share every fingerprint + // field (timestamp, model, tokens, cost, agent). When both carry an + // embedded `$.id` and the ids DIFFER, they are distinct messages -- not + // fork copies -- and must be kept separate rather than collapsed. + let dir = tempfile::tempdir().unwrap(); + let db_path = dir.path().join("opencode-next.db"); + let conn = create_opencode_v2_sqlite_db(&db_path); + + let payload = |id: &str| { + format!( + r#"{{ + "id": "{id}", + "time": {{ "created": 1783882279705, "completed": 1783882279943 }}, + "agent": "build", + "model": {{ "id": "claude-sonnet-4", "providerID": "anthropic" }}, + "cost": 0.0123, + "tokens": {{ "input": 10, "output": 5, "reasoning": 0, "cache": {{ "read": 0, "write": 0 }} }} + }}"# + ) + }; + + conn.execute( + "INSERT INTO session_message (id, session_id, type, data) VALUES (?1, ?2, ?3, ?4)", + rusqlite::params!["row_a", "ses_v2", "assistant", payload("msg_a")], + ) + .unwrap(); + conn.execute( + "INSERT INTO session_message (id, session_id, type, data) VALUES (?1, ?2, ?3, ?4)", + rusqlite::params!["row_b", "ses_v2", "assistant", payload("msg_b")], + ) + .unwrap(); + // A true fork of msg_a (same embedded id, different session/row) must + // still collapse into msg_a rather than becoming a third entry. + conn.execute( + "INSERT INTO session_message (id, session_id, type, data) VALUES (?1, ?2, ?3, ?4)", + rusqlite::params!["row_a_fork", "fork_session", "assistant", payload("msg_a")], + ) + .unwrap(); + drop(conn); + + let mut dedup_keys: Vec = parse_opencode_sqlite(&db_path) + .into_iter() + .filter_map(|m| m.dedup_key) + .collect(); + dedup_keys.sort(); + assert_eq!( + dedup_keys, + vec!["msg_a".to_string(), "msg_b".to_string()], + "distinct embedded ids stay separate; a same-id fork collapses" + ); + } + #[test] fn test_parse_opencode_structure() { let json = r#"{ @@ -629,6 +1125,42 @@ mod tests { ); } + #[test] + fn test_parse_opencode_file_requires_explicit_assistant_role() { + use std::io::Write; + // Regression: making `role` optional for the v2 SQLite path must NOT + // loosen file parsing. A file without a `role` (or a non-assistant one) + // is not assistant usage and must be skipped -- the missing-role => + // assistant shortcut applies only to the type-filtered session_message + // SQLite query, never to JSON files. + let role_less = r#"{ + "modelID": "claude-sonnet-4", + "providerID": "anthropic", + "tokens": { "input": 10, "output": 5, "reasoning": 0, "cache": { "read": 0, "write": 0 } }, + "time": { "created": 1700000000000.0 } + }"#; + let mut f1 = tempfile::Builder::new().suffix(".json").tempfile().unwrap(); + f1.write_all(role_less.as_bytes()).unwrap(); + assert!( + parse_opencode_file(f1.path()).is_none(), + "a role-less OpenCode JSON file must not be counted as assistant usage" + ); + + let user_role = r#"{ + "role": "user", + "modelID": "claude-sonnet-4", + "providerID": "anthropic", + "tokens": { "input": 10, "output": 5, "reasoning": 0, "cache": { "read": 0, "write": 0 } }, + "time": { "created": 1700000000000.0 } + }"#; + let mut f2 = tempfile::Builder::new().suffix(".json").tempfile().unwrap(); + f2.write_all(user_role.as_bytes()).unwrap(); + assert!( + parse_opencode_file(f2.path()).is_none(), + "a non-assistant OpenCode JSON file must be skipped" + ); + } + /// JSON dedup_key uses msg.id when present #[test] fn test_dedup_key_from_json_message_id() { diff --git a/vendor/tokscale-core/src/sessions/opencodereview.rs b/vendor/tokscale-core/src/sessions/opencodereview.rs new file mode 100644 index 0000000..ed8522e --- /dev/null +++ b/vendor/tokscale-core/src/sessions/opencodereview.rs @@ -0,0 +1,294 @@ +//! OpenCodeReview session parser +//! +//! OpenCodeReview stores sessions as JSONL files under +//! `~/.opencodereview/sessions//.jsonl`. + +use super::utils::{back_anchor_timestamp, file_modified_timestamp_ms, parse_timestamp_value}; +use super::UnifiedMessage; +use crate::{pricing, provider_identity, TokenBreakdown}; +use serde_json::Value; +use std::collections::HashSet; +use std::io::{BufRead, BufReader}; +use std::path::Path; + +pub fn parse_opencodereview_file(path: &Path) -> Vec { + let file = match std::fs::File::open(path) { + Ok(f) => f, + Err(_) => return Vec::new(), + }; + + let session_id = session_id_from_path(path); + let fallback_timestamp = file_modified_timestamp_ms(path); + + let mut workspace: Option = None; + let mut messages = Vec::new(); + let mut seen = HashSet::new(); + + for line in BufReader::new(file).lines() { + let Ok(line) = line else { + continue; + }; + + if !line.contains("llm_response") && !line.contains("session_start") { + continue; + } + + let Ok(value) = serde_json::from_str::(&line) else { + continue; + }; + + let record_type = value.get("type").and_then(Value::as_str).unwrap_or(""); + + if record_type == "session_start" { + if workspace.is_none() { + workspace = value.get("cwd").and_then(Value::as_str).map(str::to_string); + } + continue; + } + + if record_type != "llm_response" { + continue; + } + + let usage = match value.get("usage") { + Some(u) => u, + None => continue, + }; + + let tokens = tokens_from_usage(usage); + if tokens.total() == 0 { + continue; + } + + let explicit_timestamp = value.get("timestamp").and_then(parse_timestamp_value); + let recorded_timestamp = explicit_timestamp.unwrap_or(fallback_timestamp); + + let model_raw = value + .get("model") + .and_then(Value::as_str) + .unwrap_or("unknown"); + let model_id = pricing::aliases::resolve_alias(model_raw) + .unwrap_or(model_raw) + .to_string(); + + let provider_id = provider_identity::inferred_provider_from_model(&model_id) + .map(str::to_string) + .unwrap_or_else(|| "opencodereview".to_string()); + + let duration_ms = value + .get("duration_ms") + .and_then(Value::as_i64) + .filter(|d| *d > 0); + + let timestamp = match (explicit_timestamp, duration_ms) { + (Some(end), Some(duration)) => back_anchor_timestamp(end, duration), + _ => recorded_timestamp, + }; + + // Use the recorded end timestamp in the identity. The message timestamp + // is start-anchored for sessionization, while a replay has the same + // source record identity regardless of that derived anchor. + let dedup_key = format!( + "opencodereview:{session_id}:{recorded_timestamp}:{model_id}:{}:{}:{}:{}", + tokens.input, tokens.output, tokens.cache_read, tokens.cache_write, + ); + if !seen.insert(dedup_key.clone()) { + continue; + } + + let mut msg = UnifiedMessage::new( + "opencodereview", + model_id, + provider_id, + &session_id, + timestamp, + tokens, + 0.0, + ); + msg.dedup_key = Some(dedup_key); + msg.duration_ms = duration_ms; + + if let Some(ws) = &workspace { + if let Some(key) = super::normalize_workspace_key(ws) { + msg.workspace_label = super::workspace_label_from_key(&key); + msg.workspace_key = Some(key); + } + } + + messages.push(msg); + } + + messages +} + +fn session_id_from_path(path: &Path) -> String { + path.file_stem() + .and_then(|s| s.to_str()) + .filter(|s| !s.trim().is_empty()) + .unwrap_or("unknown") + .to_string() +} + +fn tokens_from_usage(usage: &Value) -> TokenBreakdown { + TokenBreakdown { + input: number_field(usage, "prompt_tokens"), + output: number_field(usage, "completion_tokens"), + cache_read: number_field(usage, "cache_read_tokens"), + cache_write: number_field(usage, "cache_write_tokens"), + reasoning: 0, + } +} + +fn number_field(value: &Value, field: &str) -> i64 { + value + .get(field) + .and_then(|v| { + v.as_i64() + .or_else(|| v.as_u64().map(|u| u.min(i64::MAX as u64) as i64)) + }) + .unwrap_or(0) + .max(0) +} + +#[cfg(test)] +mod tests { + use super::*; + use std::io::Write; + use tempfile::TempDir; + + fn parse_events(content: &str) -> Vec { + let dir = TempDir::new().unwrap(); + let repo_dir = dir.path().join("test-repo"); + std::fs::create_dir_all(&repo_dir).unwrap(); + let path = repo_dir.join("test-session-123.jsonl"); + let mut file = std::fs::File::create(&path).unwrap(); + file.write_all(content.as_bytes()).unwrap(); + file.flush().unwrap(); + parse_opencodereview_file(&path) + } + + fn session_start(cwd: &str) -> String { + format!( + r#"{{"type":"session_start","sessionId":"test-session-123","timestamp":"2026-01-15T10:00:00Z","cwd":"{cwd}","model":"claude-sonnet-4-20250514"}}"# + ) + } + + fn llm_response( + timestamp: &str, + model: &str, + prompt: i64, + completion: i64, + cache_read: i64, + cache_write: i64, + ) -> String { + format!( + r#"{{"type":"llm_response","sessionId":"test-session-123","timestamp":"{timestamp}","model":"{model}","duration_ms":1500,"usage":{{"prompt_tokens":{prompt},"completion_tokens":{completion},"cache_read_tokens":{cache_read},"cache_write_tokens":{cache_write}}}}}"# + ) + } + + #[test] + fn parses_single_llm_response() { + let content = format!( + "{}\n{}\n", + session_start("/home/user/project"), + llm_response( + "2026-01-15T10:00:05Z", + "claude-sonnet-4-20250514", + 1000, + 200, + 500, + 100, + ), + ); + let msgs = parse_events(&content); + + assert_eq!(msgs.len(), 1); + assert_eq!(msgs[0].client, "opencodereview"); + assert_eq!(msgs[0].tokens.input, 1000); + assert_eq!(msgs[0].tokens.output, 200); + assert_eq!(msgs[0].tokens.cache_read, 500); + assert_eq!(msgs[0].tokens.cache_write, 100); + assert_eq!(msgs[0].tokens.reasoning, 0); + assert_eq!(msgs[0].duration_ms, Some(1500)); + assert_eq!(msgs[0].session_id, "test-session-123"); + assert!(msgs[0].workspace_key.is_some()); + } + + #[test] + fn deduplicates_identical_records() { + let resp = llm_response("2026-01-15T10:00:05Z", "gpt-4o", 1000, 200, 0, 0); + let content = format!("{}\n{}\n", resp, resp); + let msgs = parse_events(&content); + assert_eq!(msgs.len(), 1); + } + + #[test] + fn parses_multiple_responses_and_skips_zero_tokens() { + let content = format!( + "{}\n{}\n{}\n", + llm_response( + "2026-01-15T10:00:05Z", + "claude-sonnet-4-20250514", + 1000, + 200, + 0, + 0 + ), + llm_response("2026-01-15T10:01:00Z", "gpt-4o", 500, 100, 0, 0), + llm_response("2026-01-15T10:02:00Z", "gpt-4o", 0, 0, 0, 0), + ); + let msgs = parse_events(&content); + assert_eq!(msgs.len(), 2); + } + + #[test] + fn oversized_unsigned_tokens_clamp() { + let usage = serde_json::json!({ "prompt_tokens": u64::MAX }); + assert_eq!(number_field(&usage, "prompt_tokens"), i64::MAX); + } + + #[test] + fn works_without_session_start() { + let content = llm_response("2026-01-15T10:00:05Z", "gpt-4o", 1000, 200, 0, 0); + let msgs = parse_events(&format!("{content}\n")); + assert_eq!(msgs.len(), 1); + assert!(msgs[0].workspace_key.is_none()); + } + + #[test] + fn start_anchor_uses_recorded_end_timestamp() { + let content = llm_response("2026-01-15T10:00:05Z", "gpt-4o", 100, 50, 0, 0); + let msgs = parse_events(&format!("{content}\n")); + + assert_eq!(msgs.len(), 1); + let expected_end = + parse_timestamp_value(&Value::String("2026-01-15T10:00:05Z".to_string())).unwrap(); + assert_eq!(msgs[0].timestamp, expected_end - 1500); + assert_eq!(msgs[0].duration_ms, Some(1500)); + } + + #[test] + fn missing_timestamp_does_not_anchor_from_file_fallback() { + let dir = TempDir::new().unwrap(); + let path = dir.path().join("missing-time.jsonl"); + let content = r#"{"type":"llm_response","model":"gpt-4o","duration_ms":1500,"usage":{"prompt_tokens":100,"completion_tokens":50}}"#; + std::fs::write(&path, format!("{content}\n")).unwrap(); + let msgs = parse_opencodereview_file(&path); + + assert_eq!(msgs.len(), 1); + assert_eq!(msgs[0].duration_ms, Some(1500)); + assert!(msgs[0].timestamp > 0); + } + + #[test] + fn session_id_derived_from_filename() { + let dir = TempDir::new().unwrap(); + let path = dir.path().join("my-unique-session.jsonl"); + let content = llm_response("2026-01-15T10:00:05Z", "gpt-4o", 100, 50, 0, 0); + std::fs::write(&path, format!("{content}\n")).unwrap(); + + let msgs = parse_opencodereview_file(&path); + assert_eq!(msgs.len(), 1); + assert_eq!(msgs[0].session_id, "my-unique-session"); + } +} diff --git a/vendor/tokscale-core/src/sessions/pi.rs b/vendor/tokscale-core/src/sessions/pi.rs index aaa1944..b4b9e60 100644 --- a/vendor/tokscale-core/src/sessions/pi.rs +++ b/vendor/tokscale-core/src/sessions/pi.rs @@ -202,12 +202,14 @@ pub fn parse_pi_file(path: &Path) -> Vec { None => continue, }; - // A missing provider field is recoverable: infer it from the model name - // (and fall back to "pi") rather than dropping a message that carries + // A missing/blank provider field is recoverable: infer it from the + // model name (e.g. a Pi "gpt-5" message with no provider maps to + // "openai"), falling back to "pi" only when inference can't + // identify the model, rather than dropping a message that carries // valid tokens. let provider = match message.provider { - Some(p) => p, - None => inferred_provider_from_model(&model) + Some(p) if !p.is_empty() => p, + _ => inferred_provider_from_model(&model) .unwrap_or("pi") .to_string(), }; @@ -278,6 +280,42 @@ mod tests { assert_eq!(messages[0].workspace_label, Some("tmp".to_string())); } + #[test] + fn test_parse_pi_infers_provider_from_model_when_absent() { + // A missing provider must be inferred from the model name (gpt-5 -> + // openai), not hardcoded to the client name. + let content = r#"{"type":"session","id":"pi_ses_005","timestamp":"2026-01-01T00:00:00.000Z","cwd":"/tmp"} +{"type":"message","id":"msg_001","parentId":null,"timestamp":"2026-01-01T00:00:01.000Z","message":{"role":"assistant","model":"gpt-5","usage":{"input":100,"output":50,"cacheRead":0,"cacheWrite":0,"totalTokens":150}}}"#; + let file = create_test_file(content); + let messages = parse_pi_file(file.path()); + + assert_eq!(messages.len(), 1); + assert_eq!(messages[0].model_id, "gpt-5"); + assert_eq!(messages[0].provider_id, "openai"); + } + + #[test] + fn test_parse_pi_infers_provider_from_model_when_blank() { + let content = r#"{"type":"session","id":"pi_ses_006","timestamp":"2026-01-01T00:00:00.000Z","cwd":"/tmp"} +{"type":"message","id":"msg_001","parentId":null,"timestamp":"2026-01-01T00:00:01.000Z","message":{"role":"assistant","model":"gpt-5","provider":"","usage":{"input":100,"output":50,"cacheRead":0,"cacheWrite":0,"totalTokens":150}}}"#; + let file = create_test_file(content); + let messages = parse_pi_file(file.path()); + + assert_eq!(messages.len(), 1); + assert_eq!(messages[0].provider_id, "openai"); + } + + #[test] + fn test_parse_pi_falls_back_to_pi_when_provider_unrecoverable() { + let content = r#"{"type":"session","id":"pi_ses_007","timestamp":"2026-01-01T00:00:00.000Z","cwd":"/tmp"} +{"type":"message","id":"msg_001","parentId":null,"timestamp":"2026-01-01T00:00:01.000Z","message":{"role":"assistant","model":"totally-unrecognized-model-xyz","usage":{"input":100,"output":50,"cacheRead":0,"cacheWrite":0,"totalTokens":150}}}"#; + let file = create_test_file(content); + let messages = parse_pi_file(file.path()); + + assert_eq!(messages.len(), 1); + assert_eq!(messages[0].provider_id, "pi"); + } + #[test] fn test_parse_pi_subagent_session_name_as_agent() { let content = r#"{"type":"session","id":"pi_subagent_001","timestamp":"2026-07-10T00:00:00.000Z","cwd":"/tmp"} diff --git a/vendor/tokscale-core/src/sessions/utils.rs b/vendor/tokscale-core/src/sessions/utils.rs index 7aeae29..6a60d32 100644 --- a/vendor/tokscale-core/src/sessions/utils.rs +++ b/vendor/tokscale-core/src/sessions/utils.rs @@ -98,6 +98,15 @@ pub(crate) fn read_file_or_none(path: &Path) -> Option> { std::fs::read(path).ok() } +/// Back-calculate a start anchor from a recorded end timestamp and elapsed +/// duration: `end - duration`. Non-positive candidates fall back to `end` so +/// sessionization does not silently discard corrupt or clock-skewed rows. +pub(crate) fn back_anchor_timestamp(end: i64, duration: i64) -> i64 { + end.checked_sub(duration) + .filter(|candidate| *candidate > 0) + .unwrap_or(end) +} + #[cfg(test)] mod tests { use super::*; diff --git a/vendor/tokscale-core/tests/streaming_snapshot.rs b/vendor/tokscale-core/tests/streaming_snapshot.rs index 06f9acd..d7ae862 100644 --- a/vendor/tokscale-core/tests/streaming_snapshot.rs +++ b/vendor/tokscale-core/tests/streaming_snapshot.rs @@ -44,6 +44,7 @@ fn snapshot_msg( message_count: 1, agent: None, dedup_key: dedup_key.map(|s| s.to_string()), + dedup_aliases: Vec::new(), is_turn_start: false, } }