From 1fc46a532b136feeff2be7960ea05e6c9f76167b Mon Sep 17 00:00:00 2001 From: Felipe Coury Date: Tue, 15 Sep 2026 18:44:57 +0000 Subject: [PATCH] Load analytics reports with server plans and account identity checks (#45762) ## Why Token plan claims can be stale, and the active account or user can change while a request is in flight. Analytics report selection and account-bound response data need to reflect the verified identity and current server plan. ## What changed - Fetch the active account's plan once per analytics session and use it to select report endpoints and supported credit breakdowns. - Add report loading with a fixed end date, account-scoped response caching, and token model filtering. Reuse payloads across grouping changes and evict invalid cached responses so requests can retry. - Prefer complete attribution for usage breakdowns within the requested range; retain legacy surface/model data when attribution is incomplete and include all features in turn-start breakdowns. - Add cancellable report-loading state with timeout and interruption errors, and preserve actionable sign-in and retry messages. - Recheck the active identity after rate-limit reads before exposing account-bound fields. ## Testing Add regression tests for server plan discovery, report routing and caching, model filtering, attribution fallback, failed-request retries, load cancellation, authentication recovery, and account or user changes during requests. GitOrigin-RevId: 45c4c09f108c9703893f3ed15613437ebd7c74a8 --- .../request_processors/account_processor.rs | 14 +- codex-rs/app-server/tests/suite/v2/mod.rs | 2 + .../suite/v2/rate_limits_identity_tests.rs | 113 ++++ .../backend-client/src/analytics_session.rs | 49 +- .../src/analytics_session_tests.rs | 113 ++++ codex-rs/backend-client/src/client.rs | 9 +- codex-rs/backend-client/src/types.rs | 6 + codex-rs/tui/src/analytics.rs | 1 + .../tui/src/analytics/account_plan_tests.rs | 167 ++++++ codex-rs/tui/src/analytics/client.rs | 183 +++++- codex-rs/tui/src/analytics/client_tests.rs | 558 +++++++++++++++++- codex-rs/tui/src/analytics/data.rs | 102 +++- codex-rs/tui/src/analytics/data_tests.rs | 59 ++ codex-rs/tui/src/analytics/models.rs | 3 +- codex-rs/tui/src/analytics/normalize.rs | 65 +- codex-rs/tui/src/analytics/normalize_tests.rs | 70 ++- codex-rs/tui/src/analytics/render.rs | 26 + 17 files changed, 1477 insertions(+), 63 deletions(-) create mode 100644 codex-rs/app-server/tests/suite/v2/rate_limits_identity_tests.rs create mode 100644 codex-rs/backend-client/src/analytics_session_tests.rs create mode 100644 codex-rs/tui/src/analytics/account_plan_tests.rs create mode 100644 codex-rs/tui/src/analytics/data_tests.rs create mode 100644 codex-rs/tui/src/analytics/render.rs diff --git a/codex-rs/app-server/src/request_processors/account_processor.rs b/codex-rs/app-server/src/request_processors/account_processor.rs index 8327cdf7ce..0f2cadabfc 100644 --- a/codex-rs/app-server/src/request_processors/account_processor.rs +++ b/codex-rs/app-server/src/request_processors/account_processor.rs @@ -1194,11 +1194,15 @@ impl AccountRequestProcessor { // Match desktop's account readiness check before exposing account-bound CTA content. // Normal rate limits remain available when older backends omit identity or banner data. - let matches_active_account = !auth.is_fedramp_account() - && response.account_id.is_some() - && response.account_id == auth.get_account_id() - && response.user_id.is_some() - && response.user_id == auth.get_chatgpt_user_id(); + // Login can change while the backend read is in flight. + let active_auth = self.auth_manager.auth().await; + let matches_active_account = active_auth.is_some_and(|auth| { + !auth.is_fedramp_account() + && response.account_id.is_some() + && response.account_id == auth.get_account_id() + && response.user_id.is_some() + && response.user_id == auth.get_chatgpt_user_id() + }); let rate_limit_upsell = response .rate_limit_upsell .filter(|_| matches_active_account); diff --git a/codex-rs/app-server/tests/suite/v2/mod.rs b/codex-rs/app-server/tests/suite/v2/mod.rs index 405d1bfbdc..5d0cbd3be5 100644 --- a/codex-rs/app-server/tests/suite/v2/mod.rs +++ b/codex-rs/app-server/tests/suite/v2/mod.rs @@ -91,6 +91,8 @@ mod process_exec; mod projects; mod rate_limit_reset_credits; mod rate_limits; +#[path = "rate_limits_identity_tests.rs"] +mod rate_limits_identity; mod realtime_conversation; mod recommended_plugins; mod remote_control; diff --git a/codex-rs/app-server/tests/suite/v2/rate_limits_identity_tests.rs b/codex-rs/app-server/tests/suite/v2/rate_limits_identity_tests.rs new file mode 100644 index 0000000000..b8356c5641 --- /dev/null +++ b/codex-rs/app-server/tests/suite/v2/rate_limits_identity_tests.rs @@ -0,0 +1,113 @@ +//! Account readiness is verified after an in-flight backend read, including concurrent login. + +use anyhow::Result; +use app_test_support::ChatGptAuthFixture; +use app_test_support::ChatGptIdTokenClaims; +use app_test_support::TestAppServer; +use app_test_support::encode_id_token; +use app_test_support::write_chatgpt_auth; +use axum::Json; +use axum::Router; +use axum::http::HeaderMap; +use axum::routing::get; +use codex_app_server_protocol::GetAccountRateLimitsResponse; +use codex_app_server_protocol::LoginAccountResponse; +use codex_config::types::AuthCredentialsStoreMode; +use pretty_assertions::assert_eq; +use serde_json::json; +use std::sync::Arc; +use std::time::Duration; +use tempfile::TempDir; +use test_case::test_case; +use tokio::net::TcpListener; +use tokio::sync::Notify; +use tokio::time::timeout; + +const READ_TIMEOUT: Duration = Duration::from_secs(/*secs*/ 10); + +#[test_case("workspace-a", "user-b", /*verified*/ false; "user_switch")] +#[test_case("workspace-b", "user-a", /*verified*/ false; "workspace_switch")] +#[test_case("workspace-a", "user-a", /*verified*/ true; "same_identity_token_refresh")] +#[tokio::test] +async fn identity_is_rechecked_after_backend_response( + account: &str, + user: &str, + verified: bool, +) -> Result<()> { + let home = TempDir::new()?; + write_chatgpt_auth( + home.path(), + ChatGptAuthFixture::new("old-token") + .account_id("workspace-a") + .chatgpt_user_id("user-a") + .plan_type("team"), + AuthCredentialsStoreMode::File, + )?; + let entered = Arc::new(Notify::new()); + let release = Arc::new(Notify::new()); + let request_entered = Arc::clone(&entered); + let response_release = Arc::clone(&release); + let listener = TcpListener::bind("127.0.0.1:0").await?; + std::fs::write( + home.path().join("config.toml"), + format!("chatgpt_base_url = \"http://{}\"\n", listener.local_addr()?), + )?; + let router = Router::new().route( + "/api/codex/usage", + get(move |headers: HeaderMap| { + let entered = Arc::clone(&request_entered); + let release = Arc::clone(&response_release); + async move { + assert_eq!(headers["authorization"], "Bearer old-token"); + assert_eq!(headers["chatgpt-account-id"], "workspace-a"); + entered.notify_one(); + release.notified().await; + Json(json!({ + "account_id": "workspace-a", "user_id": "user-a", "plan_type": "team", + "rate_limit": {"allowed": true, "limit_reached": false, + "primary_window": {"used_percent": 42, "limit_window_seconds": 3600, + "reset_after_seconds": 120, "reset_at": 2000000000}} + })) + } + }), + ); + let server = tokio::spawn(async move { axum::serve(listener, router).await }); + let mut app = TestAppServer::builder() + .with_codex_home(home.path()) + .with_env_overrides(&[("OPENAI_API_KEY", None)]) + .build_initialized_with_timeout(READ_TIMEOUT) + .await?; + let request = app + .send_request( + "account/rateLimits/read", + Some(json!({"excludeResetCreditDetails": true})), + ) + .await?; + timeout(READ_TIMEOUT, entered.notified()).await?; + let token = encode_id_token( + &ChatGptIdTokenClaims::new() + .chatgpt_account_id(account) + .chatgpt_user_id(user) + .plan_type("team"), + )?; + let login = app + .send_chatgpt_auth_tokens_login_request(token, account.into(), Some("team".into())) + .await?; + let response: LoginAccountResponse = timeout(READ_TIMEOUT, app.read_response(login)).await??; + assert_eq!(response, LoginAccountResponse::ChatgptAuthTokens {}); + release.notify_one(); + let response: GetAccountRateLimitsResponse = + timeout(READ_TIMEOUT, app.read_response(request)).await??; + server.abort(); + let snapshot = json!({ + "limitId": "codex", "planType": "team", + "primary": {"usedPercent": 42, "windowDurationMins": 60, "resetsAt": 2000000000} + }); + let expected: GetAccountRateLimitsResponse = serde_json::from_value(json!({ + "ordinaryUsageAllowed": if verified { Some(true) } else { None }, + "accountId": "workspace-a", + "rateLimits": snapshot, "rateLimitsByLimitId": {"codex": snapshot} + }))?; + assert_eq!(response, expected); + Ok(()) +} diff --git a/codex-rs/backend-client/src/analytics_session.rs b/codex-rs/backend-client/src/analytics_session.rs index 9919dfe846..d2ba698b46 100644 --- a/codex-rs/backend-client/src/analytics_session.rs +++ b/codex-rs/backend-client/src/analytics_session.rs @@ -1,4 +1,5 @@ -//! Account-scoped analytics authentication, credential recovery, and request identity checks. +//! Local analytics authentication, server plan discovery, and account/user identity checks. +//! Each new session fetches its plan once; token plan claims do not select reports. use crate::Client; use crate::RequestError; @@ -17,7 +18,7 @@ pub struct AnalyticsAccount { pub plan_type: Option, } -/// A backend client that remains bound to its initial ChatGPT account and user. +/// A backend client bound to its initial ChatGPT account and user, with a server plan snapshot. pub struct AnalyticsSession { client: Client, auth_manager: Arc, @@ -26,7 +27,7 @@ pub struct AnalyticsSession { } impl AnalyticsSession { - /// Load local ChatGPT credentials using the configured auth and HTTP policies. + /// Load local ChatGPT credentials and the current server plan using the configured HTTP policy. pub async fn from_config( config: &impl AuthManagerConfig, http_client_factory: HttpClientFactory, @@ -48,26 +49,47 @@ impl AnalyticsSession { let account = AnalyticsAccount { id, email: auth.get_account_email(), - plan_type: auth.account_plan_type(), + plan_type: None, }; let client = Client::new_without_redirects(config.chatgpt_base_url(), http_client_factory) .with_auth_provider(codex_model_provider::auth_provider_from_auth_manager( Arc::clone(&auth_manager), &auth, )); - Ok(Self { + let mut session = Self { client, auth_manager, auth, account, - }) + }; + let accounts = session + .request(|client| async move { client.get_accounts_check().await }) + .await; + // Preserve account-switch guidance even when it interrupts account discovery. + session.ensure_identity().await?; + let accounts = accounts.map_err(|error| { + if error.is_unauthorized() { + "Sign in again to load Analytics." + } else { + "Couldn't load account plan. Press R to retry Analytics." + } + })?; + session.account.plan_type = Some( + accounts + .accounts + .into_iter() + .find(|account| account.id == session.account.id) + .and_then(|account| account.plan_type) + .ok_or("Couldn't load account plan. Press R to retry Analytics.")?, + ); + Ok(session) } /// Return the account metadata captured when this session was opened. pub fn account(&self) -> &AnalyticsAccount { &self.account } - /// Reject responses or cached data after a local account or user switch. + /// Reject responses or cached data after a local account or user change. pub async fn ensure_identity(&self) -> Result<(), String> { self.auth_manager.reload().await; let current = self.auth_manager.auth().await; @@ -91,10 +113,11 @@ impl AnalyticsSession { .await .map_err(|error| RequestError::Other(anyhow::anyhow!(error)))?; let result = request(self.client.clone()).await; - if result.as_ref().is_err_and(RequestError::is_unauthorized) && recovery.has_next() { - recovery.next().await.map_err(|_| { - RequestError::Other(anyhow::anyhow!("Sign in again to view Analytics.")) - })?; + // Keep the original 401 when recovery fails so callers retain sign-in guidance. + if result.as_ref().is_err_and(RequestError::is_unauthorized) + && recovery.has_next() + && recovery.next().await.is_ok() + { continue; } self.ensure_identity() @@ -104,3 +127,7 @@ impl AnalyticsSession { } } } + +#[cfg(test)] +#[path = "analytics_session_tests.rs"] +mod tests; diff --git a/codex-rs/backend-client/src/analytics_session_tests.rs b/codex-rs/backend-client/src/analytics_session_tests.rs new file mode 100644 index 0000000000..c491cd88de --- /dev/null +++ b/codex-rs/backend-client/src/analytics_session_tests.rs @@ -0,0 +1,113 @@ +//! Authentication recovery preserves actionable status without changing the account scope. + +use super::*; +use codex_http_client::OutboundProxyPolicy; +use codex_login::ExternalAuth; +use codex_login::ExternalAuthFuture; +use codex_login::ExternalAuthRefreshContext; +use codex_login::RefreshTokenError; +use codex_protocol::auth::RefreshTokenFailedError; +use codex_protocol::auth::RefreshTokenFailedReason; +use pretty_assertions::assert_eq; +use std::sync::atomic::AtomicUsize; +use std::sync::atomic::Ordering; +use wiremock::Mock; +use wiremock::MockServer; +use wiremock::ResponseTemplate; +use wiremock::matchers::method; + +enum RefreshOutcome { + Accepted, + Rejected, +} + +struct ExternalCredentials { + auth: CodexAuth, + outcome: RefreshOutcome, + refreshes: AtomicUsize, +} + +impl ExternalAuth for ExternalCredentials { + fn resolve(&self) -> ExternalAuthFuture<'_, CodexAuth> { + Box::pin(async { Ok(self.auth.clone()) }) + } + + fn refresh(&self, _context: ExternalAuthRefreshContext) -> ExternalAuthFuture<'_, CodexAuth> { + Box::pin(async { + self.refreshes.fetch_add(/*val*/ 1, Ordering::SeqCst); + match self.outcome { + RefreshOutcome::Accepted => Ok(self.auth.clone()), + RefreshOutcome::Rejected => Err(std::io::Error::other("refresh rejected")), + } + }) + } + + fn classify_error(&self, error: std::io::Error) -> RefreshTokenError { + RefreshTokenError::Permanent(RefreshTokenFailedError::new( + RefreshTokenFailedReason::Other, + error.to_string(), + )) + } +} + +#[tokio::test] +async fn failed_recovery_preserves_unauthorized_status() { + assert_unauthorized_after_recovery(RefreshOutcome::Rejected, /*expected_requests*/ 1).await; +} + +#[tokio::test] +async fn persistent_unauthorized_stops_after_successful_recovery() { + assert_unauthorized_after_recovery(RefreshOutcome::Accepted, /*expected_requests*/ 2).await; +} + +async fn assert_unauthorized_after_recovery(outcome: RefreshOutcome, expected_requests: u64) { + let auth = CodexAuth::from_external_chatgpt_tokens( + "e30.eyJleHAiOjQxMDI0NDQ4MDAsImh0dHBzOi8vYXBpLm9wZW5haS5jb20vYXV0aCI6eyJjaGF0Z3B0X3VzZXJfaWQiOiJ1c2VyLWEifX0.test", "account-a", Some("plus"), + ).unwrap(); + let auth_manager = AuthManager::from_auth_for_testing(auth.clone()); + let credentials = Arc::new(ExternalCredentials { + auth: auth.clone(), + outcome, + refreshes: AtomicUsize::new(/*v*/ 0), + }); + auth_manager + .set_external_auth(credentials.clone()) + .await + .unwrap(); + let server = MockServer::start().await; + Mock::given(method("GET")) + .respond_with(ResponseTemplate::new(/*s*/ 401)) + .expect(expected_requests) + .mount(&server) + .await; + let client = Client::new_without_redirects( + server.uri(), + HttpClientFactory::new(OutboundProxyPolicy::ReqwestDefault), + ) + .with_auth_provider(codex_model_provider::auth_provider_from_auth_manager( + Arc::clone(&auth_manager), + &auth, + )); + let session = AnalyticsSession { + account: AnalyticsAccount { + id: "account-a".into(), + email: None, + plan_type: auth.account_plan_type(), + }, + client, + auth_manager, + auth, + }; + let error = session + .request(|client| async move { + client + .get_account_analytics(crate::AnalyticsReport::Usage, "2026-09-01", "2026-09-07") + .await + }) + .await + .unwrap_err(); + assert_eq!(error.status().map(|status| status.as_u16()), Some(401)); + session.ensure_identity().await.unwrap(); + assert_eq!(credentials.refreshes.load(Ordering::SeqCst), 1); + server.verify().await; +} diff --git a/codex-rs/backend-client/src/client.rs b/codex-rs/backend-client/src/client.rs index aeda1920cd..c60f4dd3d8 100644 --- a/codex-rs/backend-client/src/client.rs +++ b/codex-rs/backend-client/src/client.rs @@ -339,14 +339,17 @@ impl Client { Ok(self.get_rate_limits_with_reset_credits().await?.rate_limits) } - pub async fn get_accounts_check(&self) -> Result { + pub async fn get_accounts_check( + &self, + ) -> std::result::Result { let url = match self.path_style { PathStyle::CodexApi => format!("{}/api/codex/accounts/check", self.base_url), PathStyle::ChatGptApi => format!("{}/wham/accounts/check", self.base_url), }; let req = self.request(Method::GET, &url).headers(self.headers()); - let (body, ct) = self.exec_request(req, "GET", &url).await?; - self.decode_json(&url, &ct, &body) + let (body, _) = self.exec_request_detailed(req, "GET", &url).await?; + serde_json::from_str(&body) + .map_err(|_| RequestError::Other(anyhow::anyhow!("Invalid accounts response."))) } pub async fn get_token_usage_profile(&self) -> Result { diff --git a/codex-rs/backend-client/src/types.rs b/codex-rs/backend-client/src/types.rs index 6b45c6e20c..925716d67e 100644 --- a/codex-rs/backend-client/src/types.rs +++ b/codex-rs/backend-client/src/types.rs @@ -136,6 +136,9 @@ pub struct AccountsCheckResponse { #[derive(Clone, Debug, Deserialize)] pub struct AccountEntry { pub id: String, + /// Current subscription reported by the accounts endpoint, independent of token claims. + #[serde(default)] + pub plan_type: Option, pub workspace_backend_origin: Option, pub account_routing_override: Option, #[serde(default)] @@ -178,6 +181,8 @@ struct ChatGptAccountEntry { struct ChatGptAccountInfo { account_id: Option, #[serde(default)] + plan_type: Option, + #[serde(default)] name: Option, #[serde(default)] profile_picture_url: Option, @@ -200,6 +205,7 @@ impl<'de> Deserialize<'de> for AccountsCheckResponse { let account = accounts.remove(account_id)?.account; Some(AccountEntry { id: account.account_id?, + plan_type: account.plan_type, workspace_backend_origin: None, account_routing_override: None, name: account.name, diff --git a/codex-rs/tui/src/analytics.rs b/codex-rs/tui/src/analytics.rs index 8a53f2b6aa..f6dd08c536 100644 --- a/codex-rs/tui/src/analytics.rs +++ b/codex-rs/tui/src/analytics.rs @@ -4,6 +4,7 @@ mod client; mod data; mod models; mod normalize; +mod render; mod report_data; mod tokens; diff --git a/codex-rs/tui/src/analytics/account_plan_tests.rs b/codex-rs/tui/src/analytics/account_plan_tests.rs new file mode 100644 index 0000000000..8795e2e966 --- /dev/null +++ b/codex-rs/tui/src/analytics/account_plan_tests.rs @@ -0,0 +1,167 @@ +//! Server plan discovery survives stale token claims and remains bound to the local identity. + +use super::tests::live; +use super::tests::sign_in; +use super::*; +use pretty_assertions::assert_eq; +use serde_json::json; +use std::sync::atomic::AtomicUsize; +use std::sync::atomic::Ordering; +use wiremock::Mock; +use wiremock::MockServer; +use wiremock::ResponseTemplate; +use wiremock::matchers::header; +use wiremock::matchers::method; +use wiremock::matchers::path; + +#[tokio::test] +async fn refresh_uses_server_plan_and_routes_with_unchanged_local_credentials() { + let server = MockServer::start().await; + let (_home, client) = live(&server, "plus").await; + let reads = Arc::new(AtomicUsize::new(/*v*/ 0)); + let plan_reads = Arc::clone(&reads); + Mock::given(method("GET")) + .and(path("/backend-api/wham/accounts/check")) + .and(header("chatgpt-account-id", "account-a")) + .respond_with(move |_: &wiremock::Request| { + let plan = if plan_reads.fetch_add(/*val*/ 1, Ordering::SeqCst) == 0 { + "team" + } else { + "enterprise" + }; + ResponseTemplate::new(/*s*/ 200).set_body_json(json!({ + "accounts": [ + {"id": "other-account", "plan_type": "free"}, + {"id": "account-a", "plan_type": plan} + ], + "account_ordering": ["other-account", "account-a"] + })) + }) + .with_priority(/*p*/ 1) + .expect(/*r*/ 2) + .mount(&server) + .await; + for endpoint in [ + "/backend-api/wham/usage/daily-workspace-user-token-usage-breakdown", + "/backend-api/wham/usage/daily-workspace-user-credit-usage", + ] { + Mock::given(method("GET")) + .and(path(endpoint)) + .and(header("chatgpt-account-id", "account-a")) + .respond_with( + ResponseTemplate::new(/*s*/ 200) + .set_body_json(json!({"breakdown": "product", "series": [], "data": []})), + ) + .expect(/*r*/ 1) + .mount(&server) + .await; + } + let first = client + .history(Report::Credits, /*days*/ 7, Grouping::Surface) + .await + .unwrap(); + assert_eq!(client.session().await.unwrap().kind, AccountKind::Business); + assert_eq!( + client + .history(Report::Credits, /*days*/ 7, Grouping::Surface) + .await + .unwrap(), + first + ); + assert_eq!(reads.load(Ordering::SeqCst), 1); + let refreshed = Live::new(Arc::clone(&client.config), client.end_date); + refreshed + .history(Report::Credits, /*days*/ 7, Grouping::Surface) + .await + .unwrap(); + assert_eq!( + refreshed.session().await.unwrap().kind, + AccountKind::Enterprise + ); + server.verify().await; +} + +#[tokio::test] +async fn account_plan_lookup_recovers_from_unauthorized() { + let server = MockServer::start().await; + let (_home, client) = live(&server, "plus").await; + let reads = AtomicUsize::new(/*v*/ 0); + Mock::given(method("GET")) + .and(path("/backend-api/wham/accounts/check")) + .and(header("chatgpt-account-id", "account-a")) + .respond_with(move |_: &wiremock::Request| { + if reads.fetch_add(/*val*/ 1, Ordering::SeqCst) == 0 { + ResponseTemplate::new(/*s*/ 401) + } else { + ResponseTemplate::new(/*s*/ 200).set_body_json(json!({ + "accounts": {"account-a": {"account": { + "account_id": "account-a", "plan_type": "enterprise" + }}}, + "account_ordering": ["account-a"] + })) + } + }) + .with_priority(/*p*/ 1) + .expect(/*r*/ 2) + .mount(&server) + .await; + assert_eq!( + client.session().await.unwrap().kind, + AccountKind::Enterprise + ); + server.verify().await; +} + +#[tokio::test] +async fn failed_plan_discovery_can_retry_without_using_token_plan() { + for response in [ + ResponseTemplate::new(/*s*/ 503), + ResponseTemplate::new(/*s*/ 200).set_body_string("invalid account response"), + ResponseTemplate::new(/*s*/ 200).set_body_json(json!({"accounts": [{"id": "account-a"}]})), + ResponseTemplate::new(/*s*/ 200).set_body_json(json!({ + "accounts": [{"id": "other-account", "plan_type": "enterprise"}] + })), + ] { + let server = MockServer::start().await; + let (_home, client) = live(&server, "plus").await; + Mock::given(method("GET")) + .and(path("/backend-api/wham/accounts/check")) + .respond_with(response) + .with_priority(/*p*/ 1) + .expect(/*r*/ 1) + .up_to_n_times(/*n*/ 1) + .mount(&server) + .await; + assert_eq!( + client.session().await.err().unwrap(), + "Couldn't load account plan. Press R to retry Analytics." + ); + assert_eq!(client.session().await.unwrap().kind, AccountKind::Consumer); + server.verify().await; + } +} + +#[tokio::test] +async fn account_plan_lookup_rejects_account_and_user_switches() { + for (account, user) in [("account-b", "user-a"), ("account-a", "user-b")] { + let server = MockServer::start().await; + let (home, client) = live(&server, "plus").await; + Mock::given(method("GET")) + .and(path("/backend-api/wham/accounts/check")) + .respond_with(move |_: &wiremock::Request| { + sign_in(home.path(), account, user, "plus"); + ResponseTemplate::new(/*s*/ 200).set_body_json(json!({ + "accounts": [{"id": "account-a", "plan_type": "enterprise"}] + })) + }) + .with_priority(/*p*/ 1) + .expect(/*r*/ 1) + .mount(&server) + .await; + assert_eq!( + client.session().await.err().unwrap(), + "Account changed. Press R to refresh Analytics." + ); + server.verify().await; + } +} diff --git a/codex-rs/tui/src/analytics/client.rs b/codex-rs/tui/src/analytics/client.rs index 48cd7b2f7f..084f82244c 100644 --- a/codex-rs/tui/src/analytics/client.rs +++ b/codex-rs/tui/src/analytics/client.rs @@ -1,29 +1,44 @@ -//! Local account sessions and display metadata for authenticated analytics. -//! The backend client owns credentials, recovery, and request identity checks. +//! Local account report loading with backend-owned authentication and identity checks. +//! Cached responses belong to one local account and user, independently of the connected server. +//! Every report uses the view's captured end date until refresh replaces the session. use super::models::AccountAnalyticsGrouping as Grouping; +use super::models::AccountAnalyticsReport as Report; use super::models::AccountKind; +use super::report_data::AnalyticsData; use crate::legacy_core::config::Config; +use codex_backend_client::AnalyticsReport; use codex_backend_client::AnalyticsSession; +use codex_backend_client::RequestError; +use codex_protocol::account::PlanType; +use std::collections::HashMap; use std::sync::Arc; +use tokio::sync::Mutex; use tokio::sync::OnceCell; pub(super) struct Live { config: Arc, + end_date: chrono::NaiveDate, session: OnceCell, + token_models: std::sync::RwLock>, + attributed_usage: std::sync::atomic::AtomicBool, } pub(super) struct Session { pub(super) kind: AccountKind, pub(super) backend: AnalyticsSession, credit_groups: Vec, + cache: Mutex>, } impl Live { - pub(super) fn new(config: Arc) -> Self { + pub(super) fn new(config: Arc, end_date: chrono::NaiveDate) -> Self { Self { config, + end_date, session: OnceCell::new(), + token_models: std::sync::RwLock::new(Vec::new()), + attributed_usage: std::sync::atomic::AtomicBool::new(/*v*/ true), } } @@ -42,6 +57,18 @@ impl Live { .map_or(&[0], |session| &session.credit_groups) } + pub(super) fn attributed_usage(&self) -> bool { + self.attributed_usage + .load(std::sync::atomic::Ordering::Relaxed) + } + + pub(super) fn token_models(&self) -> Vec { + self.token_models + .read() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .clone() + } + pub(super) async fn session(&self) -> Result<&Session, String> { self.session .get_or_try_init(|| async { @@ -62,12 +89,162 @@ impl Live { kind: AccountKind::from(session.account().plan_type), credit_groups, backend: session, + cache: Mutex::new(HashMap::new()), }) }) .await } + + pub(super) async fn history( + &self, + report: Report, + days: u32, + grouping: Grouping, + ) -> Result, String> { + self.filtered_history(report, days, grouping, /*model_filter*/ None) + .await + } + + pub(super) async fn filtered_history( + &self, + report: Report, + days: u32, + grouping: Grouping, + model_filter: Option<&str>, + ) -> Result, String> { + let end = self.end_date; + let start = days + .checked_sub(/*rhs*/ 1) + .and_then(|offset| end.checked_sub_days(chrono::Days::new(u64::from(offset)))) + .ok_or("Invalid analytics date range.")?; + let session = self.session().await?; + session.backend.ensure_identity().await?; + let enterprise_tokens = report == Report::Usage + && matches!( + session.kind, + AccountKind::Enterprise | AccountKind::Business + ); + let grouping = + if enterprise_tokens && !matches!(grouping, Grouping::Model | Grouping::TokenType) { + Grouping::TokenType + } else { + grouping + }; + let route = + route(report, grouping, session.backend.account().plan_type).ok_or_else(|| { + "This credit breakdown is not supported for this account type.".to_string() + })?; + // Credit events have no range parameters. Other grouping changes reuse the same payload. + let key = if route == AnalyticsReport::Credits { + (route, String::new(), String::new()) + } else { + (route, start.to_string(), end.to_string()) + }; + let cached = session.cache.lock().await.get(&key).cloned(); + let response = if let Some(response) = cached { + Ok(response) + } else { + session + .backend + .request(|client| { + let (_, start, end) = key.clone(); + async move { client.get_account_analytics(route, &start, &end).await } + }) + .await + .map(AnalyticsData::from) + }; + session.backend.ensure_identity().await?; + let response = response.map_err(request_error)?; + let history = (|| { + if enterprise_tokens { + let models = super::tokens::history(response.clone(), Grouping::Model, start, end)?; + *self + .token_models + .write() + .unwrap_or_else(std::sync::PoisonError::into_inner) = + super::render::categories(&models) + .into_iter() + .map(|model| model.key) + .collect(); + super::tokens::filtered_history( + response.clone(), + grouping, + start, + end, + model_filter, + ) + .map(Some) + } else { + let grouping = if report == Report::Usage { + let attributed = + super::normalize::has_complete_attribution(&response, start, end)?; + self.attributed_usage + .store(attributed, std::sync::atomic::Ordering::Relaxed); + if !attributed && matches!(grouping, Grouping::Feature | Grouping::TaskStart) { + Grouping::Surface + } else { + grouping + } + } else { + grouping + }; + super::normalize::history(response.clone(), report, grouping, start, end) + } + })(); + let mut cache = session.cache.lock().await; + if history.is_ok() { + cache.insert(key, response); + } else { + cache.remove(&key); + } + history + } +} + +// Account type selects billing semantics, never activity-report eligibility. +fn route(report: Report, grouping: Grouping, plan: Option) -> Option { + let kind = AccountKind::from(plan); + let enterprise = kind == AccountKind::Enterprise; + Some(match report { + Report::Usage if enterprise => AnalyticsReport::EnterpriseTokens, + Report::Usage if kind == AccountKind::Business => AnalyticsReport::WorkspaceCredits, + Report::Usage => AnalyticsReport::Usage, + Report::Messages => AnalyticsReport::Messages, + Report::Credits if !Grouping::credit_groupings(plan).contains(&grouping) => return None, + Report::Credits if enterprise => AnalyticsReport::EnterpriseCredits { + breakdown: match grouping { + Grouping::Surface => "product", + Grouping::Model => "model", + Grouping::Speed => "speed", + Grouping::Reasoning => "reasoning_effort", + Grouping::Feature | Grouping::TaskStart | Grouping::TokenType => return None, + }, + }, + Report::Credits if kind == AccountKind::Business => AnalyticsReport::WorkspaceCredits, + Report::Credits => AnalyticsReport::Credits, + Report::Plugins => AnalyticsReport::Plugins { + limit: if enterprise { 8 } else { 10 }, + }, + Report::Skills => AnalyticsReport::Skills { + limit: if enterprise { 6 } else { 10 }, + }, + }) +} + +pub(super) fn request_error(error: RequestError) -> String { + match error.status().map(|status| status.as_u16()) { + Some(401) => "Sign in again to load this report.", + Some(403) => "Access denied for this report.", + Some(404) => "This report endpoint is unavailable. Press R to retry.", + _ => "Report request failed. Press R to retry.", + } + .into() } #[cfg(test)] #[path = "client_tests.rs"] pub(super) mod tests; + +#[cfg(test)] +#[path = "account_plan_tests.rs"] +mod account_plan_tests; diff --git a/codex-rs/tui/src/analytics/client_tests.rs b/codex-rs/tui/src/analytics/client_tests.rs index ce5601b811..a85f510713 100644 --- a/codex-rs/tui/src/analytics/client_tests.rs +++ b/codex-rs/tui/src/analytics/client_tests.rs @@ -13,7 +13,7 @@ use wiremock::matchers::header; use wiremock::matchers::method; use wiremock::matchers::path; -fn sign_in(home: &std::path::Path, account: &str, user: &str, plan: &str) { +pub(super) fn sign_in(home: &std::path::Path, account: &str, user: &str, plan: &str) { let claims = base64::engine::general_purpose::URL_SAFE_NO_PAD.encode( json!({ "exp": 4102444800_i64, "email": "analytics@example.test", @@ -53,7 +53,29 @@ pub(in crate::analytics) async fn live( config.chatgpt_base_url = format!("{}/backend-api", server.uri()); config.cli_auth_credentials_store_mode = codex_login::AuthCredentialsStoreMode::File; sign_in(home.path(), "account-a", "user-a", plan); - (home, Live::new(Arc::new(config))) + Mock::given(method("GET")) + .and(path("/backend-api/wham/accounts/check")) + .respond_with(ResponseTemplate::new(/*s*/ 200).set_body_json(json!({ + "accounts": [{"id": "account-a", "plan_type": plan}], + "account_ordering": ["account-a"] + }))) + .with_priority(/*p*/ 2) + .mount(server) + .await; + ( + home, + Live::new(Arc::new(config), "2026-01-15".parse().unwrap()), + ) +} + +async fn report_requests(server: &MockServer) -> Vec { + server + .received_requests() + .await + .unwrap() + .into_iter() + .filter(|request| request.url.path() != "/backend-api/wham/accounts/check") + .collect() } #[tokio::test] @@ -72,7 +94,7 @@ async fn analytics_rejects_identity_changes_during_requests() { result.unwrap_err().to_string(), "Account changed. Press R to refresh Analytics." ); - assert!(server.received_requests().await.unwrap().is_empty()); + assert!(report_requests(&server).await.is_empty()); } #[tokio::test] @@ -112,7 +134,7 @@ async fn analytics_requires_local_chatgpt_authentication() { live.session().await.err(), Some("Sign in locally with ChatGPT to view Analytics.".to_string()) ); - assert!(server.received_requests().await.unwrap().is_empty()); + assert!(report_requests(&server).await.is_empty()); } } @@ -154,8 +176,8 @@ async fn analytics_requests_use_reloaded_credentials_for_the_same_identity() { assert_eq!( result, codex_backend_client::AnalyticsResponse::Usage( - codex_backend_client::analytics_models::DailyProductSurfaceUsageResponse::default(), - ), + codex_backend_client::analytics_models::DailyProductSurfaceUsageResponse::default() + ) ); server.verify().await; } @@ -203,8 +225,8 @@ async fn analytics_retries_unauthorized_requests_after_credentials_reload() { assert_eq!( result, codex_backend_client::AnalyticsResponse::Usage( - codex_backend_client::analytics_models::DailyProductSurfaceUsageResponse::default(), - ), + codex_backend_client::analytics_models::DailyProductSurfaceUsageResponse::default() + ) ); server.verify().await; } @@ -221,3 +243,523 @@ async fn analytics_retries_session_initialization_after_sign_in() { "account-a" ); } + +#[tokio::test] +async fn analytics_matches_app_account_routes_and_parameters() { + let server = MockServer::start().await; + for (plan, report, grouping, endpoint, extra) in [ + ( + "plus", + Report::Usage, + Grouping::Model, + "usage/daily-token-usage-breakdown", + vec![("group_by", "day")], + ), + ( + "business", + Report::Usage, + Grouping::TokenType, + "usage/daily-workspace-user-token-usage-breakdown", + vec![ + ("group_by", "day"), + ("breakdown_by", "model"), + ("modes", "codex"), + ("modes", "work"), + ], + ), + ( + "business", + Report::Messages, + Grouping::Model, + "analytics/daily-workspace-usage-counts", + vec![("group_by", "day"), ("workspace_user", "true")], + ), + ( + "unknown", + Report::Usage, + Grouping::Surface, + "usage/daily-token-usage-breakdown", + vec![("group_by", "day")], + ), + ( + "unknown", + Report::Messages, + Grouping::Model, + "analytics/daily-workspace-usage-counts", + vec![("group_by", "day"), ("workspace_user", "true")], + ), + ( + "unknown", + Report::Plugins, + Grouping::Feature, + "analytics/daily-plugin-usage-metrics", + vec![ + ("group_by", "day"), + ("workspace_user", "true"), + ("top_plugin_limit", "10"), + ], + ), + ( + "unknown", + Report::Skills, + Grouping::Feature, + "analytics/daily-skill-usage-metrics", + vec![ + ("group_by", "day"), + ("workspace_user", "true"), + ("top_skill_limit", "10"), + ], + ), + ( + "edu", + Report::Credits, + Grouping::Model, + "usage/daily-workspace-user-credit-usage", + vec![("breakdown", "model")], + ), + ( + "plus", + Report::Credits, + Grouping::Surface, + "usage/credit-usage-events", + vec![], + ), + ( + "team", + Report::Usage, + Grouping::Model, + "usage/daily-workspace-user-token-usage-breakdown", + vec![("group_by", "day")], + ), + ( + "team", + Report::Credits, + Grouping::Speed, + "usage/daily-workspace-user-token-usage-breakdown", + vec![("group_by", "day")], + ), + ( + "business", + Report::Credits, + Grouping::Reasoning, + "usage/daily-workspace-user-credit-usage", + vec![("breakdown", "reasoning_effort")], + ), + ( + "plus", + Report::Messages, + Grouping::Model, + "analytics/daily-workspace-usage-counts", + vec![("group_by", "day"), ("workspace_user", "true")], + ), + ( + "plus", + Report::Plugins, + Grouping::Feature, + "analytics/daily-plugin-usage-metrics", + vec![ + ("group_by", "day"), + ("workspace_user", "true"), + ("top_plugin_limit", "10"), + ], + ), + ( + "business", + Report::Plugins, + Grouping::Feature, + "analytics/daily-plugin-usage-metrics", + vec![ + ("group_by", "day"), + ("workspace_user", "true"), + ("top_plugin_limit", "8"), + ], + ), + ( + "business", + Report::Skills, + Grouping::Feature, + "analytics/daily-skill-usage-metrics", + vec![ + ("group_by", "day"), + ("workspace_user", "true"), + ("top_skill_limit", "6"), + ], + ), + ] { + server.reset().await; + let (_home, live) = live(&server, plan).await; + let today = live.end_date; + let start = (today - chrono::Days::new(/*num*/ 6)).to_string(); + let response = if endpoint == "usage/daily-workspace-user-credit-usage" { + let breakdown = extra.iter().find(|(key, _)| *key == "breakdown").unwrap().1; + json!({"data": [], "series": [], "breakdown": breakdown}) + } else { + json!({"data": []}) + }; + Mock::given(method("GET")) + .and(path(format!("/backend-api/wham/{endpoint}"))) + .and(header("chatgpt-account-id", "account-a")) + .respond_with(ResponseTemplate::new(/*s*/ 200).set_body_json(response)) + .expect(/*r*/ 1) + .mount(&server) + .await; + assert!( + live.history(report, /*days*/ 7, grouping) + .await + .unwrap() + .is_some() + ); + let requests = report_requests(&server).await; + let mut actual = requests[0] + .url + .query_pairs() + .map(|(key, value)| (key.into_owned(), value.into_owned())) + .collect::>(); + let mut expected = extra + .into_iter() + .map(|(key, value)| (key.to_string(), value.to_string())) + .collect::>(); + if endpoint != "usage/credit-usage-events" { + expected.push(("start_date".into(), start.clone())); + expected.push(("end_date".into(), today.to_string())); + } + actual.sort(); + expected.sort(); + assert_eq!(actual, expected); + assert!(requests[0].headers.get("authorization").is_some()); + } +} + +#[tokio::test] +async fn analytics_uses_api_response_for_usage_and_turns_availability() { + let server = MockServer::start().await; + for report in [Report::Usage, Report::Messages] { + for status in [403, 404] { + server.reset().await; + let (_home, live) = live(&server, "business").await; + Mock::given(method("GET")) + .respond_with(ResponseTemplate::new(status)) + .expect(/*r*/ 1) + .mount(&server) + .await; + assert_eq!( + live.history(report, /*days*/ 7, Grouping::Surface) + .await + .unwrap_err(), + if status == 403 { + "Access denied for this report." + } else { + "This report endpoint is unavailable. Press R to retry." + } + ); + server.verify().await; + } + } +} + +#[tokio::test] +async fn analytics_reuses_payload_for_grouping_and_rejects_account_changes() { + let server = MockServer::start().await; + let (home, live) = live(&server, "team").await; + let today = live.end_date.to_string(); + Mock::given(method("GET")) + .and(path( + "/backend-api/wham/usage/daily-workspace-user-token-usage-breakdown", + )) + .respond_with( + ResponseTemplate::new(/*s*/ 200).set_body_json(json!({"data": [{ + "date": today, "product_surface_usage_values": {}, "premium_usage_values": { + "credit_usage_credits": {"cli": 12.5}, "total_usage_credits": {}, + "uncached_text_input_tokens_by_surface": {}, "cached_text_input_tokens_by_surface": {}, + "text_output_tokens_by_surface": {}, "text_total_tokens_by_surface": {} + }, + "models": [{"model": "model-a", "credits": 12.5, "speed": "fast", "on_demand_credits": 12.5}] + }]})), + ) + .expect(/*r*/ 1) + .mount(&server) + .await; + let surface = live + .history(Report::Credits, /*days*/ 7, Grouping::Surface) + .await + .unwrap() + .unwrap(); + let model = live + .history(Report::Credits, /*days*/ 7, Grouping::Model) + .await + .unwrap() + .unwrap(); + assert_eq!( + ( + surface.data.last().unwrap().total, + model.data.last().unwrap().total + ), + (12.5, 12.5) + ); + sign_in(home.path(), "account-a", "user-b", "team"); + assert_eq!( + live.history(Report::Credits, /*days*/ 7, Grouping::Model) + .await + .unwrap_err(), + "Account changed. Press R to refresh Analytics." + ); +} + +#[tokio::test] +async fn token_model_filter_reuses_the_account_scoped_response() { + let server = MockServer::start().await; + let (_home, live) = live(&server, "business").await; + let today = live.end_date.to_string(); + Mock::given(method("GET")) + .and(path("/backend-api/wham/usage/daily-workspace-user-token-usage-breakdown")) + .respond_with(ResponseTemplate::new(/*s*/ 200).set_body_json(json!({"data":[{ + "date":today, "product_surface_usage_values": {}, "groups":[ + {"dimensions":{"model":"alpha"}, "credits": 0, "uncached_text_input_tokens":10,"cached_text_input_tokens":20,"text_output_tokens":30}, + {"dimensions":{"model":"beta"}, "credits": 0, "uncached_text_input_tokens":100,"cached_text_input_tokens":200,"text_output_tokens":300} + ] + }]}))) + .expect(/*r*/ 1).mount(&server).await; + let all = live + .history(Report::Usage, /*days*/ 7, Grouping::TokenType) + .await + .unwrap() + .unwrap(); + let alpha = live + .filtered_history( + Report::Usage, + /*days*/ 7, + Grouping::TokenType, + Some("alpha"), + ) + .await + .unwrap() + .unwrap(); + let beta = live + .filtered_history( + Report::Usage, + /*days*/ 7, + Grouping::TokenType, + Some("beta"), + ) + .await + .unwrap() + .unwrap(); + assert_eq!( + (all.data[0].total, alpha.data[0].total, beta.data[0].total), + (660.0, 60.0, 600.0) + ); + assert_eq!( + live.token_models(), + vec!["beta".to_string(), "alpha".to_string()] + ); + assert_eq!(report_requests(&server).await.len(), 1); +} + +#[tokio::test] +async fn legacy_usage_falls_back_to_surface_for_attribution_only_groupings() { + let server = MockServer::start().await; + let (_home, live) = live(&server, "plus").await; + let today = live.end_date.to_string(); + Mock::given(method("GET")) + .respond_with( + ResponseTemplate::new(/*s*/ 200).set_body_json(json!({"data": [{ + "date": today, "product_surface_usage_values": {"cli": 12.5} + }]})), + ) + .expect(/*r*/ 1) + .mount(&server) + .await; + let surface = live + .history(Report::Usage, /*days*/ 7, Grouping::Surface) + .await + .unwrap(); + for grouping in [Grouping::Feature, Grouping::TaskStart] { + let history = live + .history(Report::Usage, /*days*/ 7, grouping) + .await + .unwrap(); + assert_eq!((live.attributed_usage(), history), (false, surface.clone())); + } + server.verify().await; +} + +#[tokio::test] +async fn invalid_ranges_fail_without_requesting_reports() { + let server = MockServer::start().await; + let (_home, live) = live(&server, "plus").await; + for days in [0, u32::MAX] { + assert_eq!( + live.history(Report::Usage, days, Grouping::Surface) + .await + .unwrap_err(), + "Invalid analytics date range." + ); + } + assert!(report_requests(&server).await.is_empty()); +} + +#[tokio::test] +async fn unknown_plan_exposes_no_credit_breakdowns() { + let server = MockServer::start().await; + let (_home, live) = live(&server, "unknown").await; + live.session().await.unwrap(); + assert!(live.credit_groups().is_empty()); + assert!( + live.history(Report::Credits, /*days*/ 7, Grouping::Surface) + .await + .is_err() + ); + assert!(report_requests(&server).await.is_empty()); +} + +#[tokio::test] +async fn in_flight_account_switch_retains_refresh_guidance() { + let server = MockServer::start().await; + let (home, live) = live(&server, "plus").await; + let home_path = home.path().to_path_buf(); + Mock::given(method("GET")) + .respond_with(move |_: &wiremock::Request| { + sign_in(&home_path, "account-b", "user-a", "plus"); + ResponseTemplate::new(/*s*/ 200).set_body_json(json!({"data": []})) + }) + .expect(/*r*/ 1) + .mount(&server) + .await; + assert_eq!( + live.history(Report::Usage, /*days*/ 7, Grouping::Surface) + .await + .unwrap_err(), + "Account changed. Press R to refresh Analytics." + ); + server.verify().await; +} + +#[tokio::test] +async fn missing_breakdowns_remain_unavailable() { + let server = MockServer::start().await; + let (_home, live) = live(&server, "plus").await; + let today = live.end_date.to_string(); + Mock::given(method("GET")) + .respond_with( + ResponseTemplate::new(/*s*/ 200).set_body_json( + json!({"data": [{"date": today, "product_surface_usage_values": {}}]}), + ), + ) + .expect(/*r*/ 1) + .mount(&server) + .await; + assert_eq!( + live.history(Report::Usage, /*days*/ 7, Grouping::Model) + .await, + Ok(None) + ); + server.verify().await; +} + +#[tokio::test] +async fn out_of_range_legacy_rows_do_not_disable_attribution() { + let server = MockServer::start().await; + let (_home, live) = live(&server, "plus").await; + let today = live.end_date.to_string(); + Mock::given(method("GET")) + .respond_with( + ResponseTemplate::new(/*s*/ 200).set_body_json(json!({"data": [ + {"date": "2000-01-01", "product_surface_usage_values": {}}, + {"date": today, "product_surface_usage_values": {}, "attribution": [{"thread_source": "user", "turn_trigger": "composer", "model": "example", "surface": "cli", "value": 12.0}]} + ]})), + ) + .expect(/*r*/ 1) + .mount(&server) + .await; + let history = live + .history(Report::Usage, /*days*/ 7, Grouping::Feature) + .await + .unwrap() + .unwrap(); + assert_eq!( + history.data, + vec![super::super::models::AccountAnalyticsDay { + date: live.end_date, + total: 12.0, + values: vec![super::super::models::AccountAnalyticsValue { + key: "user".into(), + label: "Tasks".into(), + value: 12.0, + }], + }] + ); + assert!(live.attributed_usage()); +} + +#[tokio::test] +async fn invalid_reports_can_retry_without_recreating_the_session() { + for plan in ["plus", "business"] { + let server = MockServer::start().await; + let (_home, live) = live(&server, plan).await; + let requests = Arc::new(std::sync::atomic::AtomicUsize::new(/*v*/ 0)); + let count = Arc::clone(&requests); + Mock::given(method("GET")) + .respond_with(move |_: &wiremock::Request| { + let body = if count.fetch_add(/*val*/ 1, std::sync::atomic::Ordering::SeqCst) == 0 { + json!({"data": [{"date": "invalid"}]}) + } else { + json!({"data": []}) + }; + ResponseTemplate::new(/*s*/ 200).set_body_json(body) + }) + .expect(/*r*/ 2) + .mount(&server) + .await; + assert!( + live.history(Report::Usage, /*days*/ 7, Grouping::Surface) + .await + .is_err() + ); + assert!( + live.history(Report::Usage, /*days*/ 7, Grouping::Surface) + .await + .is_ok() + ); + assert_eq!(requests.load(std::sync::atomic::Ordering::SeqCst), 2); + server.verify().await; + } +} + +#[tokio::test] +async fn invalid_cached_breakdowns_are_evicted_before_retry() { + let server = MockServer::start().await; + let (_home, live) = live(&server, "plus").await; + let today = live.end_date.to_string(); + let attempts = Arc::new(std::sync::atomic::AtomicUsize::new(/*v*/ 0)); + let count = Arc::clone(&attempts); + Mock::given(method("GET")) + .respond_with(move |_: &wiremock::Request| { + let model = if count.fetch_add(/*val*/ 1, std::sync::atomic::Ordering::SeqCst) == 0 { + json!({"model": "alpha", "credits": -1}) + } else { + json!({"model": "alpha", "credits": 12}) + }; + ResponseTemplate::new(/*s*/ 200).set_body_json(json!({"data": [{ + "date": today, "product_surface_usage_values": {"cli": 12}, "models": [model] + }]})) + }) + .expect(/*r*/ 2) + .mount(&server) + .await; + live.history(Report::Usage, /*days*/ 7, Grouping::Surface) + .await + .unwrap(); + assert!( + live.history(Report::Usage, /*days*/ 7, Grouping::Model) + .await + .is_err() + ); + assert!( + live.history(Report::Usage, /*days*/ 7, Grouping::Model) + .await + .unwrap() + .is_some() + ); + assert_eq!(attempts.load(std::sync::atomic::Ordering::SeqCst), 2); + server.verify().await; +} diff --git a/codex-rs/tui/src/analytics/data.rs b/codex-rs/tui/src/analytics/data.rs index d72fe15837..d1dfd70551 100644 --- a/codex-rs/tui/src/analytics/data.rs +++ b/codex-rs/tui/src/analytics/data.rs @@ -1,6 +1,12 @@ -//! Grouping metadata and numeric formatting for account analytics. +//! Report loading and grouping metadata for the account dashboard. +//! Replacing a pending load aborts it; its response cannot populate a newer range or account. use super::models::AccountAnalyticsGrouping as Grouping; +use crate::tui::FrameRequester; +use futures::FutureExt; +use std::panic::AssertUnwindSafe; +use tokio::sync::oneshot; +use tokio::task::JoinHandle; pub(super) const GROUPINGS: [Grouping; 7] = [ Grouping::Surface, @@ -12,6 +18,96 @@ pub(super) const GROUPINGS: [Grouping; 7] = [ Grouping::TokenType, ]; +pub(super) const GROUP_LABELS: [&str; 7] = [ + "Surface", + "Feature", + "Model", + "Turn start", + "Speed", + "Reasoning", + "Token type", +]; + +pub(super) struct Pending { + receiver: oneshot::Receiver, String>>, + task: JoinHandle<()>, +} + +impl Drop for Pending { + fn drop(&mut self) { + self.task.abort(); + } +} + +pub(super) enum Load { + Loading(Pending), + Ready(T), + Unavailable, + Error(String), +} + +impl Load { + pub(super) fn start( + future: impl std::future::Future, String>> + Send + 'static, + frame: FrameRequester, + ) -> Self { + let (sender, receiver) = oneshot::channel(); + let task = tokio::spawn(async move { + let result = match tokio::time::timeout( + std::time::Duration::from_secs(/*secs*/ 30), + AssertUnwindSafe(future).catch_unwind(), + ) + .await + { + Ok(Ok(result)) => result, + Ok(Err(_)) => Err("Request interrupted. Press R to retry.".into()), + Err(_) => Err("Request timed out. Press R to retry.".into()), + }; + let _ = sender.send(result); + frame.schedule_frame(); + }); + Self::Loading(Pending { receiver, task }) + } + + pub(super) fn poll(&mut self) { + if let Self::Loading(pending) = self { + *self = match pending.receiver.try_recv() { + Ok(Ok(Some(data))) => Self::Ready(data), + Ok(Ok(None)) => Self::Unavailable, + Ok(Err(message)) => Self::Error(message), + Err(oneshot::error::TryRecvError::Closed) => { + Self::Error("Request interrupted. Press R to retry.".into()) + } + Err(oneshot::error::TryRecvError::Empty) => return, + }; + } + } + + pub(super) fn ready(&self) -> Option<&T> { + if let Self::Ready(value) = self { + Some(value) + } else { + None + } + } + + pub(super) fn message(&self) -> Option<&str> { + match self { + Self::Ready(_) => None, + Self::Loading(_) => Some("Loading…"), + Self::Unavailable => Some("No history has been reported."), + Self::Error(message) => Some(message), + } + } +} + +pub(super) fn error(error: codex_app_server_client::TypedRequestError) -> String { + match error { + codex_app_server_client::TypedRequestError::Server { source, .. } => source.message, + _ => "Couldn't load analytics. Press R to retry.".into(), + } +} + /// Keep tiny refunds visible while avoiding noise on ordinary credit amounts. pub(super) fn amount(value: f64) -> String { if value == 0.0 { @@ -89,3 +185,7 @@ pub(super) fn date(date: &str) -> String { .map(|date| date.format("%b %-d").to_string()) .unwrap_or_else(|_| date.to_string()) } + +#[cfg(test)] +#[path = "data_tests.rs"] +mod tests; diff --git a/codex-rs/tui/src/analytics/data_tests.rs b/codex-rs/tui/src/analytics/data_tests.rs new file mode 100644 index 0000000000..7d1995a060 --- /dev/null +++ b/codex-rs/tui/src/analytics/data_tests.rs @@ -0,0 +1,59 @@ +//! Loading transitions and cancellation when a view replaces pending work. + +use super::*; +use pretty_assertions::assert_eq; + +#[tokio::test] +async fn completed_loads_publish_data_unavailability_and_errors() { + for (result, expected) in [ + (Ok(Some(7)), (Some(7), None)), + (Ok(None), (None, Some("No history has been reported."))), + ( + Err("temporary failure".to_string()), + (None, Some("temporary failure")), + ), + ] { + let mut load = Load::start(async move { result }, FrameRequester::test_dummy()); + let Load::Loading(pending) = &mut load else { + unreachable!() + }; + (&mut pending.task).await.unwrap(); + load.poll(); + assert_eq!((load.ready().copied(), load.message()), expected); + } +} + +#[tokio::test] +async fn replacing_a_pending_load_cancels_its_request() { + let (cancelled_tx, cancelled_rx) = oneshot::channel::<()>(); + let load = Load::<()>::start( + async move { + let _cancelled = cancelled_tx; + std::future::pending().await + }, + FrameRequester::test_dummy(), + ); + tokio::task::yield_now().await; + drop(load); + assert!(cancelled_rx.await.is_err()); +} + +#[tokio::test] +async fn panicking_loads_request_a_frame_and_report_interruption() { + let (draw, mut frames) = tokio::sync::broadcast::channel(/*capacity*/ 1); + let frame = FrameRequester::new(draw); + let mut load = Load::<()>::start(async { panic!("report task failed") }, frame.clone()); + let Load::Loading(pending) = &mut load else { + unreachable!() + }; + let _ = (&mut pending.task).await; + tokio::time::timeout(std::time::Duration::from_secs(/*secs*/ 1), frames.recv()) + .await + .unwrap() + .unwrap(); + load.poll(); + assert_eq!( + load.message(), + Some("Request interrupted. Press R to retry.") + ); +} diff --git a/codex-rs/tui/src/analytics/models.rs b/codex-rs/tui/src/analytics/models.rs index ab78ee3311..63d6435a83 100644 --- a/codex-rs/tui/src/analytics/models.rs +++ b/codex-rs/tui/src/analytics/models.rs @@ -48,7 +48,8 @@ impl AccountAnalyticsGrouping { match AccountKind::from(plan) { AccountKind::Business => &[Self::Surface, Self::Model, Self::Speed], AccountKind::Enterprise => &[Self::Surface, Self::Model, Self::Speed, Self::Reasoning], - AccountKind::Consumer | AccountKind::Unknown => &[Self::Surface], + AccountKind::Consumer => &[Self::Surface], + AccountKind::Unknown => &[], } } } diff --git a/codex-rs/tui/src/analytics/normalize.rs b/codex-rs/tui/src/analytics/normalize.rs index 8d60d9b96d..26ae206ce6 100644 --- a/codex-rs/tui/src/analytics/normalize.rs +++ b/codex-rs/tui/src/analytics/normalize.rs @@ -8,6 +8,25 @@ use chrono::DateTime; use chrono::NaiveDate; use std::collections::BTreeMap; +/// Out-of-range legacy rows do not change attribution semantics for the requested period. +pub(super) fn has_complete_attribution( + response: &AnalyticsData, + start: NaiveDate, + end: NaiveDate, +) -> Result { + response + .data + .iter() + .try_fold(/*init*/ true, |complete, record| { + let date = NaiveDate::parse_from_str( + record.date.get(..10).unwrap_or(&record.date), + "%Y-%m-%d", + ) + .map_err(|_| String::from("Analytics returned an invalid date."))?; + Ok(complete && (date < start || date > end || record.attribution.is_some())) + }) +} + pub(super) fn history( response: AnalyticsData, report: Report, @@ -27,6 +46,7 @@ pub(super) fn history( } } } + let attributed = has_complete_attribution(&response, start, end)?; let series = response.series.as_ref(); let mut days: BTreeMap)> = BTreeMap::new(); for record in response.data { @@ -39,21 +59,26 @@ pub(super) fn history( let (total, values) = days.entry(date).or_default(); match report { Report::Usage => { - let daily = match grouping { - Grouping::Surface => record.product_surface_usage_values, - Grouping::Model => record - .models - .map(|models| -> Result<_, String> { - let mut values = BTreeMap::new(); - for model in models { - let amount = - model.credits.ok_or("Plan usage amount was not reported.")?; - *values.entry(model.model).or_insert(/*default*/ 0.0) += amount; - } - Ok(values) - }) - .transpose()?, - _ => None, + let daily = if attributed { + None + } else { + match grouping { + Grouping::Surface => record.product_surface_usage_values, + Grouping::Model => record + .models + .map(|models| -> Result<_, String> { + let mut values = BTreeMap::new(); + for model in models { + let amount = model + .credits + .ok_or("Plan usage amount was not reported.")?; + *values.entry(model.model).or_insert(/*default*/ 0.0) += amount; + } + Ok(values) + }) + .transpose()?, + _ => None, + } }; if let Some(daily) = daily { for (key, amount) in daily { @@ -75,17 +100,12 @@ pub(super) fn history( *total += amount; *values.entry(key).or_default() += amount; } - } else if let Some(attribution) = record.attribution { + } else if let Some(attribution) = record.attribution.filter(|_| attributed) { for entry in attribution { if !entry.value.is_finite() || entry.value < 0.0 { return Err(String::from("Analytics returned an invalid amount.")); } *total += entry.value; - if grouping == Grouping::TaskStart - && entry.thread_source.as_deref() != Some("user") - { - continue; - } let key = match grouping { Grouping::Feature => entry.thread_source, Grouping::Model => entry.model, @@ -224,6 +244,7 @@ pub(super) fn history( } } if report == Report::Usage + && !attributed && grouping == Grouping::Model && response.units.as_deref() != Some("credits") { @@ -257,7 +278,7 @@ pub(super) fn history( } Ok(Some(AccountAnalyticsHistory { unit: match report { - Report::Usage if response.units.as_deref() == Some("credits") => { + Report::Usage if !attributed && response.units.as_deref() == Some("credits") => { AccountAnalyticsUnit::Credits } Report::Usage => AccountAnalyticsUnit::RelativeUsage, diff --git a/codex-rs/tui/src/analytics/normalize_tests.rs b/codex-rs/tui/src/analytics/normalize_tests.rs index 7191298436..140166f9e9 100644 --- a/codex-rs/tui/src/analytics/normalize_tests.rs +++ b/codex-rs/tui/src/analytics/normalize_tests.rs @@ -5,7 +5,7 @@ use pretty_assertions::assert_eq; use serde_json::json; #[test] -fn task_start_keeps_the_full_denominator_and_missing_attribution_is_unavailable() { +fn turn_start_includes_all_features_and_missing_attribution_is_unavailable() { let date = NaiveDate::from_ymd_opt(/*year*/ 2026, /*month*/ 1, /*day*/ 9).unwrap(); let response = serde_json::from_value(json!({"data": [{"date": "2026-01-09", "attribution": [ {"thread_source": "user", "turn_trigger": "composer", "value": 30}, @@ -20,11 +20,18 @@ fn task_start_keeps_the_full_denominator_and_missing_attribution_is_unavailable( data: vec![AccountAnalyticsDay { date: "2026-01-09".parse().unwrap(), total: 100.0, - values: vec![AccountAnalyticsValue { - key: "start-composer".to_string(), - label: "User messages".to_string(), - value: 30.0 - }], + values: vec![ + AccountAnalyticsValue { + key: "start-composer".to_string(), + label: "User messages".to_string(), + value: 30.0 + }, + AccountAnalyticsValue { + key: "unknown".to_string(), + label: "Unknown".to_string(), + value: 70.0 + } + ], }], }) ); @@ -190,13 +197,12 @@ fn count_reports_preserve_explicit_zero_and_omit_unreported_days() { } #[test] -fn consumer_plan_prefers_daily_products_and_models_over_attribution() { +fn legacy_consumer_history_keeps_daily_products_and_models() { let date = "2026-09-01".parse().unwrap(); let response: AnalyticsData = serde_json::from_value(json!({"units":"relative", "data":[{ "date":"2026-09-01", "product_surface_usage_values":{"work_desktop":12,"work_web":8,"cli":30,"desktop_app":50}, - "models":[{"model":"major","credits":98.5},{"model":"one-percent","credits":1},{"model":"tiny","credits":0.3},{"model":"OTHER","credits":0.2}], - "attribution":[{"surface":"wrong","model":"wrong","value":999}] + "models":[{"model":"major","credits":98.5},{"model":"one-percent","credits":1},{"model":"tiny","credits":0.3},{"model":"OTHER","credits":0.2}] }]})).unwrap(); for (grouping, expected) in [ ( @@ -243,6 +249,52 @@ fn consumer_plan_prefers_daily_products_and_models_over_attribution() { ); } +#[test] +fn complete_attribution_wins_consistently_and_incomplete_attribution_is_not_mixed() { + let date = "2026-09-01".parse().unwrap(); + let response: AnalyticsData = serde_json::from_value(json!({"units":"credits", "data":[{ + "date":"2026-09-01", "product_surface_usage_values":{"cli":999}, + "models":[{"model":"legacy","credits":999}], + "attribution":[ + {"thread_source":"user","turn_trigger":"composer","surface":"desktop_app","model":"alpha","value":30}, + {"thread_source":"subagent","turn_trigger":"goal","surface":"cli","model":"tiny","value":0.1} + ] + }]})).unwrap(); + for grouping in [ + Grouping::Feature, + Grouping::Model, + Grouping::Surface, + Grouping::TaskStart, + ] { + let normalized = history(response.clone(), Report::Usage, grouping, date, date) + .unwrap() + .unwrap(); + assert_eq!( + ( + normalized.unit, + normalized.data[0].total, + normalized.data[0].values.len() + ), + (AccountAnalyticsUnit::RelativeUsage, 30.1, 2) + ); + } + let incomplete = serde_json::from_value(json!({"data":[ + {"date":"2026-09-01", "attribution":[]}, {"date":"2026-09-02"} + ]})) + .unwrap(); + assert_eq!( + history( + incomplete, + Report::Usage, + Grouping::Feature, + date, + "2026-09-02".parse().unwrap() + ) + .unwrap(), + None + ); +} + #[test] fn duplicate_message_dates_preserve_each_records_other_remainder() { let date = NaiveDate::from_ymd_opt(/*year*/ 2026, /*month*/ 1, /*day*/ 9).unwrap(); diff --git a/codex-rs/tui/src/analytics/render.rs b/codex-rs/tui/src/analytics/render.rs new file mode 100644 index 0000000000..998a181570 --- /dev/null +++ b/codex-rs/tui/src/analytics/render.rs @@ -0,0 +1,26 @@ +//! Shared report category aggregation for token model filtering. + +use super::models::AccountAnalyticsHistory; +use super::models::AccountAnalyticsValue; +use std::collections::BTreeMap; + +pub(super) fn categories(history: &AccountAnalyticsHistory) -> Vec { + let mut totals = BTreeMap::::new(); + for value in history.data.iter().flat_map(|day| &day.values) { + totals + .entry(value.key.clone()) + .or_insert_with(|| AccountAnalyticsValue { + value: 0.0, + ..value.clone() + }) + .value += value.value; + } + let mut values = totals.into_values().collect::>(); + values.sort_by(|a, b| { + b.value + .abs() + .total_cmp(&a.value.abs()) + .then_with(|| a.key.cmp(&b.key)) + }); + values +}