diff --git a/apps/desktop-tauri/src-tauri/src/commands/providers.rs b/apps/desktop-tauri/src-tauri/src/commands/providers.rs index eaeee247537..d01451bd0df 100644 --- a/apps/desktop-tauri/src-tauri/src/commands/providers.rs +++ b/apps/desktop-tauri/src-tauri/src/commands/providers.rs @@ -6,7 +6,7 @@ use super::provider_refresh::{ use super::warning_identity::WarningIdentity; use super::*; use chrono::{Local, Utc}; -use codexbar::core::HookUsageWindow; +use codexbar::core::{HookUsageWindow, RefreshInteraction, scope_refresh_interaction}; use codexbar::notifications::{LimitResetObservation, LimitResetSample, WarningScope}; use serde::Serialize; use std::sync::Arc; @@ -223,12 +223,20 @@ async fn do_refresh_providers_with_policy( ); let enabled_count = refresh_ids.len(); + // Non-forced refreshes are automatic (timer, auto-resume, stale-gated launch and + // poll refreshes). Forced refreshes come from a user action or a follow-up to one. + let interaction = if force { + RefreshInteraction::UserInitiated + } else { + RefreshInteraction::Background + }; let handles = spawn_provider_refreshes( app, &inputs, &refresh_ids, generation, scope.refresh_account_lanes(), + interaction, ); if scope == RefreshScope::AllEnabled { super::claude_swap_usage::spawn_background_pool_read(app, &inputs.settings, false); @@ -292,6 +300,7 @@ fn spawn_provider_refreshes( provider_ids: &[ProviderId], generation: u64, refresh_account_lanes: bool, + interaction: RefreshInteraction, ) -> Vec> { let mut handles = Vec::with_capacity(provider_ids.len()); let fetch_permits = Arc::new(tokio::sync::Semaphore::new(MAX_CONCURRENT_PROVIDER_FETCHES)); @@ -322,17 +331,20 @@ fn spawn_provider_refreshes( let Ok(_permit) = fetch_permits.acquire_owned().await else { return; }; - refresh_provider( - app_handle, - id, - ctx, - generation, - ActiveTokenAccount { - id: token_account_id, - label: token_account_label, - }, - hooks_enabled, - pace_work_days, + scope_refresh_interaction( + interaction, + refresh_provider( + app_handle, + id, + ctx, + generation, + ActiveTokenAccount { + id: token_account_id, + label: token_account_label, + }, + hooks_enabled, + pace_work_days, + ), ) .await; })); diff --git a/rust/src/core/mod.rs b/rust/src/core/mod.rs index 2d600b6970f..b61eadd19c8 100755 --- a/rust/src/core/mod.rs +++ b/rust/src/core/mod.rs @@ -32,6 +32,7 @@ mod provider_state; mod quota_burndown; mod rate_window; mod redactor; +mod refresh_interaction; mod session_equivalent_forecast; mod sqlite; mod throughput; @@ -74,6 +75,7 @@ pub use quota_burndown::{ }; pub use rate_window::*; pub use redactor::*; +pub use refresh_interaction::*; pub use session_equivalent_forecast::*; pub use sqlite::*; pub use throughput::*; diff --git a/rust/src/core/refresh_interaction.rs b/rust/src/core/refresh_interaction.rs new file mode 100644 index 00000000000..66c95933c7a --- /dev/null +++ b/rust/src/core/refresh_interaction.rs @@ -0,0 +1,65 @@ +//! Whether the provider fetch running on this task was started by the user or +//! by an automatic refresh. +//! +//! Upstream scopes this per task (`ProviderInteractionContext`). The Windows +//! shell sets it around each provider fetch: non-forced refreshes (the +//! scheduled timer, the auto-resume probe, and stale-gated launch and poll +//! refreshes) are `Background`; forced refreshes (the Refresh action, tray +//! Refresh, saved settings or account changes) are `UserInitiated`. Outside a +//! scope, such as the CLI, the value is `UserInitiated`. + +use std::future::Future; + +#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)] +pub enum RefreshInteraction { + #[default] + UserInitiated, + Background, +} + +tokio::task_local! { + static CURRENT: RefreshInteraction; +} + +/// The interaction of the fetch running on this task. `UserInitiated` outside a scope. +pub fn current_refresh_interaction() -> RefreshInteraction { + CURRENT + .try_with(|interaction| *interaction) + .unwrap_or_default() +} + +/// Run `future` with `interaction` as the current refresh interaction. +pub async fn scope_refresh_interaction( + interaction: RefreshInteraction, + future: F, +) -> F::Output { + CURRENT.scope(interaction, future).await +} + +#[cfg(test)] +mod tests { + use super::*; + + #[tokio::test] + async fn defaults_to_user_initiated_outside_a_scope() { + assert_eq!( + current_refresh_interaction(), + RefreshInteraction::UserInitiated + ); + } + + #[tokio::test] + async fn background_scope_is_visible_across_awaits_and_ends_with_the_scope() { + let seen = scope_refresh_interaction(RefreshInteraction::Background, async { + tokio::task::yield_now().await; + current_refresh_interaction() + }) + .await; + + assert_eq!(seen, RefreshInteraction::Background); + assert_eq!( + current_refresh_interaction(), + RefreshInteraction::UserInitiated + ); + } +} diff --git a/rust/src/providers/kiro/mod.rs b/rust/src/providers/kiro/mod.rs index f7ad449e960..425376b55a0 100755 --- a/rust/src/providers/kiro/mod.rs +++ b/rust/src/providers/kiro/mod.rs @@ -18,8 +18,8 @@ use std::process::{Output, Stdio}; use tokio::process::Command; use crate::core::{ - FetchContext, Provider, ProviderError, ProviderFetchResult, ProviderId, RateWindow, SourceMode, - UsageSnapshot, + FetchContext, Provider, ProviderError, ProviderFetchResult, ProviderId, RateWindow, + RefreshInteraction, SourceMode, UsageSnapshot, current_refresh_interaction, }; /// Upstream `KiroProviderDescriptor.pace` (`.calendarMonthResetWindow`): the @@ -101,6 +101,12 @@ impl KiroProvider { "kiro-cli not found. Install from https://kiro.dev".to_string(), ) })?; + // Background refreshes check the usage host before any CLI command runs, + // including `whoami`. Offline, the CLI's sign-in attempt opens browser tabs. + // Manual refreshes skip the check and keep the existing path. + if current_refresh_interaction() == RefreshInteraction::Background { + usage_limits::ensure_reachable().await?; + } Self::ensure_logged_in(&cli_path).await?; // Windows intentionally uses pipe-first (stdout/stderr Stdio::piped) rather than a dual @@ -430,6 +436,14 @@ impl Provider for KiroProvider { true } + // A background refresh that fails its reachability check keeps the last good + // usage instead of showing an error. Note: the shared transport classification + // also covers timeouts and connection failures from any refresh, including manual + // ones; Kiro's CLI parse and auth errors are unaffected. + fn retains_last_good_on_transport_failure(&self) -> bool { + true + } + async fn fetch_usage(&self, ctx: &FetchContext) -> Result { tracing::debug!("Fetching Kiro usage"); diff --git a/rust/src/providers/kiro/tests.rs b/rust/src/providers/kiro/tests.rs index 3d02d964ed3..5d9edf64a47 100644 --- a/rust/src/providers/kiro/tests.rs +++ b/rust/src/providers/kiro/tests.rs @@ -279,3 +279,19 @@ fn cli_presence_maps_to_local_runtime_offline_but_state_db_stays_default() { crate::core::ProviderStateKind::NeedsAuthentication ); } + +#[test] +fn transport_failures_keep_last_good_while_sign_in_failures_replace_it() { + use crate::core::{LastGoodFailurePolicy, Provider}; + + let provider = KiroProvider::new(); + assert!(provider.retains_last_good_on_transport_failure()); + assert_eq!( + provider.last_good_failure_policy_for_error(&ProviderError::Timeout), + LastGoodFailurePolicy::Preserve + ); + assert_eq!( + provider.last_good_failure_policy_for_error(&ProviderError::AuthRequired), + LastGoodFailurePolicy::Replace + ); +} diff --git a/rust/src/providers/kiro/usage_limits.rs b/rust/src/providers/kiro/usage_limits.rs index 0d3cb256fa4..98d9718538c 100644 --- a/rust/src/providers/kiro/usage_limits.rs +++ b/rust/src/providers/kiro/usage_limits.rs @@ -11,6 +11,8 @@ const FRANKFURT_ENDPOINT: &str = "https://q.eu-central-1.amazonaws.com/"; const TARGET: &str = "AmazonCodeWhispererService.GetUsageLimits"; const TOKEN_KEY: &str = "kirocli:odic:token"; const PROFILE_KEY: &str = "api.codewhisperer.profile"; +/// Upper bound for the unauthenticated reachability check before a background CLI run. +const REACHABILITY_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(4); #[derive(Debug, Clone)] pub(super) struct KiroUsageLimits { @@ -234,6 +236,65 @@ pub(super) fn parse_usage_limits(data: &[u8]) -> Result Result<(), ProviderError> { + ensure_reachable_at(reachability_endpoint(&state_database_path())).await +} + +/// Usage host for the CLI's profile. Missing or unsupported region state +/// falls back to US East, matching upstream. +fn reachability_endpoint(database: &Path) -> &'static str { + read_profile_arn(database) + .as_deref() + .and_then(endpoint_for_profile_arn) + .unwrap_or(US_EAST_ENDPOINT) +} + +async fn ensure_reachable_at(endpoint: &str) -> Result<(), ProviderError> { + let client = crate::core::apply_app_proxy(reqwest::Client::builder()) + .timeout(REACHABILITY_TIMEOUT) + .build() + .map_err(|error| ProviderError::Other(error.to_string()))?; + // The shared classifier retains only refused connections, but offline + // Windows hosts usually fail DNS first. Every connect failure here means the + // host is unreachable, so it must keep the cached snapshot like upstream. + client.head(endpoint).send().await.map_err(|error| { + if error.is_connect() { + ProviderError::Timeout + } else { + ProviderError::Network(error) + } + })?; + Ok(()) +} + +/// The profile ARN alone, without the access token. Used only for the +/// reachability endpoint, so a signed-out CLI still gets a host choice. +fn read_profile_arn(path: &Path) -> Option { + if !path.is_file() { + return None; + } + let connection = crate::core::open_readonly_sqlite_connection( + path, + crate::core::DEFAULT_SQLITE_BUSY_TIMEOUT, + ) + .ok()?; + let profile_json: Option = connection + .query_row( + "SELECT value FROM state WHERE key = ?1", + [PROFILE_KEY], + |row| row.get(0), + ) + .optional() + .ok() + .flatten(); + json_string(profile_json.as_deref(), "arn", false) +} + fn state_database_path() -> PathBuf { if let Some(path) = std::env::var_os("KIRO_DATA_DIR").filter(|value| !value.is_empty()) { return PathBuf::from(path).join("data.sqlite3"); @@ -476,4 +537,85 @@ mod tests { ); } } + + fn state_db_with_profile(profile_arn: &str) -> (tempfile::TempDir, PathBuf) { + let dir = tempfile::tempdir().expect("temp dir"); + let path = dir.path().join("data.sqlite3"); + let connection = rusqlite::Connection::open(&path).expect("open fixture database"); + connection + .execute_batch("CREATE TABLE state (key TEXT PRIMARY KEY, value TEXT);") + .expect("create state table"); + let profile = serde_json::json!({ "arn": profile_arn }).to_string(); + connection + .execute( + "INSERT INTO state (key, value) VALUES (?1, ?2)", + [PROFILE_KEY, profile.as_str()], + ) + .expect("insert profile"); + (dir, path) + } + + #[test] + fn reachability_endpoint_follows_the_cli_profile_region() { + let (_dir, path) = + state_db_with_profile("arn:aws:codewhisperer:eu-central-1:123:profile/test"); + + assert_eq!(reachability_endpoint(&path), FRANKFURT_ENDPOINT); + } + + #[test] + fn reachability_endpoint_defaults_to_us_east_without_usable_region_state() { + let dir = tempfile::tempdir().expect("temp dir"); + assert_eq!( + reachability_endpoint(&dir.path().join("missing.sqlite3")), + US_EAST_ENDPOINT + ); + + let (_dir, path) = + state_db_with_profile("arn:aws:codewhisperer:ap-southeast-1:123:profile/test"); + assert_eq!(reachability_endpoint(&path), US_EAST_ENDPOINT); + } + + #[tokio::test] + async fn transport_failure_is_a_transport_error_that_keeps_last_good() { + // Bind and drop a listener to get a local port with nothing listening. + let port = { + let listener = std::net::TcpListener::bind(("127.0.0.1", 0)).expect("bind"); + listener.local_addr().expect("local addr").port() + }; + + let error = ensure_reachable_at(&format!("http://127.0.0.1:{port}/")) + .await + .expect_err("a closed port must fail the reachability check"); + + assert!(matches!(error, ProviderError::Timeout)); + assert!(error.is_transport_failure()); + } + + #[tokio::test] + async fn any_http_response_counts_as_reachable() { + use tokio::io::{AsyncReadExt, AsyncWriteExt}; + + let listener = tokio::net::TcpListener::bind(("127.0.0.1", 0)) + .await + .expect("bind"); + let port = listener.local_addr().expect("local addr").port(); + let server = tokio::spawn(async move { + let (mut socket, _) = listener.accept().await.expect("accept"); + let mut request = [0_u8; 1024]; + let read = socket.read(&mut request).await.expect("read request"); + assert!(read > 0, "the client must send a request before reading"); + socket + .write_all( + b"HTTP/1.1 403 Forbidden\r\nContent-Length: 0\r\nConnection: close\r\n\r\n", + ) + .await + .expect("write response"); + }); + + ensure_reachable_at(&format!("http://127.0.0.1:{port}/")) + .await + .expect("an HTTP response proves reachability"); + server.await.expect("server task"); + } }