diff --git a/apps/desktop-tauri/src-tauri/src/commands/codex_accounts.rs b/apps/desktop-tauri/src-tauri/src/commands/codex_accounts.rs index a0c1dbca34..aaa3ac5534 100644 --- a/apps/desktop-tauri/src-tauri/src/commands/codex_accounts.rs +++ b/apps/desktop-tauri/src-tauri/src/commands/codex_accounts.rs @@ -105,13 +105,19 @@ pub(crate) async fn refresh_codex_account_lanes( let api = CodexAccountApi::new(); let home_path = account.codex_home_path.clone(); let email_hint = account.email_hint.clone(); + let workspace_account_id = account.effective_workspace_account_id(); match tokio::time::timeout( std::time::Duration::from_secs(DEFAULT_FETCH_TIMEOUT_SECONDS), - api.fetch_snapshot(&home_path, email_hint.as_deref(), true), + api.fetch_snapshot_for_workspace( + &home_path, + email_hint.as_deref(), + workspace_account_id.as_deref(), + true, + ), ) .await { - Ok(Ok(snapshot)) => Some((account.id, snapshot)), + Ok(Ok(snapshot)) => Some((account, snapshot)), Ok(Err(e)) => { tracing::debug!( "codex account lane {} failed: {}", @@ -128,10 +134,28 @@ pub(crate) async fn refresh_codex_account_lanes( })); } + let current_accounts = load_codex_accounts().unwrap_or_default(); + let current_by_id: HashMap = current_accounts + .iter() + .cloned() + .map(|account| (account.id, account)) + .collect(); let mut snapshots = SnapshotStore::new().load().unwrap_or_default(); + snapshots.retain(|id, snapshot| { + current_by_id + .get(id) + .is_some_and(|account| account_snapshot_belongs_to(account, snapshot)) + }); for handle in handles { - if let Ok(Some((id, snapshot))) = handle.await { - snapshots.insert(id, snapshot); + if let Ok(Some((fetched_account, snapshot))) = handle.await + && current_by_id + .get(&fetched_account.id) + .is_some_and(|current| { + account_lane_is_current(&fetched_account, current) + && account_snapshot_belongs_to(current, &snapshot) + }) + { + snapshots.insert(fetched_account.id, snapshot); } } if let Err(e) = SnapshotStore::new().save(&snapshots) { @@ -232,16 +256,30 @@ pub async fn codex_account_fetch( let api = CodexAccountApi::new(); let home_path = target.codex_home_path.clone(); let email_hint = target.email_hint.clone(); + let workspace_account_id = target.effective_workspace_account_id(); let snapshot = tokio::time::timeout( std::time::Duration::from_secs(DEFAULT_FETCH_TIMEOUT_SECONDS), - api.fetch_snapshot(&home_path, email_hint.as_deref(), true), + api.fetch_snapshot_for_workspace( + &home_path, + email_hint.as_deref(), + workspace_account_id.as_deref(), + true, + ), ) .await .map_err(|_| "Timed out waiting for the Codex usage API.".to_string())? .map_err(into_api_message)?; // Persist snapshot to the snapshot store, keyed by account id. - if let Ok(mut snapshots) = SnapshotStore::new().load() { + if let Ok(mut snapshots) = SnapshotStore::new().load() + && load_codex_accounts().ok().is_some_and(|accounts| { + accounts.iter().any(|account| { + account.id == target.id + && account_lane_is_current(&target, account) + && account_snapshot_belongs_to(account, &snapshot) + }) + }) + { snapshots.insert(target.id, snapshot.clone()); let _ = SnapshotStore::new().save(&snapshots); } @@ -306,6 +344,58 @@ fn into_api_message(error: CodexApiError) -> String { } } +fn account_home_key(account: &CodexAccount) -> String { + std::path::absolute(&account.codex_home_path) + .unwrap_or_else(|_| account.codex_home_path.clone()) + .to_string_lossy() + .to_lowercase() +} + +/// In-flight results are only authoritative for the selected workspace and +/// managed home that started the request. +fn account_lane_is_current(started: &CodexAccount, current: &CodexAccount) -> bool { + started.id == current.id + && account_home_key(started) == account_home_key(current) + && started.effective_workspace_account_id() == current.effective_workspace_account_id() +} + +fn account_snapshot_belongs_to( + account: &CodexAccount, + snapshot: &codexbar::codex_accounts::AccountUsageSnapshot, +) -> bool { + let snapshot_workspace = snapshot + .provider_account_id + .as_deref() + .map(str::trim) + .filter(|id| !id.is_empty()) + .map(str::to_lowercase); + match (account.effective_workspace_account_id(), snapshot_workspace) { + (Some(account_workspace), Some(snapshot_workspace)) => { + account_workspace == snapshot_workspace + } + (None, None) => true, + _ => false, + } +} + +fn snapshots_for_accounts( + accounts: &[CodexAccount], + snapshots: HashMap, +) -> HashMap { + let accounts_by_id: HashMap = accounts + .iter() + .map(|account| (account.id, account)) + .collect(); + snapshots + .into_iter() + .filter(|(id, snapshot)| { + accounts_by_id + .get(id) + .is_some_and(|account| account_snapshot_belongs_to(account, snapshot)) + }) + .collect() +} + #[derive(Debug, serde::Serialize)] #[serde(rename_all = "camelCase")] pub struct CodexAccountsStateBridge { @@ -320,10 +410,12 @@ pub fn get_codex_accounts_state( ) -> Result { let _guard = state.lock().map_err(|e| e.to_string())?; let accounts = load_codex_accounts()?; + let display_names = display_names_by_id(&accounts); + let snapshots = snapshots_for_accounts(&accounts, codex_account_snapshots()?); Ok(CodexAccountsStateBridge { - display_names: display_names_by_id(&accounts), accounts, - snapshots: codex_account_snapshots()?, + display_names, + snapshots, }) } diff --git a/apps/desktop-tauri/src-tauri/src/commands/providers.rs b/apps/desktop-tauri/src-tauri/src/commands/providers.rs index 958d8002d3..237e03f88e 100644 --- a/apps/desktop-tauri/src-tauri/src/commands/providers.rs +++ b/apps/desktop-tauri/src-tauri/src/commands/providers.rs @@ -15,6 +15,7 @@ pub(crate) fn build_fetch_context( api_keys: &ApiKeys, token_accounts: &HashMap, ) -> FetchContext { + let provider = instantiate_provider(id); let cookie_source = settings.cookie_source(id); let stored_cookie = cookies.get(id.cli_name()).map(|s| s.to_string()); let stored_api_key = api_keys.get(id.cli_name()).map(|s| s.to_string()); @@ -26,6 +27,9 @@ pub(crate) fn build_fetch_context( let active_token_cookie = token_override .as_ref() .and_then(|override_data| override_data.cookie_header.clone()); + let defer_provider_browser_cookie_lookup = provider.owns_browser_cookie_resolution() + && active_token_cookie.is_none() + && stored_cookie.is_none(); let active_token_env = token_override .as_ref() .and_then(|override_data| override_data.env_override.as_ref()); @@ -78,14 +82,18 @@ pub(crate) fn build_fetch_context( } // `browser` is accepted as a legacy alias from older settings. "auto" | "browser" | "web" => { - // Try browser cookie extraction as fallback when no manual cookie is set. - // On non-Windows this is a harmless no-op that returns an error. + // Claude resolves its cached cookie and browser fallback inside + // the provider; other providers retain the shell fallback. let cookie_header = active_token_cookie.or(stored_cookie).or_else(|| { - provider_cookie_domain(id, settings).and_then(|domain| { - codexbar::browser::cookies::get_cookie_header(domain) - .ok() - .filter(|h| !h.is_empty()) - }) + if defer_provider_browser_cookie_lookup { + None + } else { + provider_cookie_domain(id, settings).and_then(|domain| { + codexbar::browser::cookies::get_cookie_header(domain) + .ok() + .filter(|h| !h.is_empty()) + }) + } }); (usage_source, cookie_header) } @@ -97,10 +105,7 @@ pub(crate) fn build_fetch_context( // historically mapped "manual + no cookie" to Cli, which surfaces as // "Source mode 'Cli' not supported". Remap to Web and try browser cookies // unless the user explicitly disabled cookies ("off"). - if source_mode == SourceMode::Cli - && cookie_source != "off" - && !instantiate_provider(id).supports_cli() - { + if source_mode == SourceMode::Cli && cookie_source != "off" && !provider.supports_cli() { if cookie_header .as_deref() .map(str::trim) @@ -505,28 +510,13 @@ pub(super) fn preserve_last_good_transient_failure( id: ProviderId, snapshot: ProviderUsageSnapshot, ) -> ProviderUsageSnapshot { - if snapshot.error.is_none() { - guard.transient_provider_failure_counts.remove(&id); - return snapshot; - } - - if id != ProviderId::Claude { + let Some(error) = snapshot.error.as_deref() else { guard.transient_provider_failure_counts.remove(&id); return snapshot; - } - - let error = snapshot.error.as_deref(); - // Hard auth loss / subscription-unavailable answers should not keep stale bars. - if is_hard_claude_auth_loss(error) { - guard.transient_provider_failure_counts.remove(&id); - return snapshot; - } + }; - let preservable = is_transient_claude_auth_error(error) - || is_claude_cli_usage_parse_failure(error) - || is_claude_cli_rate_limit_failure(error) - || is_claude_timeout_failure(error); - if !preservable { + let policy = instantiate_provider(id).last_good_failure_policy(error); + if policy == codexbar::core::LastGoodFailurePolicy::Replace { guard.transient_provider_failure_counts.remove(&id); return snapshot; } @@ -540,82 +530,53 @@ pub(super) fn preserve_last_good_transient_failure( return snapshot; }; - // Parse / rate-limit / timeout: keep last-good every time (upstream #2247). - // Transient auth (unauthorized-ish) still only preserves once so real logout surfaces. - let parse_or_rate = is_claude_cli_usage_parse_failure(error) - || is_claude_cli_rate_limit_failure(error) - || is_claude_timeout_failure(error); - let count = guard .transient_provider_failure_counts .entry(id) .or_insert(0); - if parse_or_rate || *count == 0 { - if !parse_or_rate { + match policy { + codexbar::core::LastGoodFailurePolicy::Preserve => { + tracing::warn!( + provider = id.cli_name(), + error, + "preserving last good provider snapshot after transient failure" + ); + previous + } + codexbar::core::LastGoodFailurePolicy::PreserveOnce if *count == 0 => { *count = 1; + tracing::warn!( + provider = id.cli_name(), + error, + "preserving last good provider snapshot after transient failure" + ); + previous } - tracing::warn!( - provider = id.cli_name(), - error = error.unwrap_or(""), - "preserving last good Claude snapshot after transient failure" - ); - previous - } else { - *count = count.saturating_add(1); - snapshot + codexbar::core::LastGoodFailurePolicy::PreserveOnce => { + *count = count.saturating_add(1); + snapshot + } + codexbar::core::LastGoodFailurePolicy::PreserveOnceThenSurface if *count == 0 => { + *count = 1; + tracing::warn!( + provider = id.cli_name(), + error, + "preserving last good provider snapshot after transient failure" + ); + previous + } + codexbar::core::LastGoodFailurePolicy::PreserveOnceThenSurface => { + *count = count.saturating_add(1); + let mut surfaced = previous; + surfaced.error = snapshot.error; + surfaced.error_state = snapshot.error_state; + surfaced.fetch_duration_ms = snapshot.fetch_duration_ms; + surfaced + } + codexbar::core::LastGoodFailurePolicy::Replace => snapshot, } } -fn is_transient_claude_auth_error(error: Option<&str>) -> bool { - let Some(error) = error else { - return false; - }; - let lower = error.to_ascii_lowercase(); - lower.contains("unauthorized") - || lower.contains("authentication required") - || lower.contains("auth required") -} - -fn is_hard_claude_auth_loss(error: Option<&str>) -> bool { - let Some(error) = error else { - return false; - }; - let lower = error.to_ascii_lowercase(); - // Credentials truly missing / login required — clear stale usage. - lower.contains("credentials not found") - || lower.contains("run `claude` to authenticate") - || (lower.contains("not installed") && lower.contains("claude")) - || (lower.contains("subscription") && lower.contains("unavailable")) -} - -fn is_claude_cli_usage_parse_failure(error: Option<&str>) -> bool { - let Some(error) = error else { - return false; - }; - let lower = error.to_ascii_lowercase(); - lower.contains("parse error") - || lower.contains("empty output") - || lower.contains("missing current session") - || lower.contains("treated /usage as a normal prompt") - || lower.contains("local activity stats") - || lower.contains("could not parse") -} - -fn is_claude_cli_rate_limit_failure(error: Option<&str>) -> bool { - let Some(error) = error else { - return false; - }; - let lower = error.to_ascii_lowercase(); - lower.contains("rate limit") || lower.contains("rate_limit") || lower.contains("ratelimited") -} - -fn is_claude_timeout_failure(error: Option<&str>) -> bool { - let Some(error) = error else { - return false; - }; - error.eq_ignore_ascii_case("timeout") || error.to_ascii_lowercase().contains("timed out") -} - async fn fetch_provider_snapshot( id: ProviderId, ctx: FetchContext, diff --git a/apps/desktop-tauri/src-tauri/src/commands/tests.rs b/apps/desktop-tauri/src-tauri/src/commands/tests.rs index e6e27b3eeb..bab3dc1731 100644 --- a/apps/desktop-tauri/src-tauri/src/commands/tests.rs +++ b/apps/desktop-tauri/src-tauri/src/commands/tests.rs @@ -429,6 +429,24 @@ fn fetch_context_claude_uses_oauth_without_manual_cookie() { assert!(ctx.manual_cookie_header.is_none()); } +#[test] +fn fetch_context_claude_web_source_defers_cookie_resolution_to_provider() { + let mut settings = Settings::default(); + settings.set_cookie_source(ProviderId::Claude, "browser"); + settings.set_usage_source(ProviderId::Claude, "web"); + + let ctx = super::build_fetch_context( + ProviderId::Claude, + &settings, + &ManualCookies::default(), + &ApiKeys::default(), + &HashMap::new(), + ); + + assert_eq!(ctx.source_mode, SourceMode::Web); + assert!(ctx.manual_cookie_header.is_none()); +} + #[test] fn fetch_context_claude_explicit_cli_source_still_uses_cli() { let mut settings = Settings::default(); @@ -1066,6 +1084,96 @@ fn claude_repeated_auth_failure_surfaces_error() { assert!(surfaced.error.is_some()); } +#[test] +fn claude_cloudflare_challenge_retains_prior_usage_while_surfaceing_guidance() { + let metadata = instantiate_provider(ProviderId::Claude).metadata().clone(); + let result = ProviderFetchResult { + usage: codexbar::core::UsageSnapshot::new(codexbar::core::RateWindow::new(42.0)), + cost: None, + wayfinder_usage: None, + source_label: "OAuth".to_string(), + pace_authoritative: true, + }; + let good = + ProviderUsageSnapshot::from_fetch_result(ProviderId::Claude, &metadata, &result, None); + let challenge = codexbar::providers::claude::CLOUDFLARE_CHALLENGE_MESSAGE; + let error = ProviderUsageSnapshot::from_error( + ProviderId::Claude, + &metadata, + challenge.to_string(), + codexbar::core::ProviderStateKind::Unknown, + ); + let mut state = crate::state::AppState::new(); + state.provider_cache.push(good); + + let surfaced = super::providers::preserve_last_good_transient_failure( + &mut state, + ProviderId::Claude, + error, + ); + + assert_eq!(surfaced.error, None); + assert_eq!(surfaced.primary.used_percent, 42.0); + assert_eq!( + super::providers::preserve_last_good_transient_failure( + &mut state, + ProviderId::Claude, + ProviderUsageSnapshot::from_error( + ProviderId::Claude, + &metadata, + challenge.to_string(), + codexbar::core::ProviderStateKind::Unknown, + ), + ) + .error + .as_deref(), + Some(challenge) + ); +} + +#[test] +fn claude_cloudflare_challenge_keeps_prior_usage_when_guidance_surfaces() { + let metadata = instantiate_provider(ProviderId::Claude).metadata().clone(); + let result = ProviderFetchResult { + usage: codexbar::core::UsageSnapshot::new(codexbar::core::RateWindow::new(42.0)), + cost: None, + wayfinder_usage: None, + source_label: "Web".to_string(), + pace_authoritative: true, + }; + let mut good = + ProviderUsageSnapshot::from_fetch_result(ProviderId::Claude, &metadata, &result, None); + good.updated_at = "2026-09-01T00:00:00Z".to_string(); + let error = ProviderUsageSnapshot::from_error( + ProviderId::Claude, + &metadata, + codexbar::providers::claude::CLOUDFLARE_CHALLENGE_MESSAGE.to_string(), + codexbar::core::ProviderStateKind::Unknown, + ); + let mut state = crate::state::AppState::new(); + state.provider_cache.push(good.clone()); + + let first = super::providers::preserve_last_good_transient_failure( + &mut state, + ProviderId::Claude, + error.clone(), + ); + let second = super::providers::preserve_last_good_transient_failure( + &mut state, + ProviderId::Claude, + error, + ); + + assert_eq!(first.error, None); + assert_eq!(first.primary.used_percent, 42.0); + assert_eq!( + second.error.as_deref(), + Some(codexbar::providers::claude::CLOUDFLARE_CHALLENGE_MESSAGE,) + ); + assert_eq!(second.primary.used_percent, 42.0); + assert_eq!(second.updated_at, good.updated_at); +} + #[test] fn claude_cli_parse_failure_keeps_last_good_every_time() { let metadata = instantiate_provider(ProviderId::Claude).metadata().clone(); @@ -1159,6 +1267,16 @@ fn claude_error_message_explains_missing_sign_in() { ); } +#[test] +fn claude_cloudflare_error_preserves_distinct_recovery_guidance() { + let challenge = codexbar::providers::claude::CLOUDFLARE_CHALLENGE_MESSAGE; + let message = super::friendly_provider_error(ProviderId::Claude, challenge); + + assert_eq!(message, challenge); + assert!(message.contains("OAuth")); + assert!(message.contains("different network")); +} + #[test] fn non_claude_error_message_is_preserved() { let message = super::friendly_provider_error( diff --git a/apps/desktop-tauri/src/components/MenuCard.test.tsx b/apps/desktop-tauri/src/components/MenuCard.test.tsx index 8f16ed05ee..1432854ab1 100644 --- a/apps/desktop-tauri/src/components/MenuCard.test.tsx +++ b/apps/desktop-tauri/src/components/MenuCard.test.tsx @@ -1,3 +1,4 @@ +import { readFileSync } from "node:fs"; import { fireEvent, render, screen, waitFor } from "@testing-library/react"; import { beforeEach, describe, expect, it, vi } from "vitest"; @@ -684,3 +685,62 @@ describe("MenuCard", () => { expect(await screen.findByText("3分前")).toBeInTheDocument(); }); }); + +// The SwiftUI fix this regression came from protecting a cached native +// measurement. The Windows card has no cached measurement layer: its live +// forecast is a normal flex row whose width is recomputed by WebView2. +if (!import.meta.dirname) { + throw new Error("import.meta.dirname unavailable to vitest runner"); +} +const stylesSource = readFileSync(import.meta.dirname + "/../styles.css", "utf8"); + +function ruleBlock(source: string, selector: string): string { + const escaped = selector.replace(/[^\w-]/g, "\\$&"); + const match = source.match( + new RegExp("(?:^|\\r?\\n)" + escaped + "\\s*\\{([^}]*)\\}"), + ); + expect(match).not.toBeNull(); + return match![1]; +} + +describe("MenuCard live forecast layout", () => { + it("renders the changing forecast in the current full-width flex row", async () => { + const snapshot = provider(null, 20); + snapshot.secondary = rateWindow(35, { windowMinutes: 7 * 24 * 60 }); + snapshot.secondaryLabel = "Weekly"; + snapshot.sessionEquivalentForecast = { + estimatedWindowsToExhaustWeekly: 123, + windowsUntilReset: 4, + availableWindowsUntilReset: 4, + sampleCount: 8, + weeklyResetsAt: "2026-06-01T00:00:00Z", + weeklyUsedPercent: 35, + }; + + const { container } = renderCard(snapshot); + const forecast = await screen.findByText("Estimated: 123 session quotas left"); + const row = forecast.closest(".menu-metric__forecast"); + + expect(row).toBeInTheDocument(); + expect(row).toHaveClass("menu-metric__row"); + expect(row?.parentElement).toHaveClass("menu-metric"); + expect(container.querySelector(".menu-card__content")).toBeInTheDocument(); + + const card = ruleBlock(stylesSource, ".menu-card"); + expect(card).toContain("align-items: stretch"); + const content = ruleBlock(stylesSource, ".menu-card__content"); + expect(content).toContain("display: flex"); + expect(content).toContain("flex-direction: column"); + const metricRow = ruleBlock(stylesSource, ".menu-metric__row"); + expect(metricRow).toContain("min-width: 0"); + const forecastLabel = ruleBlock( + stylesSource, + ".menu-metric__forecast .menu-metric__pct", + ); + expect(forecastLabel).toContain("flex: 1 1 auto"); + expect(forecastLabel).toContain("min-width: 0"); + expect(forecastLabel).toContain("overflow: hidden"); + expect(forecastLabel).toContain("text-overflow: ellipsis"); + expect(forecastLabel).toContain("white-space: nowrap"); + }); +}); diff --git a/apps/desktop-tauri/src/styles.css b/apps/desktop-tauri/src/styles.css index 1818e53865..80e96b4f23 100644 --- a/apps/desktop-tauri/src/styles.css +++ b/apps/desktop-tauri/src/styles.css @@ -4304,6 +4304,17 @@ html:has(.menu-surface--tray) { flex: 0 0 auto; } +/* A live forecast is its own full-width row. Release the base percent label's + no-shrink rule so refreshed text is constrained by the row and tail-truncated + inside the fixed tray width instead of painting past it. */ +.menu-metric__forecast .menu-metric__pct { + flex: 1 1 auto; + min-width: 0; + overflow: hidden; + text-overflow: ellipsis; + white-space: nowrap; +} + .menu-metric__reset { font-size: 11px; color: var(--text-secondary); diff --git a/docs/PROVIDERS.md b/docs/PROVIDERS.md index 206969d599..edff732cfd 100644 --- a/docs/PROVIDERS.md +++ b/docs/PROVIDERS.md @@ -63,6 +63,14 @@ Optional status polling (provider status pages) is available via CLI `--status` Desktop tab id: `usageSpend`. The desktop and Overview consume one shared spend catalog. Codex and Claude local logs are first-class; routed OpenCodex usage enriches the matching Codex, OpenCode Go, Kimi, or DeepSeek subscription instead of appearing as a second fake provider. xAI and OpenRouter can publish exact provider-metered daily USD spend when their management credentials are configured, while Grok local sessions contribute tokens only. Missing spend sources remain unknown rather than becoming a false `$0`. Do not invent cross-currency totals. +### AWS Bedrock monitoring + +AWS Bedrock is a Windows provider backed by signed Cost Explorer requests and optional CloudWatch activity. It is disabled by default, and monitoring requests can add charges to your AWS bill. AWS currently charges $0.01 per Cost Explorer API request; paginated monthly-spend reads can therefore use more than one billed request, while optional CloudWatch activity is billed under CloudWatch pricing. + +The shared refresh interval controls automatic provider polling. `0` / Manual disables the recurring timer, but explicit refreshes and **Refresh when the menu opens** can still fetch Bedrock data. Disable Bedrock itself to stop its app refreshes. + +`CODEXBAR_BEDROCK_BUDGET` changes only the displayed monthly progress. It does **not** cap AWS charges, stop polling, or enforce a billing limit. + Custom pricing overlays are exact-match overrides used only where the local spend contract has matching provider/model token evidence. Explicit zero rates mean free; omitted rate fields stay unknown. The Usage & Spend surface keeps provenance/coverage visible, preserves cost-only model rows when token coverage is partial, and can Copy JSON or save the same JSON contract through the native file picker. ### OpenCode, Codex quota, and local cost boundaries diff --git a/rust/src/codex_accounts/account_manager.rs b/rust/src/codex_accounts/account_manager.rs index 9cb16121c4..4c28129c73 100644 --- a/rust/src/codex_accounts/account_manager.rs +++ b/rust/src/codex_accounts/account_manager.rs @@ -228,7 +228,7 @@ impl CodexAccountManager { self.sync_ambient_global_state( ambient_account .as_ref() - .and_then(|account| account.provider_account_id.clone()), + .and_then(CodexAccount::effective_workspace_account_id), self.target_account_id(target)?, ); @@ -261,7 +261,7 @@ impl CodexAccountManager { fs::copy(&source_auth_path, destination_home.join("auth.json"))?; let now = utc_now(); - Ok(CodexAccount::new( + let mut materialized = CodexAccount::new( account.id, account.nickname.clone(), account.email_hint.clone(), @@ -272,7 +272,9 @@ impl CodexAccountManager { account.created_at, now, Some(account.last_authenticated_at.unwrap_or(now)), - )) + ); + materialized.workspace_account_id = account.workspace_account_id.clone(); + Ok(materialized) } fn backup_ambient_auth(&self) -> Result, CodexAccountManagerError> { @@ -288,8 +290,8 @@ impl CodexAccountManager { } fn target_account_id(&self, target: &CodexAccount) -> Result, CodexApiError> { - if let Some(account_id) = &target.provider_account_id { - return Ok(Some(account_id.clone())); + if let Some(account_id) = target.effective_workspace_account_id() { + return Ok(Some(account_id)); } let identity = load_identity(&target.codex_home_path)?; Ok(identity.provider_account_id) @@ -449,26 +451,29 @@ impl CodexAccountManager { } let now = utc_now(); - Ok(CodexAccount::new( + let mut authenticated = CodexAccount::new( existing .map(|account| account.id) .unwrap_or_else(Uuid::new_v4), existing.and_then(|account| account.nickname.clone()), identity .email + .clone() .or_else(|| existing.and_then(|account| account.email_hint.clone())), identity .auth_subject + .clone() .or_else(|| existing.and_then(|account| account.auth_subject.clone())), - identity - .provider_account_id - .or_else(|| existing.and_then(|account| account.provider_account_id.clone())), + provider_account_id_after_auth(&identity, existing), home_path.to_path_buf(), source, existing.map(|account| account.created_at).unwrap_or(now), now, Some(now), - )) + ); + authenticated.workspace_account_id = + existing.and_then(|account| account.workspace_account_id.clone()); + Ok(authenticated) } fn discovered_managed_account( @@ -524,6 +529,34 @@ fn candidate_account( ) } +/// Keep a v0.56.3 provider id when it is the legacy selected workspace. A +/// fresh auth read may report the auth-file default instead; that value must +/// not silently replace the app-owned selection. +fn provider_account_id_after_auth( + identity: &AuthBackedIdentity, + existing: Option<&CodexAccount>, +) -> Option { + if let Some(existing) = existing + && existing.workspace_account_id.is_none() + && existing.provider_account_id.is_some() + && identity + .provider_account_id + .as_deref() + .map(str::trim) + .is_none_or(|auth_id| { + existing + .normalized_provider_account_id() + .is_some_and(|selected_id| selected_id != auth_id.to_lowercase()) + }) + { + return existing.provider_account_id.clone(); + } + identity + .provider_account_id + .clone() + .or_else(|| existing.and_then(|account| account.provider_account_id.clone())) +} + fn build_discovered_account( matched: Option<&CodexAccount>, identity: AuthBackedIdentity, @@ -531,20 +564,20 @@ fn build_discovered_account( source: CodexAccountSource, discovered_at: DateTime, ) -> CodexAccount { - CodexAccount::new( + let mut discovered = CodexAccount::new( matched .map(|account| account.id) .unwrap_or_else(Uuid::new_v4), matched.and_then(|account| account.nickname.clone()), identity .email + .clone() .or_else(|| matched.and_then(|account| account.email_hint.clone())), identity .auth_subject + .clone() .or_else(|| matched.and_then(|account| account.auth_subject.clone())), - identity - .provider_account_id - .or_else(|| matched.and_then(|account| account.provider_account_id.clone())), + provider_account_id_after_auth(&identity, matched), home_path, source, matched @@ -556,7 +589,10 @@ fn build_discovered_account( matched .and_then(|account| account.last_authenticated_at) .or(Some(discovered_at)), - ) + ); + discovered.workspace_account_id = + matched.and_then(|account| account.workspace_account_id.clone()); + discovered } fn directory_timestamp(path: &Path) -> DateTime { diff --git a/rust/src/codex_accounts/api.rs b/rust/src/codex_accounts/api.rs index 5482dfa323..b67eff2d7b 100644 --- a/rust/src/codex_accounts/api.rs +++ b/rust/src/codex_accounts/api.rs @@ -282,6 +282,21 @@ impl CodexAccountApi { email_hint: Option<&str>, verify_live_data: bool, ) -> Result { + self.fetch_snapshot_for_workspace(codex_home_path, email_hint, None, verify_live_data) + .await + } + + /// Fetch a snapshot while scoping every usage/credits request to the + /// app-selected workspace. The selected id is request metadata only: the + /// auth file remains untouched and may retain a different default. + pub async fn fetch_snapshot_for_workspace( + &self, + codex_home_path: &Path, + email_hint: Option<&str>, + workspace_account_id: Option<&str>, + verify_live_data: bool, + ) -> Result { + let workspace_account_id = workspace_account_id.and_then(|id| normalize_string(Some(id))); let mut credentials = load_credentials(codex_home_path)?; if credentials.needs_refresh() @@ -295,7 +310,13 @@ impl CodexAccountApi { } let result = self - .fetch_once(codex_home_path, &credentials, email_hint, verify_live_data) + .fetch_once( + codex_home_path, + &credentials, + email_hint, + workspace_account_id.as_deref(), + verify_live_data, + ) .await; if !matches!(&result, Err(CodexApiError::Message(msg)) if msg == UNAUTHORIZED_MESSAGE) || credentials.refresh_token.is_empty() @@ -308,7 +329,13 @@ impl CodexAccountApi { // cannot block the fetch already in progress. let _saved_retry = save_credentials(codex_home_path, &refreshed); return self - .fetch_once(codex_home_path, &refreshed, email_hint, verify_live_data) + .fetch_once( + codex_home_path, + &refreshed, + email_hint, + workspace_account_id.as_deref(), + verify_live_data, + ) .await; } result @@ -319,14 +346,25 @@ impl CodexAccountApi { codex_home_path: &Path, credentials: &AuthCredentials, email_hint: Option<&str>, + workspace_account_id: Option<&str>, verify_live_data: bool, ) -> Result { if verify_live_data { - self.fetch_verified(codex_home_path, credentials, email_hint) - .await + self.fetch_verified( + codex_home_path, + credentials, + email_hint, + workspace_account_id, + ) + .await } else { - self.fetch_single(codex_home_path, credentials, email_hint) - .await + self.fetch_single( + codex_home_path, + credentials, + email_hint, + workspace_account_id, + ) + .await } } @@ -336,18 +374,34 @@ impl CodexAccountApi { codex_home_path: &Path, credentials: &AuthCredentials, email_hint: Option<&str>, + workspace_account_id: Option<&str>, ) -> Result { let first = self - .fetch_single(codex_home_path, credentials, email_hint) + .fetch_single( + codex_home_path, + credentials, + email_hint, + workspace_account_id, + ) .await?; let second = self - .fetch_single(codex_home_path, credentials, email_hint) + .fetch_single( + codex_home_path, + credentials, + email_hint, + workspace_account_id, + ) .await?; if is_equivalent(&first, &second) { return Ok(second); } let third = self - .fetch_single(codex_home_path, credentials, email_hint) + .fetch_single( + codex_home_path, + credentials, + email_hint, + workspace_account_id, + ) .await?; if is_equivalent(&first, &third) || is_equivalent(&second, &third) { return Ok(third); @@ -362,13 +416,18 @@ impl CodexAccountApi { codex_home_path: &Path, credentials: &AuthCredentials, fallback_email: Option<&str>, + workspace_account_id: Option<&str>, ) -> Result { let identity = identity_from_credentials(credentials); + let remote_account_id = workspace_account_id + .and_then(|id| normalize_string(Some(id))) + .or_else(|| identity.provider_account_id.clone()) + .or_else(|| credentials.account_id.clone()); let response = self .fetch_usage( codex_home_path, &credentials.access_token, - credentials.account_id.as_deref(), + remote_account_id.as_deref(), ) .await?; let rate_limit = response.get("rate_limit").and_then(|v| v.as_object()); @@ -380,9 +439,7 @@ impl CodexAccountApi { Ok(AccountUsageSnapshot { email: identity.email.or_else(|| normalize_string(fallback_email)), - provider_account_id: identity - .provider_account_id - .or_else(|| credentials.account_id.clone()), + provider_account_id: remote_account_id, plan: normalize_string(response.get("plan_type").and_then(|v| v.as_str())) .or(identity.plan), allowed: rate_limit diff --git a/rust/src/codex_accounts/models.rs b/rust/src/codex_accounts/models.rs index f69029965e..edbc0e9e1e 100644 --- a/rust/src/codex_accounts/models.rs +++ b/rust/src/codex_accounts/models.rs @@ -100,7 +100,14 @@ pub struct CodexAccount { pub nickname: Option, pub email_hint: Option, pub auth_subject: Option, + /// Legacy persisted workspace selection. New records should prefer + /// `workspace_account_id`, but this remains a valid selected-workspace + /// fallback for v0.56.3 accounts. pub provider_account_id: Option, + /// App-owned remote workspace selection. This deliberately is not copied + /// into the Codex auth file, whose account id may name another default. + #[serde(default)] + pub workspace_account_id: Option, pub codex_home_path: PathBuf, pub source: CodexAccountSource, pub created_at: DateTime, @@ -131,6 +138,7 @@ impl CodexAccount { email_hint, auth_subject, provider_account_id, + workspace_account_id: None, codex_home_path, source, created_at, @@ -178,6 +186,35 @@ impl CodexAccount { normalize_identifier(self.provider_account_id.as_deref()) } + pub fn normalized_workspace_account_id(&self) -> Option { + normalize_identifier(self.workspace_account_id.as_deref()) + } + + /// The remote workspace owned by the app for this account. + /// + /// `provider_account_id` was the selected workspace field in the local + /// v0.56.3 store. Keep it as the compatibility fallback, while an explicit + /// selection always wins over the auth file's default account id. + pub fn effective_workspace_account_id(&self) -> Option { + self.normalized_workspace_account_id() + .or_else(|| self.normalized_provider_account_id()) + } + + /// Whether the app-selected workspace differs from the auth file default. + /// A missing side is not a proven mismatch, matching the upstream guard. + pub fn selected_workspace_differs_from_auth_default( + &self, + auth_default_account_id: Option<&str>, + ) -> bool { + match ( + self.effective_workspace_account_id(), + normalize_identifier(auth_default_account_id), + ) { + (Some(selected), Some(default_id)) => selected != default_id, + _ => false, + } + } + pub fn standardized_home_path(&self) -> String { std::path::absolute(&self.codex_home_path) .unwrap_or_else(|_| self.codex_home_path.clone()) @@ -186,7 +223,7 @@ impl CodexAccount { } fn display_identity(&self) -> String { - self.normalized_provider_account_id() + self.effective_workspace_account_id() .unwrap_or_else(|| self.id.to_string().to_lowercase()) } @@ -204,14 +241,14 @@ impl CodexAccount { return true; } if let (Some(a), Some(b)) = ( - self.normalized_provider_account_id(), - other.normalized_provider_account_id(), + self.effective_workspace_account_id(), + other.effective_workspace_account_id(), ) && a == b { return true; } - if self.normalized_provider_account_id().is_some() - || other.normalized_provider_account_id().is_some() + if self.effective_workspace_account_id().is_some() + || other.effective_workspace_account_id().is_some() { return false; } @@ -253,10 +290,26 @@ impl CodexAccount { }; pick(&mut self.email_hint, other.email_hint.as_ref()); pick(&mut self.auth_subject, other.auth_subject.as_ref()); - pick( - &mut self.provider_account_id, - other.provider_account_id.as_ref(), - ); + if self.workspace_account_id.is_none() { + if other.workspace_account_id.is_some() { + self.workspace_account_id = other.workspace_account_id.clone(); + } else if self.provider_account_id.is_none() + || self.normalized_provider_account_id() == other.normalized_provider_account_id() + { + pick( + &mut self.provider_account_id, + other.provider_account_id.as_ref(), + ); + } + } else { + // The explicit app-owned selection is authoritative. The legacy + // provider field may still refresh as auth metadata, but must never + // replace the selected workspace above. + pick( + &mut self.provider_account_id, + other.provider_account_id.as_ref(), + ); + } if prefer_other { self.source = other.source; @@ -330,6 +383,8 @@ pub struct RemovedAccountIdentity { pub email_hint: Option, pub auth_subject: Option, pub provider_account_id: Option, + #[serde(default)] + pub workspace_account_id: Option, pub codex_home_path: PathBuf, pub source: CodexAccountSource, pub removed_at: DateTime, @@ -342,6 +397,7 @@ impl RemovedAccountIdentity { email_hint: account.email_hint.clone(), auth_subject: account.auth_subject.clone(), provider_account_id: account.provider_account_id.clone(), + workspace_account_id: account.workspace_account_id.clone(), codex_home_path: account.codex_home_path.clone(), source: account.source, removed_at: utc_now(), @@ -353,18 +409,14 @@ impl RemovedAccountIdentity { return true; } if let (Some(a), Some(b)) = ( - normalize_identifier(self.provider_account_id.as_deref()), - account.normalized_provider_account_id(), + self.effective_workspace_account_id(), + account.effective_workspace_account_id(), ) && a == b { return true; } - if self - .provider_account_id - .as_ref() - .map(|v| !v.trim().is_empty()) - .unwrap_or(false) - || account.provider_account_id.as_ref().is_some() + if self.effective_workspace_account_id().is_some() + || account.effective_workspace_account_id().is_some() { return false; } @@ -391,6 +443,11 @@ impl RemovedAccountIdentity { .to_string_lossy() .to_lowercase() } + + fn effective_workspace_account_id(&self) -> Option { + normalize_identifier(self.workspace_account_id.as_deref()) + .or_else(|| normalize_identifier(self.provider_account_id.as_deref())) + } } /// A single quota window (session or weekly). @@ -629,6 +686,48 @@ mod tests { assert!(!a.matches(&b)); } + #[test] + fn explicit_workspace_beats_auth_default_and_survives_discovery_merge() { + let mut selected = account( + "11111111-1111-1111-1111-111111111111", + "/managed/selected", + CodexAccountSource::ManagedByApp, + Some("auth-default-a"), + ); + selected.workspace_account_id = Some("selected-workspace-b".to_string()); + let discovered = account( + "22222222-2222-2222-2222-222222222222", + "/managed/selected", + CodexAccountSource::ManagedByApp, + Some("auth-default-a"), + ); + + assert_eq!( + selected.effective_workspace_account_id().as_deref(), + Some("selected-workspace-b") + ); + assert!(selected.selected_workspace_differs_from_auth_default(Some("auth-default-a"))); + selected.merge_from(&discovered); + assert_eq!( + selected.effective_workspace_account_id().as_deref(), + Some("selected-workspace-b") + ); + } + + #[test] + fn legacy_provider_account_id_is_selected_workspace_fallback() { + let account = account( + "11111111-1111-1111-1111-111111111111", + "/managed/selected", + CodexAccountSource::ManagedByApp, + Some("Selected-Workspace-B"), + ); + assert_eq!( + account.effective_workspace_account_id().as_deref(), + Some("selected-workspace-b") + ); + } + #[test] fn source_displays_and_ownership() { assert_eq!(CodexAccountSource::Ambient.display_name(), "System"); diff --git a/rust/src/core/cost_cache_budget.rs b/rust/src/core/cost_cache_budget.rs index 4edd0e02f4..2495fd0957 100644 --- a/rust/src/core/cost_cache_budget.rs +++ b/rust/src/core/cost_cache_budget.rs @@ -323,6 +323,7 @@ mod tests { size, days: day_map, parsed_bytes: parsed, + codex_scan_target_size: None, last_model: None, last_totals: None, codex_token_timestamps_monotonic: None, diff --git a/rust/src/core/jsonl_scanner.rs b/rust/src/core/jsonl_scanner.rs index ef0dd20e47..7689ea0236 100755 --- a/rust/src/core/jsonl_scanner.rs +++ b/rust/src/core/jsonl_scanner.rs @@ -224,6 +224,11 @@ pub struct CostUsageFileUsage { pub days: HashMap>>, /// Bytes parsed so far (for incremental parsing) pub parsed_bytes: Option, + /// Frozen logical end of the scan target. A growing rollout may have a + /// physical tail beyond this boundary; that tail remains queued until a + /// later pass can consume complete records from it. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub codex_scan_target_size: Option, /// Last model seen (for delta calculations) pub last_model: Option, /// Last token totals (for delta calculations) @@ -305,6 +310,9 @@ pub struct CodexParseResult { pub records: Vec, /// Bytes parsed pub parsed_bytes: i64, + /// Stable logical target reached by this parse. This may be behind the + /// physical EOF when the tail ended inside an incomplete JSONL record. + pub scan_target_size: i64, /// Last model seen pub last_model: Option, /// Last totals seen diff --git a/rust/src/core/jsonl_scanner/codex.rs b/rust/src/core/jsonl_scanner/codex.rs index 349bc5bb7c..e4c2186359 100644 --- a/rust/src/core/jsonl_scanner/codex.rs +++ b/rust/src/core/jsonl_scanner/codex.rs @@ -5,7 +5,7 @@ mod parser; use helpers::{ BoundedJsonlLine, CODEX_JSONL_MAX_LINE_BYTES, nonempty_json_string, parse_rfc3339_timestamp, - read_bounded_jsonl_line, session_meta_field, + read_bounded_jsonl_line, read_bounded_jsonl_line_until, session_meta_field, }; use parser::CodexParserState; @@ -109,8 +109,10 @@ impl JsonlScanner { break; }; let (line_bytes, consumed) = match line { - BoundedJsonlLine::Retained { bytes, consumed } => (bytes, consumed), - BoundedJsonlLine::Discarded { consumed } => { + BoundedJsonlLine::Retained { + bytes, consumed, .. + } => (bytes, consumed), + BoundedJsonlLine::Discarded { consumed, .. } => { bytes_examined = bytes_examined.saturating_add(consumed); continue; } @@ -242,6 +244,41 @@ impl JsonlScanner { token_timestamps_monotonic, cancel, false, + None, + max_bytes_to_read, + ) + } + + /// Parse a Codex file against a caller-owned frozen target. The target is + /// intentionally separate from the current physical EOF so an active + /// rollout cannot make a bounded catch-up pass chase its own growth. + #[allow( + clippy::too_many_arguments, + reason = "resume state mirrors the persisted parser cache" + )] + pub(crate) fn parse_codex_file_with_state_bounded_target( + file_path: &Path, + range: &CostUsageDayRange, + start_offset: i64, + initial_model: Option, + initial_totals: Option, + previous_token_timestamp: Option, + token_timestamps_monotonic: Option, + cancel: Option<&AtomicBool>, + scan_target_size: Option, + max_bytes_to_read: Option, + ) -> std::io::Result { + Self::parse_codex_file_with_state_bounded_internal( + file_path, + range, + start_offset, + initial_model, + initial_totals, + previous_token_timestamp, + token_timestamps_monotonic, + cancel, + false, + scan_target_size, max_bytes_to_read, ) } @@ -270,6 +307,35 @@ impl JsonlScanner { None, cancel, true, + None, + max_bytes_to_read, + ) + } + + /// Fork equivalent of [`Self::parse_codex_file_with_state_bounded_target`]. + #[allow( + clippy::too_many_arguments, + reason = "fork parse state mirrors the persisted parser cache" + )] + pub(crate) fn parse_codex_file_with_state_bounded_fork_target( + file_path: &Path, + range: &CostUsageDayRange, + initial_totals: CodexTotals, + cancel: Option<&AtomicBool>, + scan_target_size: Option, + max_bytes_to_read: Option, + ) -> std::io::Result { + Self::parse_codex_file_with_state_bounded_internal( + file_path, + range, + 0, + None, + Some(initial_totals), + None, + None, + cancel, + true, + scan_target_size, max_bytes_to_read, ) } @@ -288,6 +354,7 @@ impl JsonlScanner { token_timestamps_monotonic: Option, cancel: Option<&AtomicBool>, fork_baseline_mode: bool, + scan_target_size: Option, max_bytes_to_read: Option, ) -> std::io::Result { let file = File::open(file_path)?; @@ -298,9 +365,15 @@ impl JsonlScanner { )] let file_size = file.metadata()?.len() as i64; + let safe_start_offset = start_offset.clamp(0, file_size); + let requested_target_size = scan_target_size + .unwrap_or(file_size) + .max(safe_start_offset) + .min(file_size); + let mut reader = BufReader::new(file); - if start_offset > 0 { - reader.seek(SeekFrom::Start(start_offset as u64))?; + if safe_start_offset > 0 { + reader.seek(SeekFrom::Start(safe_start_offset as u64))?; } let mut parser = CodexParserState::with_timestamp_state_and_fork_mode( @@ -310,14 +383,16 @@ impl JsonlScanner { token_timestamps_monotonic, fork_baseline_mode, ); - let mut parsed_bytes = start_offset; + let mut parsed_bytes = safe_start_offset; + let mut committed_bytes = safe_start_offset; let mut cancelled = false; let mut budget_exhausted = false; + let mut incomplete_tail = false; loop { if max_bytes_to_read.is_some_and(|limit| { - parsed_bytes.saturating_sub(start_offset) >= limit.max(0) - && parsed_bytes < file_size + parsed_bytes.saturating_sub(safe_start_offset) >= limit.max(0) + && parsed_bytes < requested_target_size }) { budget_exhausted = true; break; @@ -326,7 +401,16 @@ impl JsonlScanner { cancelled = true; break; } - let Some(line) = read_bounded_jsonl_line(&mut reader, CODEX_JSONL_MAX_LINE_BYTES)? + let remaining_to_target = requested_target_size.saturating_sub(parsed_bytes); + if remaining_to_target == 0 { + break; + } + let max_total_bytes = usize::try_from(remaining_to_target).ok(); + let Some(line) = read_bounded_jsonl_line_until( + &mut reader, + CODEX_JSONL_MAX_LINE_BYTES, + max_total_bytes, + )? else { break; }; @@ -334,33 +418,66 @@ impl JsonlScanner { cancelled = true; break; } - let (line_bytes, consumed) = match line { - BoundedJsonlLine::Retained { bytes, consumed } => (Some(bytes), consumed), - BoundedJsonlLine::Discarded { consumed } => (None, consumed), + let (line_bytes, consumed, terminated_by_newline) = match line { + BoundedJsonlLine::Retained { + bytes, + consumed, + terminated_by_newline, + } => (Some(bytes), consumed, terminated_by_newline), + BoundedJsonlLine::Discarded { + consumed, + terminated_by_newline, + } => (None, consumed, terminated_by_newline), }; let consumed_i64 = i64::try_from(consumed).unwrap_or(i64::MAX); parsed_bytes = parsed_bytes.saturating_add(consumed_i64); let Some(line_bytes) = line_bytes else { + if terminated_by_newline { + committed_bytes = parsed_bytes; + } else { + incomplete_tail = true; + parsed_bytes = committed_bytes; + break; + } continue; }; if line_bytes.is_empty() { + committed_bytes = parsed_bytes; continue; } let Ok(line) = std::str::from_utf8(&line_bytes) else { - continue; + if terminated_by_newline { + committed_bytes = parsed_bytes; + continue; + } + incomplete_tail = true; + parsed_bytes = committed_bytes; + break; }; let line = line.strip_suffix('\r').unwrap_or(line); + if !terminated_by_newline && serde_json::from_str::(line).is_err() { + incomplete_tail = true; + parsed_bytes = committed_bytes; + break; + } parser.process_line(line, range); + committed_bytes = parsed_bytes; } - let bytes_read = parsed_bytes.saturating_sub(start_offset).max(0); - let is_complete = !cancelled && !budget_exhausted && parsed_bytes >= file_size; + let effective_target_size = if incomplete_tail && !cancelled && !budget_exhausted { + committed_bytes + } else { + requested_target_size + }; + let is_complete = !cancelled && !budget_exhausted && parsed_bytes >= effective_target_size; + let bytes_read = parsed_bytes.saturating_sub(safe_start_offset).max(0); Ok(CodexParseResult { records: parser.records, - parsed_bytes: if is_complete { - file_size.max(parsed_bytes) + parsed_bytes, + scan_target_size: if is_complete { + effective_target_size } else { - parsed_bytes + requested_target_size }, last_model: parser.current_model, last_totals: parser.previous_totals, diff --git a/rust/src/core/jsonl_scanner/codex/helpers.rs b/rust/src/core/jsonl_scanner/codex/helpers.rs index fd69046e61..e01a4764ee 100644 --- a/rust/src/core/jsonl_scanner/codex/helpers.rs +++ b/rust/src/core/jsonl_scanner/codex/helpers.rs @@ -149,37 +149,94 @@ pub(super) fn cumulative_reasoning_delta( /// Keeping the discarded case separate prevents callers from accidentally /// treating an oversized prefix as a parseable empty line. pub(super) enum BoundedJsonlLine { - Retained { bytes: Vec, consumed: usize }, - Discarded { consumed: usize }, + Retained { + bytes: Vec, + consumed: usize, + terminated_by_newline: bool, + }, + Discarded { + consumed: usize, + terminated_by_newline: bool, + }, } /// Read one JSONL line, discarding content when it exceeds `max_bytes`. pub(super) fn read_bounded_jsonl_line( reader: &mut R, max_bytes: usize, +) -> std::io::Result> { + read_bounded_jsonl_line_until(reader, max_bytes, None) +} + +/// Read one JSONL line without consuming past a frozen logical target. +/// +/// A target can end in the middle of a record while Codex is writing it. The +/// caller can then retain the last committed line boundary and retry the tail +/// after a later append, instead of publishing a partial record. +pub(super) fn read_bounded_jsonl_line_until( + reader: &mut R, + max_bytes: usize, + max_total_bytes: Option, ) -> std::io::Result> { let mut line = Vec::new(); let mut saw_bytes = false; let mut discarding = false; let mut consumed_total = 0; + let mut remaining_total = max_total_bytes; loop { + if remaining_total == Some(0) { + return Ok(saw_bytes.then_some(if discarding { + BoundedJsonlLine::Discarded { + consumed: consumed_total, + terminated_by_newline: false, + } + } else { + BoundedJsonlLine::Retained { + bytes: line, + consumed: consumed_total, + terminated_by_newline: false, + } + })); + } + let chunk = reader.fill_buf()?; if chunk.is_empty() { return Ok(saw_bytes.then_some(if discarding { BoundedJsonlLine::Discarded { consumed: consumed_total, + terminated_by_newline: false, + } + } else { + BoundedJsonlLine::Retained { + bytes: line, + consumed: consumed_total, + terminated_by_newline: false, + } + })); + } + + let visible_len = + remaining_total.map_or(chunk.len(), |remaining| remaining.min(chunk.len())); + if visible_len == 0 { + return Ok(saw_bytes.then_some(if discarding { + BoundedJsonlLine::Discarded { + consumed: consumed_total, + terminated_by_newline: false, } } else { BoundedJsonlLine::Retained { bytes: line, consumed: consumed_total, + terminated_by_newline: false, } })); } - let newline = chunk.iter().position(|byte| *byte == b'\n'); - let segment_end = newline.unwrap_or(chunk.len()); - let segment = &chunk[..segment_end]; + + let visible_chunk = &chunk[..visible_len]; + let newline = visible_chunk.iter().position(|byte| *byte == b'\n'); + let segment_end = newline.unwrap_or(visible_len); + let segment = &visible_chunk[..segment_end]; saw_bytes = true; if !discarding { @@ -195,15 +252,35 @@ pub(super) fn read_bounded_jsonl_line( let consumed = segment_end + usize::from(newline.is_some()); reader.consume(consumed); consumed_total += consumed; + if let Some(remaining) = remaining_total.as_mut() { + *remaining = remaining.saturating_sub(consumed); + } if newline.is_some() { return Ok(Some(if discarding { BoundedJsonlLine::Discarded { consumed: consumed_total, + terminated_by_newline: true, + } + } else { + BoundedJsonlLine::Retained { + bytes: line, + consumed: consumed_total, + terminated_by_newline: true, + } + })); + } + + if remaining_total == Some(0) { + return Ok(Some(if discarding { + BoundedJsonlLine::Discarded { + consumed: consumed_total, + terminated_by_newline: false, } } else { BoundedJsonlLine::Retained { bytes: line, consumed: consumed_total, + terminated_by_newline: false, } })); } diff --git a/rust/src/core/jsonl_scanner/tests.rs b/rust/src/core/jsonl_scanner/tests.rs index cb08f37b57..733d1a361c 100644 --- a/rust/src/core/jsonl_scanner/tests.rs +++ b/rust/src/core/jsonl_scanner/tests.rs @@ -501,6 +501,52 @@ fn codex_append_timestamp_state_is_output_equivalent_and_boundary_only() { assert_eq!(full_input, 30); } +#[test] +fn codex_parse_publishes_only_the_committed_prefix_before_an_incomplete_tail() { + let mut file = tempfile::NamedTempFile::new().expect("temp file"); + let day = NaiveDate::from_ymd_opt(2026, 5, 31).unwrap(); + let range = CostUsageDayRange::new(day, day); + let committed_line = r#"{"timestamp":"2026-05-31T10:00:00Z","type":"event_msg","payload":{"type":"token_count","info":{"model":"gpt-5.5","total_token_usage":{"input_tokens":10,"cached_input_tokens":0,"output_tokens":1}}}}"#; + writeln!(file, "{committed_line}").unwrap(); + let committed_bytes = + i64::try_from(committed_line.len() + 1).expect("fixture line length fits i64"); + + let complete_tail = r#"{"timestamp":"2026-05-31T10:00:01Z","type":"event_msg","payload":{"type":"token_count","info":{"model":"gpt-5.5","total_token_usage":{"input_tokens":20,"cached_input_tokens":0,"output_tokens":2}}}}"#; + let split = complete_tail.len() / 2; + write!(file, "{}", &complete_tail[..split]).unwrap(); + file.flush().unwrap(); + + let partial = JsonlScanner::parse_codex_file(file.path(), &range, 0, None, None) + .expect("parse committed prefix"); + assert_eq!(partial.records.len(), 1); + assert_eq!(partial.records[0].input, 10); + assert_eq!(partial.parsed_bytes, committed_bytes); + assert_eq!(partial.scan_target_size, committed_bytes); + assert!(partial.is_complete, "the logical prefix is complete"); + + writeln!(file, "{}", &complete_tail[split..]).unwrap(); + let resumed = JsonlScanner::parse_codex_file_with_state( + file.path(), + &range, + partial.parsed_bytes, + partial.last_model, + partial.last_totals, + partial.last_token_timestamp, + partial.token_timestamps_monotonic, + None, + ) + .expect("resume completed tail"); + assert_eq!(resumed.records.len(), 1); + assert_eq!(resumed.records[0].input, 10); + assert_eq!( + resumed.parsed_bytes, + i64::try_from(std::fs::metadata(file.path()).unwrap().len()) + .expect("fixture file length fits i64") + ); + assert_eq!(resumed.scan_target_size, resumed.parsed_bytes); + assert!(resumed.is_complete); +} + #[test] fn codex_parser_discards_oversized_line_and_recovers_next_record() { let mut file = tempfile::NamedTempFile::new().expect("temp file"); @@ -1006,6 +1052,7 @@ fn catch_up_snapshot_preserves_established_codex_cost_and_tokens() { HashMap::from([("gpt-5.6-sol".to_string(), vec![1_000, 250, 100])]), )]), parsed_bytes: Some(100), + codex_scan_target_size: None, last_model: Some("gpt-5.6-sol".to_string()), last_totals: None, codex_token_timestamps_monotonic: Some(true), @@ -1023,6 +1070,7 @@ fn catch_up_snapshot_preserves_established_codex_cost_and_tokens() { size: 10, days: HashMap::new(), parsed_bytes: Some(10), + codex_scan_target_size: None, last_model: None, last_totals: None, codex_token_timestamps_monotonic: None, @@ -1077,6 +1125,7 @@ fn save_cache_persists_small_codex_artifact() { HashMap::from([("gpt-5.6-sol".to_string(), vec![10, 0, 1])]), )]), parsed_bytes: None, + codex_scan_target_size: None, last_model: None, last_totals: None, codex_token_timestamps_monotonic: None, @@ -1143,6 +1192,7 @@ fn save_cache_refuses_non_bounded_provider_oversize() { size: 100, days: HashMap::new(), parsed_bytes: None, + codex_scan_target_size: None, last_model: None, last_totals: None, codex_token_timestamps_monotonic: None, @@ -1178,6 +1228,7 @@ fn save_cache_refusal_removes_preexisting_destination_artifact() { HashMap::from([("gpt-5.6-sol".to_string(), vec![10, 0, 1])]), )]), parsed_bytes: None, + codex_scan_target_size: None, last_model: None, last_totals: None, codex_token_timestamps_monotonic: None, diff --git a/rust/src/core/provider.rs b/rust/src/core/provider.rs index cff8512331..9d97116ca4 100755 --- a/rust/src/core/provider.rs +++ b/rust/src/core/provider.rs @@ -641,6 +641,15 @@ impl Default for FetchContext { } } +/// How the shell should treat a failed refresh when a prior good snapshot exists. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum LastGoodFailurePolicy { + Replace, + Preserve, + PreserveOnce, + PreserveOnceThenSurface, +} + /// Trait that all providers must implement #[async_trait] pub trait Provider: Send + Sync { @@ -678,6 +687,16 @@ pub trait Provider: Send + Sync { None } + /// Whether browser-cookie discovery/recovery is owned by the provider. + fn owns_browser_cookie_resolution(&self) -> bool { + false + } + + /// How the shell should treat a failed refresh when a prior good snapshot exists. + fn last_good_failure_policy(&self, _error: &str) -> LastGoodFailurePolicy { + LastGoodFailurePolicy::Replace + } + /// Presentation-safe availability state for a refresh error. The default /// maps `ProviderError` variants, treating `NotInstalled` as a missing /// credential (most providers raise it for a missing API key or auth diff --git a/rust/src/cost_scanner/codex.rs b/rust/src/cost_scanner/codex.rs index 175c6aad43..360024c44b 100644 --- a/rust/src/cost_scanner/codex.rs +++ b/rust/src/cost_scanner/codex.rs @@ -1,6 +1,8 @@ use super::*; +mod logical_target; mod reconciliation; +use logical_target::*; use reconciliation::*; fn rebuild_cache_days(cache: &mut CostUsageCache) { @@ -89,26 +91,6 @@ fn summary_from_cached_report( } } -fn cached_codex_file_is_complete_for_range( - cache: &CostUsageCache, - path_key: &str, - range: &CostUsageDayRange, -) -> bool { - JsonlScanner::cache_covers_range(cache, range) - && cache.files.get(path_key).is_some_and(|usage| { - let Ok(metadata) = fs::metadata(path_key) else { - return false; - }; - #[allow(clippy::cast_possible_wrap, reason = "session file sizes fit i64")] - let size = metadata.len().min(i64::MAX as u64) as i64; - usage.mtime_unix_ms == system_time_to_unix_ms(metadata.modified().ok()) - && usage.size == size - && usage.parsed_bytes.unwrap_or(0) >= size - && !usage.codex_unresolved_fork_parent - && codex_fork_parent_is_safe(cache, usage) - }) -} - fn codex_fork_parent_is_safe(cache: &CostUsageCache, usage: &CostUsageFileUsage) -> bool { usage.codex_forked_from_id.as_deref().is_none() || codex_parent_baseline( @@ -306,7 +288,7 @@ impl CostScanner { || !cache.files.is_empty())) .then(|| JsonlScanner::cached_cost_report_from_days(&cache)); - let (candidates, discovery_complete) = + let (mut candidates, discovery_complete) = self.collect_codex_candidates(&sessions_dirs, &range, &cache, cancel, &mut stats); let candidate_limit = if self.options.codex_candidate_limit == 0 { usize::MAX @@ -326,11 +308,13 @@ impl CostScanner { let mut bytes_read_this_refresh = 0_i64; let mut pending_next = cache.codex_pending_paths.clone(); let pending_paths_before_pass = cache.codex_pending_paths.clone(); + prioritize_codex_pending_candidates(&mut candidates, &pending_paths_before_pass); if discovery_complete && !is_cancelled(cancel) { pending_next .retain(|path| !cached_codex_file_is_complete_for_range(&cache, path, &range)); } + let mut incomplete_processed = Vec::new(); for (index, candidate) in candidates.iter().enumerate() { if is_cancelled(cancel) || index >= candidate_limit @@ -381,11 +365,26 @@ impl CostScanner { .saturating_add(u64::try_from(outcome.bytes_read.max(0)).unwrap_or(u64::MAX)); let key = candidate.path.to_string_lossy().to_string(); pending_next.retain(|pending| pending != &key); - if !outcome.is_complete { - pending_next.push(key); + let observed_size = fs::metadata(&candidate.path) + .ok() + .map(|metadata| { + #[allow( + clippy::cast_possible_wrap, + reason = "file sizes are clamped to i64::MAX" + )] + let size = metadata.len().min(i64::MAX as u64) as i64; + size + }) + .unwrap_or(0); + let has_unconsumed_tail = cache.files.get(&key).is_some_and(|usage| { + codex_logical_target_has_unconsumed_tail(observed_size, usage) + }); + if !outcome.is_complete || has_unconsumed_tail { + incomplete_processed.push(key); stats.files_deferred = stats.files_deferred.saturating_add(1); } } + pending_next.extend(incomplete_processed); let mut pruned_paths_pending = Vec::new(); if discovery_complete && !is_cancelled(cancel) { @@ -720,9 +719,10 @@ impl CostScanner { let path_key = path.to_string_lossy().to_string(); let cached = cache.files.get(&path_key).cloned(); let cache_covers_range = JsonlScanner::cache_covers_range(cache, range); - let trace_was_pruned = cached - .as_ref() - .is_some_and(|entry| entry.size > size && entry.parsed_bytes.unwrap_or(0) > size); + let trace_was_pruned = cached.as_ref().is_some_and(|entry| { + entry.size > size + && (entry.parsed_bytes.unwrap_or(0) > size || codex_scan_target_size(entry) > size) + }); if trace_was_pruned && !self.options.is_app_driven() { // A shrinking trace invalidates the append cursor. Preserve the // validated cache and queue the path for an explicit cold refresh @@ -751,6 +751,18 @@ impl CostScanner { .then(|| cached.as_ref()?.codex_fork_timestamp.clone()) .flatten() }); + let cached_identity_changed = cached.as_ref().is_some_and(|entry| { + session_metadata + .session_id + .as_ref() + .zip(entry.codex_session_id.as_ref()) + .is_some_and(|(current, previous)| current != previous) + || session_metadata + .forked_from_id + .as_ref() + .zip(entry.codex_forked_from_id.as_ref()) + .is_some_and(|(current, previous)| current != previous) + }); let is_fork = codex_forked_from_id.is_some(); let fork_baseline = codex_forked_from_id.as_deref().and_then(|parent_id| { codex_parent_baseline(cache, parent_id, codex_fork_timestamp.as_deref()) @@ -764,6 +776,7 @@ impl CostScanner { size, days: HashMap::new(), parsed_bytes: Some(0), + codex_scan_target_size: None, last_model: None, last_totals: None, codex_token_timestamps_monotonic: None, @@ -786,6 +799,7 @@ impl CostScanner { && !entry.codex_unresolved_fork_parent && entry.mtime_unix_ms == mtime_ms && entry.size == size + && codex_scan_target_size(entry) == size && entry.parsed_bytes.unwrap_or(0) >= size { let (session_cost, has_tokens) = @@ -801,7 +815,10 @@ impl CostScanner { }; } - if !is_fork && let Some(entry) = &cached { + if !is_fork + && !cached_identity_changed + && let Some(entry) = &cached + { let start_offset = entry.parsed_bytes.unwrap_or(0); let same_partial = size == entry.size && mtime_ms == entry.mtime_unix_ms && start_offset < size; @@ -814,7 +831,8 @@ impl CostScanner { && parser_state_safe && JsonlScanner::is_line_boundary_offset(path, start_offset) { - let parse_result = match JsonlScanner::parse_codex_file_with_state_bounded( + let resumable_target_size = codex_resumable_scan_target_size(size, entry); + let parse_result = match JsonlScanner::parse_codex_file_with_state_bounded_target( path, range, start_offset, @@ -823,6 +841,7 @@ impl CostScanner { entry.codex_last_token_timestamp.clone(), entry.codex_token_timestamps_monotonic, cancel, + resumable_target_size, max_bytes_to_read, ) { Ok(result) => result, @@ -850,6 +869,7 @@ impl CostScanner { size, days, parsed_bytes: Some(parse_result.parsed_bytes), + codex_scan_target_size: Some(parse_result.scan_target_size), last_model: parse_result.last_model.or_else(|| entry.last_model.clone()), last_totals: parse_result .last_totals @@ -871,12 +891,16 @@ impl CostScanner { } } + let parse_target_size = cached + .as_ref() + .and_then(|entry| codex_resumable_scan_target_size(size, entry)); let parse_result = match if let Some(baseline) = fork_baseline.clone() { - JsonlScanner::parse_codex_file_with_state_bounded_fork( + JsonlScanner::parse_codex_file_with_state_bounded_fork_target( path, range, baseline, cancel, + parse_target_size, max_bytes_to_read, ) } else { @@ -906,6 +930,7 @@ impl CostScanner { size, days: HashMap::new(), parsed_bytes: Some(0), + codex_scan_target_size: None, last_model: None, last_totals: None, codex_token_timestamps_monotonic: None, @@ -941,6 +966,7 @@ impl CostScanner { size, days, parsed_bytes: Some(parse_result.parsed_bytes), + codex_scan_target_size: Some(parse_result.scan_target_size), last_model: parse_result.last_model, last_totals: parse_result.last_totals, codex_token_timestamps_monotonic: parse_result.token_timestamps_monotonic, diff --git a/rust/src/cost_scanner/codex/logical_target.rs b/rust/src/cost_scanner/codex/logical_target.rs new file mode 100644 index 0000000000..90200d4b8c --- /dev/null +++ b/rust/src/cost_scanner/codex/logical_target.rs @@ -0,0 +1,80 @@ +use super::*; + +pub(super) fn cached_codex_file_is_complete_for_range( + cache: &CostUsageCache, + path_key: &str, + range: &CostUsageDayRange, +) -> bool { + JsonlScanner::cache_covers_range(cache, range) + && cache.files.get(path_key).is_some_and(|usage| { + let Ok(metadata) = fs::metadata(path_key) else { + return false; + }; + #[allow(clippy::cast_possible_wrap, reason = "session file sizes fit i64")] + let size = metadata.len().min(i64::MAX as u64) as i64; + usage.mtime_unix_ms == system_time_to_unix_ms(metadata.modified().ok()) + && usage.size == size + && codex_scan_target_size(usage) == size + && usage.parsed_bytes.unwrap_or(0) >= size + && !usage.codex_unresolved_fork_parent + && super::codex_fork_parent_is_safe(cache, usage) + }) +} + +/// Give paths already in the durable queue their saved turn before newly +/// discovered dirty paths. The scanner appends unfinished paths after this +/// pass, making the queue a round-robin cursor instead of a newest-first loop. +pub(super) fn prioritize_codex_pending_candidates( + candidates: &mut Vec, + pending_paths: &[String], +) { + if pending_paths.is_empty() || candidates.len() < 2 { + return; + } + + let mut pending = Vec::with_capacity(candidates.len()); + let mut fresh = Vec::with_capacity(candidates.len()); + for candidate in candidates.drain(..) { + let key = candidate.path.to_string_lossy(); + if pending_paths + .iter() + .any(|path| path.as_str() == key.as_ref()) + { + pending.push(candidate); + } else { + fresh.push(candidate); + } + } + pending.extend(fresh); + candidates.extend(pending); +} + +/// Return the persisted logical end of a Codex parse. Older cache entries did +/// not have a frozen target, so their physical size remains the safe fallback. +pub(super) fn codex_scan_target_size(usage: &CostUsageFileUsage) -> i64 { + usage.codex_scan_target_size.unwrap_or(usage.size).max(0) +} + +/// A cached prefix is resumable toward its original target when the target is +/// still present in the current file. The caller separately validates the +/// byte-boundary/parser-state invariants before using the cursor. +pub(super) fn codex_resumable_scan_target_size( + metadata_size: i64, + usage: &CostUsageFileUsage, +) -> Option { + let parsed_bytes = usage.parsed_bytes.unwrap_or(usage.size).max(0); + let target_size = codex_scan_target_size(usage); + (parsed_bytes < target_size && target_size <= metadata_size).then_some(target_size) +} + +/// Whether a logically complete prefix still has physical bytes that must be +/// revisited. This is the catch-up cursor for a growing rollout or a retained +/// incomplete tail. +pub(super) fn codex_logical_target_has_unconsumed_tail( + metadata_size: i64, + usage: &CostUsageFileUsage, +) -> bool { + let parsed_bytes = usage.parsed_bytes.unwrap_or(usage.size).max(0); + let target_size = codex_scan_target_size(usage); + parsed_bytes < metadata_size || target_size < metadata_size +} diff --git a/rust/src/cost_scanner/tests.rs b/rust/src/cost_scanner/tests.rs index e4eaa1fde1..65da4bcdf8 100644 --- a/rust/src/cost_scanner/tests.rs +++ b/rust/src/cost_scanner/tests.rs @@ -517,6 +517,7 @@ fn cached_usage_with_packed(day: &str, model: &str, packed: Vec) -> CostUsa HashMap::from([(model.to_string(), packed)]), )]), parsed_bytes: Some(1), + codex_scan_target_size: None, last_model: None, last_totals: None, codex_token_timestamps_monotonic: None, @@ -1212,6 +1213,7 @@ fn cancelled_fresh_cache_hit_is_not_authoritative() { size: 100, days: usage.clone(), parsed_bytes: Some(100), + codex_scan_target_size: None, last_model: Some("gpt-5.6-sol".to_string()), last_totals: None, codex_token_timestamps_monotonic: Some(true), @@ -1425,6 +1427,147 @@ fn cost_scan_resumes_appended_bytes() { assert!(cached_file.codex_last_token_timestamp.is_some()); } +#[test] +fn bounded_growing_rollout_freezes_target_and_resumes_a_retained_tail() { + let root = tempfile::tempdir().unwrap(); + let sessions = root.path().join("sessions"); + let cache_root = root.path().join("cache"); + let path = + write_codex_session_fixture_with_inputs(&sessions, "growing.jsonl", &[100, 200, 300]); + let initial_size = i64::try_from(std::fs::metadata(&path).unwrap().len()) + .expect("fixture file length fits i64"); + let first_line_bytes = i64::try_from( + std::fs::read(&path) + .unwrap() + .split(|byte| *byte == b'\n') + .next() + .unwrap() + .len(), + ) + .expect("fixture line length fits i64") + + 1; + + let mut options = CostScanOptions::app_driven(); + options.codex_max_session_file_bytes = first_line_bytes; + options.codex_max_scan_bytes_per_refresh = first_line_bytes; + let scanner = CostScanner::new(7) + .with_options(options) + .with_cache_root(&cache_root) + .with_sessions_dirs(vec![sessions.clone()]); + + let (first, _, first_cache) = scanner.scan_codex_detailed_with_cache(None); + assert_eq!(first.input_tokens, 100); + let first_usage = first_cache + .files + .get(&path.to_string_lossy().to_string()) + .unwrap(); + assert_eq!(first_usage.codex_scan_target_size, Some(initial_size)); + assert_eq!(first_usage.parsed_bytes, Some(first_line_bytes)); + assert!(first_cache.codex_scan_incomplete); + + let timestamp = (Utc::now() - Duration::minutes(10)) + .format("%Y-%m-%dT%H:%M:%S%.3fZ") + .to_string(); + let append_line = |total: u64| { + format!( + r#"{{"timestamp":"{timestamp}","type":"event_msg","payload":{{"type":"token_count","info":{{"model":"gpt-5","total_token_usage":{{"input_tokens":{total},"cached_input_tokens":0,"output_tokens":{}}}}}}}}} +"#, + total / 10 + ) + }; + let mut next_total = 400_u64; + let mut bounded_summary = first; + let mut bounded_cache = first_cache; + for _ in 0..8 { + let line = append_line(next_total); + next_total += 100; + let mut file = std::fs::OpenOptions::new() + .append(true) + .open(&path) + .unwrap(); + file.write_all(line.as_bytes()).unwrap(); + drop(file); + + (bounded_summary, _, bounded_cache) = scanner.scan_codex_detailed_with_cache(None); + let usage = bounded_cache + .files + .get(&path.to_string_lossy().to_string()) + .unwrap(); + assert_eq!(usage.codex_scan_target_size, Some(initial_size)); + assert!(usage.parsed_bytes.unwrap_or_default() <= initial_size); + if usage.parsed_bytes == Some(initial_size) { + break; + } + } + + assert_eq!(bounded_summary.input_tokens, 300); + let bounded_usage = bounded_cache + .files + .get(&path.to_string_lossy().to_string()) + .unwrap(); + assert_eq!(bounded_usage.parsed_bytes, Some(initial_size)); + assert_eq!(bounded_usage.codex_scan_target_size, Some(initial_size)); + assert!( + bounded_cache.codex_scan_incomplete, + "the appended tail stays queued" + ); + + for _ in 0..32 { + if !bounded_cache.codex_scan_incomplete { + break; + } + (bounded_summary, _, bounded_cache) = scanner.scan_codex_detailed_with_cache(None); + } + assert!(!bounded_cache.codex_scan_incomplete); + let stable_summary = bounded_summary.clone(); + let stable_size = i64::try_from(std::fs::metadata(&path).unwrap().len()) + .expect("fixture file length fits i64"); + + let partial_line = append_line(next_total); + let split = partial_line.len() / 2; + let mut file = std::fs::OpenOptions::new() + .append(true) + .open(&path) + .unwrap(); + file.write_all(&partial_line.as_bytes()[..split]).unwrap(); + drop(file); + let (partial_summary, _, partial_cache) = scanner.scan_codex_detailed_with_cache(None); + let partial_usage = partial_cache + .files + .get(&path.to_string_lossy().to_string()) + .unwrap(); + assert_eq!(partial_summary.input_tokens, stable_summary.input_tokens); + assert_eq!(partial_usage.parsed_bytes, Some(stable_size)); + assert_eq!(partial_usage.codex_scan_target_size, Some(stable_size)); + assert!(partial_cache.codex_scan_incomplete); + + let mut file = std::fs::OpenOptions::new() + .append(true) + .open(&path) + .unwrap(); + file.write_all(&partial_line.as_bytes()[split..]).unwrap(); + drop(file); + next_total += 100; + let (resumed_summary, _, resumed_cache) = scanner.scan_codex_detailed_with_cache(None); + assert!(!resumed_cache.codex_scan_incomplete); + assert_eq!(resumed_summary.input_tokens, next_total - 100); + let resumed_usage = resumed_cache + .files + .get(&path.to_string_lossy().to_string()) + .unwrap(); + assert_eq!( + resumed_usage.parsed_bytes, + Some( + i64::try_from(std::fs::metadata(&path).unwrap().len()) + .expect("fixture file length fits i64"), + ) + ); + assert_eq!( + resumed_usage.codex_scan_target_size, + resumed_usage.parsed_bytes + ); +} + #[test] fn cost_scan_midline_rewrite_forces_full_parse_not_resume() { // F2 (upstream 0.48.0 #2648): when a file is rewritten/truncated so the diff --git a/rust/src/providers/antigravity/local_proto.rs b/rust/src/providers/antigravity/local_proto.rs index 6fa92e4df6..099f25e07b 100644 --- a/rust/src/providers/antigravity/local_proto.rs +++ b/rust/src/providers/antigravity/local_proto.rs @@ -14,6 +14,7 @@ pub(super) struct ParsedTurn { pub timestamp_ms: Option, pub model: Option, pub label: Option, + pub step_uuid: Option, } #[derive(Clone, Copy)] @@ -113,6 +114,25 @@ fn text(field: Field<'_>) -> Option> { Some((!value.is_empty()).then(|| value.to_string())) } +fn identity_text(field: Field<'_>) -> Option> { + let value = std::str::from_utf8(message(field)?).ok()?; + Some((!value.trim().is_empty()).then(|| value.to_string())) +} + +fn timestamp_millis(seconds: Option, nanos: u64) -> Option> { + let Some(seconds) = seconds else { + return Some(None); + }; + if seconds == 0 || seconds > 253_402_300_799 || nanos > 999_999_999 { + return None; + } + let seconds = i64::try_from(seconds).ok()?; + let nanos = i64::try_from(nanos).ok()?; + Some(Some( + seconds.checked_mul(1000)?.checked_add(nanos / 1_000_000)?, + )) +} + pub(super) fn parse_turn(root: &[u8]) -> Option { let mut turn = ParsedTurn::default(); let mut seconds = None; @@ -120,6 +140,9 @@ pub(super) fn parse_turn(root: &[u8]) -> Option { let mut found_chat = false; fields(root, |field| { if field.number != 1 { + if field.number == 4 { + turn.step_uuid = identity_text(field)?; + } return Some(()); } found_chat = true; @@ -128,16 +151,37 @@ pub(super) fn parse_turn(root: &[u8]) -> Option { if !found_chat { return None; } - turn.timestamp_ms = match seconds { - Some(value) if value > 0 && value <= 253_402_300_799 && nanos <= 999_999_999 => { - let seconds = i64::try_from(value).ok()?; - let nanos = i64::try_from(nanos).ok()?; - seconds.checked_mul(1000)?.checked_add(nanos / 1_000_000) + turn.timestamp_ms = timestamp_millis(seconds, nanos)?; + Some(turn) +} + +pub(super) fn parse_step_metadata(root: &[u8]) -> Option<(Option, Option)> { + let mut step_uuid = None; + let mut seconds = None; + let mut nanos = 0_u64; + let mut timestamp_is_valid = true; + fields(root, |field| { + match field.number { + 1 => { + let Some(timestamp) = message(field) else { + timestamp_is_valid = false; + return Some(()); + }; + if parse_timestamp_field(timestamp, &mut seconds, &mut nanos).is_none() { + timestamp_is_valid = false; + } + } + 12 => step_uuid = identity_text(field)?, + _ => {} } - Some(_) => return None, - None => None, + Some(()) + })?; + let timestamp_ms = if timestamp_is_valid { + timestamp_millis(seconds, nanos)? + } else { + None }; - Some(turn) + Some((step_uuid, timestamp_ms)) } fn parse_chat( @@ -180,26 +224,30 @@ fn parse_generation(bytes: &[u8], seconds: &mut Option, nanos: &mut u64) -> if field.number != 4 { return Some(()); } - fields(message(field)?, |stamp| { - match stamp.number { - 1 => { - let value = integer(stamp)?; - if value == 0 || value > 253_402_300_799 { - return None; - } - *seconds = Some(value); + parse_timestamp_field(message(field)?, seconds, nanos) + }) +} + +fn parse_timestamp_field(bytes: &[u8], seconds: &mut Option, nanos: &mut u64) -> Option<()> { + fields(bytes, |stamp| { + match stamp.number { + 1 => { + let value = integer(stamp)?; + if value == 0 || value > 253_402_300_799 { + return None; } - 2 => { - let value = integer(stamp)?; - if value > 999_999_999 { - return None; - } - *nanos = value; + *seconds = Some(value); + } + 2 => { + let value = integer(stamp)?; + if value > 999_999_999 { + return None; } - _ => {} + *nanos = value; } - Some(()) - }) + _ => {} + } + Some(()) }) } @@ -264,4 +312,47 @@ mod tests { fn rejects_malformed_varint() { assert!(parse_turn(&[0x0a, 0x80]).is_none()); } + + #[test] + fn decodes_new_root_envelope_step_uuid_and_step_metadata() { + let step_uuid = "step-uuid-1"; + let mut usage = Vec::new(); + usage.extend(field_varint(1, 10)); + usage.extend(field_varint(2, 20)); + usage.extend(field_varint(9, 40)); + let chat = field_bytes(4, &usage); + let mut root = field_bytes(2, &[1, 2]); + root.extend(field_bytes(4, step_uuid.as_bytes())); + root.extend(field_bytes(1, &chat)); + + let turn = parse_turn(&root).unwrap(); + assert_eq!(turn.step_uuid.as_deref(), Some(step_uuid)); + assert_eq!(turn.timestamp_ms, None); + + let timestamp = field_varint(1, 1_787_572_800); + let metadata = field_bytes( + 1, + ×tamp + .into_iter() + .chain(field_varint(2, 123_000_000)) + .collect::>(), + ); + let mut metadata = metadata; + metadata.extend(field_bytes(12, step_uuid.as_bytes())); + assert_eq!( + parse_step_metadata(&metadata), + Some((Some(step_uuid.to_string()), Some(1_787_572_800_123))) + ); + } + + #[test] + fn invalid_step_timestamp_is_evidence_without_a_timestamp() { + let mut metadata = field_bytes(1, &field_varint(2, 1_000_000_000)); + metadata.extend(field_bytes(12, b"step-uuid-1")); + + assert_eq!( + parse_step_metadata(&metadata), + Some((Some("step-uuid-1".to_string()), None)) + ); + } } diff --git a/rust/src/providers/antigravity/local_sqlite.rs b/rust/src/providers/antigravity/local_sqlite.rs index 52cf904a91..1149988b99 100644 --- a/rust/src/providers/antigravity/local_sqlite.rs +++ b/rust/src/providers/antigravity/local_sqlite.rs @@ -6,8 +6,9 @@ use std::time::{Duration as StdDuration, Instant}; use chrono::{DateTime, Duration, Local, TimeZone, Utc}; use rusqlite::{Connection, OpenFlags, TransactionBehavior, types::ValueRef}; -use super::local_proto::{ParsedTurn, parse_turn}; +use super::local_proto::{ParsedTurn, parse_step_metadata, parse_turn}; use super::local_sessions::{LocalHistoryCoverage, LocalSessionSummary}; +use super::local_step_resolver::{StepOccurrence, StepTimestamp, resolve_step_timestamps}; const MAX_DATABASES: usize = 500; const MAX_DIRECTORY_ENTRIES: usize = 10_000; @@ -73,6 +74,29 @@ struct Event { total: u64, } +#[derive(Debug)] +struct PendingTimestampRow { + row: i64, + step_uuid: String, + turn: ParsedTurn, + total: u64, +} + +#[derive(Debug)] +struct ParsedRows { + events: Vec, + pending: Vec, + occurrences: HashMap>, + database_bytes: usize, + complete: bool, +} + +#[derive(Debug)] +struct StepTimestampScan { + timestamps: HashMap>, + complete: bool, +} + pub(super) fn database_roots(gemini_base: &Path) -> [PathBuf; 3] { [ gemini_base.join("antigravity-cli").join("conversations"), @@ -281,7 +305,65 @@ fn read_database(path: &Path, budget: &mut Budget) -> rusqlite::Result<(Vec>() + .into_iter() + .map(|step_uuid| { + let occurrences = rows + .occurrences + .get(&step_uuid) + .cloned() + .unwrap_or_default(); + (step_uuid, occurrences) + }) + .collect::>(); + let step_scan = + match read_step_timestamps(&tx, &needed_occurrences, budget, &mut rows.database_bytes) { + Ok(scan) => scan, + Err(_) => return Ok((rows.events, false)), + }; + if !step_scan.complete { + return Ok((rows.events, false)); + } + + let resolved = resolve_step_timestamps(&step_scan.timestamps, &needed_occurrences); + let recovered = append_recovered_events( + &mut rows.events, + &session, + &rows.pending, + &resolved, + &rows.occurrences, + ); + if recovered < rows.pending.len() { + rows.complete = false; + } + Ok((rows.events, rows.complete)) +} + +fn read_generation_rows( + conn: &Connection, + session: &str, + budget: &mut Budget, +) -> rusqlite::Result { + let mut statement = conn.prepare( "SELECT idx, CASE WHEN typeof(data) = 'blob' THEN length(data) END, CASE WHEN typeof(data) = 'blob' AND length(data) <= ?2 THEN data END FROM main.gen_metadata NOT INDEXED LIMIT ?1", )?; let row_limit = i64::try_from(MAX_ROWS_PER_DATABASE + 1).unwrap_or(i64::MAX); @@ -291,6 +373,8 @@ fn read_database(path: &Path, budget: &mut Budget) -> rusqlite::Result<(Vec> = HashMap::new(); while let Some(row) = query.next()? { if !budget.check() { @@ -350,31 +434,198 @@ fn read_database(path: &Path, budget: &mut Budget) -> rusqlite::Result<(Vec events.push(Event { + session: session.to_string(), + row: idx, + turn, + total, + }), + (None, Some(step_uuid)) => pending.push(PendingTimestampRow { + row: idx, + step_uuid, + turn, + total, + }), + (None, None) => complete = false, } - let Some(input) = usage.system_prompt.checked_add(usage.new_input) else { + } + + Ok(ParsedRows { + events, + pending, + occurrences, + database_bytes, + complete, + }) +} + +fn token_total(usage: &super::local_proto::ParsedUsage) -> Option { + usage + .system_prompt + .checked_add(usage.new_input) + .and_then(|value| value.checked_add(usage.output)) + .and_then(|value| value.checked_add(usage.cache_read)) + .and_then(|value| value.checked_add(usage.reasoning)) +} + +fn read_step_timestamps( + conn: &Connection, + needed_occurrences: &HashMap>, + budget: &mut Budget, + database_bytes: &mut usize, +) -> rusqlite::Result { + let mut statement = conn.prepare( + "SELECT idx, CASE WHEN typeof(metadata) = 'blob' THEN length(metadata) END, CASE WHEN typeof(metadata) = 'blob' AND length(metadata) <= ?1 THEN metadata END FROM main.steps NOT INDEXED", + )?; + let blob_limit = i64::try_from(MAX_BLOB_BYTES).unwrap_or(i64::MAX); + let mut query = statement.query([blob_limit])?; + let mut timestamps = HashMap::>::new(); + let mut complete = true; + let mut rows_are_valid = true; + + while let Some(row) = query.next()? { + if !budget.check() { + complete = false; + break; + } + budget.rows += 1; + if budget.rows > MAX_ROWS { complete = false; + break; + } + + let idx: i64 = match row.get(0) { + Ok(value) if value >= 0 => value, + _ => { + rows_are_valid = false; + continue; + } + }; + let declared: Option = row.get(1).ok(); + let Some(declared) = declared.and_then(|value| usize::try_from(value).ok()) else { + rows_are_valid = false; + continue; + }; + *database_bytes = match (*database_bytes).checked_add(declared) { + Some(value) if value <= MAX_DATABASE_BYTES => value, + _ => { + complete = false; + break; + } + }; + budget.bytes = match budget.bytes.checked_add(declared) { + Some(value) if value <= MAX_TOTAL_BYTES => value, + _ => { + complete = false; + break; + } + }; + if declared == 0 || declared > MAX_BLOB_BYTES { + rows_are_valid = false; + continue; + } + + let blob = match row.get_ref(2)? { + ValueRef::Blob(bytes) if bytes.len() == declared => bytes, + _ => { + rows_are_valid = false; + continue; + } + }; + let Some((step_uuid, timestamp_ms)) = parse_step_metadata(blob) else { + rows_are_valid = false; + continue; + }; + let Some(step_uuid) = step_uuid.filter(|value| !value.is_empty()) else { + rows_are_valid = false; + continue; + }; + if needed_occurrences.contains_key(&step_uuid) { + timestamps + .entry(step_uuid) + .or_default() + .push(StepTimestamp { + row: idx, + timestamp_ms, + }); + } + } + + Ok(StepTimestampScan { + timestamps, + complete: complete && rows_are_valid, + }) +} + +fn append_recovered_events( + events: &mut Vec, + session: &str, + pending: &[PendingTimestampRow], + resolved: &HashMap>, + occurrences: &HashMap>, +) -> usize { + let mut occurrence_offsets = HashMap::>::new(); + for (step_uuid, occurrences) in occurrences { + let mut rows = occurrences + .iter() + .map(|occurrence| occurrence.row) + .collect::>(); + rows.sort_unstable(); + if rows.windows(2).any(|pair| pair[0] == pair[1]) { + continue; + } + occurrence_offsets.insert( + step_uuid.clone(), + rows.into_iter() + .enumerate() + .map(|(offset, row)| (row, offset)) + .collect(), + ); + } + + let mut recovered = 0; + let mut pending_rows = pending.iter().collect::>(); + pending_rows.sort_by_key(|pending| pending.row); + for pending in pending_rows { + let Some(timestamps) = resolved.get(&pending.step_uuid) else { continue; }; - let Some(total) = input - .checked_add(usage.output) - .and_then(|value| value.checked_add(usage.cache_read)) - .and_then(|value| value.checked_add(usage.reasoning)) + let Some(offset) = occurrence_offsets + .get(&pending.step_uuid) + .and_then(|rows| rows.get(&pending.row)) else { - complete = false; continue; }; + let Some(timestamp_ms) = timestamps.get(*offset) else { + continue; + }; + let mut turn = pending.turn.clone(); + turn.timestamp_ms = Some(*timestamp_ms); events.push(Event { - session: session.clone(), - row: idx, + session: session.to_string(), + row: pending.row, turn, - total, + total: pending.total, }); + recovered += 1; } - - Ok((events, complete)) + events.sort_by_key(|event| event.row); + recovered } fn supported_schema(conn: &Connection, budget: &mut Budget) -> rusqlite::Result { @@ -409,9 +660,59 @@ fn supported_schema(conn: &Connection, budget: &mut Budget) -> rusqlite::Result< return Ok(false); } + has_stored_columns(conn, "gen_metadata", &["idx", "data"], budget) +} + +fn supported_steps_schema(conn: &Connection, budget: &mut Budget) -> rusqlite::Result { + let mut statement = + conn.prepare("SELECT name, type, rootpage FROM main.sqlite_master LIMIT ?1")?; + let mut rows = statement.query([i64::try_from(MAX_SCHEMA_ENTRIES + 1).unwrap_or(i64::MAX)])?; + let mut found = false; + let mut schema_entries = 0usize; + while let Some(row) = rows.next()? { + if !budget.check() { + return Ok(false); + } + schema_entries += 1; + if schema_entries > MAX_SCHEMA_ENTRIES { + return Ok(false); + } + let name: String = row.get(0)?; + let kind: String = row.get(1)?; + if !budget.charge_schema_text(&name) || !budget.charge_schema_text(&kind) { + return Ok(false); + } + if !name.eq_ignore_ascii_case("steps") { + continue; + } + let rootpage: i64 = row.get(2)?; + if kind != "table" || rootpage <= 0 || found { + return Ok(false); + } + found = true; + } + if !found { + return Ok(false); + } + + has_stored_columns(conn, "steps", &["idx", "metadata"], budget) +} + +fn has_stored_columns( + conn: &Connection, + table: &str, + required: &[&str], + budget: &mut Budget, +) -> rusqlite::Result { + let table = match table { + "gen_metadata" | "steps" => table, + _ => return Ok(false), + }; + let mut columns = HashSet::new(); let mut schema_columns = 0usize; - let mut info = conn.prepare("PRAGMA main.table_xinfo('gen_metadata')")?; + let query = format!("PRAGMA main.table_xinfo('{table}')"); + let mut info = conn.prepare(&query)?; let mut rows = info.query([])?; while let Some(row) = rows.next()? { if !budget.check() { @@ -438,9 +739,13 @@ fn supported_schema(conn: &Connection, budget: &mut Budget) -> rusqlite::Result< } columns.insert(name.to_ascii_lowercase()); } - Ok(columns.contains("idx") && columns.contains("data")) + Ok(required.iter().all(|column| columns.contains(*column))) } +#[cfg(test)] +#[path = "local_sqlite_synthetic_tests.rs"] +mod synthetic_tests; + #[cfg(test)] mod tests { use super::*; diff --git a/rust/src/providers/antigravity/local_sqlite_synthetic_tests.rs b/rust/src/providers/antigravity/local_sqlite_synthetic_tests.rs new file mode 100644 index 0000000000..6b8d3a361a --- /dev/null +++ b/rust/src/providers/antigravity/local_sqlite_synthetic_tests.rs @@ -0,0 +1,297 @@ +use std::fs; +use std::path::PathBuf; + +use chrono::{DateTime, TimeZone, Utc}; +use rusqlite::{Connection, params}; +use tempfile::TempDir; + +use super::*; + +const NOW_SECONDS: u64 = 1_800_000_000; + +fn now() -> DateTime { + Utc.timestamp_opt(i64::try_from(NOW_SECONDS).unwrap(), 0) + .single() + .unwrap() +} + +fn varint(mut value: u64) -> Vec { + let mut bytes = Vec::new(); + loop { + let mut byte = (value & 0x7f) as u8; + value >>= 7; + if value != 0 { + byte |= 0x80; + } + bytes.push(byte); + if value == 0 { + return bytes; + } + } +} + +fn field_varint(number: u64, value: u64) -> Vec { + let mut bytes = varint(number << 3); + bytes.extend(varint(value)); + bytes +} + +fn field_bytes(number: u64, value: &[u8]) -> Vec { + let mut bytes = varint((number << 3) | 2); + bytes.extend(varint(value.len() as u64)); + bytes.extend(value); + bytes +} + +fn turn_blob(step_uuid: Option<&str>, input: u64, timestamp: Option) -> Vec { + let mut usage = field_varint(1, 11); + usage.extend(field_varint(2, input)); + usage.extend(field_varint(5, 50)); + usage.extend(field_varint(9, 30)); + usage.extend(field_varint(10, 7)); + + let mut chat = field_bytes(4, &usage); + if let Some(seconds) = timestamp { + let mut stamp = field_varint(1, seconds); + stamp.extend(field_varint(2, 250_000_000)); + chat.extend(field_bytes(9, &field_bytes(4, &stamp))); + } + + let mut root = Vec::new(); + if let Some(step_uuid) = step_uuid { + // Newer records can begin with the turn-coordination envelope (field 2). + root.extend(field_bytes(2, &[1, 2])); + root.extend(field_bytes(4, step_uuid.as_bytes())); + } + root.extend(field_bytes(1, &chat)); + root +} + +fn step_metadata(step_uuid: Option<&str>, timestamp: Option) -> Vec { + let mut metadata = Vec::new(); + if let Some(seconds) = timestamp { + let mut stamp = field_varint(1, seconds); + stamp.extend(field_varint(2, 0)); + metadata.extend(field_bytes(1, &stamp)); + } + if let Some(step_uuid) = step_uuid { + metadata.extend(field_bytes(12, step_uuid.as_bytes())); + } + metadata +} + +fn malformed_step_metadata(step_uuid: &str) -> Vec { + let mut metadata = field_bytes(12, step_uuid.as_bytes()); + metadata.extend([0x0a, 0x80]); + metadata +} + +type SyntheticStepRows<'a> = &'a [(i64, Option>)]; + +fn database( + dir: &TempDir, + session: &str, + generation_rows: &[(i64, Vec)], + step_rows: Option>, +) -> PathBuf { + let root = dir.path().join("conversations"); + fs::create_dir_all(&root).unwrap(); + let path = root.join(format!("{session}.db")); + let conn = Connection::open(&path).unwrap(); + conn.execute("CREATE TABLE gen_metadata (idx INTEGER, data BLOB)", []) + .unwrap(); + for (row, blob) in generation_rows { + conn.execute( + "INSERT INTO gen_metadata (idx, data) VALUES (?1, ?2)", + params![*row, blob.as_slice()], + ) + .unwrap(); + } + if let Some(step_rows) = step_rows { + conn.execute("CREATE TABLE steps (idx INTEGER, metadata BLOB)", []) + .unwrap(); + for (row, blob) in step_rows { + conn.execute( + "INSERT INTO steps (idx, metadata) VALUES (?1, ?2)", + params![*row, blob.as_deref()], + ) + .unwrap(); + } + } + drop(conn); + path +} + +fn summary(dir: &TempDir) -> LocalSessionSummary { + let root = dir.path().join("conversations"); + let SQLiteScan::Summary(summary) = summarize(&[root], now(), 30) else { + panic!("database should be attempted"); + }; + summary +} + +#[test] +fn newer_step_timestamp_recovery_preserves_legacy_and_new_totals() { + let dir = tempfile::tempdir().unwrap(); + let uuid = "new-step"; + database( + &dir, + "mixed", + &[ + (0, turn_blob(None, 100, Some(NOW_SECONDS - 120))), + (1, turn_blob(Some(uuid), 200, None)), + ], + Some(&[(10, Some(step_metadata(Some(uuid), Some(NOW_SECONDS - 60))))]), + ); + + let summary = summary(&dir); + + assert_eq!(summary.coverage, LocalHistoryCoverage::Complete); + assert_eq!(summary.total_tokens, 496); + assert_eq!(summary.session_count, 1); +} + +#[test] +fn absent_steps_table_withholds_newer_rows_but_keeps_embedded_totals() { + let dir = tempfile::tempdir().unwrap(); + database( + &dir, + "missing-table", + &[ + (0, turn_blob(None, 100, Some(NOW_SECONDS - 120))), + (1, turn_blob(Some("missing-table"), 200, None)), + ], + None, + ); + + let summary = summary(&dir); + + assert_eq!(summary.coverage, LocalHistoryCoverage::Partial); + assert_eq!(summary.total_tokens, 198); +} + +#[test] +fn reused_step_uuid_follows_stored_idx_order() { + let dir = tempfile::tempdir().unwrap(); + let uuid = "reused-step"; + database( + &dir, + "ordered", + &[ + (0, turn_blob(Some(uuid), 100, None)), + (1, turn_blob(Some(uuid), 200, None)), + ], + Some(&[ + (20, Some(step_metadata(Some(uuid), Some(NOW_SECONDS - 60)))), + (10, Some(step_metadata(Some(uuid), Some(NOW_SECONDS - 120)))), + ]), + ); + + let summary = summary(&dir); + + assert_eq!(summary.coverage, LocalHistoryCoverage::Complete); + assert_eq!(summary.total_tokens, 496); +} + +#[test] +fn missing_lowest_step_timestamp_withholds_pending_tokens() { + let dir = tempfile::tempdir().unwrap(); + let uuid = "missing-lowest"; + database( + &dir, + "missing", + &[(0, turn_blob(Some(uuid), 100, None))], + Some(&[ + (10, Some(step_metadata(Some(uuid), None))), + (20, Some(step_metadata(Some(uuid), Some(NOW_SECONDS - 60)))), + ]), + ); + + let summary = summary(&dir); + + assert_eq!(summary.coverage, LocalHistoryCoverage::Partial); + assert_eq!(summary.total_tokens, 0); +} + +#[test] +fn duplicate_or_malformed_step_rows_fail_closed() { + for (session, step_rows) in [ + ( + "duplicate", + vec![ + ( + 10, + Some(step_metadata(Some("duplicate"), Some(NOW_SECONDS - 120))), + ), + ( + 10, + Some(step_metadata(Some("duplicate"), Some(NOW_SECONDS - 60))), + ), + ], + ), + ( + "malformed", + vec![(10, Some(malformed_step_metadata("malformed")))], + ), + ] { + let dir = tempfile::tempdir().unwrap(); + database( + &dir, + session, + &[(0, turn_blob(Some(session), 100, None))], + Some(&step_rows), + ); + + let summary = summary(&dir); + + assert_eq!(summary.coverage, LocalHistoryCoverage::Partial, "{session}"); + assert_eq!(summary.total_tokens, 0, "{session}"); + } +} + +#[test] +fn null_or_unidentified_step_rows_fail_closed() { + for (session, step_row) in [ + ("null-step", (10, None)), + ( + "unidentified-step", + (10, Some(step_metadata(None, Some(NOW_SECONDS - 60)))), + ), + ] { + let dir = tempfile::tempdir().unwrap(); + database( + &dir, + session, + &[(0, turn_blob(Some(session), 100, None))], + Some(&[step_row]), + ); + + let summary = summary(&dir); + + assert_eq!(summary.coverage, LocalHistoryCoverage::Partial, "{session}"); + assert_eq!(summary.total_tokens, 0, "{session}"); + } +} + +#[test] +fn embedded_and_step_timestamps_must_agree() { + let dir = tempfile::tempdir().unwrap(); + let uuid = "conflicting"; + database( + &dir, + "conflict", + &[ + (0, turn_blob(Some(uuid), 100, Some(NOW_SECONDS - 120))), + (1, turn_blob(Some(uuid), 200, None)), + ], + Some(&[ + (10, Some(step_metadata(Some(uuid), Some(NOW_SECONDS - 60)))), + (20, Some(step_metadata(Some(uuid), Some(NOW_SECONDS - 30)))), + ]), + ); + + let summary = summary(&dir); + + assert_eq!(summary.coverage, LocalHistoryCoverage::Partial); + assert_eq!(summary.total_tokens, 198); +} diff --git a/rust/src/providers/antigravity/local_step_resolver.rs b/rust/src/providers/antigravity/local_step_resolver.rs new file mode 100644 index 0000000000..12bfb2e42c --- /dev/null +++ b/rust/src/providers/antigravity/local_step_resolver.rs @@ -0,0 +1,148 @@ +use std::collections::HashMap; + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(super) struct StepOccurrence { + pub row: i64, + pub timestamp_ms: Option, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(super) struct StepTimestamp { + pub row: i64, + pub timestamp_ms: Option, +} + +pub(super) fn resolve_step_timestamps( + step_timestamps: &HashMap>, + needed_occurrences: &HashMap>, +) -> HashMap> { + let mut resolved = HashMap::new(); + for (step_uuid, timestamps) in step_timestamps { + let Some(occurrences) = needed_occurrences.get(step_uuid) else { + continue; + }; + + let mut sorted_occurrences = occurrences.clone(); + sorted_occurrences.sort_by_key(|occurrence| occurrence.row); + if has_duplicate_rows(sorted_occurrences.iter().map(|occurrence| occurrence.row)) { + continue; + } + + let mut sorted_timestamps = timestamps.clone(); + sorted_timestamps.sort_by_key(|timestamp| timestamp.row); + if has_duplicate_rows(sorted_timestamps.iter().map(|timestamp| timestamp.row)) { + continue; + } + + let needed_count = sorted_occurrences.len(); + let ordered_timestamps = sorted_timestamps + .iter() + .map(|timestamp| timestamp.timestamp_ms) + .collect::>(); + let selected = if ordered_timestamps.len() == 1 { + let Some(timestamp) = ordered_timestamps[0] else { + continue; + }; + vec![timestamp; needed_count] + } else { + if ordered_timestamps.len() < needed_count { + continue; + } + let candidates = &ordered_timestamps[..needed_count]; + let Some(selected) = candidates.iter().copied().collect::>>() else { + continue; + }; + selected + }; + + if sorted_occurrences + .iter() + .zip(selected.iter()) + .any(|(occurrence, timestamp)| { + occurrence + .timestamp_ms + .is_some_and(|embedded| embedded != *timestamp) + }) + { + continue; + } + resolved.insert(step_uuid.clone(), selected); + } + resolved +} + +fn has_duplicate_rows(rows: impl Iterator) -> bool { + let mut prior = None; + for row in rows { + if prior == Some(row) { + return true; + } + prior = Some(row); + } + false +} + +#[cfg(test)] +mod tests { + use super::*; + + fn occurrences(rows: &[(i64, Option)]) -> Vec { + rows.iter() + .map(|&(row, timestamp_ms)| StepOccurrence { row, timestamp_ms }) + .collect() + } + + fn timestamps(rows: &[(i64, Option)]) -> Vec { + rows.iter() + .map(|&(row, timestamp_ms)| StepTimestamp { row, timestamp_ms }) + .collect() + } + + #[test] + fn orders_rows_and_repeats_one_shared_timestamp() { + let occurrences = + HashMap::from([("step".to_string(), occurrences(&[(0, None), (1, None)]))]); + let timestamp_map = HashMap::from([( + "step".to_string(), + timestamps(&[(20, Some(200)), (10, Some(100))]), + )]); + + assert_eq!( + resolve_step_timestamps(×tamp_map, &occurrences)["step"], + vec![100, 200] + ); + + let timestamps = HashMap::from([("step".to_string(), timestamps(&[(10, Some(100))]))]); + assert_eq!( + resolve_step_timestamps(×tamps, &occurrences)["step"], + vec![100, 100] + ); + } + + #[test] + fn withholds_missing_duplicate_and_conflicting_evidence() { + let occurrences = + HashMap::from([("step".to_string(), occurrences(&[(0, None), (1, None)]))]); + for timestamps in [ + timestamps(&[(10, None), (20, Some(200))]), + timestamps(&[(10, Some(100)), (10, Some(200))]), + ] { + let evidence = HashMap::from([("step".to_string(), timestamps)]); + assert!(!resolve_step_timestamps(&evidence, &occurrences).contains_key("step")); + } + } + + #[test] + fn embedded_timestamp_must_agree_with_aligned_step() { + let occurrences = HashMap::from([( + "step".to_string(), + occurrences(&[(0, Some(100)), (1, None)]), + )]); + let timestamps = HashMap::from([( + "step".to_string(), + timestamps(&[(10, Some(101)), (20, Some(200))]), + )]); + + assert!(!resolve_step_timestamps(×tamps, &occurrences).contains_key("step")); + } +} diff --git a/rust/src/providers/antigravity/mod.rs b/rust/src/providers/antigravity/mod.rs index 359b235b32..4ab5bbaada 100755 --- a/rust/src/providers/antigravity/mod.rs +++ b/rust/src/providers/antigravity/mod.rs @@ -6,6 +6,7 @@ mod local_proto; pub mod local_sessions; mod local_sqlite; +mod local_step_resolver; mod quota_summary; use async_trait::async_trait; diff --git a/rust/src/providers/claude/cloudflare_tests.rs b/rust/src/providers/claude/cloudflare_tests.rs new file mode 100644 index 0000000000..a636105192 --- /dev/null +++ b/rust/src/providers/claude/cloudflare_tests.rs @@ -0,0 +1,76 @@ +use super::{CLOUDFLARE_BODY_PREFIX_BYTES, classify_web_http_error}; +use crate::core::ProviderError; +use reqwest::{StatusCode, header}; + +#[test] +fn tagged_forbidden_response_returns_oauth_and_network_guidance() { + let mut headers = header::HeaderMap::new(); + headers.insert( + "cf-mitigated", + header::HeaderValue::from_static(" ChAlLeNgE "), + ); + + let error = + classify_web_http_error("usage", StatusCode::FORBIDDEN, &headers, b"challenge page"); + + assert_eq!( + error.to_string(), + crate::providers::claude::CLOUDFLARE_CHALLENGE_MESSAGE + ); + assert!(error.to_string().contains("OAuth")); + assert!(error.to_string().contains("different network")); +} + +#[test] +fn just_a_moment_fixture_is_detected_only_in_the_bounded_prefix() { + let error = classify_web_http_error( + "organizations", + StatusCode::FORBIDDEN, + &header::HeaderMap::new(), + br#"Just a moment..."#, + ); + assert_eq!( + error.to_string(), + crate::providers::claude::CLOUDFLARE_CHALLENGE_MESSAGE + ); + + let mut late_marker = vec![b'x'; CLOUDFLARE_BODY_PREFIX_BYTES]; + late_marker.extend_from_slice(b"Just a moment"); + assert!(matches!( + classify_web_http_error( + "usage", + StatusCode::FORBIDDEN, + &header::HeaderMap::new(), + &late_marker, + ), + ProviderError::AuthRequired + )); +} + +#[test] +fn ordinary_forbidden_and_all_unauthorized_responses_stay_auth_failures() { + let mut challenge_headers = header::HeaderMap::new(); + challenge_headers.insert( + "cf-mitigated", + header::HeaderValue::from_static("challenge"), + ); + + assert!(matches!( + classify_web_http_error( + "usage", + StatusCode::UNAUTHORIZED, + &challenge_headers, + b"Just a moment", + ), + ProviderError::AuthRequired + )); + assert!(matches!( + classify_web_http_error( + "usage", + StatusCode::FORBIDDEN, + &header::HeaderMap::new(), + b"permission denied", + ), + ProviderError::AuthRequired + )); +} diff --git a/rust/src/providers/claude/mod.rs b/rust/src/providers/claude/mod.rs index 0a01617703..700f21fbcd 100755 --- a/rust/src/providers/claude/mod.rs +++ b/rust/src/providers/claude/mod.rs @@ -19,8 +19,8 @@ use std::time::{Duration, Instant}; use crate::cli::tty_runner::{TtyCommandOptions, TtyCommandRunner}; use crate::core::{ - FetchContext, Provider, ProviderError, ProviderFetchResult, ProviderId, ProviderMetadata, - RateWindow, SourceMode, UsageSnapshot, + FetchContext, LastGoodFailurePolicy, Provider, ProviderError, ProviderFetchResult, ProviderId, + ProviderMetadata, RateWindow, SourceMode, UsageSnapshot, }; use admin_api::ClaudeAdminApiFetcher; @@ -79,6 +79,13 @@ fn is_oauth_revoked_error(error: &ProviderError) -> bool { pub use oauth::ClaudeOAuthFetcher; pub use web_api::ClaudeWebApiFetcher; +/// Recovery guidance for a Claude web request blocked by a Cloudflare challenge. +pub const CLOUDFLARE_CHALLENGE_MESSAGE: &str = concat!( + "claude.ai is behind a Cloudflare challenge, often caused by VPN or datacenter networks. ", + "Re-authenticating will not help. Switch Claude Usage source to OAuth in Settings ", + "(Usage credits balance will be unavailable), or try a different network." +); + /// Whether the user explicitly consented to reading (and refreshing) Claude /// Code's own credentials. Upstream #2634/#2745: without consent the /// file/keyring sources stay closed and refreshed tokens are never rotated @@ -383,6 +390,41 @@ async fn run_claude_pty_probe( }) } +fn last_good_failure_policy_for_error(error: &str) -> LastGoodFailurePolicy { + let lower = error.to_ascii_lowercase(); + if lower.contains("credentials not found") + || (lower.contains("run") && lower.contains("claude") && lower.contains("authenticate")) + || (lower.contains("not installed") && lower.contains("claude")) + || (lower.contains("subscription") && lower.contains("unavailable")) + { + return LastGoodFailurePolicy::Replace; + } + if lower.contains(&CLOUDFLARE_CHALLENGE_MESSAGE.to_ascii_lowercase()) { + return LastGoodFailurePolicy::PreserveOnceThenSurface; + } + if lower.contains("parse error") + || lower.contains("empty output") + || lower.contains("missing current session") + || lower.contains("treated /usage as a normal prompt") + || lower.contains("local activity stats") + || lower.contains("could not parse") + || lower.contains("rate limit") + || lower.contains("rate_limit") + || lower.contains("ratelimited") + || error.eq_ignore_ascii_case("timeout") + || lower.contains("timed out") + { + return LastGoodFailurePolicy::Preserve; + } + if lower.contains("unauthorized") + || lower.contains("authentication required") + || lower.contains("auth required") + { + return LastGoodFailurePolicy::PreserveOnce; + } + LastGoodFailurePolicy::Replace +} + #[async_trait] impl Provider for ClaudeProvider { fn id(&self) -> ProviderId { @@ -433,6 +475,14 @@ impl Provider for ClaudeProvider { true } + fn owns_browser_cookie_resolution(&self) -> bool { + true + } + + fn last_good_failure_policy(&self, error: &str) -> LastGoodFailurePolicy { + last_good_failure_policy_for_error(error) + } + fn detect_version(&self) -> Option { detect_claude_version() } diff --git a/rust/src/providers/claude/web_api.rs b/rust/src/providers/claude/web_api.rs index 3fd53d7ee7..ddebee622a 100755 --- a/rust/src/providers/claude/web_api.rs +++ b/rust/src/providers/claude/web_api.rs @@ -1,14 +1,59 @@ //! Claude Web API fetcher - uses browser cookies to fetch usage from claude.ai use chrono::{DateTime, Utc}; -use reqwest::{Client, header}; +use reqwest::{Client, StatusCode, header}; use serde::Deserialize; -use crate::browser::cookies::get_cookie_header; use crate::core::{ CostSnapshot, NamedRateWindow, ProviderError, ProviderFetchResult, RateWindow, UsageSnapshot, }; +use super::CLOUDFLARE_CHALLENGE_MESSAGE; + +const CLOUDFLARE_BODY_PREFIX_BYTES: usize = 64 * 1024; + +fn is_cloudflare_challenge_response( + status: StatusCode, + headers: &header::HeaderMap, + body: &[u8], +) -> bool { + if status != StatusCode::FORBIDDEN { + return false; + } + + if headers + .get("cf-mitigated") + .and_then(|value| value.to_str().ok()) + .map(str::trim) + .is_some_and(|value| value.eq_ignore_ascii_case("challenge")) + { + return true; + } + + let prefix = &body[..body.len().min(CLOUDFLARE_BODY_PREFIX_BYTES)]; + std::str::from_utf8(prefix) + .ok() + .is_some_and(|text| text.to_ascii_lowercase().contains("just a moment")) +} + +fn classify_web_http_error( + label: &str, + status: StatusCode, + headers: &header::HeaderMap, + body: &[u8], +) -> ProviderError { + if status == StatusCode::UNAUTHORIZED { + return ProviderError::AuthRequired; + } + if status == StatusCode::FORBIDDEN { + if is_cloudflare_challenge_response(status, headers, body) { + return ProviderError::Other(CLOUDFLARE_CHALLENGE_MESSAGE.to_string()); + } + return ProviderError::AuthRequired; + } + ProviderError::Other(format!("Failed to get {label}: {status}")) +} + /// Read the response body as text, then deserialize as JSON. On failure, include /// non-sensitive shape metadata so auth redirects, error envelopes, and schema /// changes are distinguishable without exposing account data in UI/log output. @@ -261,7 +306,6 @@ impl ClaudeWebApiFetcher { return self.fetch_with_cookie_header(&cookie_header).await; } - // Try multiple domains - Claude uses different domains for different services let domains = [ "claude.ai", "claude.com", @@ -269,22 +313,25 @@ impl ClaudeWebApiFetcher { "anthropic.com", ]; - for domain in domains { - match get_cookie_header(domain) { - Ok(cookie_header) if !cookie_header.is_empty() => { - tracing::debug!("Found cookies for {}", domain); - return self.fetch_with_cookie_header(&cookie_header).await; - } - Ok(_) => { - tracing::debug!("No cookies found for {}", domain); - } - Err(e) => { - tracing::debug!("Failed to get cookies for {}: {}", domain, e); + // A challenge is a network-path failure, not evidence that a cached + // session is invalid. Keep the last validated cookie for the next + // refresh and only invalidate it for an ordinary auth response. + use crate::browser::cookie_cache::CookieHeaderCache; + if let Some(cached) = CookieHeaderCache::load(crate::core::ProviderId::Claude) { + match self.fetch_with_cookie_header(&cached.cookie_header).await { + Ok(result) => return Ok(result), + Err(error) if is_cookie_authentication_failure(&error) => { + CookieHeaderCache::clear(crate::core::ProviderId::Claude); } + Err(error) => return Err(error), } } - Err(ProviderError::NoCookies) + let cookie_header = crate::providers::browser_cookie_header(&domains)?; + let result = self.fetch_with_cookie_header(&cookie_header).await?; + let _stored = + CookieHeaderCache::store(crate::core::ProviderId::Claude, &cookie_header, "browser"); + Ok(result) } /// Fetch usage with a provided cookie header @@ -470,11 +517,16 @@ impl ClaudeWebApiFetcher { .send() .await?; - if !response.status().is_success() { - return Err(ProviderError::Other(format!( - "Failed to get organizations: {}", - response.status() - ))); + let status = response.status(); + if !status.is_success() { + let response_headers = response.headers().clone(); + let body = response.bytes().await?; + return Err(classify_web_http_error( + "organizations", + status, + &response_headers, + &body, + )); } let orgs: Vec = parse_json_with_body(response, "organizations").await?; @@ -500,11 +552,16 @@ impl ClaudeWebApiFetcher { .send() .await?; - if !response.status().is_success() { - return Err(ProviderError::Other(format!( - "Failed to get usage: {}", - response.status() - ))); + let status = response.status(); + if !status.is_success() { + let response_headers = response.headers().clone(); + let body = response.bytes().await?; + return Err(classify_web_http_error( + "usage", + status, + &response_headers, + &body, + )); } parse_json_with_body(response, "usage").await @@ -679,6 +736,10 @@ impl Default for ClaudeWebApiFetcher { } } +fn is_cookie_authentication_failure(error: &ProviderError) -> bool { + matches!(error, ProviderError::AuthRequired) +} + fn cookie_value(cookie_header: &str, name: &str) -> Option { cookie_header.split(';').find_map(|part| { let (key, value) = part.trim().split_once('=')?; @@ -1234,3 +1295,7 @@ mod tests { assert_eq!(snapshot.extra_rate_windows.len(), 2); } } + +#[cfg(test)] +#[path = "cloudflare_tests.rs"] +mod cloudflare_tests; diff --git a/rust/src/providers/fixtures/claude/cloudflare-challenge.html b/rust/src/providers/fixtures/claude/cloudflare-challenge.html new file mode 100644 index 0000000000..08c6a47476 --- /dev/null +++ b/rust/src/providers/fixtures/claude/cloudflare-challenge.html @@ -0,0 +1,5 @@ + + + Just a moment... +

Checking your browser before accessing claude.ai.

+