From fe4cbc7b59b2f8414ce2dbcd71fb39fe2acdecd1 Mon Sep 17 00:00:00 2001 From: RCD <90105158+Finesssee@users.noreply.github.com> Date: Wed, 30 Sep 2026 14:40:21 +0700 Subject: [PATCH] Port upstream 0.60.4: keep Chrome DeepSeek balance through same-session transport failures --- .../src/commands/last_good_owner_tests.rs | 160 ++++++++ .../src-tauri/src/commands/mod.rs | 2 + .../src-tauri/src/commands/providers.rs | 83 +++- .../src-tauri/src/commands/tests.rs | 12 + apps/desktop-tauri/src-tauri/src/state.rs | 4 + rust/src/cli/diagnose.rs | 1 + rust/src/cli/hooks.rs | 1 + rust/src/core/last_good_owner.rs | 112 ++++++ rust/src/core/last_good_owner_tests.rs | 86 ++++ rust/src/core/mod.rs | 2 + rust/src/core/provider.rs | 11 + rust/src/core/provider_state.rs | 1 + rust/src/core/usage_snapshot.rs | 13 + rust/src/providers/deepseek/chrome_session.rs | 160 ++++++++ .../deepseek/chrome_session_tests.rs | 78 ++++ rust/src/providers/deepseek/mod.rs | 93 ++++- .../providers/deepseek/platform_balance.rs | 160 ++++++++ .../deepseek/platform_balance_tests.rs | 93 +++++ .../providers/deepseek/session_resolver.rs | 311 +++++++++++++++ .../deepseek/session_resolver_tests.rs | 366 ++++++++++++++++++ rust/src/providers/deepseek/tests.rs | 31 ++ 21 files changed, 1762 insertions(+), 18 deletions(-) create mode 100644 apps/desktop-tauri/src-tauri/src/commands/last_good_owner_tests.rs create mode 100644 rust/src/core/last_good_owner.rs create mode 100644 rust/src/core/last_good_owner_tests.rs create mode 100644 rust/src/providers/deepseek/chrome_session.rs create mode 100644 rust/src/providers/deepseek/chrome_session_tests.rs create mode 100644 rust/src/providers/deepseek/platform_balance.rs create mode 100644 rust/src/providers/deepseek/platform_balance_tests.rs create mode 100644 rust/src/providers/deepseek/session_resolver.rs create mode 100644 rust/src/providers/deepseek/session_resolver_tests.rs diff --git a/apps/desktop-tauri/src-tauri/src/commands/last_good_owner_tests.rs b/apps/desktop-tauri/src-tauri/src/commands/last_good_owner_tests.rs new file mode 100644 index 0000000000..e9c7753d8c --- /dev/null +++ b/apps/desktop-tauri/src-tauri/src/commands/last_good_owner_tests.rs @@ -0,0 +1,160 @@ +//! Owner-checked last-good retention in the provider refresh shell. +//! +//! A failure tied to a live session may keep the cached snapshot only when the +//! same session produced that snapshot. Scenarios use DeepSeek's Chrome +//! session balance, the provider that declares owned transport failures. + +use super::ProviderUsageSnapshot; +use super::providers::{preserve_last_good_transient_failure, record_last_good_owner}; +use crate::state::AppState; +use codexbar::core::{ + LastGoodOwner, ProviderError, ProviderFetchResult, ProviderId, ProviderStateKind, RateWindow, + UsageSnapshot, instantiate_provider, +}; + +const MEASURED_AT: &str = "2026-09-01T00:00:00Z"; + +fn owner(profile: &str, token: &str) -> LastGoodOwner { + LastGoodOwner::derive("deepseek-platform-balance", profile, token).expect("owner") +} + +fn cached_balance() -> ProviderUsageSnapshot { + let metadata = instantiate_provider(ProviderId::DeepSeek) + .metadata() + .clone(); + let result = ProviderFetchResult::new(UsageSnapshot::new(RateWindow::new(0.0)), "web"); + let mut snapshot = + ProviderUsageSnapshot::from_fetch_result(ProviderId::DeepSeek, &metadata, &result, None); + snapshot.updated_at = MEASURED_AT.to_string(); + snapshot +} + +fn failed_refresh() -> ProviderUsageSnapshot { + let metadata = instantiate_provider(ProviderId::DeepSeek) + .metadata() + .clone(); + ProviderUsageSnapshot::from_error( + ProviderId::DeepSeek, + &metadata, + "Timeout".to_string(), + ProviderStateKind::Unknown, + ) +} + +fn state_with_cached_balance(cached_owner: Option) -> AppState { + let mut state = AppState::new(); + state.provider_cache.push(cached_balance()); + record_last_good_owner(&mut state, ProviderId::DeepSeek, cached_owner); + state +} + +fn refresh_failure(state: &mut AppState, error: &ProviderError) -> ProviderUsageSnapshot { + preserve_last_good_transient_failure(state, ProviderId::DeepSeek, failed_refresh(), error) +} + +#[test] +fn matching_session_keeps_the_balance_and_its_original_time() { + let session = owner("chrome:Default", "token-a"); + let mut state = state_with_cached_balance(Some(session.clone())); + + let kept = refresh_failure( + &mut state, + &ProviderError::Timeout.with_failure_owner(Some(session)), + ); + + assert_eq!(kept.error, None); + assert_eq!(kept.updated_at, MEASURED_AT); +} + +#[test] +fn different_profile_or_token_shows_the_error() { + for failed_session in [ + owner("chrome:Profile 1", "token-a"), + owner("chrome:Default", "token-b"), + ] { + let mut state = state_with_cached_balance(Some(owner("chrome:Default", "token-a"))); + let shown = refresh_failure( + &mut state, + &ProviderError::Timeout.with_failure_owner(Some(failed_session)), + ); + assert_eq!(shown.error.as_deref(), Some("Timeout")); + } +} + +#[test] +fn failure_without_a_session_owner_fails_closed() { + let mut state = state_with_cached_balance(Some(owner("chrome:Default", "token-a"))); + + let shown = refresh_failure(&mut state, &ProviderError::Timeout.with_failure_owner(None)); + + assert_eq!(shown.error.as_deref(), Some("Timeout")); +} + +#[test] +fn balance_without_an_owner_is_never_retained() { + // An API-key balance, a decoded cache, or a seeded snapshot has no owner. + let session = owner("chrome:Default", "token-a"); + let mut state = state_with_cached_balance(None); + + let shown = refresh_failure( + &mut state, + &ProviderError::Timeout.with_failure_owner(Some(session)), + ); + + assert_eq!(shown.error.as_deref(), Some("Timeout")); +} + +#[test] +fn unattributed_transport_failure_does_not_retain_deepseek_balances() { + let mut state = state_with_cached_balance(Some(owner("chrome:Default", "token-a"))); + + let shown = refresh_failure(&mut state, &ProviderError::Timeout); + + assert_eq!(shown.error.as_deref(), Some("Timeout")); +} + +#[test] +fn fresh_snapshot_replaces_or_clears_the_owner() { + let first = owner("chrome:Default", "token-a"); + let second = owner("chrome:Default", "token-b"); + let mut state = state_with_cached_balance(Some(first.clone())); + + record_last_good_owner(&mut state, ProviderId::DeepSeek, Some(second.clone())); + assert_eq!( + state.last_good_owners.get(&ProviderId::DeepSeek), + Some(&second) + ); + let shown = refresh_failure( + &mut state, + &ProviderError::Timeout.with_failure_owner(Some(first)), + ); + assert_eq!(shown.error.as_deref(), Some("Timeout")); + + record_last_good_owner(&mut state, ProviderId::DeepSeek, None); + assert!(state.last_good_owners.is_empty()); +} + +#[test] +fn other_providers_keep_their_existing_failure_policy() { + let metadata = instantiate_provider(ProviderId::Codex).metadata().clone(); + let result = ProviderFetchResult::new(UsageSnapshot::new(RateWindow::new(42.0)), "OAuth"); + let good = + ProviderUsageSnapshot::from_fetch_result(ProviderId::Codex, &metadata, &result, None); + let failed = ProviderUsageSnapshot::from_error( + ProviderId::Codex, + &metadata, + "Timeout".to_string(), + ProviderStateKind::Unknown, + ); + let mut state = AppState::new(); + state.provider_cache.push(good); + + let kept = preserve_last_good_transient_failure( + &mut state, + ProviderId::Codex, + failed, + &ProviderError::Timeout, + ); + + assert_eq!(kept.error, None); +} diff --git a/apps/desktop-tauri/src-tauri/src/commands/mod.rs b/apps/desktop-tauri/src-tauri/src/commands/mod.rs index fdfaeda967..e1798ee433 100644 --- a/apps/desktop-tauri/src-tauri/src/commands/mod.rs +++ b/apps/desktop-tauri/src-tauri/src/commands/mod.rs @@ -75,6 +75,8 @@ pub use surface::*; pub use system::*; pub(crate) use usage_items::*; +#[cfg(test)] +mod last_good_owner_tests; #[cfg(test)] mod tests; diff --git a/apps/desktop-tauri/src-tauri/src/commands/providers.rs b/apps/desktop-tauri/src-tauri/src/commands/providers.rs index 103d47b869..7c55e07d07 100644 --- a/apps/desktop-tauri/src-tauri/src/commands/providers.rs +++ b/apps/desktop-tauri/src-tauri/src/commands/providers.rs @@ -39,6 +39,7 @@ pub(crate) fn invalidate_account_usage( state.is_refreshing = false; state.provider_refresh_started_at = None; state.transient_provider_failure_counts.remove(&id); + state.last_good_owners.remove(&id); state.auto_resume.clear_provider(id); state .provider_cache @@ -334,6 +335,9 @@ pub(crate) fn invalidate_provider_refresh_and_prune_disabled( guard .transient_provider_failure_counts .retain(|id, _| enabled_ids.contains(id)); + guard + .last_good_owners + .retain(|id, _| enabled_ids.contains(id)); guard .provider_cache_updated_at_by_provider .retain(|id, _| enabled_ids.contains(id)); @@ -565,7 +569,7 @@ async fn refresh_provider( token_account_id: Option, hooks_enabled: bool, ) { - let (snapshot, account_identity, failure_policy) = + let (snapshot, account_identity, retention) = fetch_provider_snapshot(id, ctx, token_account_id).await; let fresh_snapshot = snapshot.error.is_none(); @@ -584,8 +588,12 @@ async fn refresh_provider( &mut guard, id, snapshot, - failure_policy, + retention.policy, + &retention.failure_ownership, ); + if fresh_snapshot { + record_last_good_owner(&mut guard, id, retention.fresh_owner); + } // F6 (upstream 0.48.0): backfill missing reset timestamps from the // cached snapshot before persisting and publishing. let cached = guard @@ -762,7 +770,30 @@ pub(super) fn preserve_last_good_transient_failure( error: &codexbar::core::ProviderError, ) -> ProviderUsageSnapshot { let policy = instantiate_provider(id).last_good_failure_policy_for_error(error); - preserve_last_good_transient_failure_with_policy(guard, id, snapshot, Some(policy)) + preserve_last_good_transient_failure_with_policy( + guard, + id, + snapshot, + Some(policy), + &error.failure_ownership(), + ) +} + +/// Remember which live session produced the fresh snapshot now cached for +/// `id`, or forget any earlier owner when the fresh snapshot has none. +pub(super) fn record_last_good_owner( + guard: &mut AppState, + id: ProviderId, + owner: Option, +) { + match owner { + Some(owner) => { + guard.last_good_owners.insert(id, owner); + } + None => { + guard.last_good_owners.remove(&id); + } + } } fn preserve_last_good_transient_failure_with_policy( @@ -770,6 +801,7 @@ fn preserve_last_good_transient_failure_with_policy( id: ProviderId, snapshot: ProviderUsageSnapshot, policy: Option, + ownership: &codexbar::core::FailureOwnership, ) -> ProviderUsageSnapshot { let Some(error) = snapshot.error.as_deref() else { guard.transient_provider_failure_counts.remove(&id); @@ -781,6 +813,12 @@ fn preserve_last_good_transient_failure_with_policy( guard.transient_provider_failure_counts.remove(&id); return snapshot; } + // A failure tied to a live session may keep only a snapshot that the same + // session produced. Anything else shows the error. + if !ownership.allows_retention(guard.last_good_owners.get(&id)) { + guard.transient_provider_failure_counts.remove(&id); + return snapshot; + } let Some(mut previous) = guard .provider_cache @@ -841,20 +879,28 @@ fn preserve_last_good_transient_failure_with_policy( } } +/// What a refresh tells the shell about keeping or replacing the last good +/// snapshot. +#[derive(Default)] +struct RefreshRetention { + /// Failure handling when a prior good snapshot exists. + policy: Option, + /// Session the failed request belonged to, when the provider can tell. + failure_ownership: codexbar::core::FailureOwnership, + /// Session that produced a fresh snapshot. + fresh_owner: Option, +} + async fn fetch_provider_snapshot( id: ProviderId, ctx: FetchContext, token_account_id: Option, -) -> ( - ProviderUsageSnapshot, - Option, - Option, -) { +) -> (ProviderUsageSnapshot, Option, RefreshRetention) { let provider = instantiate_provider(id); let metadata = provider.metadata().clone(); let started = std::time::Instant::now(); - let (mut snapshot, account_identity, failure_policy) = + let (mut snapshot, account_identity, retention) = match tokio::time::timeout(provider_fetch_timeout(id, &ctx), provider.fetch_usage(&ctx)) .await { @@ -868,7 +914,10 @@ async fn fetch_provider_snapshot( token_account_id, ), account_identity, - None, + RefreshRetention { + fresh_owner: result.last_good_owner.clone(), + ..RefreshRetention::default() + }, ) } Ok(Err(e)) => { @@ -881,7 +930,11 @@ async fn fetch_provider_snapshot( provider.error_state_kind(&e), ), None, - Some(policy), + RefreshRetention { + policy: Some(policy), + failure_ownership: e.failure_ownership(), + fresh_owner: None, + }, ) } Err(_) => { @@ -895,13 +948,17 @@ async fn fetch_provider_snapshot( provider.error_state_kind(&error), ), None, - Some(policy), + RefreshRetention { + policy: Some(policy), + failure_ownership: error.failure_ownership(), + fresh_owner: None, + }, ) } }; record_provider_fetch_duration(id, &mut snapshot, started); - (snapshot, account_identity, failure_policy) + (snapshot, account_identity, retention) } fn record_provider_fetch_duration( diff --git a/apps/desktop-tauri/src-tauri/src/commands/tests.rs b/apps/desktop-tauri/src-tauri/src/commands/tests.rs index d25a0a064a..433efc9474 100644 --- a/apps/desktop-tauri/src-tauri/src/commands/tests.rs +++ b/apps/desktop-tauri/src-tauri/src/commands/tests.rs @@ -948,6 +948,7 @@ fn usage_item_descriptors_keep_raw_ids_and_redact_titles() { has_successful_claude_cli_quota: false, pace_authoritative: true, account_identity: None, + last_good_owner: None, }; let mut snapshot = ProviderUsageSnapshot::from_fetch_result(ProviderId::Codex, &metadata, &result, None); @@ -1172,6 +1173,7 @@ fn provider_cache_upsert_replaces_existing_provider() { has_successful_claude_cli_quota: false, pace_authoritative: true, account_identity: None, + last_good_owner: None, }; let mut first = ProviderUsageSnapshot::from_fetch_result(ProviderId::Codex, &metadata, &result, None); @@ -1200,6 +1202,7 @@ fn provider_cache_prunes_disabled_providers() { has_successful_claude_cli_quota: false, pace_authoritative: true, account_identity: None, + last_good_owner: None, }; let codex = ProviderUsageSnapshot::from_fetch_result(ProviderId::Codex, &metadata, &result, None); @@ -1236,6 +1239,7 @@ fn claude_transient_auth_failure_preserves_first_last_good_snapshot() { has_successful_claude_cli_quota: false, pace_authoritative: true, account_identity: None, + last_good_owner: None, }; let good = ProviderUsageSnapshot::from_fetch_result(ProviderId::Claude, &metadata, &result, None); @@ -1273,6 +1277,7 @@ fn codex_transient_transport_failure_helper_uses_typed_policy() { has_successful_claude_cli_quota: false, pace_authoritative: true, account_identity: None, + last_good_owner: None, }; let good = ProviderUsageSnapshot::from_fetch_result(ProviderId::Codex, &metadata, &result, None); @@ -1309,6 +1314,7 @@ fn claude_repeated_auth_failure_surfaces_error() { has_successful_claude_cli_quota: false, pace_authoritative: true, account_identity: None, + last_good_owner: None, }; let good = ProviderUsageSnapshot::from_fetch_result(ProviderId::Claude, &metadata, &result, None); @@ -1352,6 +1358,7 @@ fn claude_cloudflare_challenge_retains_prior_usage_while_surfaceing_guidance() { has_successful_claude_cli_quota: false, pace_authoritative: true, account_identity: None, + last_good_owner: None, }; let good = ProviderUsageSnapshot::from_fetch_result(ProviderId::Claude, &metadata, &result, None); @@ -1406,6 +1413,7 @@ fn claude_cloudflare_challenge_keeps_prior_usage_when_guidance_surfaces() { has_successful_claude_cli_quota: false, pace_authoritative: true, account_identity: None, + last_good_owner: None, }; let mut good = ProviderUsageSnapshot::from_fetch_result(ProviderId::Claude, &metadata, &result, None); @@ -1457,6 +1465,7 @@ fn claude_cli_parse_failure_keeps_last_good_every_time() { has_successful_claude_cli_quota: true, pace_authoritative: true, account_identity: None, + last_good_owner: None, }; let good = ProviderUsageSnapshot::from_fetch_result(ProviderId::Claude, &metadata, &result, None); @@ -1504,6 +1513,7 @@ fn claude_hard_credentials_missing_does_not_preserve_stale() { has_successful_claude_cli_quota: false, pace_authoritative: true, account_identity: None, + last_good_owner: None, }; let good = ProviderUsageSnapshot::from_fetch_result(ProviderId::Claude, &metadata, &result, None); @@ -1708,6 +1718,7 @@ fn japanese_provider_snapshot_localizes_weekly_label() { has_successful_claude_cli_quota: false, pace_authoritative: true, account_identity: None, + last_good_owner: None, }; let snapshot = @@ -1742,6 +1753,7 @@ fn japanese_provider_snapshot_localizes_pace_reserve_description() { has_successful_claude_cli_quota: false, pace_authoritative: true, account_identity: None, + last_good_owner: None, }; let snapshot = diff --git a/apps/desktop-tauri/src-tauri/src/state.rs b/apps/desktop-tauri/src-tauri/src/state.rs index 10d626ec43..289ab89d55 100644 --- a/apps/desktop-tauri/src-tauri/src/state.rs +++ b/apps/desktop-tauri/src-tauri/src/state.rs @@ -121,6 +121,9 @@ pub struct AppState { pub tray_anchor: Option, pub provider_cache: Vec, pub transient_provider_failure_counts: HashMap, + /// Live session behind each provider's cached good snapshot, for + /// owner-checked last-good retention. In memory only. + pub last_good_owners: HashMap, pub provider_cache_updated_at: Option, /// Per-provider freshness for scoped background refreshes. The aggregate /// timestamp cannot tell a dedicated auto-resume watcher whether its own @@ -193,6 +196,7 @@ impl AppState { tray_anchor: None, provider_cache: Vec::new(), transient_provider_failure_counts: HashMap::new(), + last_good_owners: HashMap::new(), provider_cache_updated_at: None, provider_cache_updated_at_by_provider: HashMap::new(), provider_refresh_started_at: None, diff --git a/rust/src/cli/diagnose.rs b/rust/src/cli/diagnose.rs index 64d0b54d59..f26a261f1a 100644 --- a/rust/src/cli/diagnose.rs +++ b/rust/src/cli/diagnose.rs @@ -400,6 +400,7 @@ fn error_category(err: &ProviderError) -> &'static str { | ProviderError::NoCookies => "auth", ProviderError::OAuthTransient(_) => "api", ProviderError::Network(_) | ProviderError::Timeout => "network", + ProviderError::OwnedTransport { source, .. } => error_category(source), ProviderError::NotInstalled(_) | ProviderError::UnsupportedSource(_) => "config", ProviderError::Parse(_) => "parse", ProviderError::Other(message) => { diff --git a/rust/src/cli/hooks.rs b/rust/src/cli/hooks.rs index 2f26498cf1..6c0d1f616c 100644 --- a/rust/src/cli/hooks.rs +++ b/rust/src/cli/hooks.rs @@ -409,6 +409,7 @@ fn hook_refresh_failure_status(error: &ProviderError) -> String { "network_error".into() } } + ProviderError::OwnedTransport { source, .. } => hook_refresh_failure_status(source), ProviderError::NotInstalled(_) => "error".into(), ProviderError::Parse(_) | ProviderError::UnsupportedSource(_) | ProviderError::Other(_) => { "error".into() diff --git a/rust/src/core/last_good_owner.rs b/rust/src/core/last_good_owner.rs new file mode 100644 index 0000000000..dbc0433dcb --- /dev/null +++ b/rust/src/core/last_good_owner.rs @@ -0,0 +1,112 @@ +//! Proof of which live session supplied a cached provider snapshot. +//! +//! A provider that reads a balance through a browser session can attach a +//! [`LastGoodOwner`] to its fetch result. The shell keeps that owner in memory +//! next to the cached snapshot and retains the snapshot through a transport +//! failure only when the failed request came from the same owner. The owner is +//! a salted digest of the session scope and secret, is never serialized, and +//! never appears in `Debug` output. + +use sha2::{Digest, Sha256}; + +use super::ProviderError; + +const OWNER_NAMESPACE: &str = "com.codexbar.last-good-owner.v1"; + +/// Opaque, in-memory identity of the session that produced a snapshot. +#[derive(Clone, PartialEq, Eq)] +pub struct LastGoodOwner(String); + +impl LastGoodOwner { + /// Derive an owner from a provider namespace, a scope such as a browser + /// profile id, and the session secret. Returns `None` when the scope or the + /// secret is empty after trimming, so an unidentified session cannot own a + /// snapshot. + pub fn derive(namespace: &str, scope: &str, secret: &str) -> Option { + let scope = scope.trim(); + let secret = secret.trim(); + if scope.is_empty() || secret.is_empty() { + return None; + } + let mut hasher = Sha256::new(); + for part in [OWNER_NAMESPACE, namespace, scope, secret] { + hasher.update(part.as_bytes()); + hasher.update([0u8]); + } + let digest = hasher.finalize(); + Some(Self( + digest.iter().map(|byte| format!("{byte:02x}")).collect(), + )) + } +} + +impl std::fmt::Debug for LastGoodOwner { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.write_str("LastGoodOwner()") + } +} + +/// Which session a failed refresh came from, as far as the provider can prove. +#[derive(Debug, Clone, Default, PartialEq, Eq)] +pub enum FailureOwnership { + /// The error carries no session identity; the provider's policy decides. + #[default] + Unchecked, + /// The failed request used exactly this session. + Owned(LastGoodOwner), + /// The failure has no attributable session (for example a resolution + /// deadline, or a session already rejected). Retention must fail closed. + Unowned, +} + +impl FailureOwnership { + /// Whether a cached snapshot owned by `cached` may be shown in place of + /// this failure. + pub fn allows_retention(&self, cached: Option<&LastGoodOwner>) -> bool { + match self { + Self::Unchecked => true, + Self::Owned(owner) => cached == Some(owner), + Self::Unowned => false, + } + } +} + +impl ProviderError { + /// Wrap a transport failure with the session that produced it. + /// + /// A failure that is not a transport failure is returned unchanged, because + /// only transport failures are eligible for owner-checked retention. + pub fn with_failure_owner(self, owner: Option) -> Self { + if self.is_transport_failure() { + Self::OwnedTransport { + owner, + source: Box::new(self), + } + } else { + self + } + } + + /// Session ownership of this failure for last-good retention. + pub fn failure_ownership(&self) -> FailureOwnership { + match self { + Self::OwnedTransport { + owner: Some(owner), .. + } => FailureOwnership::Owned(owner.clone()), + Self::OwnedTransport { owner: None, .. } => FailureOwnership::Unowned, + _ => FailureOwnership::Unchecked, + } + } + + /// The underlying error with any session-ownership wrapper removed. + pub fn without_failure_owner(&self) -> &Self { + match self { + Self::OwnedTransport { source, .. } => source.without_failure_owner(), + other => other, + } + } +} + +#[cfg(test)] +#[path = "last_good_owner_tests.rs"] +mod tests; diff --git a/rust/src/core/last_good_owner_tests.rs b/rust/src/core/last_good_owner_tests.rs new file mode 100644 index 0000000000..525899f018 --- /dev/null +++ b/rust/src/core/last_good_owner_tests.rs @@ -0,0 +1,86 @@ +use super::*; + +fn owner(scope: &str, secret: &str) -> LastGoodOwner { + LastGoodOwner::derive("test", scope, secret).expect("owner") +} + +#[test] +fn owner_requires_scope_and_secret() { + assert!(LastGoodOwner::derive("test", " ", "secret").is_none()); + assert!(LastGoodOwner::derive("test", "chrome:Default", " \n").is_none()); +} + +#[test] +fn owner_trims_and_separates_scope_from_secret() { + assert_eq!( + owner(" chrome:Default ", " token "), + owner("chrome:Default", "token") + ); + assert_ne!( + owner("chrome:Default", "token"), + owner("chrome:Profile 1", "token") + ); + assert_ne!( + owner("chrome:Default", "token"), + owner("chrome:Default", "token2") + ); + assert_ne!(owner("ab", "c"), owner("a", "bc")); +} + +#[test] +fn owner_debug_output_never_contains_digest_or_secret() { + let rendered = format!("{:?}", owner("chrome:Default", "super-secret-token")); + assert_eq!(rendered, "LastGoodOwner()"); +} + +#[test] +fn only_transport_failures_are_wrapped() { + let wrapped = ProviderError::Timeout.with_failure_owner(Some(owner("p", "t"))); + assert!(matches!(wrapped, ProviderError::OwnedTransport { .. })); + assert!(wrapped.is_transport_failure()); + + let parse = ProviderError::Parse("bad".into()).with_failure_owner(Some(owner("p", "t"))); + assert!(matches!(parse, ProviderError::Parse(_))); +} + +#[test] +fn ownership_of_wrapped_and_plain_errors() { + let a = owner("p", "t"); + assert_eq!( + ProviderError::Timeout + .with_failure_owner(Some(a.clone())) + .failure_ownership(), + FailureOwnership::Owned(a) + ); + assert_eq!( + ProviderError::Timeout + .with_failure_owner(None) + .failure_ownership(), + FailureOwnership::Unowned + ); + assert_eq!( + ProviderError::Timeout.failure_ownership(), + FailureOwnership::Unchecked + ); +} + +#[test] +fn retention_requires_matching_owner_unless_unchecked() { + let a = owner("p", "t"); + let b = owner("p", "other"); + assert!(FailureOwnership::Unchecked.allows_retention(None)); + assert!(FailureOwnership::Owned(a.clone()).allows_retention(Some(&a))); + assert!(!FailureOwnership::Owned(a.clone()).allows_retention(Some(&b))); + assert!(!FailureOwnership::Owned(a).allows_retention(None)); + assert!(!FailureOwnership::Unowned.allows_retention(Some(&b))); +} + +#[test] +fn wrapper_displays_the_underlying_message() { + let wrapped = ProviderError::Timeout.with_failure_owner(None); + assert_eq!(wrapped.to_string(), "Timeout"); + assert!(matches!( + wrapped.without_failure_owner(), + ProviderError::Timeout + )); +} diff --git a/rust/src/core/mod.rs b/rust/src/core/mod.rs index 436c7d5e0a..510501ad4b 100755 --- a/rust/src/core/mod.rs +++ b/rust/src/core/mod.rs @@ -14,6 +14,7 @@ mod hooks; mod http; mod http_proxy; mod jsonl_scanner; +mod last_good_owner; mod models_dev_pricing; mod openai_dashboard; mod provider; @@ -42,6 +43,7 @@ pub use hooks::*; pub use http::*; pub use http_proxy::*; pub use jsonl_scanner::*; +pub use last_good_owner::{FailureOwnership, LastGoodOwner}; pub use models_dev_pricing::*; pub use openai_dashboard::*; pub use provider::*; diff --git a/rust/src/core/provider.rs b/rust/src/core/provider.rs index cc954839f7..f9b052f086 100755 --- a/rust/src/core/provider.rs +++ b/rust/src/core/provider.rs @@ -6,6 +6,7 @@ use std::collections::HashMap; use thiserror::Error; use super::ProviderFetchResult; +use super::last_good_owner::LastGoodOwner; use super::provider_state::ProviderStateKind; /// Unique identifier for a provider @@ -653,6 +654,15 @@ pub enum ProviderError { #[error("{0}")] Other(String), + + /// A transport failure tagged with the session that produced it, so the + /// shell can retain a cached snapshot only for the same session. Build it + /// with [`ProviderError::with_failure_owner`]. + #[error("{source}")] + OwnedTransport { + owner: Option, + source: Box, + }, } impl ProviderError { @@ -664,6 +674,7 @@ impl ProviderError { ReqwestFailureClass::Timeout | ReqwestFailureClass::Connect ), ProviderError::Timeout => true, + ProviderError::OwnedTransport { source, .. } => source.is_transport_failure(), _ => false, } } diff --git a/rust/src/core/provider_state.rs b/rust/src/core/provider_state.rs index a4087e1ca5..7c6cd97a15 100644 --- a/rust/src/core/provider_state.rs +++ b/rust/src/core/provider_state.rs @@ -72,6 +72,7 @@ impl ProviderError { | ProviderError::Parse(_) | ProviderError::UnsupportedSource(_) | ProviderError::Other(_) => ProviderStateKind::Unknown, + ProviderError::OwnedTransport { source, .. } => source.state_kind(), } } } diff --git a/rust/src/core/usage_snapshot.rs b/rust/src/core/usage_snapshot.rs index ebbdcb8c20..9612832ea6 100755 --- a/rust/src/core/usage_snapshot.rs +++ b/rust/src/core/usage_snapshot.rs @@ -3,6 +3,7 @@ use chrono::{DateTime, Utc}; use serde::{Deserialize, Serialize}; +use super::LastGoodOwner; use super::ProviderDisplayDetail; use super::RateWindow; /// Subscription dates explicitly reported by an authenticated provider @@ -643,6 +644,11 @@ pub struct ProviderFetchResult { /// actions. It never crosses the frontend bridge. #[serde(skip)] pub account_identity: Option, + + /// Live session that supplied this result, used only for owner-checked + /// last-good retention. It is in memory only and never crosses the bridge. + #[serde(skip)] + pub last_good_owner: Option, } fn default_pace_authoritative() -> bool { @@ -662,9 +668,16 @@ impl ProviderFetchResult { has_successful_claude_cli_quota: false, pace_authoritative: true, account_identity: None, + last_good_owner: None, } } + /// Record which live session supplied this result. + pub fn with_last_good_owner(mut self, owner: Option) -> Self { + self.last_good_owner = owner; + self + } + /// Attach the provider's stable account identity without exposing it to /// serialized UI payloads. pub fn with_account_identity(mut self, account_identity: impl Into) -> Self { diff --git a/rust/src/providers/deepseek/chrome_session.rs b/rust/src/providers/deepseek/chrome_session.rs new file mode 100644 index 0000000000..098ee91be7 --- /dev/null +++ b/rust/src/providers/deepseek/chrome_session.rs @@ -0,0 +1,160 @@ +//! Import DeepSeek Platform sessions from Chrome's local storage. +//! +//! The platform dashboard keeps the signed-in session as a `userToken` item in +//! `localStorage` for `https://platform.deepseek.com`. Each Chrome profile holds +//! its own copy, so a candidate is identified by `chrome:`. + +use std::fmt; + +use serde_json::Value; + +use crate::browser::detection::{BrowserDetector, BrowserType}; +use crate::browser::leveldb::local_storage::{local_storage_dir, read_local_storage_entries}; + +const PLATFORM_ORIGIN: &str = "https://platform.deepseek.com"; +const TOKEN_KEY: &str = "userToken"; +const TOKEN_FIELDS: [&str; 5] = ["value", "token", "access_token", "accessToken", "userToken"]; +const MIN_TOKEN_CHARS: usize = 20; + +/// Environment variable that pins the Chrome profile to read when several +/// profiles hold a DeepSeek session. It accepts `chrome:` or a +/// profile directory path. +pub(super) const PROFILE_ID_ENV: &str = "CODEXBAR_DEEPSEEK_PROFILE_ID"; + +/// A session token found in one Chrome profile. +#[derive(Clone, PartialEq, Eq)] +pub(super) struct TokenInfo { + pub(super) id: String, + pub(super) token: String, + pub(super) label: String, +} + +impl fmt::Debug for TokenInfo { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("TokenInfo") + .field("id", &self.id) + .field("token", &"") + .field("label", &self.label) + .finish() + } +} + +/// Which profile the user asked for, if any. +#[derive(Debug, Clone, Default, PartialEq, Eq)] +pub(super) struct ProfileSelection { + pub(super) profile_id: Option, +} + +impl ProfileSelection { + pub(super) fn from_env() -> Self { + Self::from_value(std::env::var(PROFILE_ID_ENV).ok().as_deref()) + } + + pub(super) fn from_value(value: Option<&str>) -> Self { + Self { + profile_id: value + .map(canonical_profile_id) + .filter(|profile| !profile.is_empty()), + } + } +} + +/// Normalize a profile setting: a directory path becomes `chrome:`. +pub(super) fn canonical_profile_id(raw: &str) -> String { + let value = raw.trim(); + let is_path = value.starts_with('/') || value.contains('\\') || value.contains(":/"); + if !is_path { + return value.to_string(); + } + let leaf = value + .trim_end_matches(['/', '\\']) + .rsplit(['/', '\\']) + .next() + .unwrap_or_default(); + if leaf.is_empty() { + value.to_string() + } else { + format!("chrome:{leaf}") + } +} + +/// Read the DeepSeek session token of every Chrome profile that has one. +/// +/// Blocking: reads LevelDB files from disk. Never logs token values. +pub(super) fn import_tokens() -> Vec { + let Some(chrome) = BrowserDetector::detect(BrowserType::Chrome) else { + tracing::debug!("deepseek chrome session: Chrome profiles not found"); + return Vec::new(); + }; + + let mut tokens = Vec::new(); + for profile in &chrome.profiles { + let id = format!("chrome:{}", profile.name); + let dir = local_storage_dir(&profile.path); + let entries = match read_local_storage_entries(&dir, PLATFORM_ORIGIN) { + Ok(entries) => entries, + Err(error) => { + tracing::debug!(profile = %id, %error, "deepseek chrome session: local storage unreadable"); + continue; + } + }; + let Some(token) = entries + .iter() + .find(|entry| entry.key == TOKEN_KEY) + .and_then(|entry| extract_user_token(&entry.value)) + else { + tracing::debug!(profile = %id, "deepseek chrome session: no userToken"); + continue; + }; + tracing::debug!(profile = %id, "deepseek chrome session: found userToken"); + tokens.push(TokenInfo { + label: format!("Google Chrome {}", profile.name), + id, + token, + }); + } + tokens.sort_by(|a, b| a.id.cmp(&b.id)); + tokens +} + +/// Extract the token from a `userToken` value: a JSON object carrying the +/// token in one of the known fields, or a bare (optionally quoted) token. +pub(super) fn extract_user_token(raw: &str) -> Option { + let trimmed = raw.trim(); + if trimmed.is_empty() { + return None; + } + + if let Ok(value @ (Value::Object(_) | Value::Array(_))) = serde_json::from_str::(trimmed) + { + return token_from_json(&value); + } + + let unquoted = if (trimmed.starts_with('"') && trimmed.ends_with('"')) + || (trimmed.starts_with('\'') && trimmed.ends_with('\'')) + { + trimmed + .get(1..trimmed.len().saturating_sub(1)) + .unwrap_or_default() + } else { + trimmed + }; + is_plausible_token(unquoted).then(|| unquoted.to_string()) +} + +fn token_from_json(value: &Value) -> Option { + let object = value.as_object()?; + TOKEN_FIELDS.iter().find_map(|field| { + let token = object.get(*field)?.as_str()?; + is_plausible_token(token).then(|| token.to_string()) + }) +} + +fn is_plausible_token(value: &str) -> bool { + let trimmed = value.trim(); + trimmed.chars().count() >= MIN_TOKEN_CHARS && !trimmed.chars().any(char::is_whitespace) +} + +#[cfg(test)] +#[path = "chrome_session_tests.rs"] +mod tests; diff --git a/rust/src/providers/deepseek/chrome_session_tests.rs b/rust/src/providers/deepseek/chrome_session_tests.rs new file mode 100644 index 0000000000..98f4b4ba9b --- /dev/null +++ b/rust/src/providers/deepseek/chrome_session_tests.rs @@ -0,0 +1,78 @@ +use super::*; + +const TOKEN: &str = "abcdefghijklmnopqrstuvwxyz0123456789"; + +#[test] +fn extracts_token_from_json_object_fields() { + for field in ["value", "token", "access_token", "accessToken", "userToken"] { + let raw = format!(r#"{{"{field}":"{TOKEN}","__version":"0"}}"#); + assert_eq!(extract_user_token(&raw).as_deref(), Some(TOKEN), "{field}"); + } +} + +#[test] +fn json_object_skips_implausible_fields_and_prefers_field_order() { + let raw = format!(r#"{{"value":"short","token":"{TOKEN}"}}"#); + assert_eq!(extract_user_token(&raw).as_deref(), Some(TOKEN)); + assert_eq!(extract_user_token(r#"{"value":"short"}"#), None); + assert_eq!(extract_user_token(r#"{"value":42}"#), None); +} + +#[test] +fn extracts_bare_and_quoted_tokens() { + assert_eq!(extract_user_token(TOKEN).as_deref(), Some(TOKEN)); + assert_eq!( + extract_user_token(&format!(" \"{TOKEN}\"\n")).as_deref(), + Some(TOKEN) + ); + assert_eq!( + extract_user_token(&format!("'{TOKEN}'")).as_deref(), + Some(TOKEN) + ); +} + +#[test] +fn rejects_empty_short_and_spaced_values() { + assert_eq!(extract_user_token(""), None); + assert_eq!(extract_user_token(" "), None); + assert_eq!(extract_user_token("short-token"), None); + assert_eq!(extract_user_token("abcdefghij klmnopqrstuvwxyz"), None); + assert_eq!(extract_user_token("[\"abcdefghijklmnopqrstuvwxyz\"]"), None); +} + +#[test] +fn canonical_profile_id_normalizes_paths_only() { + assert_eq!(canonical_profile_id(" chrome:Default "), "chrome:Default"); + assert_eq!( + canonical_profile_id("C:\\Users\\me\\AppData\\Local\\Google\\Chrome\\User Data\\Profile 1"), + "chrome:Profile 1" + ); + assert_eq!( + canonical_profile_id("/home/me/.config/google-chrome/Default/"), + "chrome:Default" + ); +} + +#[test] +fn profile_selection_ignores_blank_values() { + assert_eq!(ProfileSelection::from_value(None).profile_id, None); + assert_eq!(ProfileSelection::from_value(Some(" ")).profile_id, None); + assert_eq!( + ProfileSelection::from_value(Some("chrome:Profile 2")) + .profile_id + .as_deref(), + Some("chrome:Profile 2") + ); +} + +#[test] +fn token_info_debug_redacts_the_token() { + let info = TokenInfo { + id: "chrome:Default".into(), + token: TOKEN.into(), + label: "Google Chrome Default".into(), + }; + let rendered = format!("{info:?}"); + assert!(!rendered.contains(TOKEN)); + assert!(rendered.contains("chrome:Default")); +} diff --git a/rust/src/providers/deepseek/mod.rs b/rust/src/providers/deepseek/mod.rs index 4461944d77..1a684a2071 100644 --- a/rust/src/providers/deepseek/mod.rs +++ b/rust/src/providers/deepseek/mod.rs @@ -1,6 +1,9 @@ //! DeepSeek provider implementation. //! -//! Fetches API account balance from DeepSeek's `/user/balance` endpoint. +//! Fetches API account balance from DeepSeek's `/user/balance` endpoint. When +//! no API key is configured, Auto mode reads the Platform balance through a +//! signed-in Chrome session and keeps that balance visible through temporary +//! connection failures of the same session. use async_trait::async_trait; use reqwest::Client; @@ -8,16 +11,31 @@ use std::collections::{HashMap, HashSet}; use serde::Deserialize; +mod chrome_session; +mod platform_balance; pub mod pricing; +mod session_resolver; use crate::core::{ - CostSnapshot, FetchContext, Provider, ProviderError, ProviderFetchResult, ProviderId, - ProviderMetadata, RateWindow, SourceMode, UsageSnapshot, + CostSnapshot, FetchContext, LastGoodFailurePolicy, Provider, ProviderError, + ProviderFetchResult, ProviderId, ProviderMetadata, RateWindow, SourceMode, UsageSnapshot, }; const DEEPSEEK_API_BASE: &str = "https://api.deepseek.com"; const DEEPSEEK_CREDENTIAL_TARGET: &str = "codexbar-deepseek"; +/// Upper bound for importing Chrome sessions and validating them. Staying +/// under the shell's fetch timeout keeps a slow resolution attributable: it +/// fails with no session owner instead of a shell-level timeout. +const CHROME_SESSION_DEADLINE: std::time::Duration = std::time::Duration::from_secs(28); + +/// Validation results outlive a single provider instance, which the shell +/// recreates for every refresh. +static CHROME_SESSION_CACHE: std::sync::LazyLock = + std::sync::LazyLock::new(|| { + session_resolver::ValidationCache::new(session_resolver::VALIDITY_TTL) + }); + #[derive(Debug, Deserialize)] struct BalanceResponse { #[serde(default)] @@ -203,7 +221,57 @@ impl DeepSeekProvider { ctx: &FetchContext, ) -> Result { let api_key = Self::get_api_key(ctx.api_key.as_deref())?; + self.fetch_usage_with_api_key(&api_key).await + } + /// Balance from a signed-in Chrome session on platform.deepseek.com, used + /// when no API key is configured. Reads the balance only. + async fn fetch_usage_chrome_session( + &self, + missing_key: ProviderError, + ) -> Result { + let selection = chrome_session::ProfileSelection::from_env(); + let client = self.client.clone(); + let resolution = tokio::time::timeout(CHROME_SESSION_DEADLINE, async { + let candidates = tokio::task::spawn_blocking(chrome_session::import_tokens) + .await + .unwrap_or_default(); + session_resolver::resolve(&candidates, &selection, &CHROME_SESSION_CACHE, |token| { + let client = client.clone(); + async move { platform_balance::fetch_platform_balance(&client, &token).await } + }) + .await + }) + .await; + + match resolution { + // A deadline has no attributable session, so it never retains a balance. + Err(_) => Err(ProviderError::Timeout.with_failure_owner(None)), + Ok(session_resolver::Resolution::Balance { balance, owner }) => Ok( + ProviderFetchResult::new(Self::snapshot_from_balance(*balance), "web") + .with_last_good_owner(owner), + ), + Ok(session_resolver::Resolution::SessionRequired) => Err(match missing_key { + ProviderError::NotInstalled(message) => ProviderError::NotInstalled(format!( + "{message} Or sign in to platform.deepseek.com in Chrome." + )), + other => other, + }), + Ok(session_resolver::Resolution::SelectionRequired(profiles)) => { + Err(ProviderError::Other(format!( + "DeepSeek is signed in on several Chrome profiles ({}). Set {} to choose one.", + profiles.join(", "), + chrome_session::PROFILE_ID_ENV, + ))) + } + Ok(session_resolver::Resolution::Failed(error)) => Err(error), + } + } + + async fn fetch_usage_with_api_key( + &self, + api_key: &str, + ) -> Result { let resp = self .client .get(format!("{DEEPSEEK_API_BASE}/user/balance")) @@ -228,7 +296,7 @@ impl DeepSeekProvider { let mut usage = Self::snapshot_from_balance(balance); let mut result = ProviderFetchResult::new(usage.clone(), "api"); - if let Ok(summary) = self.fetch_usage_summary(&api_key).await { + if let Ok(summary) = self.fetch_usage_summary(api_key).await { usage = Self::apply_usage_summary(usage, &summary); result = ProviderFetchResult::new(usage, "api"); if summary.month_cost > 0.0 || !summary.model_costs.is_empty() { @@ -501,7 +569,11 @@ impl Provider for DeepSeekProvider { async fn fetch_usage(&self, ctx: &FetchContext) -> Result { match ctx.source_mode { - SourceMode::Auto | SourceMode::OAuth => self.fetch_usage_api(ctx).await, + SourceMode::Auto => match Self::get_api_key(ctx.api_key.as_deref()) { + Ok(api_key) => self.fetch_usage_with_api_key(&api_key).await, + Err(missing_key) => self.fetch_usage_chrome_session(missing_key).await, + }, + SourceMode::OAuth => self.fetch_usage_api(ctx).await, SourceMode::Web | SourceMode::Cli => { Err(ProviderError::UnsupportedSource(ctx.source_mode)) } @@ -511,6 +583,17 @@ impl Provider for DeepSeekProvider { fn available_sources(&self) -> Vec { vec![SourceMode::Auto, SourceMode::OAuth] } + + /// Only a transport failure attributed to a Chrome session keeps the last + /// balance; the shell then checks that the session matches the cached one. + /// API-key balances and every other failure surface normally. + fn last_good_failure_policy_for_error(&self, error: &ProviderError) -> LastGoodFailurePolicy { + if matches!(error, ProviderError::OwnedTransport { .. }) { + LastGoodFailurePolicy::Preserve + } else { + LastGoodFailurePolicy::Replace + } + } } fn parse_money(value: &str) -> f64 { diff --git a/rust/src/providers/deepseek/platform_balance.rs b/rust/src/providers/deepseek/platform_balance.rs new file mode 100644 index 0000000000..27e0ca0ec7 --- /dev/null +++ b/rust/src/providers/deepseek/platform_balance.rs @@ -0,0 +1,160 @@ +//! DeepSeek Platform balance read through a signed-in web session. +//! +//! The platform dashboard exposes the account wallets at +//! `GET /api/v0/users/get_user_summary`, authorized with the session's +//! `userToken`. This module only reads the balance; it never reads token-level +//! usage or cost. + +use std::collections::BTreeMap; + +use reqwest::{Client, StatusCode}; +use serde::Deserialize; + +use super::{BalanceInfo, BalanceResponse, FlexibleF64, FlexibleI64, select_balance_info}; +use crate::core::ProviderError; + +const PLATFORM_USER_SUMMARY_URL: &str = + "https://platform.deepseek.com/api/v0/users/get_user_summary"; + +/// Envelope error codes the platform uses for a missing or expired session. +const AUTH_ERROR_CODES: [i64; 2] = [40002, 40003]; + +/// Top-level envelope. Error envelopes are not schema-stable, so `data` stays +/// raw until the envelope code has been checked. +#[derive(Debug, Deserialize)] +struct UserSummaryResponse { + code: Option, + data: Option, +} + +#[derive(Debug, Deserialize)] +struct UserSummaryData { + biz_code: Option, + biz_data: Option, +} + +#[derive(Debug, Deserialize)] +struct UserSummary { + #[serde(default)] + normal_wallets: Vec, + #[serde(default)] + bonus_wallets: Vec, +} + +#[derive(Debug, Deserialize)] +struct Wallet { + currency: String, + balance: FlexibleF64, +} + +/// Fetch the wallet balance for one platform session token. +/// +/// `ProviderError::AuthRequired` means the platform rejected the session +/// (HTTP 401/403 or an auth envelope code). Transport failures keep their +/// reqwest classification so callers can tell an outage from a rejection. +pub(super) async fn fetch_platform_balance( + client: &Client, + token: &str, +) -> Result { + let token = token.trim(); + if token.is_empty() { + return Err(ProviderError::AuthRequired); + } + let response = client + .get(PLATFORM_USER_SUMMARY_URL) + .header("Authorization", format!("Bearer {token}")) + .header("Accept", "application/json") + .header("x-client-platform", "web") + .send() + .await?; + let status = response.status(); + if matches!(status, StatusCode::UNAUTHORIZED | StatusCode::FORBIDDEN) { + return Err(ProviderError::AuthRequired); + } + if status != StatusCode::OK { + return Err(ProviderError::Other(format!( + "DeepSeek Platform returned status {status}" + ))); + } + let body = response.bytes().await?; + parse_platform_balance(&body) +} + +pub(super) fn parse_platform_balance(body: &[u8]) -> Result { + let parse_error = |error: serde_json::Error| { + ProviderError::Parse(format!( + "Failed to parse DeepSeek Platform balance: {error}" + )) + }; + let envelope: UserSummaryResponse = serde_json::from_slice(body).map_err(parse_error)?; + check_code(envelope.code, "user summary code")?; + + let data: UserSummaryData = envelope + .data + .map(serde_json::from_value) + .transpose() + .map_err(parse_error)? + .ok_or_else(|| ProviderError::Parse("Missing DeepSeek Platform user summary".into()))?; + check_code(data.biz_code, "user summary biz_code")?; + + let summary: UserSummary = data + .biz_data + .map(serde_json::from_value) + .transpose() + .map_err(parse_error)? + .ok_or_else(|| ProviderError::Parse("Missing DeepSeek Platform biz_data".into()))?; + Ok(balance_response(&summary)) +} + +fn check_code(code: Option, label: &str) -> Result<(), ProviderError> { + match code.map(|code| code.0) { + None | Some(0) => Ok(()), + Some(code) if AUTH_ERROR_CODES.contains(&code) => Err(ProviderError::AuthRequired), + Some(code) => Err(ProviderError::Other(format!( + "DeepSeek Platform {label} {code}" + ))), + } +} + +/// Sum wallets per currency and express them as the API-key balance shape, so +/// both lanes render through the same snapshot code. +fn balance_response(summary: &UserSummary) -> BalanceResponse { + let mut totals: BTreeMap<&str, (f64, f64)> = BTreeMap::new(); + for wallet in &summary.normal_wallets { + totals.entry(wallet.currency.as_str()).or_default().0 += wallet.balance.0; + } + for wallet in &summary.bonus_wallets { + totals.entry(wallet.currency.as_str()).or_default().1 += wallet.balance.0; + } + + let balance_infos: Vec = totals + .into_iter() + .map(|(currency, (topped_up, granted))| BalanceInfo { + currency: currency.to_string(), + total_balance: (topped_up + granted).to_string(), + granted_balance: granted.to_string(), + topped_up_balance: topped_up.to_string(), + }) + .collect(); + + let Some(selected) = select_balance_info(&balance_infos) else { + return BalanceResponse { + is_available: false, + balance_infos: vec![BalanceInfo { + currency: "USD".into(), + total_balance: "0".into(), + granted_balance: "0".into(), + topped_up_balance: "0".into(), + }], + }; + }; + let is_available = selected.total_balance.parse::().unwrap_or(0.0) > 0.0; + BalanceResponse { + is_available, + balance_infos, + } +} + +#[cfg(test)] +#[path = "platform_balance_tests.rs"] +mod tests; diff --git a/rust/src/providers/deepseek/platform_balance_tests.rs b/rust/src/providers/deepseek/platform_balance_tests.rs new file mode 100644 index 0000000000..eee4c68c4f --- /dev/null +++ b/rust/src/providers/deepseek/platform_balance_tests.rs @@ -0,0 +1,93 @@ +use super::*; + +fn parse(body: &str) -> Result { + parse_platform_balance(body.as_bytes()) +} + +fn summary_body(normal: &str, bonus: &str) -> String { + format!( + r#"{{"code":0,"msg":"","data":{{"biz_code":0,"biz_msg":"","biz_data":{{"normal_wallets":{normal},"bonus_wallets":{bonus}}}}}}}"# + ) +} + +#[test] +fn sums_wallets_per_currency_and_prefers_funded_usd() { + let body = summary_body( + r#"[{"currency":"USD","balance":"2.50"},{"currency":"CNY","balance":"9"}]"#, + r#"[{"currency":"USD","balance":1.25}]"#, + ); + let response = parse(&body).unwrap(); + assert!(response.is_available); + + let snapshot = crate::providers::deepseek::DeepSeekProvider::snapshot_from_balance(response); + assert_eq!( + snapshot.primary.reset_description.as_deref(), + Some("$3.75 (Paid: $2.50 / Granted: $1.25)") + ); +} + +#[test] +fn falls_back_to_funded_cny_when_usd_is_empty() { + let body = summary_body( + r#"[{"currency":"USD","balance":"0"},{"currency":"CNY","balance":"40"}]"#, + r#"[{"currency":"CNY","balance":"2.25"}]"#, + ); + let snapshot = + crate::providers::deepseek::DeepSeekProvider::snapshot_from_balance(parse(&body).unwrap()); + assert_eq!( + snapshot.login_method.as_deref(), + Some("CNY balance: ¥42.25") + ); +} + +#[test] +fn empty_wallets_report_an_unavailable_usd_balance() { + let response = parse(&summary_body("[]", "[]")).unwrap(); + assert!(!response.is_available); + assert_eq!(response.balance_infos.len(), 1); + assert_eq!(response.balance_infos[0].currency, "USD"); +} + +#[test] +fn zero_balance_is_not_available() { + let response = parse(&summary_body(r#"[{"currency":"USD","balance":"0"}]"#, "[]")).unwrap(); + assert!(!response.is_available); +} + +#[test] +fn auth_envelope_codes_are_session_rejections() { + for code in [40002, 40003] { + let body = format!(r#"{{"code":{code},"msg":"auth","data":{{"unexpected":true}}}}"#); + assert!(matches!(parse(&body), Err(ProviderError::AuthRequired))); + + let body = format!( + r#"{{"code":0,"data":{{"biz_code":{code},"biz_msg":"auth","biz_data":"not a summary"}}}}"# + ); + assert!(matches!(parse(&body), Err(ProviderError::AuthRequired))); + } +} + +#[test] +fn other_envelope_codes_are_not_session_rejections() { + let error = parse(r#"{"code":50000,"msg":"busy","data":null}"#).unwrap_err(); + assert!(matches!(error, ProviderError::Other(_)), "{error:?}"); + assert!(!error.is_transport_failure()); + + let error = parse(r#"{"code":0,"data":{"biz_code":1,"biz_data":null}}"#).unwrap_err(); + assert!(matches!(error, ProviderError::Other(_)), "{error:?}"); +} + +#[test] +fn malformed_or_missing_payloads_are_parse_errors() { + for body in [ + "not json", + r#"{"code":0}"#, + r#"{"code":0,"data":{"biz_code":0}}"#, + r#"{"code":0,"data":{"biz_code":0,"biz_data":{"normal_wallets":[{"currency":"USD","balance":"abc"}]}}}"#, + ] { + assert!( + matches!(parse(body), Err(ProviderError::Parse(_))), + "{body}" + ); + } +} diff --git a/rust/src/providers/deepseek/session_resolver.rs b/rust/src/providers/deepseek/session_resolver.rs new file mode 100644 index 0000000000..f38ffc2d34 --- /dev/null +++ b/rust/src/providers/deepseek/session_resolver.rs @@ -0,0 +1,311 @@ +//! Choose the Chrome session that supplies the DeepSeek Platform balance, and +//! decide which session a failed request belongs to. +//! +//! Every candidate is validated against the platform (a rejected session is +//! remembered for [`VALIDITY_TTL`]). The selected session's balance is always +//! fetched live. A transport failure keeps the identity of the session it came +//! from, so the shell can keep showing the last balance only when that same +//! session (same profile, same token) produced it. A session the platform has +//! rejected never keeps that authority. + +use std::collections::HashMap; +use std::future::Future; +use std::sync::Mutex; +use std::time::{Duration, Instant}; + +use futures::future::join_all; + +use super::BalanceResponse; +use super::chrome_session::{ProfileSelection, TokenInfo}; +use crate::core::{LastGoodOwner, ProviderError}; + +/// How long a validation result is trusted before the platform is asked again. +pub(super) const VALIDITY_TTL: Duration = Duration::from_secs(30 * 60); + +const OWNER_NAMESPACE: &str = "deepseek-platform-balance"; + +/// Owner of a balance fetched with `token` from Chrome profile `profile_id`. +pub(super) fn balance_owner(profile_id: &str, token: &str) -> Option { + LastGoodOwner::derive(OWNER_NAMESPACE, profile_id, token) +} + +/// Outcome of asking the platform about one candidate. +enum Validation { + Valid(Box), + Rejected, + Unavailable(Option), +} + +struct Outcome { + index: usize, + validation: Validation, +} + +/// Result of resolving the selected Chrome session. +#[derive(Debug)] +pub(super) enum Resolution { + Balance { + balance: Box, + owner: Option, + }, + /// No profile holds a usable DeepSeek session. + SessionRequired, + /// Several profiles hold a session and none was selected. + SelectionRequired(Vec), + Failed(ProviderError), +} + +#[derive(Clone, Copy)] +struct Lookup { + fresh: Option, + last_known: Option, +} + +struct CacheEntry { + proof: Option, + status: bool, + checked_at: Instant, +} + +/// Remembers which sessions the platform accepted or rejected. Holds a digest +/// of each token, never the token. +pub(super) struct ValidationCache { + ttl: Duration, + entries: Mutex>, +} + +impl ValidationCache { + pub(super) fn new(ttl: Duration) -> Self { + Self { + ttl, + entries: Mutex::new(HashMap::new()), + } + } + + fn lookup(&self, candidate: &TokenInfo, now: Instant) -> Lookup { + let proof = balance_owner(&candidate.id, &candidate.token); + let entries = self.entries.lock().unwrap_or_else(|e| e.into_inner()); + match entries.get(&candidate.id) { + Some(entry) if entry.proof == proof => Lookup { + fresh: (now.saturating_duration_since(entry.checked_at) < self.ttl) + .then_some(entry.status), + last_known: Some(entry.status), + }, + _ => Lookup { + fresh: None, + last_known: None, + }, + } + } + + fn record(&self, candidate: &TokenInfo, status: bool, now: Instant) { + let mut entries = self.entries.lock().unwrap_or_else(|e| e.into_inner()); + entries.insert( + candidate.id.clone(), + CacheEntry { + proof: balance_owner(&candidate.id, &candidate.token), + status, + checked_at: now, + }, + ); + } +} + +/// Resolve the session to use and fetch its balance through `validate`. +pub(super) async fn resolve( + candidates: &[TokenInfo], + selection: &ProfileSelection, + cache: &ValidationCache, + validate: F, +) -> Resolution +where + F: Fn(String) -> Fut, + Fut: Future>, +{ + if candidates.is_empty() { + return Resolution::SessionRequired; + } + + let now = Instant::now(); + let lookups: Vec = candidates.iter().map(|c| cache.lookup(c, now)).collect(); + let selected_id = selection.profile_id.as_deref(); + let to_validate: Vec = (0..candidates.len()) + .filter(|&i| lookups[i].fresh.is_none() || Some(candidates[i].id.as_str()) == selected_id) + .collect(); + + let mut outcomes = run_validation(candidates, &to_validate, &validate).await; + record(cache, candidates, &outcomes, now); + let mut statuses = resolved_statuses(&lookups, &outcomes); + let mut valid = valid_indices(&statuses); + + // A status that came only from the cache carries no balance. + if let Some(index) = selected_candidate(candidates, &valid, selection) + && !has_balance(&outcomes, index) + { + outcomes.extend(run_validation(candidates, &[index], &validate).await); + record(cache, candidates, &outcomes, now); + statuses = resolved_statuses(&lookups, &outcomes); + valid = valid_indices(&statuses); + } + + let all: Vec = (0..candidates.len()).collect(); + let unresolved = selected_candidate(candidates, &all, selection); + + if valid.is_empty() { + let failure = + unresolved.and_then(|index| take_failure(candidates, &mut outcomes, &statuses, index)); + if let Some(error) = failure { + return Resolution::Failed(error); + } + if outcomes + .iter() + .any(|o| matches!(o.validation, Validation::Unavailable(_))) + { + return Resolution::Failed(validation_unavailable()); + } + return Resolution::SessionRequired; + } + + let Some(selected) = selected_candidate(candidates, &valid, selection) else { + let rejected = unresolved.is_some_and(|index| statuses[index] == Some(false)); + let failure = if rejected { + None + } else { + unresolved.and_then(|index| take_failure(candidates, &mut outcomes, &statuses, index)) + }; + return match failure { + Some(error) => Resolution::Failed(error), + None => Resolution::SelectionRequired( + valid.iter().map(|&i| candidates[i].id.clone()).collect(), + ), + }; + }; + + if let Some(balance) = take_balance(&mut outcomes, selected) { + let candidate = &candidates[selected]; + return Resolution::Balance { + balance, + owner: balance_owner(&candidate.id, &candidate.token), + }; + } + Resolution::Failed( + take_failure(candidates, &mut outcomes, &statuses, selected) + .unwrap_or_else(validation_unavailable), + ) +} + +fn validation_unavailable() -> ProviderError { + ProviderError::Other("Chrome DeepSeek session could not be validated right now".to_string()) +} + +async fn run_validation( + candidates: &[TokenInfo], + indices: &[usize], + validate: &F, +) -> Vec +where + F: Fn(String) -> Fut, + Fut: Future>, +{ + join_all(indices.iter().map(|&index| async move { + let validation = match validate(candidates[index].token.clone()).await { + Ok(balance) => Validation::Valid(Box::new(balance)), + Err(ProviderError::AuthRequired) => Validation::Rejected, + Err(error) => Validation::Unavailable(Some(error)), + }; + Outcome { index, validation } + })) + .await +} + +fn record(cache: &ValidationCache, candidates: &[TokenInfo], outcomes: &[Outcome], now: Instant) { + for outcome in outcomes { + match outcome.validation { + Validation::Valid(_) => cache.record(&candidates[outcome.index], true, now), + Validation::Rejected => cache.record(&candidates[outcome.index], false, now), + Validation::Unavailable(_) => {} + } + } +} + +/// Known status per candidate: a fresh cache entry, overridden by this round's +/// outcomes. An unavailable check falls back to the last known status. +fn resolved_statuses(lookups: &[Lookup], outcomes: &[Outcome]) -> Vec> { + let mut statuses: Vec> = lookups.iter().map(|l| l.fresh).collect(); + for outcome in outcomes { + match outcome.validation { + Validation::Valid(_) => statuses[outcome.index] = Some(true), + Validation::Rejected => statuses[outcome.index] = Some(false), + Validation::Unavailable(_) => { + if let Some(known) = lookups[outcome.index].last_known { + statuses[outcome.index] = Some(known); + } + } + } + } + statuses +} + +fn valid_indices(statuses: &[Option]) -> Vec { + statuses + .iter() + .enumerate() + .filter_map(|(index, status)| (*status == Some(true)).then_some(index)) + .collect() +} + +/// The candidate the user's selection points at among `pool`. Without an +/// explicit selection only a single candidate is unambiguous. +fn selected_candidate( + candidates: &[TokenInfo], + pool: &[usize], + selection: &ProfileSelection, +) -> Option { + match selection.profile_id.as_deref() { + Some(id) => pool.iter().copied().find(|&i| candidates[i].id == id), + None => (pool.len() == 1).then(|| pool[0]), + } +} + +fn has_balance(outcomes: &[Outcome], index: usize) -> bool { + outcomes + .iter() + .any(|o| o.index == index && matches!(o.validation, Validation::Valid(_))) +} + +fn take_balance(outcomes: &mut [Outcome], index: usize) -> Option> { + let outcome = outcomes + .iter_mut() + .rev() + .find(|o| o.index == index && matches!(o.validation, Validation::Valid(_)))?; + match std::mem::replace(&mut outcome.validation, Validation::Unavailable(None)) { + Validation::Valid(balance) => Some(balance), + _ => None, + } +} + +/// The latest failed check of `index` as an error. A transport failure carries +/// the session that produced it, except when the platform already rejected +/// that session: an outage after a rejection must not revive the old balance. +fn take_failure( + candidates: &[TokenInfo], + outcomes: &mut [Outcome], + statuses: &[Option], + index: usize, +) -> Option { + let latest = outcomes.iter_mut().rev().find(|o| o.index == index)?; + let Validation::Unavailable(error) = &mut latest.validation else { + return None; + }; + let error = error.take()?; + let owner = if statuses[index] == Some(false) { + None + } else { + balance_owner(&candidates[index].id, &candidates[index].token) + }; + Some(error.with_failure_owner(owner)) +} + +#[cfg(test)] +#[path = "session_resolver_tests.rs"] +mod tests; diff --git a/rust/src/providers/deepseek/session_resolver_tests.rs b/rust/src/providers/deepseek/session_resolver_tests.rs new file mode 100644 index 0000000000..d0da44ebfe --- /dev/null +++ b/rust/src/providers/deepseek/session_resolver_tests.rs @@ -0,0 +1,366 @@ +use std::sync::Arc; +use std::sync::atomic::{AtomicUsize, Ordering}; + +use super::*; +use crate::core::FailureOwnership; +use crate::providers::deepseek::BalanceInfo; + +fn candidate(id: &str, token: &str) -> TokenInfo { + TokenInfo { + id: id.to_string(), + token: token.to_string(), + label: format!("Google Chrome {id}"), + } +} + +fn select(id: &str) -> ProfileSelection { + ProfileSelection::from_value(Some(id)) +} + +fn balance(total: &str) -> BalanceResponse { + BalanceResponse { + is_available: true, + balance_infos: vec![BalanceInfo { + currency: "USD".into(), + total_balance: total.into(), + granted_balance: "0".into(), + topped_up_balance: total.into(), + }], + } +} + +fn fresh_cache() -> ValidationCache { + ValidationCache::new(VALIDITY_TTL) +} + +fn owner_of(id: &str, token: &str) -> LastGoodOwner { + balance_owner(id, token).expect("owner") +} + +/// Token-keyed outcomes: `ok-*` succeed, `rejected-*` are refused by the +/// platform, `down-*` hit a transport failure, `broken-*` a non-transport one. +async fn answer(token: String) -> Result { + if token.starts_with("rejected") { + Err(ProviderError::AuthRequired) + } else if token.starts_with("down") { + Err(ProviderError::Timeout) + } else if token.starts_with("broken") { + Err(ProviderError::Parse("bad body".into())) + } else { + Ok(balance("5")) + } +} + +fn failed(resolution: Resolution) -> ProviderError { + match resolution { + Resolution::Failed(error) => error, + other => panic!("expected failure, got {other:?}"), + } +} + +#[tokio::test] +async fn no_candidates_requires_a_session() { + let resolution = resolve(&[], &ProfileSelection::default(), &fresh_cache(), answer).await; + assert!(matches!(resolution, Resolution::SessionRequired)); +} + +#[tokio::test] +async fn single_valid_session_supplies_balance_and_owner() { + let candidates = [candidate("chrome:Default", "ok-token")]; + let resolution = resolve( + &candidates, + &ProfileSelection::default(), + &fresh_cache(), + answer, + ) + .await; + let Resolution::Balance { balance, owner } = resolution else { + panic!("expected balance"); + }; + assert!(balance.is_available); + assert_eq!(owner, Some(owner_of("chrome:Default", "ok-token"))); +} + +#[tokio::test] +async fn owner_follows_profile_and_token() { + assert_ne!( + owner_of("chrome:Default", "ok-one"), + owner_of("chrome:Default", "ok-two") + ); + assert_ne!( + owner_of("chrome:Default", "ok-one"), + owner_of("chrome:Profile 1", "ok-one") + ); + assert_eq!( + owner_of("chrome:Default", "ok-one"), + owner_of(" chrome:Default", "ok-one ") + ); +} + +#[tokio::test] +async fn rejected_session_requires_a_new_sign_in() { + let candidates = [candidate("chrome:Default", "rejected-token")]; + let resolution = resolve( + &candidates, + &ProfileSelection::default(), + &fresh_cache(), + answer, + ) + .await; + assert!(matches!(resolution, Resolution::SessionRequired)); +} + +#[tokio::test] +async fn transport_failure_keeps_the_session_that_failed() { + let candidates = [candidate("chrome:Default", "down-token")]; + let error = failed( + resolve( + &candidates, + &ProfileSelection::default(), + &fresh_cache(), + answer, + ) + .await, + ); + assert!(error.is_transport_failure()); + assert_eq!( + error.failure_ownership(), + FailureOwnership::Owned(owner_of("chrome:Default", "down-token")) + ); +} + +#[tokio::test] +async fn non_transport_failure_is_not_owned() { + let candidates = [candidate("chrome:Default", "broken-token")]; + let error = failed( + resolve( + &candidates, + &ProfileSelection::default(), + &fresh_cache(), + answer, + ) + .await, + ); + assert!(matches!(error, ProviderError::Parse(_))); + assert_eq!(error.failure_ownership(), FailureOwnership::Unchecked); +} + +#[tokio::test] +async fn outage_after_rejection_cannot_revive_the_old_balance() { + let cache = fresh_cache(); + let candidates = [candidate("chrome:Default", "token-flaky")]; + let calls = Arc::new(AtomicUsize::new(0)); + let flaky = |token: String| { + let calls = Arc::clone(&calls); + async move { + let _ = token; + match calls.fetch_add(1, Ordering::SeqCst) { + 0 => Err(ProviderError::AuthRequired), + _ => Err(ProviderError::Timeout), + } + } + }; + let selection = select("chrome:Default"); + + assert!(matches!( + resolve(&candidates, &selection, &cache, &flaky).await, + Resolution::SessionRequired + )); + let error = failed(resolve(&candidates, &selection, &cache, &flaky).await); + assert!(error.is_transport_failure()); + assert_eq!(error.failure_ownership(), FailureOwnership::Unowned); +} + +#[tokio::test] +async fn several_sessions_without_a_selection_ask_for_one() { + let candidates = [ + candidate("chrome:Default", "ok-a"), + candidate("chrome:Profile 1", "ok-b"), + ]; + let resolution = resolve( + &candidates, + &ProfileSelection::default(), + &fresh_cache(), + answer, + ) + .await; + let Resolution::SelectionRequired(ids) = resolution else { + panic!("expected selection"); + }; + assert_eq!(ids, vec!["chrome:Default", "chrome:Profile 1"]); +} + +#[tokio::test] +async fn explicit_selection_supplies_that_profiles_balance() { + let candidates = [ + candidate("chrome:Default", "ok-a"), + candidate("chrome:Profile 1", "ok-b"), + ]; + let resolution = resolve( + &candidates, + &select("chrome:Profile 1"), + &fresh_cache(), + answer, + ) + .await; + let Resolution::Balance { owner, .. } = resolution else { + panic!("expected balance"); + }; + assert_eq!(owner, Some(owner_of("chrome:Profile 1", "ok-b"))); +} + +#[tokio::test] +async fn selected_profile_keeps_its_own_outage_beside_a_working_profile() { + let candidates = [ + candidate("chrome:Default", "ok-a"), + candidate("chrome:Profile 1", "down-b"), + ]; + let error = failed( + resolve( + &candidates, + &select("chrome:Profile 1"), + &fresh_cache(), + answer, + ) + .await, + ); + assert_eq!( + error.failure_ownership(), + FailureOwnership::Owned(owner_of("chrome:Profile 1", "down-b")) + ); +} + +#[tokio::test] +async fn ambiguous_outage_does_not_borrow_another_profiles_failure() { + let candidates = [ + candidate("chrome:Default", "ok-a"), + candidate("chrome:Profile 1", "down-b"), + ]; + let resolution = resolve( + &candidates, + &ProfileSelection::default(), + &fresh_cache(), + answer, + ) + .await; + // Only Default is known valid; with no explicit selection the single valid + // profile is selected and supplies its own balance. + let Resolution::Balance { owner, .. } = resolution else { + panic!("expected balance"); + }; + assert_eq!(owner, Some(owner_of("chrome:Default", "ok-a"))); +} + +#[tokio::test] +async fn rejected_selection_beside_a_valid_profile_asks_for_a_choice() { + let candidates = [ + candidate("chrome:Default", "ok-a"), + candidate("chrome:Profile 1", "rejected-b"), + ]; + let resolution = resolve( + &candidates, + &select("chrome:Profile 1"), + &fresh_cache(), + answer, + ) + .await; + let Resolution::SelectionRequired(ids) = resolution else { + panic!("expected selection, got {resolution:?}"); + }; + assert_eq!(ids, vec!["chrome:Default"]); +} + +#[tokio::test] +async fn unknown_selection_asks_for_a_choice() { + let candidates = [ + candidate("chrome:Default", "ok-a"), + candidate("chrome:Profile 1", "ok-b"), + ]; + let resolution = resolve( + &candidates, + &select("chrome:Missing"), + &fresh_cache(), + answer, + ) + .await; + assert!(matches!(resolution, Resolution::SelectionRequired(_))); +} + +#[tokio::test] +async fn fresh_cache_skips_revalidating_unselected_profiles() { + let cache = fresh_cache(); + let calls = Arc::new(AtomicUsize::new(0)); + let counting = |token: String| { + let calls = Arc::clone(&calls); + async move { + calls.fetch_add(1, Ordering::SeqCst); + answer(token).await + } + }; + let candidates = [ + candidate("chrome:Default", "ok-a"), + candidate("chrome:Profile 1", "ok-b"), + ]; + let selection = select("chrome:Default"); + + resolve(&candidates, &selection, &cache, &counting).await; + assert_eq!(calls.load(Ordering::SeqCst), 2); + resolve(&candidates, &selection, &cache, &counting).await; + // Only the selected profile is asked again: its balance is always live. + assert_eq!(calls.load(Ordering::SeqCst), 3); +} + +#[tokio::test] +async fn changed_token_in_the_same_profile_is_validated_again() { + let cache = fresh_cache(); + let selection = select("chrome:Default"); + let first = [candidate("chrome:Default", "rejected-old")]; + assert!(matches!( + resolve(&first, &selection, &cache, answer).await, + Resolution::SessionRequired + )); + + let second = [candidate("chrome:Default", "ok-new")]; + let resolution = resolve(&second, &selection, &cache, answer).await; + let Resolution::Balance { owner, .. } = resolution else { + panic!("expected balance"); + }; + assert_eq!(owner, Some(owner_of("chrome:Default", "ok-new"))); +} + +#[tokio::test] +async fn cached_status_survives_a_later_validation_outage() { + let cache = ValidationCache::new(Duration::ZERO); + let candidates = [candidate("chrome:Default", "token-x")]; + let calls = Arc::new(AtomicUsize::new(0)); + let script = |_token: String| { + let calls = Arc::clone(&calls); + async move { + match calls.fetch_add(1, Ordering::SeqCst) { + 0 => Ok(balance("5")), + _ => Err(ProviderError::Timeout), + } + } + }; + let selection = ProfileSelection::default(); + + assert!(matches!( + resolve(&candidates, &selection, &cache, &script).await, + Resolution::Balance { .. } + )); + let error = failed(resolve(&candidates, &selection, &cache, &script).await); + assert_eq!( + error.failure_ownership(), + FailureOwnership::Owned(owner_of("chrome:Default", "token-x")) + ); +} + +#[test] +fn validation_cache_stores_no_plain_token() { + let cache = fresh_cache(); + let candidate = candidate("chrome:Default", "super-secret-token-value"); + cache.record(&candidate, true, Instant::now()); + let entries = cache.entries.lock().unwrap(); + let rendered = format!("{:?}", entries.get("chrome:Default").unwrap().proof); + assert!(!rendered.contains("super-secret")); +} diff --git a/rust/src/providers/deepseek/tests.rs b/rust/src/providers/deepseek/tests.rs index badaaf05e1..4acac37f8d 100644 --- a/rust/src/providers/deepseek/tests.rs +++ b/rust/src/providers/deepseek/tests.rs @@ -239,3 +239,34 @@ fn applies_deepseek_summary_as_extra_windows() { Some("¥0.0000 · Current month") ); } + +#[test] +fn only_owned_transport_failures_retain_the_last_balance() { + let provider = DeepSeekProvider::new(); + let owner = session_resolver::balance_owner("chrome:Default", "token-value"); + + assert_eq!( + provider + .last_good_failure_policy_for_error(&ProviderError::Timeout.with_failure_owner(owner)), + LastGoodFailurePolicy::Preserve + ); + // A deadline has no owner but is still an owned failure; the shell then + // fails closed because no session can match. + assert_eq!( + provider + .last_good_failure_policy_for_error(&ProviderError::Timeout.with_failure_owner(None)), + LastGoodFailurePolicy::Preserve + ); + for error in [ + ProviderError::Timeout, + ProviderError::AuthRequired, + ProviderError::Other("HTTP 500".into()), + ProviderError::Parse("bad".into()), + ] { + assert_eq!( + provider.last_good_failure_policy_for_error(&error), + LastGoodFailurePolicy::Replace, + "{error:?}" + ); + } +}