Skip to content

fix: scope session-equivalent forecast history by account #276

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
zaszlo wants to merge 1 commit into
base: main
Choose a base branch
from
Open

fix: scope session-equivalent forecast history by account #276

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
63 changes: 58 additions & 5 deletions apps/desktop-tauri/src-tauri/src/commands/bridge.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -155,6 +155,7 @@ impl ProviderUsageSnapshot {
id: ProviderId,
metadata: &ProviderMetadata,
result: &ProviderFetchResult,
token_account_id: Option<uuid::Uuid>,
) -> Self {
let usage = &result.usage;

Expand Down Expand Up @@ -194,8 +195,17 @@ impl ProviderUsageSnapshot {
s
});

let session_equivalent_forecast =
session_equivalent_forecast_for(id, &usage.primary, usage.secondary.as_ref());
// Scope forecast history to the signed-in account so switching accounts on one
// provider does not blend burn samples across plans. Codex publishes no email or
// organization (ADR 0003 ambient/managed lanes), so its discriminator is the
// managed token-account id.
let account_key = forecast_account_key(usage, token_account_id);
let session_equivalent_forecast = session_equivalent_forecast_for(
id,
account_key.as_deref(),
&usage.primary,
usage.secondary.as_ref(),
);

Self {
provider_id: id.cli_name().to_string(),
Expand Down Expand Up @@ -302,8 +312,45 @@ impl ProviderUsageSnapshot {
}
}

/// Account discriminator that forecast history is scoped to.
///
/// Deliberately mirrors `quota_notification_account_identity` precedence
/// (token account -> email -> organization) so a single account is never seen as two
/// different identities by the notification and forecast subsystems. Kept as a separate
/// function because that one consumes an already-built `ProviderUsageSnapshot`, while the
/// forecast needs the key *while* the snapshot is being built.
///
/// `providers::tests::forecast_account_key_matches_notification_identity` pins them
/// together.
pub(super) fn forecast_account_key(
usage: &codexbar::core::UsageSnapshot,
token_account_id: Option<uuid::Uuid>,
) -> Option<String> {
if let Some(id) = token_account_id {
return Some(format!("token-account:{}", id.as_hyphenated()));
}
if let Some(email) = usage
.account_email
.as_deref()
.map(str::trim)
.filter(|s| !s.is_empty())
{
return Some(email.to_ascii_lowercase());
}
if let Some(org) = usage
.account_organization
.as_deref()
.map(str::trim)
.filter(|s| !s.is_empty())
{
return Some(format!("org:{}", org.to_ascii_lowercase()));
}
None
}

