Port upstream 0.60.4: keep Chrome DeepSeek balance through same-session transport failures (stacked on #646)

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<LastGoodOwner>) -> 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);
+}

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;
 

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<uuid::Uuid>,
     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,14 +770,38 @@ 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<codexbar::core::LastGoodOwner>,
+) {
+    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(
     guard: &mut AppState,
     id: ProviderId,
     snapshot: ProviderUsageSnapshot,
     policy: Option<codexbar::core::LastGoodFailurePolicy>,
+    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<codexbar::core::LastGoodFailurePolicy>,
+    /// 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<codexbar::core::LastGoodOwner>,
+}
+
 async fn fetch_provider_snapshot(
     id: ProviderId,
     ctx: FetchContext,
     token_account_id: Option<uuid::Uuid>,
-) -> (
-    ProviderUsageSnapshot,
-    Option<String>,
-    Option<codexbar::core::LastGoodFailurePolicy>,
-) {
+) -> (ProviderUsageSnapshot, Option<String>, 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(