fn session_equivalent_forecast_for(
id: ProviderId,
account_key: Option<&str>,
session: &RateWindow,
weekly: Option<&RateWindow>,
) -> Option<SessionEquivalentForecastSnapshot> {
Expand All @@ -313,10 +360,16 @@ fn session_equivalent_forecast_for(
let weekly = weekly?;
let now = chrono::Utc::now();
let provider_id = id.cli_name();
codexbar::core::record_provider_windows(provider_id, session, Some(weekly), now);
codexbar::core::record_provider_windows(provider_id, account_key, session, Some(weekly), now);
let work_days = Settings::load().weekly_progress_work_days;
let forecast =
codexbar::core::forecast_for_provider(provider_id, session, weekly, now, work_days)?;
let forecast = codexbar::core::forecast_for_provider(
provider_id,
account_key,
session,
weekly,
now,
work_days,
)?;
Some(SessionEquivalentForecastSnapshot {
estimated_windows_to_exhaust_weekly: forecast.estimated_windows_to_exhaust_weekly,
windows_until_reset: forecast.windows_until_reset,
Expand Down
54 changes: 50 additions & 4 deletions apps/desktop-tauri/src-tauri/src/commands/providers.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -343,12 +343,19 @@ fn spawn_provider_refreshes(
&inputs.api_keys,
&inputs.token_accounts,
);
// Resolved here rather than inside the fetch: forecast history is keyed by
// account, and the managed-account id is the only discriminator Codex exposes.
let token_account_id = inputs
.token_accounts
.get(&id)
.and_then(ProviderAccountData::active_account)
.map(|account| account.id);

handles.push(tokio::spawn(async move {
let Ok(_permit) = fetch_permits.acquire_owned().await else {
return;
};
refresh_provider(app_handle, id, ctx, generation).await;
refresh_provider(app_handle, id, ctx, generation, token_account_id).await;
}));
}

Expand All @@ -373,8 +380,9 @@ async fn refresh_provider(
id: ProviderId,
ctx: FetchContext,
generation: u64,
token_account_id: Option<uuid::Uuid>,
) {
let snapshot = fetch_provider_snapshot(id, ctx).await;
let snapshot = fetch_provider_snapshot(id, ctx, token_account_id).await;

let state = app.state::<Mutex<AppState>>();
let published = if let Ok(mut guard) = state.lock() {
Expand Down Expand Up @@ -585,7 +593,11 @@ fn is_claude_timeout_failure(error: Option<&str>) -> bool {
error.eq_ignore_ascii_case("timeout") || error.to_ascii_lowercase().contains("timed out")
}

async fn fetch_provider_snapshot(id: ProviderId, ctx: FetchContext) -> ProviderUsageSnapshot {
async fn fetch_provider_snapshot(
id: ProviderId,
ctx: FetchContext,
token_account_id: Option<uuid::Uuid>,
) -> ProviderUsageSnapshot {
let provider = instantiate_provider(id);
let metadata = provider.metadata().clone();
let started = std::time::Instant::now();
Expand All @@ -594,7 +606,9 @@ async fn fetch_provider_snapshot(id: ProviderId, ctx: FetchContext) -> ProviderU
match tokio::time::timeout(provider_fetch_timeout(id, &ctx), provider.fetch_usage(&ctx))
.await
{
Ok(Ok(result)) => ProviderUsageSnapshot::from_fetch_result(id, &metadata, &result),
Ok(Ok(result)) => {
ProviderUsageSnapshot::from_fetch_result(id, &metadata, &result, token_account_id)
}
Ok(Err(e)) => ProviderUsageSnapshot::from_error(
id,
&metadata,
Expand Down Expand Up @@ -1006,6 +1020,38 @@ mod predictive_warning_tests {
snapshot.plan_name = None;
assert_eq!(quota_notification_account_identity(&snapshot, None), "");
}

/// The forecast scope key and the notification identity must never disagree.
/// If they did, one account would be seen as two identities and its burn history
/// would be split, silently halving the sample count behind every forecast.
#[test]
fn forecast_account_key_matches_notification_identity() {
use crate::commands::bridge::forecast_account_key;

let token = uuid::Uuid::parse_str("aaaaaaaa-aaaa-aaaa-aaaa-aaaaaaaaaaaa").unwrap();
let mut usage = codexbar::core::UsageSnapshot::new(codexbar::core::RateWindow::new(1.0));
let mut snapshot = empty_snapshot();

for (email, org) in [
(Some("Person@Example.com"), Some("Acme Org")),
(Some("Person@Example.com"), None),
(None, Some("Acme Org")),
(None, None),
] {
usage.account_email = email.map(str::to_string);
usage.account_organization = org.map(str::to_string);
snapshot.account_email = usage.account_email.clone();
snapshot.account_organization = usage.account_organization.clone();

for tok in [Some(token), None] {
assert_eq!(
forecast_account_key(&usage, tok).unwrap_or_default(),
quota_notification_account_identity(&snapshot, tok),
"identity drift for email={email:?} org={org:?} token={tok:?}"
);
}
}
}
}

#[cfg(test)]
Expand Down
28 changes: 18 additions & 10 deletions apps/desktop-tauri/src-tauri/src/commands/tests.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -774,7 +774,8 @@ fn provider_cache_upsert_replaces_existing_provider() {
wayfinder_usage: None,
source_label: "CLI".to_string(),
};
let mut first = ProviderUsageSnapshot::from_fetch_result(ProviderId::Codex, &metadata, &result);
let mut first =
ProviderUsageSnapshot::from_fetch_result(ProviderId::Codex, &metadata, &result, None);
let mut second = first.clone();
first.error = Some("old".to_string());
second.error = Some("new".to_string());
Expand All @@ -796,10 +797,11 @@ fn provider_cache_prunes_disabled_providers() {
wayfinder_usage: None,
source_label: "CLI".to_string(),
};
let codex = ProviderUsageSnapshot::from_fetch_result(ProviderId::Codex, &metadata, &result);
let codex =
ProviderUsageSnapshot::from_fetch_result(ProviderId::Codex, &metadata, &result, None);
let claude_meta = instantiate_provider(ProviderId::Claude).metadata().clone();
let claude =
ProviderUsageSnapshot::from_fetch_result(ProviderId::Claude, &claude_meta, &result);
ProviderUsageSnapshot::from_fetch_result(ProviderId::Claude, &claude_meta, &result, None);

let mut cache = vec![codex, claude];
super::prune_provider_cache_to_enabled(&mut cache, &[ProviderId::Codex]);
Expand All @@ -826,7 +828,7 @@ fn hiding_codex_spark_rows_preserves_other_extra_usage() {
source_label: "CLI".to_string(),
};
let mut snapshot =
ProviderUsageSnapshot::from_fetch_result(ProviderId::Codex, &metadata, &result);
ProviderUsageSnapshot::from_fetch_result(ProviderId::Codex, &metadata, &result, None);
snapshot.extra_rate_windows = vec![
NamedRateWindowSnapshot {
id: "codex-spark".to_string(),
Expand Down Expand Up @@ -855,7 +857,8 @@ fn claude_transient_auth_failure_preserves_first_last_good_snapshot() {
wayfinder_usage: None,
source_label: "OAuth".to_string(),
};
let good = ProviderUsageSnapshot::from_fetch_result(ProviderId::Claude, &metadata, &result);
let good =
ProviderUsageSnapshot::from_fetch_result(ProviderId::Claude, &metadata, &result, None);
let error = ProviderUsageSnapshot::from_error(
ProviderId::Claude,
&metadata,
Expand Down Expand Up @@ -883,7 +886,8 @@ fn claude_repeated_auth_failure_surfaces_error() {
wayfinder_usage: None,
source_label: "OAuth".to_string(),
};
let good = ProviderUsageSnapshot::from_fetch_result(ProviderId::Claude, &metadata, &result);
let good =
ProviderUsageSnapshot::from_fetch_result(ProviderId::Claude, &metadata, &result, None);
let first_error = ProviderUsageSnapshot::from_error(
ProviderId::Claude,
&metadata,
Expand Down Expand Up @@ -916,7 +920,8 @@ fn claude_cli_parse_failure_keeps_last_good_every_time() {
wayfinder_usage: None,
source_label: "CLI".to_string(),
};
let good = ProviderUsageSnapshot::from_fetch_result(ProviderId::Claude, &metadata, &result);
let good =
ProviderUsageSnapshot::from_fetch_result(ProviderId::Claude, &metadata, &result, None);
let err = ProviderUsageSnapshot::from_error(
ProviderId::Claude,
&metadata,
Expand Down Expand Up @@ -949,7 +954,8 @@ fn claude_hard_credentials_missing_does_not_preserve_stale() {
wayfinder_usage: None,
source_label: "OAuth".to_string(),
};
let good = ProviderUsageSnapshot::from_fetch_result(ProviderId::Claude, &metadata, &result);
let good =
ProviderUsageSnapshot::from_fetch_result(ProviderId::Claude, &metadata, &result, None);
let err = ProviderUsageSnapshot::from_error(
ProviderId::Claude,
&metadata,
Expand Down Expand Up @@ -1081,7 +1087,8 @@ fn japanese_provider_snapshot_localizes_weekly_label() {
source_label: "OAuth".to_string(),
};

let snapshot = ProviderUsageSnapshot::from_fetch_result(ProviderId::Claude, &metadata, &result);
let snapshot =
ProviderUsageSnapshot::from_fetch_result(ProviderId::Claude, &metadata, &result, None);

// Secondary label stays raw; localization happens at render time.
assert_eq!(snapshot.secondary_label, Some("Weekly".to_string()));
Expand Down Expand Up @@ -1109,7 +1116,8 @@ fn japanese_provider_snapshot_localizes_pace_reserve_description() {
source_label: "OAuth".to_string(),
};

let snapshot = ProviderUsageSnapshot::from_fetch_result(ProviderId::Claude, &metadata, &result);
let snapshot =
ProviderUsageSnapshot::from_fetch_result(ProviderId::Claude, &metadata, &result, None);

// Reserve data stays raw; localization happens at render time.
let secondary = snapshot.secondary.as_ref().expect("secondary window");
Expand Down
Loading
Loading