From 91ca2d80bded5f35aefb836ce18b81bdfae6cb20 Mon Sep 17 00:00:00 2001 From: Jeff Lloyd Date: Mon, 27 Jul 2026 15:22:02 -0400 Subject: [PATCH 1/8] [REV-1714] Use server-authoritative AI credit availability in the client Consume the new User.aiCreditAvailability GraphQL field (warp-server - Mirror the AICreditAvailability schema types and add cynic bindings, a focused GetAICreditAvailability query, and a piggybacked selection on the workspace metadata query. - Hold the shared server decision on AIRequestUsageModel with last-known-good semantics on refresh failures; has_any_ai_remaining now returns the server decision once one exists and only falls back to the legacy local derivation before the first successful fetch. - Feed the state from exactly two paths: the workspace metadata refresh piggyback and coalesced targeted fetches on auth completion, API-key/credential changes, workspace selection changes, and add-on credit/overage changes. Reset on logout. - Map server denial reasons to prompt alert presentation and gate the AI Assistant zero-state/prepared prompts on the shared availability. Co-Authored-By: Oz --- .../blocklist/agent_view/agent_message_bar.rs | 1 + app/src/ai/blocklist/prompt/prompt_alert.rs | 47 ++++- .../ai/blocklist/prompt/prompt_alert_tests.rs | 123 ++++++++++++ app/src/ai/credit_availability.rs | 111 +++++++++++ app/src/ai/credit_availability_tests.rs | 98 ++++++++++ app/src/ai/mod.rs | 2 + app/src/ai/request_usage_model.rs | 112 ++++++++++- app/src/ai/request_usage_model_tests.rs | 178 ++++++++++++++++++ app/src/ai_assistant/panel.rs | 7 +- app/src/ai_assistant/transcript.rs | 5 +- app/src/auth/auth_manager.rs | 1 + app/src/auth/mod.rs | 4 + app/src/code_review/comment_list_view.rs | 6 +- app/src/lib.rs | 28 +++ app/src/server/server_api/ai.rs | 38 +++- app/src/settings_view/ai_page.rs | 1 + app/src/settings_view/teams_page.rs | 4 + app/src/workspaces/gql_convert.rs | 1 + app/src/workspaces/update_manager.rs | 14 ++ app/src/workspaces/update_manager_tests.rs | 45 +++++ app/src/workspaces/user_workspaces.rs | 17 ++ app/src/workspaces/user_workspaces_tests.rs | 4 + crates/graphql/src/api/ai.rs | 30 +++ .../api/queries/get_ai_credit_availability.rs | 61 ++++++ .../get_workspaces_metadata_for_user.rs | 7 + crates/graphql/src/api/queries/mod.rs | 1 + crates/warp_graphql_schema/api/schema.graphql | 24 +++ 27 files changed, 957 insertions(+), 13 deletions(-) create mode 100644 app/src/ai/blocklist/prompt/prompt_alert_tests.rs create mode 100644 app/src/ai/credit_availability.rs create mode 100644 app/src/ai/credit_availability_tests.rs create mode 100644 crates/graphql/src/api/queries/get_ai_credit_availability.rs diff --git a/app/src/ai/blocklist/agent_view/agent_message_bar.rs b/app/src/ai/blocklist/agent_view/agent_message_bar.rs index 633a6e78602..7ccfe586aae 100644 --- a/app/src/ai/blocklist/agent_view/agent_message_bar.rs +++ b/app/src/ai/blocklist/agent_view/agent_message_bar.rs @@ -231,6 +231,7 @@ impl AgentMessageBar { if matches!( event, AIRequestUsageModelEvent::RequestUsageUpdated + | AIRequestUsageModelEvent::CreditAvailabilityUpdated | AIRequestUsageModelEvent::AmbientCreditsBannerDismissed ) { ctx.notify(); diff --git a/app/src/ai/blocklist/prompt/prompt_alert.rs b/app/src/ai/blocklist/prompt/prompt_alert.rs index 2161072c247..06ab9d06d77 100644 --- a/app/src/ai/blocklist/prompt/prompt_alert.rs +++ b/app/src/ai/blocklist/prompt/prompt_alert.rs @@ -12,6 +12,7 @@ use warpui::{ use crate::ai::AIRequestUsageModel; use crate::ai::blocklist::error_color; +use crate::ai::credit_availability::{AICreditAvailability, AICreditDenialReason}; use crate::auth::AuthStateProvider; use crate::network::NetworkStatus; use crate::server::ids::ServerId; @@ -142,6 +143,16 @@ impl PromptAlertView { } } + // The server-authoritative availability decision drives the alert once + // it has been fetched; local data below is only a pre-fetch fallback. + if let Some(availability) = request_usage_model.server_availability() { + return Self::state_from_server_availability(availability, app); + } + + // Legacy locally derived fallback, used only before the first + // successful availability fetch (e.g. right after startup or against + // servers that don't support the availability field yet). + // Next, make sure the user isn't delinquent in their plan. let workspace = UserWorkspaces::as_ref(app).current_workspace(); if workspace.is_some_and(|w| w.billing_metadata.is_delinquent_due_to_payment_issue()) { @@ -153,8 +164,38 @@ impl PromptAlertView { return PromptAlertState::NoAlert; } + Self::out_of_credits_presentation(app) + } + + /// Maps the server-authoritative availability decision to presentation + /// state. The server decides *whether* AI is available; workspace policy + /// only shapes the call-to-action copy. + fn state_from_server_availability( + availability: AICreditAvailability, + app: &AppContext, + ) -> PromptAlertState { + if availability.available { + return PromptAlertState::NoAlert; + } + + match availability.denial_reason { + AICreditDenialReason::Delinquent => PromptAlertState::DelinquentDueToPaymentIssue, + AICreditDenialReason::EnterpriseTeamSpendLimitHit + | AICreditDenialReason::EnterprisePerUserSpendLimitHit + | AICreditDenialReason::EnterpriseWorkspaceSpendLimitHit => { + PromptAlertState::MonthlyOveragesSpendLimitReached + } + AICreditDenialReason::None + | AICreditDenialReason::OutOfCredits + | AICreditDenialReason::Unknown => Self::out_of_credits_presentation(app), + } + } + + /// Picks the most actionable presentation for an out-of-credits denial + /// based on the current workspace's overage policy. + fn out_of_credits_presentation(app: &AppContext) -> PromptAlertState { // Check if overages are available. - if let Some(workspace) = workspace { + if let Some(workspace) = UserWorkspaces::as_ref(app).current_workspace() { let are_overages_toggleable = workspace.are_overages_toggleable(); let are_overages_enabled = workspace.are_overages_enabled(); @@ -478,3 +519,7 @@ impl TypedActionView for PromptAlertView { } } } + +#[cfg(test)] +#[path = "prompt_alert_tests.rs"] +mod tests; diff --git a/app/src/ai/blocklist/prompt/prompt_alert_tests.rs b/app/src/ai/blocklist/prompt/prompt_alert_tests.rs new file mode 100644 index 00000000000..c8e7da9c933 --- /dev/null +++ b/app/src/ai/blocklist/prompt/prompt_alert_tests.rs @@ -0,0 +1,123 @@ +use std::sync::Arc; + +use warpui::App; + +use super::*; +use crate::ai::credit_availability::AICreditSource; +use crate::server::server_api::ServerApiProvider; +use crate::server::server_api::team::MockTeamClient; +use crate::server::server_api::workspace::MockWorkspaceClient; + +fn initialize_app(app: &mut App) { + app.add_singleton_model(|_| NetworkStatus::new()); + app.add_singleton_model(|_| AuthStateProvider::new_for_test()); + app.add_singleton_model(|_| ServerApiProvider::new_for_test()); + app.add_singleton_model(|ctx| { + UserWorkspaces::mock( + Arc::new(MockTeamClient::new()), + Arc::new(MockWorkspaceClient::new()), + vec![], + ctx, + ) + }); + if app + .models_of_type::() + .is_empty() + { + app.update(crate::settings::init_and_register_user_preferences); + } + app.update(|ctx| { + warpui_extras::secure_storage::register_noop("test", ctx); + ctx.add_singleton_model(ApiKeyManager::new); + }); + app.add_singleton_model(|_| crate::pricing::PricingInfoModel::new()); + app.add_singleton_model(|ctx| { + AIRequestUsageModel::new_for_test(ServerApiProvider::as_ref(ctx).get_ai_client(), ctx) + }); +} + +fn apply_server_availability(app: &mut App, availability: AICreditAvailability) { + AIRequestUsageModel::handle(app).update(app, |model, ctx| { + model.apply_server_availability(Ok(availability), ctx); + }); +} + +fn determine_state(app: &mut App) -> PromptAlertState { + app.read(PromptAlertView::determine_state) +} + +#[test] +fn test_server_available_maps_to_no_alert() { + App::test((), |mut app| async move { + initialize_app(&mut app); + apply_server_availability( + &mut app, + AICreditAvailability::available_with_source(Some(AICreditSource::BaseLimit)), + ); + assert_eq!(determine_state(&mut app), PromptAlertState::NoAlert); + }); +} + +#[test] +fn test_server_delinquent_maps_to_delinquency_alert() { + App::test((), |mut app| async move { + initialize_app(&mut app); + apply_server_availability( + &mut app, + AICreditAvailability::unavailable(AICreditDenialReason::Delinquent), + ); + assert_eq!( + determine_state(&mut app), + PromptAlertState::DelinquentDueToPaymentIssue + ); + }); +} + +#[test] +fn test_server_spend_limit_reasons_map_to_spend_limit_alert() { + App::test((), |mut app| async move { + initialize_app(&mut app); + for reason in [ + AICreditDenialReason::EnterpriseTeamSpendLimitHit, + AICreditDenialReason::EnterprisePerUserSpendLimitHit, + AICreditDenialReason::EnterpriseWorkspaceSpendLimitHit, + ] { + apply_server_availability(&mut app, AICreditAvailability::unavailable(reason)); + assert_eq!( + determine_state(&mut app), + PromptAlertState::MonthlyOveragesSpendLimitReached, + "unexpected alert state for {reason:?}", + ); + } + }); +} + +#[test] +fn test_server_out_of_credits_maps_to_request_limit_reached() { + App::test((), |mut app| async move { + initialize_app(&mut app); + // With no workspace overage policy in play, an out-of-credits denial + // falls through to the generic request limit alert. + for reason in [ + AICreditDenialReason::OutOfCredits, + AICreditDenialReason::Unknown, + ] { + apply_server_availability(&mut app, AICreditAvailability::unavailable(reason)); + assert_eq!( + determine_state(&mut app), + PromptAlertState::RequestLimitReached, + "unexpected alert state for {reason:?}", + ); + } + }); +} + +#[test] +fn test_legacy_fallback_used_before_first_server_response() { + App::test((), |mut app| async move { + initialize_app(&mut app); + // No server availability applied: the default request limit info has + // requests remaining, so the legacy derivation reports no alert. + assert_eq!(determine_state(&mut app), PromptAlertState::NoAlert); + }); +} diff --git a/app/src/ai/credit_availability.rs b/app/src/ai/credit_availability.rs new file mode 100644 index 00000000000..50845cd05fb --- /dev/null +++ b/app/src/ai/credit_availability.rs @@ -0,0 +1,111 @@ +//! Domain types for the server-authoritative AI credit availability decision +//! (`User.aiCreditAvailability`). The server evaluates the same credit +//! waterfall used to authorize AI requests, so these values are the source of +//! truth for whether the user can start an interactive AI request. +use warp_graphql::ai::{ + AICreditAvailability as GqlAICreditAvailability, + AICreditAvailabilityDenialReason as GqlAICreditAvailabilityDenialReason, + AICreditAvailabilitySource as GqlAICreditAvailabilitySource, +}; + +/// Stable, client-safe reason the server reports when no inference access +/// exists. `None` is reported when the user is available. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum AICreditDenialReason { + None, + OutOfCredits, + Delinquent, + EnterpriseTeamSpendLimitHit, + EnterprisePerUserSpendLimitHit, + EnterpriseWorkspaceSpendLimitHit, + /// A reason from a newer server that this client version doesn't know. + /// Treated as a generic denial for presentation purposes. + Unknown, +} + +/// The credit source the server selected when inference access exists. +/// Capability-only access (e.g. BYO API key) has no credit source. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum AICreditSource { + BaseLimit, + BonusGrant, + Payg, + Overage, + AmbientBonusGrant, + /// A source from a newer server that this client version doesn't know. + Unknown, +} + +/// The server-authoritative answer to "can this user start an interactive AI +/// request right now". +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub struct AICreditAvailability { + pub available: bool, + pub denial_reason: AICreditDenialReason, + pub credit_source: Option, +} + +impl AICreditAvailability { + pub fn available_with_source(credit_source: Option) -> Self { + Self { + available: true, + denial_reason: AICreditDenialReason::None, + credit_source, + } + } + + pub fn unavailable(denial_reason: AICreditDenialReason) -> Self { + Self { + available: false, + denial_reason, + credit_source: None, + } + } +} + +impl From for AICreditDenialReason { + fn from(reason: GqlAICreditAvailabilityDenialReason) -> Self { + match reason { + GqlAICreditAvailabilityDenialReason::None => Self::None, + GqlAICreditAvailabilityDenialReason::OutOfCredits => Self::OutOfCredits, + GqlAICreditAvailabilityDenialReason::Delinquent => Self::Delinquent, + GqlAICreditAvailabilityDenialReason::EnterpriseTeamSpendLimitHit => { + Self::EnterpriseTeamSpendLimitHit + } + GqlAICreditAvailabilityDenialReason::EnterprisePerUserSpendLimitHit => { + Self::EnterprisePerUserSpendLimitHit + } + GqlAICreditAvailabilityDenialReason::EnterpriseWorkspaceSpendLimitHit => { + Self::EnterpriseWorkspaceSpendLimitHit + } + GqlAICreditAvailabilityDenialReason::Other(_) => Self::Unknown, + } + } +} + +impl From for AICreditSource { + fn from(source: GqlAICreditAvailabilitySource) -> Self { + match source { + GqlAICreditAvailabilitySource::BaseLimit => Self::BaseLimit, + GqlAICreditAvailabilitySource::BonusGrant => Self::BonusGrant, + GqlAICreditAvailabilitySource::Payg => Self::Payg, + GqlAICreditAvailabilitySource::Overage => Self::Overage, + GqlAICreditAvailabilitySource::AmbientBonusGrant => Self::AmbientBonusGrant, + GqlAICreditAvailabilitySource::Other(_) => Self::Unknown, + } + } +} + +impl From for AICreditAvailability { + fn from(availability: GqlAICreditAvailability) -> Self { + Self { + available: availability.available, + denial_reason: availability.denial_reason.into(), + credit_source: availability.credit_source.map(Into::into), + } + } +} + +#[cfg(test)] +#[path = "credit_availability_tests.rs"] +mod tests; diff --git a/app/src/ai/credit_availability_tests.rs b/app/src/ai/credit_availability_tests.rs new file mode 100644 index 00000000000..bbf61bc1095 --- /dev/null +++ b/app/src/ai/credit_availability_tests.rs @@ -0,0 +1,98 @@ +use warp_graphql::ai::{ + AICreditAvailability as GqlAICreditAvailability, + AICreditAvailabilityDenialReason as GqlDenialReason, AICreditAvailabilitySource as GqlSource, +}; + +use super::{AICreditAvailability, AICreditDenialReason, AICreditSource}; + +#[test] +fn converts_every_documented_denial_reason() { + let cases = [ + (GqlDenialReason::None, AICreditDenialReason::None), + ( + GqlDenialReason::OutOfCredits, + AICreditDenialReason::OutOfCredits, + ), + ( + GqlDenialReason::Delinquent, + AICreditDenialReason::Delinquent, + ), + ( + GqlDenialReason::EnterpriseTeamSpendLimitHit, + AICreditDenialReason::EnterpriseTeamSpendLimitHit, + ), + ( + GqlDenialReason::EnterprisePerUserSpendLimitHit, + AICreditDenialReason::EnterprisePerUserSpendLimitHit, + ), + ( + GqlDenialReason::EnterpriseWorkspaceSpendLimitHit, + AICreditDenialReason::EnterpriseWorkspaceSpendLimitHit, + ), + ]; + for (gql, expected) in cases { + assert_eq!(AICreditDenialReason::from(gql), expected); + } +} + +#[test] +fn converts_every_documented_credit_source() { + let cases = [ + (GqlSource::BaseLimit, AICreditSource::BaseLimit), + (GqlSource::BonusGrant, AICreditSource::BonusGrant), + (GqlSource::Payg, AICreditSource::Payg), + (GqlSource::Overage, AICreditSource::Overage), + ( + GqlSource::AmbientBonusGrant, + AICreditSource::AmbientBonusGrant, + ), + ]; + for (gql, expected) in cases { + assert_eq!(AICreditSource::from(gql), expected); + } +} + +#[test] +fn converts_unknown_enum_values_to_unknown() { + assert_eq!( + AICreditDenialReason::from(GqlDenialReason::Other("FUTURE_REASON".to_string())), + AICreditDenialReason::Unknown + ); + assert_eq!( + AICreditSource::from(GqlSource::Other("FUTURE_SOURCE".to_string())), + AICreditSource::Unknown + ); +} + +#[test] +fn converts_full_availability_payload() { + let available = AICreditAvailability::from(GqlAICreditAvailability { + available: true, + denial_reason: GqlDenialReason::None, + credit_source: Some(GqlSource::BaseLimit), + }); + assert_eq!( + available, + AICreditAvailability::available_with_source(Some(AICreditSource::BaseLimit)) + ); + + let capability_only = AICreditAvailability::from(GqlAICreditAvailability { + available: true, + denial_reason: GqlDenialReason::None, + credit_source: None, + }); + assert_eq!( + capability_only, + AICreditAvailability::available_with_source(None) + ); + + let denied = AICreditAvailability::from(GqlAICreditAvailability { + available: false, + denial_reason: GqlDenialReason::OutOfCredits, + credit_source: None, + }); + assert_eq!( + denied, + AICreditAvailability::unavailable(AICreditDenialReason::OutOfCredits) + ); +} diff --git a/app/src/ai/mod.rs b/app/src/ai/mod.rs index 90a52d71cbe..28f6ac22553 100644 --- a/app/src/ai/mod.rs +++ b/app/src/ai/mod.rs @@ -30,6 +30,7 @@ pub(crate) mod conversation_navigation; pub(crate) mod conversation_rename; pub(crate) mod conversation_status_ui; pub(crate) mod conversation_utils; +pub mod credit_availability; pub(crate) mod custom_model_router_editor; pub(crate) mod custom_model_routers; pub(crate) mod document; @@ -55,6 +56,7 @@ pub(crate) mod skills; pub(crate) mod tui_api_keys; pub(crate) mod voice; pub use agent_tips::*; +pub use credit_availability::*; pub use request_usage_model::*; use warpui::AppContext; #[cfg(not(target_family = "wasm"))] diff --git a/app/src/ai/request_usage_model.rs b/app/src/ai/request_usage_model.rs index 1aa50cec424..5fcf09dce5d 100644 --- a/app/src/ai/request_usage_model.rs +++ b/app/src/ai/request_usage_model.rs @@ -15,6 +15,7 @@ use warpui::{AppContext, Entity, ModelContext, SingletonEntity}; use crate::BlocklistAIHistoryModel; use crate::ai::agent::AIAgentExchangeId; use crate::ai::agent::conversation::AIConversationId; +use crate::ai::credit_availability::AICreditAvailability; use crate::auth::AuthStateProvider; use crate::pricing::PricingInfoModel; use crate::server::server_api::ai::AIClient; @@ -184,6 +185,24 @@ fn get_cached_ambient_credits_banner_dismissed(app_mut: &mut AppContext) -> bool .unwrap_or_default() } +/// The server-authoritative AI credit availability state shared by all AI +/// surfaces. It is fed from exactly two paths: the workspace metadata refresh +/// piggyback (primary cadence) and targeted fetches after meaningful state +/// changes (auth completion, plan/billing, workspace selection, credentials). +#[derive(Default)] +struct ServerAvailabilityState { + /// The last successfully fetched decision. Retained as last-known-good + /// when a refresh fails; never cleared by transient errors. + latest: Option, + /// When `latest` was last updated from a successful response. + last_success_time: Option, + /// Whether a targeted availability fetch is currently in flight. Used to + /// coalesce coincident event-triggered fetches. + refresh_in_flight: bool, + /// The most recent refresh failure, kept until the next success. + last_error: Option, +} + pub struct AIRequestUsageModel { ai_client: Arc, @@ -194,6 +213,8 @@ pub struct AIRequestUsageModel { bonus_grants: Vec, + server_availability: ServerAvailabilityState, + /// Whether the buy credits banner has been dismissed by the user. buy_addon_credits_banner_dismissed: bool, @@ -207,6 +228,8 @@ impl Entity for AIRequestUsageModel { pub enum AIRequestUsageModelEvent { RequestUsageUpdated, + /// The server-authoritative credit availability decision was updated. + CreditAvailabilityUpdated, AmbientCreditsBannerDismissed, RequestBonusRefunded { requests_refunded: i32, @@ -228,6 +251,7 @@ impl AIRequestUsageModel { request_limit_info, last_update_time: None, bonus_grants: vec![], + server_availability: ServerAvailabilityState::default(), buy_addon_credits_banner_dismissed: false, ambient_credits_banner_dismissed, } @@ -240,6 +264,7 @@ impl AIRequestUsageModel { last_update_time: None, request_limit_info: RequestLimitInfo::default(), bonus_grants: vec![], + server_availability: ServerAvailabilityState::default(), buy_addon_credits_banner_dismissed: false, ambient_credits_banner_dismissed: get_cached_ambient_credits_banner_dismissed(ctx), } @@ -310,6 +335,72 @@ impl AIRequestUsageModel { ctx.emit(AIRequestUsageModelEvent::RequestUsageUpdated); } + /// The server-authoritative availability decision, if one has been + /// successfully fetched this session (last-known-good on refresh failure). + pub fn server_availability(&self) -> Option { + self.server_availability.latest + } + + /// Records the outcome of an availability fetch, whether piggybacked on a + /// workspace metadata refresh or from a targeted fetch. + /// + /// A failure keeps the last-known-good decision: a transport or resolver + /// error must never flip availability in either direction, and legacy + /// locally derived availability must not be re-enabled once a valid server + /// decision has been received. + pub fn apply_server_availability( + &mut self, + result: Result, + ctx: &mut ModelContext, + ) { + match result { + Ok(availability) => { + self.server_availability.latest = Some(availability); + self.server_availability.last_success_time = Some(Instant::now()); + self.server_availability.last_error = None; + ctx.emit(AIRequestUsageModelEvent::CreditAvailabilityUpdated); + ctx.notify(); + } + Err(e) => { + log::warn!("Failed to refresh AI credit availability: {e:#}"); + self.server_availability.last_error = Some(format!("{e:#}")); + } + } + } + + /// Fetches the server-authoritative availability decision in response to a + /// meaningful state change (auth completion, plan/billing change, + /// workspace selection change, or API-key/credential change). + /// + /// Coincident triggers are coalesced: if a fetch is already in flight this + /// is a no-op. There is intentionally no retry loop — the next qualifying + /// trigger or workspace metadata refresh serves as the retry. + pub fn request_availability_refresh(&mut self, ctx: &mut ModelContext) { + if !AuthStateProvider::as_ref(ctx).get().is_logged_in() { + return; + } + if self.server_availability.refresh_in_flight { + return; + } + self.server_availability.refresh_in_flight = true; + + let ai_client = self.ai_client.clone(); + ctx.spawn( + async move { ai_client.get_ai_credit_availability().await }, + |model, result, ctx| { + model.server_availability.refresh_in_flight = false; + model.apply_server_availability(result, ctx); + }, + ); + } + + /// Clears the server-authoritative availability state, e.g. on logout. + pub fn reset_server_availability(&mut self, ctx: &mut ModelContext) { + self.server_availability = ServerAvailabilityState::default(); + ctx.emit(AIRequestUsageModelEvent::CreditAvailabilityUpdated); + ctx.notify(); + } + pub fn provide_negative_feedback_response_for_ai_conversation( &mut self, client_conversation_id: AIConversationId, @@ -421,7 +512,23 @@ impl AIRequestUsageModel { self.requests_remaining() > 0 } - /// Returns `true` if the user meets one of the following conditions: + /// Returns `true` if the user can start an interactive AI request. + /// Use this method as the starting point for AI availability checking. + /// + /// Once a server-authoritative availability decision has been received + /// this session, that decision (last-known-good on transient refresh + /// failures) is the only authority. The locally derived fallback below is + /// used solely before the first successful fetch, e.g. right after startup + /// or against servers that don't support the availability field yet. + pub fn has_any_ai_remaining(&self, ctx: &AppContext) -> bool { + if let Some(availability) = self.server_availability.latest { + return availability.available; + } + self.has_any_ai_remaining_from_local_state(ctx) + } + + /// Legacy locally derived availability check. Returns `true` if the user + /// meets one of the following conditions: /// 1. user has ai credits from the plan base limit /// 2. user has overage enabled /// 3. user has bonus grants (either team grants or user grants) @@ -430,8 +537,7 @@ impl AIRequestUsageModel { /// 6. user's team has self-serve auto-reload enabled within its monthly spend limit /// 7. user has BYOK enabled and has either provided at least one API key or /// connected a Grok subscription - /// Use this method as the starting point for AI availability checking. - pub fn has_any_ai_remaining(&self, ctx: &AppContext) -> bool { + fn has_any_ai_remaining_from_local_state(&self, ctx: &AppContext) -> bool { let current_workspace = UserWorkspaces::as_ref(ctx).current_workspace(); let has_base_plan_ai_requests = self.has_requests_remaining(); diff --git a/app/src/ai/request_usage_model_tests.rs b/app/src/ai/request_usage_model_tests.rs index f9b08827b45..8b3642692e7 100644 --- a/app/src/ai/request_usage_model_tests.rs +++ b/app/src/ai/request_usage_model_tests.rs @@ -8,9 +8,11 @@ use warp_graphql::billing::{AddonCreditsOption, OveragesPricing, PricingInfo}; use warpui::{App, ModelHandle}; use super::*; +use crate::ai::credit_availability::{AICreditDenialReason, AICreditSource}; use crate::auth::AuthStateProvider; use crate::pricing::PricingInfoModel; use crate::server::server_api::ServerApiProvider; +use crate::server::server_api::ai::MockAIClient; use crate::server::server_api::team::MockTeamClient; use crate::server::server_api::workspace::MockWorkspaceClient; use crate::workspaces::user_workspaces::UserWorkspaces; @@ -73,6 +75,16 @@ fn add_request_usage_model_without_auth(app: &mut App) -> ModelHandle, +) -> ModelHandle { + register_user_preferences_for_tests(app); + app.add_singleton_model(|ctx| AIRequestUsageModel::new_for_test(ai_client, ctx)) +} + fn set_addon_credits_pricing_info(app: &mut App) { PricingInfoModel::handle(app).update(app, |model, ctx| { model.update_pricing_info( @@ -904,3 +916,169 @@ fn test_has_any_ai_remaining_false_with_only_ambient_bonus_credits() { }); }); } + +#[test] +fn test_server_availability_overrides_locally_derived_state() { + App::test((), |mut app| async move { + app.add_singleton_model(UserWorkspaces::default_mock); + let request_usage_model = add_request_usage_model(&mut app); + + request_usage_model.update(&mut app, |model, ctx| { + // Local state says AI is available. + model.request_limit_info = RequestLimitInfo::new_for_test(10, 5); + assert!(model.has_any_ai_remaining(ctx)); + + // The server-authoritative decision wins over local state. + model.apply_server_availability( + Ok(AICreditAvailability::unavailable( + AICreditDenialReason::OutOfCredits, + )), + ctx, + ); + assert!(!model.has_any_ai_remaining(ctx)); + + // And in the other direction: local state is exhausted, but the + // server reports a usable fallback source. + model.request_limit_info = RequestLimitInfo::new_for_test(10, 10); + model.apply_server_availability( + Ok(AICreditAvailability::available_with_source(Some( + AICreditSource::BonusGrant, + ))), + ctx, + ); + assert!(model.has_any_ai_remaining(ctx)); + }); + }); +} + +#[test] +fn test_availability_refresh_failure_keeps_last_known_good() { + App::test((), |mut app| async move { + app.add_singleton_model(UserWorkspaces::default_mock); + let request_usage_model = add_request_usage_model(&mut app); + + request_usage_model.update(&mut app, |model, ctx| { + // Local state says AI is available, so any legacy fallback would + // report `true`. + model.request_limit_info = RequestLimitInfo::new_for_test(10, 5); + + let denied = AICreditAvailability::unavailable(AICreditDenialReason::Delinquent); + model.apply_server_availability(Ok(denied), ctx); + model.apply_server_availability(Err(anyhow::anyhow!("transient failure")), ctx); + + // The last-known-good server decision is retained: the error is + // recorded but neither flips availability nor re-enables the + // legacy locally derived decision. + assert_eq!(model.server_availability(), Some(denied)); + assert!(!model.has_any_ai_remaining(ctx)); + assert_eq!( + model.server_availability.last_error.as_deref(), + Some("transient failure") + ); + }); + }); +} + +#[test] +fn test_availability_refresh_failure_before_first_success_uses_legacy_fallback() { + App::test((), |mut app| async move { + app.add_singleton_model(UserWorkspaces::default_mock); + let request_usage_model = add_request_usage_model(&mut app); + + request_usage_model.update(&mut app, |model, ctx| { + model.request_limit_info = RequestLimitInfo::new_for_test(10, 5); + model.apply_server_availability(Err(anyhow::anyhow!("unsupported operation")), ctx); + + // Without any successful fetch (e.g. server doesn't support the + // field yet), the legacy locally derived decision still applies. + assert_eq!(model.server_availability(), None); + assert!(model.has_any_ai_remaining(ctx)); + }); + }); +} + +#[test] +fn test_reset_server_availability_restores_legacy_fallback() { + App::test((), |mut app| async move { + app.add_singleton_model(UserWorkspaces::default_mock); + let request_usage_model = add_request_usage_model(&mut app); + + request_usage_model.update(&mut app, |model, ctx| { + model.request_limit_info = RequestLimitInfo::new_for_test(10, 5); + model.apply_server_availability( + Ok(AICreditAvailability::unavailable( + AICreditDenialReason::OutOfCredits, + )), + ctx, + ); + assert!(!model.has_any_ai_remaining(ctx)); + + // On logout the server decision is cleared and pre-fetch behavior + // is restored for the next principal. + model.reset_server_availability(ctx); + assert_eq!(model.server_availability(), None); + assert!(model.has_any_ai_remaining(ctx)); + }); + }); +} + +#[test] +fn test_availability_refresh_coalesces_concurrent_fetches() { + App::test((), |mut app| async move { + app.add_singleton_model(|_| AuthStateProvider::new_for_test()); + + let mut ai_client = MockAIClient::new(); + // Exactly one fetch may go out even though two triggers fire while the + // first request is still in flight. + ai_client + .expect_get_ai_credit_availability() + .times(1) + .returning(|| { + Ok(AICreditAvailability::available_with_source(Some( + AICreditSource::BaseLimit, + ))) + }); + let request_usage_model = + add_request_usage_model_with_client(&mut app, Arc::new(ai_client)); + + request_usage_model.update(&mut app, |model, ctx| { + model.request_availability_refresh(ctx); + model.request_availability_refresh(ctx); + assert!(model.server_availability.refresh_in_flight); + }); + + // Let the spawned fetch complete. + warpui::r#async::Timer::after(std::time::Duration::from_millis(100)).await; + + request_usage_model.read(&app, |model, _| { + assert!(!model.server_availability.refresh_in_flight); + assert_eq!( + model.server_availability(), + Some(AICreditAvailability::available_with_source(Some( + AICreditSource::BaseLimit, + ))) + ); + }); + }); +} + +#[test] +fn test_availability_refresh_skipped_when_logged_out() { + App::test((), |mut app| async move { + app.add_singleton_model(|_| AuthStateProvider::new_logged_out_for_test()); + + // No expectations: any fetch would panic the test. + let ai_client = MockAIClient::new(); + let request_usage_model = + add_request_usage_model_with_client(&mut app, Arc::new(ai_client)); + + request_usage_model.update(&mut app, |model, ctx| { + model.request_availability_refresh(ctx); + assert!(!model.server_availability.refresh_in_flight); + }); + + request_usage_model.read(&app, |model, _| { + assert_eq!(model.server_availability(), None); + }); + }); +} diff --git a/app/src/ai_assistant/panel.rs b/app/src/ai_assistant/panel.rs index 7bf9f67d923..5a9af183c7d 100644 --- a/app/src/ai_assistant/panel.rs +++ b/app/src/ai_assistant/panel.rs @@ -33,6 +33,7 @@ use super::{ AI_ASSISTANT_FEATURE_NAME, AI_ASSISTANT_LOGO_COLOR, AI_ASSISTANT_SVG_PATH, ASK_AI_ASSISTANT_TEXT, AskAIType, PROMPT_CHARACTER_LIMIT, }; +use crate::ai::AIRequestUsageModel; use crate::appearance::Appearance; use crate::editor::{ EditorOptions, EditorView, Event as EditorEvent, PropagateAndNoOpNavigationKeys, TextOptions, @@ -688,10 +689,6 @@ impl AIAssistantPanelView { self.requests_model.as_ref(app).request_status() } - fn num_remaining_reqs(&self, app: &AppContext) -> usize { - self.requests_model.as_ref(app).num_remaining_reqs() - } - #[cfg(feature = "integration_tests")] pub fn editor(&self) -> &ViewHandle { &self.editor @@ -916,7 +913,7 @@ impl AIAssistantPanelView { .finish(), ); - if self.num_remaining_reqs(app) > 0 { + if AIRequestUsageModel::as_ref(app).has_any_ai_remaining(app) { column.add_children([ Container::new(render_prepared_response_button( appearance, diff --git a/app/src/ai_assistant/transcript.rs b/app/src/ai_assistant/transcript.rs index c4388a92fd4..ed9fc6b16d8 100644 --- a/app/src/ai_assistant/transcript.rs +++ b/app/src/ai_assistant/transcript.rs @@ -28,6 +28,7 @@ use super::utils::{ TranscriptPartSubType, code_block_position_id, markdown_segments_from_text, render_prepared_response_button, render_request_limit_info, save_as_workflow_position_id, }; +use crate::ai::AIRequestUsageModel; use crate::appearance::Appearance; use crate::send_telemetry_from_ctx; use crate::server::telemetry::{SaveAsWorkflowModalSource, TelemetryEvent, WarpAIActionType}; @@ -816,7 +817,7 @@ impl View for Transcript { let theme = appearance.theme(); let transcript = self.requests_model.as_ref(app).transcript(); let request_status = self.requests_model.as_ref(app).request_status(); - let num_remaining_reqs = self.requests_model.as_ref(app).num_remaining_reqs(); + let has_ai_available = AIRequestUsageModel::as_ref(app).has_any_ai_remaining(app); let mut blocks = Flex::column(); for (index, part) in transcript.iter().enumerate() { @@ -850,7 +851,7 @@ impl View for Transcript { if !transcript.is_empty() && matches!(request_status, RequestStatus::NotInFlight) { // Only show the prepared responses if the last response wasn't an error // and the user still has remaining requests. - if !transcript.last().is_none_or(|p| p.assistant.is_error) && num_remaining_reqs > 0 { + if !transcript.last().is_none_or(|p| p.assistant.is_error) && has_ai_available { blocks.add_child( Container::new(self.render_prepared_responses(appearance)) .with_margin_top(15.) diff --git a/app/src/auth/auth_manager.rs b/app/src/auth/auth_manager.rs index 2ae8969344b..7e2f4e02856 100644 --- a/app/src/auth/auth_manager.rs +++ b/app/src/auth/auth_manager.rs @@ -432,6 +432,7 @@ impl AuthManager { AIRequestUsageModel::handle(ctx).update(ctx, |usage_model, ctx| { usage_model.refresh_request_usage_async(ctx); + usage_model.request_availability_refresh(ctx); }); LLMPreferences::handle(ctx).update(ctx, |prefs, ctx| { diff --git a/app/src/auth/mod.rs b/app/src/auth/mod.rs index 175c62c92b3..f4ce22281f4 100644 --- a/app/src/auth/mod.rs +++ b/app/src/auth/mod.rs @@ -33,6 +33,7 @@ use crate::ai::blocklist::agent_view::orchestration_pill_bar_model::Orchestratio use crate::ai::execution_profiles::profiles::AIExecutionProfilesModel; #[cfg(not(target_family = "wasm"))] use crate::ai::mcp::TemplatableMCPServerManager; +use crate::ai::request_usage_model::AIRequestUsageModel; use crate::ai_assistant::requests::REQUEST_LIMIT_INFO_CACHE_KEY; use crate::cloud_object::model::persistence::CloudModel; use crate::code::editor_management::{CodeEditorStatus, CodeEditorSummary}; @@ -254,6 +255,9 @@ pub fn log_out(app: &mut AppContext) { TemplatableMCPServerManager::handle(app).update(app, |manager, ctx| { manager.sync_builtin_servers(false, ctx); }); + AIRequestUsageModel::handle(app).update(app, |usage_model, ctx| { + usage_model.reset_server_availability(ctx); + }); BlocklistAIHistoryModel::handle(app).update(app, |history_model, _| { history_model.reset(); }); diff --git a/app/src/code_review/comment_list_view.rs b/app/src/code_review/comment_list_view.rs index 1db6d8cdaef..e5276f57470 100644 --- a/app/src/code_review/comment_list_view.rs +++ b/app/src/code_review/comment_list_view.rs @@ -235,7 +235,11 @@ impl CommentListView { // Keep the stored button state in sync when AI availability changes. ctx.subscribe_to_model(&AIRequestUsageModel::handle(ctx), |me, _, event, ctx| { - if let AIRequestUsageModelEvent::RequestUsageUpdated = event { + if matches!( + event, + AIRequestUsageModelEvent::RequestUsageUpdated + | AIRequestUsageModelEvent::CreditAvailabilityUpdated + ) { me.sync_send_button(ctx); } }); diff --git a/app/src/lib.rs b/app/src/lib.rs index 61419492656..f1f2dfd5e96 100644 --- a/app/src/lib.rs +++ b/app/src/lib.rs @@ -1680,6 +1680,34 @@ pub(crate) fn initialize_app( manager }); + // Keep the server-authoritative AI credit availability fresh across + // meaningful state changes. The primary cadence is the workspace metadata + // refresh piggyback; these targeted triggers cover changes that don't + // immediately produce a metadata response. `TeamsChanged` is deliberately + // not a trigger: it fires on every metadata poll, whose response already + // carries the piggybacked availability. + ctx.subscribe_to_model(&UserWorkspaces::handle(ctx), |_, event, ctx| { + if matches!( + event, + UserWorkspacesEvent::CurrentWorkspaceChanged + | UserWorkspacesEvent::AiOveragesUpdated + | UserWorkspacesEvent::PurchaseAddonCreditsSuccess + ) { + AIRequestUsageModel::handle(ctx).update(ctx, |usage_model, ctx| { + usage_model.request_availability_refresh(ctx); + }); + } + }); + ctx.subscribe_to_model( + &::ai::api_keys::ApiKeyManager::handle(ctx), + |_, event, ctx| { + let ::ai::api_keys::ApiKeyManagerEvent::KeysUpdated = event; + AIRequestUsageModel::handle(ctx).update(ctx, |usage_model, ctx| { + usage_model.request_availability_refresh(ctx); + }); + }, + ); + ctx.add_singleton_model(AntivirusInfo::new); cfg_if::cfg_if! { diff --git a/app/src/server/server_api/ai.rs b/app/src/server/server_api/ai.rs index e4853cd18fe..af29e050b51 100644 --- a/app/src/server/server_api/ai.rs +++ b/app/src/server/server_api/ai.rs @@ -76,6 +76,10 @@ use warp_graphql::queries::free_available_models::{ FreeAvailableModels, FreeAvailableModelsInput, FreeAvailableModelsResult, FreeAvailableModelsVariables, }; +#[cfg(not(feature = "agent_mode_evals"))] +use warp_graphql::queries::get_ai_credit_availability::{ + GetAICreditAvailability, GetAICreditAvailabilityVariables, +}; use warp_graphql::queries::get_available_harnesses::{ GetAvailableHarnesses, GetAvailableHarnessesVariables, }; @@ -117,7 +121,6 @@ use super::download::write_response_body_to_path; use super::harness_support::{UploadField, UploadFieldValue, UploadTarget}; #[cfg(not(feature = "agent_mode_evals"))] use crate::ai::BonusGrant; -use crate::ai::RequestUsageInfo; pub use crate::ai::agent::UserQueryMode; use crate::ai::agent::api::ServerConversationToken; use crate::ai::agent::conversation::{ @@ -141,6 +144,7 @@ use crate::ai::llms::{ }; #[cfg(feature = "agent_mode_evals")] use crate::ai::request_usage_model::RequestLimitInfo; +use crate::ai::{AICreditAvailability, RequestUsageInfo}; use crate::ai_assistant::execution_context::WarpAiExecutionContext; use crate::ai_assistant::requests::GenerateDialogueResult; use crate::ai_assistant::utils::TranscriptPart; @@ -1157,6 +1161,10 @@ pub trait AIClient: 'static + Send + Sync { async fn get_request_limit_info(&self) -> Result; + /// Fetches the server-authoritative decision on whether the authenticated + /// user can start an interactive AI request. + async fn get_ai_credit_availability(&self) -> Result; + /// Returns conversation usage history for the current user over the requested number of days. /// /// If `last_updated_end_timestamp` is provided, only conversations updated before that timestamp are returned. @@ -1810,6 +1818,34 @@ impl AIClient for ServerApi { } } + #[cfg(feature = "agent_mode_evals")] + async fn get_ai_credit_availability(&self) -> Result { + Ok(AICreditAvailability::available_with_source(Some( + crate::ai::AICreditSource::BaseLimit, + ))) + } + + #[cfg(not(feature = "agent_mode_evals"))] + async fn get_ai_credit_availability(&self) -> Result { + let variables = GetAICreditAvailabilityVariables { + request_context: get_request_context(), + }; + let operation = GetAICreditAvailability::build(variables); + let response = self.send_graphql_request(operation, None).await?; + + match response.user { + warp_graphql::queries::get_ai_credit_availability::UserResult::UserOutput(output) => { + Ok(output.user.ai_credit_availability.into()) + } + warp_graphql::queries::get_ai_credit_availability::UserResult::UserFacingError(e) => { + Err(anyhow!(get_user_facing_error_message(e))) + } + warp_graphql::queries::get_ai_credit_availability::UserResult::Unknown => { + Err(anyhow!("failed to get AI credit availability")) + } + } + } + async fn get_conversation_usage_history( &self, days: Option, diff --git a/app/src/settings_view/ai_page.rs b/app/src/settings_view/ai_page.rs index 1d235f75284..25c20af1831 100644 --- a/app/src/settings_view/ai_page.rs +++ b/app/src/settings_view/ai_page.rs @@ -1796,6 +1796,7 @@ impl AISettingsPageView { ctx.subscribe_to_model(&ai_request_model, |me, _, event, ctx| { match event { AIRequestUsageModelEvent::RequestUsageUpdated => ctx.notify(), + AIRequestUsageModelEvent::CreditAvailabilityUpdated => ctx.notify(), AIRequestUsageModelEvent::RequestBonusRefunded { .. } => ctx.notify(), AIRequestUsageModelEvent::AmbientCreditsBannerDismissed => {} } diff --git a/app/src/settings_view/teams_page.rs b/app/src/settings_view/teams_page.rs index 6f1b74dca38..a7182abea6a 100644 --- a/app/src/settings_view/teams_page.rs +++ b/app/src/settings_view/teams_page.rs @@ -933,6 +933,10 @@ impl TeamsPageView { ctx.emit(TeamsPageViewEvent::TeamsChanged); } + UserWorkspacesEvent::CurrentWorkspaceChanged => { + // A workspace selection change always emits `TeamsChanged` too, + // which already refreshes this page. + } UserWorkspacesEvent::ToggleInviteLinksSuccess => { self.show_success("Toggled invite links", ctx); ctx.notify(); diff --git a/app/src/workspaces/gql_convert.rs b/app/src/workspaces/gql_convert.rs index 192307f4101..6514f66a20d 100644 --- a/app/src/workspaces/gql_convert.rs +++ b/app/src/workspaces/gql_convert.rs @@ -1167,6 +1167,7 @@ impl From for WorkspacesMetadataResponse { joinable_teams, experiments, feature_model_choices, + ai_credit_availability: Some(gql_user.ai_credit_availability.into()), } } } diff --git a/app/src/workspaces/update_manager.rs b/app/src/workspaces/update_manager.rs index c41d9d109c6..ae407918874 100644 --- a/app/src/workspaces/update_manager.rs +++ b/app/src/workspaces/update_manager.rs @@ -16,6 +16,7 @@ use super::user_workspaces::{ }; use super::workspace::WorkspaceUid; use crate::ai::llms::LLMPreferences; +use crate::ai::request_usage_model::AIRequestUsageModel; use crate::auth::AuthStateProvider; use crate::cloud_object::CloudObjectEventEntrypoint; use crate::network::{NetworkStatus, NetworkStatusEvent, NetworkStatusKind}; @@ -125,6 +126,7 @@ impl TeamUpdateManager { joinable_teams: vec![], experiments: None, feature_model_choices: None, + ai_credit_availability: None, }, pricing_info: None, }) @@ -347,6 +349,12 @@ impl TeamUpdateManager { }); } + if let Some(availability) = response.metadata.ai_credit_availability { + AIRequestUsageModel::handle(ctx).update(ctx, |usage_model, ctx| { + usage_model.apply_server_availability(Ok(availability), ctx); + }); + } + let workspaces = response.metadata.workspaces; let joinable_teams = response.metadata.joinable_teams; @@ -469,6 +477,12 @@ impl TeamUpdateManager { let joinable_teams = user_workspaces_access.joinable_teams; let experiments = user_workspaces_access.experiments; + if let Some(availability) = user_workspaces_access.ai_credit_availability { + AIRequestUsageModel::handle(ctx).update(ctx, |usage_model, ctx| { + usage_model.apply_server_availability(Ok(availability), ctx); + }); + } + UserWorkspaces::handle(ctx).update(ctx, |user_workspaces, ctx| { user_workspaces.update_workspaces(workspaces.clone(), ctx); user_workspaces.update_joinable_teams(joinable_teams.clone(), ctx); diff --git a/app/src/workspaces/update_manager_tests.rs b/app/src/workspaces/update_manager_tests.rs index 308a73721a5..e0d60c28496 100644 --- a/app/src/workspaces/update_manager_tests.rs +++ b/app/src/workspaces/update_manager_tests.rs @@ -4,6 +4,7 @@ use itertools::Itertools; use warpui::{AddSingletonModel, App}; use super::*; +use crate::ai::credit_availability::{AICreditAvailability, AICreditDenialReason}; use crate::auth::AuthManager; use crate::cloud_object::model::actions::ObjectActions; use crate::cloud_object::model::persistence::CloudModel; @@ -96,6 +97,7 @@ fn test_leaving_team_removes_objects() { joinable_teams: vec![], experiments: None, feature_model_choices: None, + ai_credit_availability: None, }, pricing_info: None, }) @@ -165,6 +167,7 @@ fn test_leaving_team_removes_objects() { joinable_teams: vec![], experiments: None, feature_model_choices: None, + ai_credit_availability: None, }, pricing_info: None, }), @@ -207,3 +210,45 @@ fn test_leaving_team_removes_objects() { }); }); } + +#[test] +fn test_workspace_metadata_piggyback_feeds_ai_credit_availability() { + App::test((), |mut app| async move { + let team_client = Arc::new(MockTeamClient::new()); + initialize_app( + team_client.clone(), + Arc::new(MockWorkspaceClient::new()), + vec![], + &mut app, + ); + if app + .models_of_type::() + .is_empty() + { + app.update(crate::settings::init_and_register_user_preferences); + } + app.add_singleton_model(|ctx| { + AIRequestUsageModel::new_for_test(ServerApiProvider::as_ref(ctx).get_ai_client(), ctx) + }); + let team_update_manager = + app.add_singleton_model(|ctx| TeamUpdateManager::new(team_client, None, ctx)); + + let availability = AICreditAvailability::unavailable(AICreditDenialReason::OutOfCredits); + team_update_manager.update(&mut app, |manager, ctx| { + manager.on_workspaces_updated( + Ok(WorkspacesMetadataResponse { + workspaces: vec![], + joinable_teams: vec![], + experiments: None, + feature_model_choices: None, + ai_credit_availability: Some(availability), + }), + ctx, + ); + }); + + AIRequestUsageModel::handle(&app).read(&app, |model, _| { + assert_eq!(model.server_availability(), Some(availability)); + }); + }); +} diff --git a/app/src/workspaces/user_workspaces.rs b/app/src/workspaces/user_workspaces.rs index c15aced929e..b54a401ca70 100644 --- a/app/src/workspaces/user_workspaces.rs +++ b/app/src/workspaces/user_workspaces.rs @@ -19,7 +19,9 @@ use super::workspace::{ AdminEnablementSetting, BillingMetadata, CustomerType, EnterpriseSecretRegex, HostEnablementSetting, UgcCollectionEnablementSetting, Workspace, WorkspaceUid, }; +use crate::ai::credit_availability::AICreditAvailability; use crate::ai::llms::LLMModelHost; +use crate::ai::request_usage_model::AIRequestUsageModel; use crate::auth::{AuthStateProvider, UserUid}; use crate::channel::ChannelState; use crate::cloud_object::model::persistence::CloudModel; @@ -77,6 +79,8 @@ pub enum UserWorkspacesEvent { PurchaseAddonCreditsRejected(anyhow::Error), /// Fired whenever the set of teams the user is on changes. TeamsChanged, + /// Fired when the selected workspace actually changes to a different one. + CurrentWorkspaceChanged, CodebaseContextEnablementChanged, /// Fired when a service agreement's sunsetted_to_build_ts field is updated. SunsettedToBuildDataUpdated, @@ -109,6 +113,9 @@ pub struct WorkspacesMetadataResponse { /// It makes most sense to fetch this in workspaces which is queried every 10 minutes. /// This is list of available LLM models for the user. pub feature_model_choices: Option, + /// The server-authoritative AI credit availability decision, piggybacked + /// on the metadata query so every refresh keeps the shared state fresh. + pub ai_credit_availability: Option, } // A representation of all data we fetch at a single time via our 10 minute poll. @@ -459,9 +466,13 @@ impl UserWorkspaces { workspace_uid: WorkspaceUid, ctx: &mut ModelContext, ) { + let changed = *self.current_workspace_uid != Some(workspace_uid); *self.current_workspace_uid = Some(workspace_uid); self.reconcile_window_team_assignments(); self.notify_and_emit_teams_changed(ctx); + if changed { + ctx.emit(UserWorkspacesEvent::CurrentWorkspaceChanged); + } } /// Returns `true` if active AI is allowed for the current workspace, based on billing config. @@ -980,6 +991,12 @@ impl UserWorkspaces { }); } + if let Some(availability) = response.metadata.ai_credit_availability { + AIRequestUsageModel::handle(ctx).update(ctx, |usage_model, ctx| { + usage_model.apply_server_availability(Ok(availability), ctx); + }); + } + let workspaces = response.metadata.workspaces; let joinable_teams = response.metadata.joinable_teams; diff --git a/app/src/workspaces/user_workspaces_tests.rs b/app/src/workspaces/user_workspaces_tests.rs index 5254ad57c7c..58e4d85f127 100644 --- a/app/src/workspaces/user_workspaces_tests.rs +++ b/app/src/workspaces/user_workspaces_tests.rs @@ -167,6 +167,7 @@ fn test_loading_all_spaces_after_switching_from_offline() { joinable_teams: vec![], experiments: None, feature_model_choices: None, + ai_credit_availability: None, }, pricing_info: None, }) @@ -184,6 +185,7 @@ fn test_loading_all_spaces_after_switching_from_offline() { joinable_teams: vec![], experiments: None, feature_model_choices: None, + ai_credit_availability: None, }, pricing_info: None, }) @@ -321,6 +323,7 @@ fn test_aws_bedrock_credentials_respect_user_setting() { joinable_teams: vec![], experiments: None, feature_model_choices: None, + ai_credit_availability: None, }, pricing_info: None, }) @@ -377,6 +380,7 @@ fn test_aws_bedrock_credentials_enforced_by_admin() { joinable_teams: vec![], experiments: None, feature_model_choices: None, + ai_credit_availability: None, }, pricing_info: None, }) diff --git a/crates/graphql/src/api/ai.rs b/crates/graphql/src/api/ai.rs index 1e0df2be2aa..07192e102c1 100644 --- a/crates/graphql/src/api/ai.rs +++ b/crates/graphql/src/api/ai.rs @@ -12,6 +12,36 @@ pub enum RequestLimitRefreshDuration { EveryTwoWeeks, } +#[derive(cynic::Enum, Clone, Debug, PartialEq, Eq)] +pub enum AICreditAvailabilityDenialReason { + None, + OutOfCredits, + Delinquent, + EnterpriseTeamSpendLimitHit, + EnterprisePerUserSpendLimitHit, + EnterpriseWorkspaceSpendLimitHit, + #[cynic(fallback)] + Other(String), +} + +#[derive(cynic::Enum, Clone, Debug, PartialEq, Eq)] +pub enum AICreditAvailabilitySource { + BaseLimit, + BonusGrant, + Payg, + Overage, + AmbientBonusGrant, + #[cynic(fallback)] + Other(String), +} + +#[derive(cynic::QueryFragment, Debug, Clone)] +pub struct AICreditAvailability { + pub available: bool, + pub denial_reason: AICreditAvailabilityDenialReason, + pub credit_source: Option, +} + #[derive(cynic::QueryFragment, Debug)] pub struct RequestLimitInfo { pub is_unlimited: bool, diff --git a/crates/graphql/src/api/queries/get_ai_credit_availability.rs b/crates/graphql/src/api/queries/get_ai_credit_availability.rs new file mode 100644 index 00000000000..f0b63269556 --- /dev/null +++ b/crates/graphql/src/api/queries/get_ai_credit_availability.rs @@ -0,0 +1,61 @@ +use crate::ai::AICreditAvailability; +use crate::error::UserFacingError; +use crate::request_context::RequestContext; +use crate::schema; + +/* +query GetAICreditAvailability($requestContext: RequestContext!) { + user(requestContext: $requestContext) { + ... on UserOutput { + user { + aiCreditAvailability { + available + denialReason + creditSource + } + } + } + ... on UserFacingError { + error { + message + } + } + } +} +*/ + +#[derive(cynic::QueryVariables, Debug)] +pub struct GetAICreditAvailabilityVariables { + pub request_context: RequestContext, +} + +#[derive(cynic::QueryFragment, Debug)] +pub struct UserOutput { + pub user: User, +} + +#[derive(cynic::QueryFragment, Debug)] +pub struct User { + pub ai_credit_availability: AICreditAvailability, +} + +#[derive(cynic::QueryFragment, Debug)] +#[cynic( + graphql_type = "RootQuery", + variables = "GetAICreditAvailabilityVariables" +)] +pub struct GetAICreditAvailability { + #[arguments(requestContext: $request_context)] + pub user: UserResult, +} +crate::client::define_operation! { + get_ai_credit_availability(GetAICreditAvailabilityVariables) -> GetAICreditAvailability; +} + +#[derive(cynic::InlineFragments, Debug)] +pub enum UserResult { + UserOutput(UserOutput), + UserFacingError(UserFacingError), + #[cynic(fallback)] + Unknown, +} diff --git a/crates/graphql/src/api/queries/get_workspaces_metadata_for_user.rs b/crates/graphql/src/api/queries/get_workspaces_metadata_for_user.rs index e803455b7ed..a91a7631308 100644 --- a/crates/graphql/src/api/queries/get_workspaces_metadata_for_user.rs +++ b/crates/graphql/src/api/queries/get_workspaces_metadata_for_user.rs @@ -1,3 +1,4 @@ +use crate::ai::AICreditAvailability; use crate::billing::PricingInfo; use crate::experiment::Experiment; use crate::request_context::RequestContext; @@ -13,6 +14,11 @@ query GetWorkspacesMetadataForUser($requestContext: RequestContext!) { profile { uid } + aiCreditAvailability { + available + denialReason + creditSource + } workspaces { uid name @@ -226,6 +232,7 @@ pub enum PricingInfoResult { #[derive(cynic::QueryFragment, Debug)] pub struct User { pub profile: UserProfile, + pub ai_credit_availability: AICreditAvailability, pub workspaces: Vec, pub experiments: Option>, pub discoverable_teams: Vec, diff --git a/crates/graphql/src/api/queries/mod.rs b/crates/graphql/src/api/queries/mod.rs index a84407cd5e3..4a6771ba11b 100644 --- a/crates/graphql/src/api/queries/mod.rs +++ b/crates/graphql/src/api/queries/mod.rs @@ -2,6 +2,7 @@ pub mod api_keys; pub mod codebase_context_config; pub mod free_available_models; pub mod get_ai_conversation_format; +pub mod get_ai_credit_availability; pub mod get_ai_overages_for_workspace; pub mod get_available_harnesses; pub mod get_blocks_for_user; diff --git a/crates/warp_graphql_schema/api/schema.graphql b/crates/warp_graphql_schema/api/schema.graphql index 1dff585bc2d..4e228f14c64 100644 --- a/crates/warp_graphql_schema/api/schema.graphql +++ b/crates/warp_graphql_schema/api/schema.graphql @@ -79,6 +79,29 @@ type AIConversationFormat { hasTaskList: Boolean! } +type AICreditAvailability { + available: Boolean! + creditSource: AICreditAvailabilitySource + denialReason: AICreditAvailabilityDenialReason! +} + +enum AICreditAvailabilityDenialReason { + DELINQUENT + ENTERPRISE_PER_USER_SPEND_LIMIT_HIT + ENTERPRISE_TEAM_SPEND_LIMIT_HIT + ENTERPRISE_WORKSPACE_SPEND_LIMIT_HIT + NONE + OUT_OF_CREDITS +} + +enum AICreditAvailabilitySource { + AMBIENT_BONUS_GRANT + BASE_LIMIT + BONUS_GRANT + OVERAGE + PAYG +} + enum AICreditsUsageAndCostSubjectType { SERVICE_ACCOUNT TEAM @@ -4322,6 +4345,7 @@ The User type is the primary way of interacting with our graph. Nearly all data is scoped to the currently logged-in user. """ type User { + aiCreditAvailability: AICreditAvailability! anonymousUserInfo: AnonymousUserInfo availableHarnesses: AvailableHarnesses! billingMetadata: BillingMetadata From 3edd41c935496ba13a1c508510c1332840b646d1 Mon Sep 17 00:00:00 2001 From: Jeff Lloyd Date: Mon, 27 Jul 2026 16:01:30 -0400 Subject: [PATCH 2/8] [REV-1714] Treat capability-only availability without a local key as out of credits The server reports BYO capability at the policy level because personal API keys are stored only on the client. Since nearly every tier allows BYOK, capability-only availability (available with no credit source) would otherwise suppress out-of-credits messaging for everyone. Pair the server's capability-only answer with the one fact only the client knows: whether a usable BYO path actually exists (a stored key/custom endpoint/Grok subscription permitted by policy, or a team-managed custom LLM). Without one, capability-only availability is treated as out of credits for gating and prompt-alert presentation. An explicit server denial is never overridden by local key presence. Pairs with the server-side denial precedence fix in warpdotdev/warp-server#13369. Co-Authored-By: Oz --- app/src/ai/blocklist/prompt/prompt_alert.rs | 9 +- .../ai/blocklist/prompt/prompt_alert_tests.rs | 39 +++++++- app/src/ai/request_usage_model.rs | 40 +++++++- app/src/ai/request_usage_model_tests.rs | 94 +++++++++++++++++++ 4 files changed, 176 insertions(+), 6 deletions(-) diff --git a/app/src/ai/blocklist/prompt/prompt_alert.rs b/app/src/ai/blocklist/prompt/prompt_alert.rs index 06ab9d06d77..a27ccaefa64 100644 --- a/app/src/ai/blocklist/prompt/prompt_alert.rs +++ b/app/src/ai/blocklist/prompt/prompt_alert.rs @@ -175,7 +175,14 @@ impl PromptAlertView { app: &AppContext, ) -> PromptAlertState { if availability.available { - return PromptAlertState::NoAlert; + // Capability-only availability (no credit source) is refined with + // local key knowledge by `has_any_ai_remaining`: the server cannot + // see locally stored API keys, so without a usable BYO path the + // user is effectively out of credits. + if AIRequestUsageModel::as_ref(app).has_any_ai_remaining(app) { + return PromptAlertState::NoAlert; + } + return Self::out_of_credits_presentation(app); } match availability.denial_reason { diff --git a/app/src/ai/blocklist/prompt/prompt_alert_tests.rs b/app/src/ai/blocklist/prompt/prompt_alert_tests.rs index c8e7da9c933..e6a62975a8c 100644 --- a/app/src/ai/blocklist/prompt/prompt_alert_tests.rs +++ b/app/src/ai/blocklist/prompt/prompt_alert_tests.rs @@ -7,8 +7,13 @@ use crate::ai::credit_availability::AICreditSource; use crate::server::server_api::ServerApiProvider; use crate::server::server_api::team::MockTeamClient; use crate::server::server_api::workspace::MockWorkspaceClient; +use crate::workspaces::workspace::{ByoApiKeyPolicy, Workspace, WorkspaceUid}; fn initialize_app(app: &mut App) { + initialize_app_with_workspaces(app, vec![]); +} + +fn initialize_app_with_workspaces(app: &mut App, workspaces: Vec) { app.add_singleton_model(|_| NetworkStatus::new()); app.add_singleton_model(|_| AuthStateProvider::new_for_test()); app.add_singleton_model(|_| ServerApiProvider::new_for_test()); @@ -16,7 +21,7 @@ fn initialize_app(app: &mut App) { UserWorkspaces::mock( Arc::new(MockTeamClient::new()), Arc::new(MockWorkspaceClient::new()), - vec![], + workspaces, ctx, ) }); @@ -121,3 +126,35 @@ fn test_legacy_fallback_used_before_first_server_response() { assert_eq!(determine_state(&mut app), PromptAlertState::NoAlert); }); } + +#[test] +fn test_capability_only_without_local_key_maps_to_out_of_credits() { + App::test((), |mut app| async move { + initialize_app(&mut app); + // The server allows BYO by policy but found no credit source; with no + // locally stored key this is effectively out of credits. + apply_server_availability(&mut app, AICreditAvailability::available_with_source(None)); + assert_eq!( + determine_state(&mut app), + PromptAlertState::RequestLimitReached + ); + }); +} + +#[test] +fn test_capability_only_with_local_key_maps_to_no_alert() { + App::test((), |mut app| async move { + let uid = WorkspaceUid::from(crate::server::ids::ServerId::from(1_i64)); + let mut workspace = Workspace::from_local_cache(uid, "Test Workspace".to_string(), None); + workspace.billing_metadata.tier.byo_api_key_policy = + Some(ByoApiKeyPolicy { enabled: true }); + initialize_app_with_workspaces(&mut app, vec![workspace]); + + ApiKeyManager::handle(&app).update(&mut app, |manager, ctx| { + manager.set_openai_key(Some("test-key".to_string()), ctx); + }); + + apply_server_availability(&mut app, AICreditAvailability::available_with_source(None)); + assert_eq!(determine_state(&mut app), PromptAlertState::NoAlert); + }); +} diff --git a/app/src/ai/request_usage_model.rs b/app/src/ai/request_usage_model.rs index 5fcf09dce5d..d6db6f86b1e 100644 --- a/app/src/ai/request_usage_model.rs +++ b/app/src/ai/request_usage_model.rs @@ -517,16 +517,48 @@ impl AIRequestUsageModel { /// /// Once a server-authoritative availability decision has been received /// this session, that decision (last-known-good on transient refresh - /// failures) is the only authority. The locally derived fallback below is - /// used solely before the first successful fetch, e.g. right after startup - /// or against servers that don't support the availability field yet. + /// failures) is the authority. The locally derived fallback below is used + /// solely before the first successful fetch, e.g. right after startup or + /// against servers that don't support the availability field yet. pub fn has_any_ai_remaining(&self, ctx: &AppContext) -> bool { if let Some(availability) = self.server_availability.latest { - return availability.available; + return Self::server_availability_permits_ai(availability, ctx); } self.has_any_ai_remaining_from_local_state(ctx) } + /// Interprets the server-authoritative decision. Capability-only + /// availability (available with no credit source) means "no Warp credits, + /// but BYO inference is allowed by policy" — the server cannot see locally + /// stored API keys, so the client contributes that one fact and treats + /// capability-only availability without a usable BYO path as out of + /// credits. + fn server_availability_permits_ai( + availability: AICreditAvailability, + ctx: &AppContext, + ) -> bool { + if !availability.available { + return false; + } + if availability.credit_source.is_some() { + return true; + } + Self::has_usable_byo_inference_path(ctx) + } + + /// Whether a BYO inference path is actually usable from this client: a + /// locally stored API key, custom endpoint, or connected Grok subscription + /// (when the BYOK policy allows it), or a team-managed custom LLM + /// configuration. + fn has_usable_byo_inference_path(ctx: &AppContext) -> bool { + let has_byo_credentials = UserWorkspaces::as_ref(ctx).is_byo_api_key_enabled(ctx) + && ApiKeyManager::as_ref(ctx).has_any_key(); + let has_team_custom_llm = UserWorkspaces::as_ref(ctx) + .current_team() + .is_some_and(|team| team.is_custom_llm_enabled()); + has_byo_credentials || has_team_custom_llm + } + /// Legacy locally derived availability check. Returns `true` if the user /// meets one of the following conditions: /// 1. user has ai credits from the plan base limit diff --git a/app/src/ai/request_usage_model_tests.rs b/app/src/ai/request_usage_model_tests.rs index 8b3642692e7..394323a1740 100644 --- a/app/src/ai/request_usage_model_tests.rs +++ b/app/src/ai/request_usage_model_tests.rs @@ -1022,6 +1022,100 @@ fn test_reset_server_availability_restores_legacy_fallback() { }); } +#[test] +fn test_capability_only_availability_requires_local_byo_path() { + App::test((), |mut app| async move { + // BYOK is allowed by policy, but no key has been stored locally. + let (_uid, mut workspace) = create_test_workspace(); + workspace.billing_metadata.tier.byo_api_key_policy = + Some(ByoApiKeyPolicy { enabled: true }); + add_user_workspaces_with_workspace(&mut app, workspace); + let request_usage_model = add_request_usage_model(&mut app); + + // Capability-only: the server allows BYO by policy but found no Warp + // credit source. It cannot see locally stored keys. + request_usage_model.update(&mut app, |model, ctx| { + model.request_limit_info = RequestLimitInfo::new_for_test(10, 10); + model.apply_server_availability( + Ok(AICreditAvailability::available_with_source(None)), + ctx, + ); + assert!( + !model.has_any_ai_remaining(ctx), + "capability-only availability without a stored key should be treated as out of credits", + ); + }); + + // Storing a key makes the capability usable. + ApiKeyManager::handle(&app).update(&mut app, |manager, ctx| { + manager.set_openai_key(Some("test-key".to_string()), ctx); + }); + request_usage_model.read(&app, |model, ctx| { + assert!( + model.has_any_ai_remaining(ctx), + "capability-only availability with a stored key should permit AI", + ); + }); + }); +} + +#[test] +fn test_capability_only_availability_with_team_custom_llm() { + App::test((), |mut app| async move { + // A team-managed custom LLM configuration is a usable BYO path even + // without any locally stored API key. + let mut team = crate::workspaces::team::Team::from_local_cache( + 123.into(), + "Test Team".to_string(), + None, + None, + None, + ); + team.organization_settings.llm_settings.enabled = true; + let (uid, _workspace) = create_test_workspace(); + let workspace = + Workspace::from_local_cache(uid, "Test Workspace".to_string(), Some(vec![team])); + add_user_workspaces_with_workspace(&mut app, workspace); + let request_usage_model = add_request_usage_model(&mut app); + + request_usage_model.update(&mut app, |model, ctx| { + model.request_limit_info = RequestLimitInfo::new_for_test(10, 10); + model.apply_server_availability( + Ok(AICreditAvailability::available_with_source(None)), + ctx, + ); + assert!(model.has_any_ai_remaining(ctx)); + }); + }); +} + +#[test] +fn test_server_unavailable_overrides_local_byo_key() { + App::test((), |mut app| async move { + // A locally stored key never overrides an explicit server denial — + // the refinement only applies to capability-only availability. + let (_uid, mut workspace) = create_test_workspace(); + workspace.billing_metadata.tier.byo_api_key_policy = + Some(ByoApiKeyPolicy { enabled: true }); + add_user_workspaces_with_workspace(&mut app, workspace); + let request_usage_model = add_request_usage_model(&mut app); + + ApiKeyManager::handle(&app).update(&mut app, |manager, ctx| { + manager.set_openai_key(Some("test-key".to_string()), ctx); + }); + + request_usage_model.update(&mut app, |model, ctx| { + model.apply_server_availability( + Ok(AICreditAvailability::unavailable( + AICreditDenialReason::Delinquent, + )), + ctx, + ); + assert!(!model.has_any_ai_remaining(ctx)); + }); + }); +} + #[test] fn test_availability_refresh_coalesces_concurrent_fetches() { App::test((), |mut app| async move { From 1e6a6f7a7c9ea643c1a42998c67809f6a3366463 Mon Sep 17 00:00:00 2001 From: Jeff Lloyd Date: Tue, 28 Jul 2026 14:46:29 -0400 Subject: [PATCH 3/8] [REV-1714] Trust definite server availability; refine OUT_OF_CREDITS with local keys Follows the server-side semantics rework (warp-server#13369): available now means the server knows for a fact requests can run (a Warp credit source or a configured server-managed BYO path), so the client trusts it outright. An OUT_OF_CREDITS denial means the server found no path it can see - locally stored API keys are request-level parameters invisible to it - so the client supplies that one fact and permits AI when a usable local key exists (policy permitting). Delinquency and spend-limit denials are hard rejections that local keys never bypass. The team-custom-LLM client check is dropped: server-managed BYO paths are now reflected directly in the server's availability decision. Co-Authored-By: Oz --- app/src/ai/blocklist/prompt/prompt_alert.rs | 19 ++++----- .../ai/blocklist/prompt/prompt_alert_tests.rs | 20 +++++----- app/src/ai/request_usage_model.rs | 38 +++++++++--------- app/src/ai/request_usage_model_tests.rs | 40 ++++++++----------- 4 files changed, 55 insertions(+), 62 deletions(-) diff --git a/app/src/ai/blocklist/prompt/prompt_alert.rs b/app/src/ai/blocklist/prompt/prompt_alert.rs index a27ccaefa64..a0556eea738 100644 --- a/app/src/ai/blocklist/prompt/prompt_alert.rs +++ b/app/src/ai/blocklist/prompt/prompt_alert.rs @@ -175,14 +175,7 @@ impl PromptAlertView { app: &AppContext, ) -> PromptAlertState { if availability.available { - // Capability-only availability (no credit source) is refined with - // local key knowledge by `has_any_ai_remaining`: the server cannot - // see locally stored API keys, so without a usable BYO path the - // user is effectively out of credits. - if AIRequestUsageModel::as_ref(app).has_any_ai_remaining(app) { - return PromptAlertState::NoAlert; - } - return Self::out_of_credits_presentation(app); + return PromptAlertState::NoAlert; } match availability.denial_reason { @@ -194,7 +187,15 @@ impl PromptAlertView { } AICreditDenialReason::None | AICreditDenialReason::OutOfCredits - | AICreditDenialReason::Unknown => Self::out_of_credits_presentation(app), + | AICreditDenialReason::Unknown => { + // An out-of-credits denial only means the server found no path + // it can see; a locally stored API key still permits requests, + // which `has_any_ai_remaining` accounts for. + if AIRequestUsageModel::as_ref(app).has_any_ai_remaining(app) { + return PromptAlertState::NoAlert; + } + Self::out_of_credits_presentation(app) + } } } diff --git a/app/src/ai/blocklist/prompt/prompt_alert_tests.rs b/app/src/ai/blocklist/prompt/prompt_alert_tests.rs index e6a62975a8c..2ae765f59b6 100644 --- a/app/src/ai/blocklist/prompt/prompt_alert_tests.rs +++ b/app/src/ai/blocklist/prompt/prompt_alert_tests.rs @@ -128,21 +128,18 @@ fn test_legacy_fallback_used_before_first_server_response() { } #[test] -fn test_capability_only_without_local_key_maps_to_out_of_credits() { +fn test_server_managed_availability_maps_to_no_alert() { App::test((), |mut app| async move { initialize_app(&mut app); - // The server allows BYO by policy but found no credit source; with no - // locally stored key this is effectively out of credits. + // `available` with no credit source means a server-managed BYO path + // is configured — definite availability, no local key required. apply_server_availability(&mut app, AICreditAvailability::available_with_source(None)); - assert_eq!( - determine_state(&mut app), - PromptAlertState::RequestLimitReached - ); + assert_eq!(determine_state(&mut app), PromptAlertState::NoAlert); }); } #[test] -fn test_capability_only_with_local_key_maps_to_no_alert() { +fn test_out_of_credits_with_local_key_maps_to_no_alert() { App::test((), |mut app| async move { let uid = WorkspaceUid::from(crate::server::ids::ServerId::from(1_i64)); let mut workspace = Workspace::from_local_cache(uid, "Test Workspace".to_string(), None); @@ -154,7 +151,12 @@ fn test_capability_only_with_local_key_maps_to_no_alert() { manager.set_openai_key(Some("test-key".to_string()), ctx); }); - apply_server_availability(&mut app, AICreditAvailability::available_with_source(None)); + // The server cannot see the locally stored key; the client refines + // its OUT_OF_CREDITS answer. + apply_server_availability( + &mut app, + AICreditAvailability::unavailable(AICreditDenialReason::OutOfCredits), + ); assert_eq!(determine_state(&mut app), PromptAlertState::NoAlert); }); } diff --git a/app/src/ai/request_usage_model.rs b/app/src/ai/request_usage_model.rs index d6db6f86b1e..816f4535d82 100644 --- a/app/src/ai/request_usage_model.rs +++ b/app/src/ai/request_usage_model.rs @@ -15,7 +15,7 @@ use warpui::{AppContext, Entity, ModelContext, SingletonEntity}; use crate::BlocklistAIHistoryModel; use crate::ai::agent::AIAgentExchangeId; use crate::ai::agent::conversation::AIConversationId; -use crate::ai::credit_availability::AICreditAvailability; +use crate::ai::credit_availability::{AICreditAvailability, AICreditDenialReason}; use crate::auth::AuthStateProvider; use crate::pricing::PricingInfoModel; use crate::server::server_api::ai::AIClient; @@ -527,36 +527,34 @@ impl AIRequestUsageModel { self.has_any_ai_remaining_from_local_state(ctx) } - /// Interprets the server-authoritative decision. Capability-only - /// availability (available with no credit source) means "no Warp credits, - /// but BYO inference is allowed by policy" — the server cannot see locally - /// stored API keys, so the client contributes that one fact and treats - /// capability-only availability without a usable BYO path as out of - /// credits. + /// Interprets the server-authoritative decision. `available` means the + /// server knows for a fact that requests can run (a Warp credit source or + /// a server-managed BYO path) and is trusted outright. An `OutOfCredits` + /// denial means the server found no path *it can see* — locally stored + /// API keys are request-level parameters invisible to it — so the client + /// contributes that one fact. Every other denial is a hard rejection that + /// local keys cannot bypass. fn server_availability_permits_ai( availability: AICreditAvailability, ctx: &AppContext, ) -> bool { - if !availability.available { - return false; - } - if availability.credit_source.is_some() { + if availability.available { return true; } - Self::has_usable_byo_inference_path(ctx) + matches!( + availability.denial_reason, + AICreditDenialReason::OutOfCredits + ) && Self::has_usable_byo_inference_path(ctx) } /// Whether a BYO inference path is actually usable from this client: a /// locally stored API key, custom endpoint, or connected Grok subscription - /// (when the BYOK policy allows it), or a team-managed custom LLM - /// configuration. + /// (when the BYOK policy allows it). Server-managed paths (team-managed + /// keys/endpoints, enterprise custom LLM) are already reflected in the + /// server's availability decision. fn has_usable_byo_inference_path(ctx: &AppContext) -> bool { - let has_byo_credentials = UserWorkspaces::as_ref(ctx).is_byo_api_key_enabled(ctx) - && ApiKeyManager::as_ref(ctx).has_any_key(); - let has_team_custom_llm = UserWorkspaces::as_ref(ctx) - .current_team() - .is_some_and(|team| team.is_custom_llm_enabled()); - has_byo_credentials || has_team_custom_llm + UserWorkspaces::as_ref(ctx).is_byo_api_key_enabled(ctx) + && ApiKeyManager::as_ref(ctx).has_any_key() } /// Legacy locally derived availability check. Returns `true` if the user diff --git a/app/src/ai/request_usage_model_tests.rs b/app/src/ai/request_usage_model_tests.rs index 394323a1740..74a24f6309b 100644 --- a/app/src/ai/request_usage_model_tests.rs +++ b/app/src/ai/request_usage_model_tests.rs @@ -1023,7 +1023,7 @@ fn test_reset_server_availability_restores_legacy_fallback() { } #[test] -fn test_capability_only_availability_requires_local_byo_path() { +fn test_out_of_credits_refined_by_local_byo_key() { App::test((), |mut app| async move { // BYOK is allowed by policy, but no key has been stored locally. let (_uid, mut workspace) = create_test_workspace(); @@ -1032,50 +1032,42 @@ fn test_capability_only_availability_requires_local_byo_path() { add_user_workspaces_with_workspace(&mut app, workspace); let request_usage_model = add_request_usage_model(&mut app); - // Capability-only: the server allows BYO by policy but found no Warp - // credit source. It cannot see locally stored keys. + // OUT_OF_CREDITS means the server found no path it can see; locally + // stored keys are request-level parameters invisible to it. request_usage_model.update(&mut app, |model, ctx| { model.request_limit_info = RequestLimitInfo::new_for_test(10, 10); model.apply_server_availability( - Ok(AICreditAvailability::available_with_source(None)), + Ok(AICreditAvailability::unavailable( + AICreditDenialReason::OutOfCredits, + )), ctx, ); assert!( !model.has_any_ai_remaining(ctx), - "capability-only availability without a stored key should be treated as out of credits", + "out of credits without a stored key should gate AI", ); }); - // Storing a key makes the capability usable. + // Storing a key supplies the one fact the server cannot know. ApiKeyManager::handle(&app).update(&mut app, |manager, ctx| { manager.set_openai_key(Some("test-key".to_string()), ctx); }); request_usage_model.read(&app, |model, ctx| { assert!( model.has_any_ai_remaining(ctx), - "capability-only availability with a stored key should permit AI", + "out of credits with a stored key should permit AI", ); }); }); } #[test] -fn test_capability_only_availability_with_team_custom_llm() { +fn test_server_managed_availability_trusted_without_local_keys() { App::test((), |mut app| async move { - // A team-managed custom LLM configuration is a usable BYO path even - // without any locally stored API key. - let mut team = crate::workspaces::team::Team::from_local_cache( - 123.into(), - "Test Team".to_string(), - None, - None, - None, - ); - team.organization_settings.llm_settings.enabled = true; - let (uid, _workspace) = create_test_workspace(); - let workspace = - Workspace::from_local_cache(uid, "Test Workspace".to_string(), Some(vec![team])); - add_user_workspaces_with_workspace(&mut app, workspace); + // `available` with no credit source now means a server-managed BYO + // path (team keys/endpoints or enterprise custom LLM) is configured; + // the server knows this for a fact, so no local key is required. + app.add_singleton_model(UserWorkspaces::default_mock); let request_usage_model = add_request_usage_model(&mut app); request_usage_model.update(&mut app, |model, ctx| { @@ -1092,8 +1084,8 @@ fn test_capability_only_availability_with_team_custom_llm() { #[test] fn test_server_unavailable_overrides_local_byo_key() { App::test((), |mut app| async move { - // A locally stored key never overrides an explicit server denial — - // the refinement only applies to capability-only availability. + // A locally stored key never overrides a hard server denial — the + // local refinement applies only to OUT_OF_CREDITS. let (_uid, mut workspace) = create_test_workspace(); workspace.billing_metadata.tier.byo_api_key_policy = Some(ByoApiKeyPolicy { enabled: true }); From 97a665abaf15c936ef2bfb4c55df0e2d6385c37f Mon Sep 17 00:00:00 2001 From: Jeff Lloyd Date: Wed, 29 Jul 2026 16:47:25 -0400 Subject: [PATCH 4/8] Count loaded local-chain Bedrock credentials as a usable BYO path The server now only vouches for BYO configs it can operate itself, so Bedrock local-chain setups (region-only, credentials resolved on the client) surface as OUT_OF_CREDITS. Refine that denial locally by treating loaded AWS credentials with Bedrock enabled as a usable inference path. Co-Authored-By: Oz --- .../ai/blocklist/prompt/prompt_alert_tests.rs | 3 +- app/src/ai/request_usage_model.rs | 24 +++++--- app/src/ai/request_usage_model_tests.rs | 56 ++++++++++++++++++- 3 files changed, 71 insertions(+), 12 deletions(-) diff --git a/app/src/ai/blocklist/prompt/prompt_alert_tests.rs b/app/src/ai/blocklist/prompt/prompt_alert_tests.rs index 2ae765f59b6..e2c5862648c 100644 --- a/app/src/ai/blocklist/prompt/prompt_alert_tests.rs +++ b/app/src/ai/blocklist/prompt/prompt_alert_tests.rs @@ -1,5 +1,6 @@ use std::sync::Arc; +use ai::LLMProvider; use warpui::App; use super::*; @@ -148,7 +149,7 @@ fn test_out_of_credits_with_local_key_maps_to_no_alert() { initialize_app_with_workspaces(&mut app, vec![workspace]); ApiKeyManager::handle(&app).update(&mut app, |manager, ctx| { - manager.set_openai_key(Some("test-key".to_string()), ctx); + manager.set_provider_key(LLMProvider::OpenAI, Some("test-key".to_string()), ctx); }); // The server cannot see the locally stored key; the client refines diff --git a/app/src/ai/request_usage_model.rs b/app/src/ai/request_usage_model.rs index 816f4535d82..b3f66e891c8 100644 --- a/app/src/ai/request_usage_model.rs +++ b/app/src/ai/request_usage_model.rs @@ -1,6 +1,6 @@ use std::sync::Arc; -use ai::api_keys::ApiKeyManager; +use ai::api_keys::{ApiKeyManager, AwsCredentialsState}; use anyhow::Context as _; use chrono::{DateTime, Local, Utc}; use futures::channel::oneshot::{self, Receiver}; @@ -547,14 +547,22 @@ impl AIRequestUsageModel { ) && Self::has_usable_byo_inference_path(ctx) } - /// Whether a BYO inference path is actually usable from this client: a - /// locally stored API key, custom endpoint, or connected Grok subscription - /// (when the BYOK policy allows it). Server-managed paths (team-managed - /// keys/endpoints, enterprise custom LLM) are already reflected in the - /// server's availability decision. + /// Whether a BYO inference path is usable with credentials held on this + /// machine: a stored API key, custom endpoint, or Grok subscription (when + /// the BYOK policy allows it), or loaded local-chain AWS credentials for + /// an enabled Bedrock host. Server-managed paths are already reflected in + /// the server's availability decision. fn has_usable_byo_inference_path(ctx: &AppContext) -> bool { - UserWorkspaces::as_ref(ctx).is_byo_api_key_enabled(ctx) - && ApiKeyManager::as_ref(ctx).has_any_key() + let user_workspaces = UserWorkspaces::as_ref(ctx); + let api_keys = ApiKeyManager::as_ref(ctx); + if user_workspaces.is_byo_api_key_enabled(ctx) && api_keys.has_any_key() { + return true; + } + user_workspaces.is_aws_bedrock_credentials_enabled(ctx) + && matches!( + api_keys.aws_credentials_state(), + AwsCredentialsState::Loaded { .. } + ) } /// Legacy locally derived availability check. Returns `true` if the user diff --git a/app/src/ai/request_usage_model_tests.rs b/app/src/ai/request_usage_model_tests.rs index 74a24f6309b..81bc2538e83 100644 --- a/app/src/ai/request_usage_model_tests.rs +++ b/app/src/ai/request_usage_model_tests.rs @@ -1,7 +1,8 @@ use std::sync::Arc; +use std::time::SystemTime; use ai::LLMProvider; -use ai::api_keys::{ApiKeyManager, GrokTokens}; +use ai::api_keys::{ApiKeyManager, AwsCredentials, AwsCredentialsState, GrokTokens}; use chrono::Duration; use warp_core::features::FeatureFlag; use warp_graphql::billing::{AddonCreditsOption, OveragesPricing, PricingInfo}; @@ -1050,7 +1051,7 @@ fn test_out_of_credits_refined_by_local_byo_key() { // Storing a key supplies the one fact the server cannot know. ApiKeyManager::handle(&app).update(&mut app, |manager, ctx| { - manager.set_openai_key(Some("test-key".to_string()), ctx); + manager.set_provider_key(LLMProvider::OpenAI, Some("test-key".to_string()), ctx); }); request_usage_model.read(&app, |model, ctx| { assert!( @@ -1061,6 +1062,55 @@ fn test_out_of_credits_refined_by_local_byo_key() { }); } +#[test] +fn test_out_of_credits_refined_by_local_bedrock_credentials() { + App::test((), |mut app| async move { + // Bedrock via the local AWS chain: the org enables the host, but the + // credentials live on this machine. + let (_uid, mut workspace) = create_test_workspace(); + workspace.settings.llm_settings.enabled = true; + workspace.settings.llm_settings.host_configs.insert( + crate::ai::llms::LLMModelHost::AwsBedrock, + crate::workspaces::workspace::LlmHostSettings { + enabled: true, + enablement_setting: crate::workspaces::workspace::HostEnablementSetting::Enforce, + ..Default::default() + }, + ); + add_user_workspaces_with_workspace(&mut app, workspace); + let request_usage_model = add_request_usage_model(&mut app); + + request_usage_model.update(&mut app, |model, ctx| { + model.request_limit_info = RequestLimitInfo::new_for_test(10, 10); + model.apply_server_availability( + Ok(AICreditAvailability::unavailable( + AICreditDenialReason::OutOfCredits, + )), + ctx, + ); + assert!(!model.has_any_ai_remaining(ctx)); + }); + + ApiKeyManager::handle(&app).update(&mut app, |manager, ctx| { + manager.set_aws_credentials_state( + AwsCredentialsState::Loaded { + credentials: AwsCredentials::new( + "access".to_string(), + "secret".to_string(), + None, + None, + ), + loaded_at: SystemTime::now(), + }, + ctx, + ); + }); + request_usage_model.read(&app, |model, ctx| { + assert!(model.has_any_ai_remaining(ctx)); + }); + }); +} + #[test] fn test_server_managed_availability_trusted_without_local_keys() { App::test((), |mut app| async move { @@ -1093,7 +1143,7 @@ fn test_server_unavailable_overrides_local_byo_key() { let request_usage_model = add_request_usage_model(&mut app); ApiKeyManager::handle(&app).update(&mut app, |manager, ctx| { - manager.set_openai_key(Some("test-key".to_string()), ctx); + manager.set_provider_key(LLMProvider::OpenAI, Some("test-key".to_string()), ctx); }); request_usage_model.update(&mut app, |model, ctx| { From 66d7564454b4a8cb953a8f8e8f1eef7f1b6b34cc Mon Sep 17 00:00:00 2001 From: Oz Date: Sun, 2 Aug 2026 14:54:35 +0000 Subject: [PATCH 5/8] Fix semantic merge gaps: register AI usage model + telemetry in tests Master added user-level purchase-policy tests (constructing GqlUser and driving on_workspaces_updated) and provider-credential telemetry on ApiKeyManager::set_provider_key. After merging the server-authoritative credit-availability change, those test paths now exercise the availability apply path and telemetry, so register the required singletons in the affected test harnesses and add the new ai_credit_availability field to the gql_user helper. Co-Authored-By: Oz --- .../ai/blocklist/prompt/prompt_alert_tests.rs | 2 ++ app/src/workspaces/user_workspaces_tests.rs | 21 +++++++++++++++++++ 2 files changed, 23 insertions(+) diff --git a/app/src/ai/blocklist/prompt/prompt_alert_tests.rs b/app/src/ai/blocklist/prompt/prompt_alert_tests.rs index e2c5862648c..4fe7931f2ce 100644 --- a/app/src/ai/blocklist/prompt/prompt_alert_tests.rs +++ b/app/src/ai/blocklist/prompt/prompt_alert_tests.rs @@ -8,6 +8,7 @@ use crate::ai::credit_availability::AICreditSource; use crate::server::server_api::ServerApiProvider; use crate::server::server_api::team::MockTeamClient; use crate::server::server_api::workspace::MockWorkspaceClient; +use crate::server::telemetry::context_provider::AppTelemetryContextProvider; use crate::workspaces::workspace::{ByoApiKeyPolicy, Workspace, WorkspaceUid}; fn initialize_app(app: &mut App) { @@ -18,6 +19,7 @@ fn initialize_app_with_workspaces(app: &mut App, workspaces: Vec) { app.add_singleton_model(|_| NetworkStatus::new()); app.add_singleton_model(|_| AuthStateProvider::new_for_test()); app.add_singleton_model(|_| ServerApiProvider::new_for_test()); + app.add_singleton_model(AppTelemetryContextProvider::new_context_provider); app.add_singleton_model(|ctx| { UserWorkspaces::mock( Arc::new(MockTeamClient::new()), diff --git a/app/src/workspaces/user_workspaces_tests.rs b/app/src/workspaces/user_workspaces_tests.rs index c1e6c8dc240..cee7310e917 100644 --- a/app/src/workspaces/user_workspaces_tests.rs +++ b/app/src/workspaces/user_workspaces_tests.rs @@ -135,6 +135,20 @@ fn initialize_window_team_test_app(app: &mut App, workspaces: Vec) { }); } +// Registers the shared AI usage model (and its dependencies) so that the +// `on_workspaces_updated` metadata-apply path — which now piggybacks the +// server-authoritative AI credit availability onto every refresh — can update +// it without panicking on a missing singleton. +fn register_ai_usage_model(app: &mut App) { + app.add_singleton_model(|_| ServerApiProvider::new_for_test()); + if app.models_of_type::().is_empty() { + app.update(crate::settings::init_and_register_user_preferences); + } + app.add_singleton_model(|ctx| { + AIRequestUsageModel::new_for_test(ServerApiProvider::as_ref(ctx).get_ai_client(), ctx) + }); +} + #[test] fn test_loading_all_spaces_after_switching_from_offline() { let _flag = FeatureFlag::KnowledgeSidebar.override_enabled(true); @@ -1724,6 +1738,11 @@ fn gql_user( profile: GqlUserProfile { uid: "test-user".to_string(), }, + ai_credit_availability: warp_graphql::ai::AICreditAvailability { + available: true, + denial_reason: warp_graphql::ai::AICreditAvailabilityDenialReason::None, + credit_source: None, + }, billing_metadata: user_purchase_policy.map(|policy| UserPurchasePolicyBillingMetadata { tier: UserPurchasePolicyTier { purchase_add_on_credits_policy: Some(policy), @@ -1739,6 +1758,7 @@ fn gql_user( fn test_user_level_policy_survives_placeholder_filtering_for_teamless_users() { App::test((), |mut app| async move { initialize_window_team_test_app(&mut app, vec![]); + register_ai_usage_model(&mut app); // The real conversion path: a teamless user's ONLY workspace is the // placeholder, which must stay filtered out of `workspaces`, while @@ -1794,6 +1814,7 @@ fn test_user_level_policy_survives_placeholder_filtering_for_teamless_users() { fn test_workspace_policy_wins_over_user_level_policy() { App::test((), |mut app| async move { initialize_window_team_test_app(&mut app, vec![]); + register_ai_usage_model(&mut app); let standard_policy = GqlPurchaseAddOnCreditsPolicy { enabled: true, From c89cf8cb4113f97550cfdebf54f014fd9ebf55c4 Mon Sep 17 00:00:00 2001 From: Oz Date: Sun, 2 Aug 2026 15:50:54 +0000 Subject: [PATCH 6/8] Fix CI: register AIRequestUsageModel in transcript test harness The transcript render path now reads AIRequestUsageModel::has_any_ai_remaining (server-authoritative availability), which panicked in transcript_tests because the harness never registered that singleton (nor its AuthStateProvider / ServerApiProvider deps). Production always registers AIRequestUsageModel during app init (app/src/lib.rs), so no production path can hit this; register the singleton (and deps) in the shared transcript-test app builder. Co-Authored-By: Oz --- app/src/ai_assistant/transcript_tests.rs | 13 ++++++++++++- 1 file changed, 12 insertions(+), 1 deletion(-) diff --git a/app/src/ai_assistant/transcript_tests.rs b/app/src/ai_assistant/transcript_tests.rs index bcb6a316287..d5af8386d58 100644 --- a/app/src/ai_assistant/transcript_tests.rs +++ b/app/src/ai_assistant/transcript_tests.rs @@ -1,7 +1,8 @@ -use warpui::App; use warpui::platform::WindowStyle; +use warpui::{App, SingletonEntity}; use super::Transcript; +use crate::ai::AIRequestUsageModel; use crate::ai_assistant::requests::Requests; use crate::ai_assistant::test_util::{ default_assistant_transcript_part, default_code_block_segment, default_formatted_message, @@ -9,6 +10,8 @@ use crate::ai_assistant::test_util::{ }; use crate::ai_assistant::utils::{CodeBlockIndex, TranscriptPart, TranscriptPartSubType}; use crate::appearance; +use crate::auth::AuthStateProvider; +use crate::server::server_api::ServerApiProvider; use crate::test_util::settings::initialize_settings_for_tests; use crate::workspaces::user_workspaces::UserWorkspaces; @@ -45,6 +48,14 @@ fn initialize_app(app: &mut App) { initialize_settings_for_tests(app); appearance::register(app); app.add_singleton_model(UserWorkspaces::default_mock); + // The transcript render path reads the shared AI usage model to decide + // whether to show prepared responses, so it must be registered (as it + // always is in the real app). + app.add_singleton_model(|_| AuthStateProvider::new_for_test()); + app.add_singleton_model(|_| ServerApiProvider::new_for_test()); + app.add_singleton_model(|ctx| { + AIRequestUsageModel::new_for_test(ServerApiProvider::as_ref(ctx).get_ai_client(), ctx) + }); } #[test] From 051b06dc6d1ed15f336e4cc19d8ab76ddd082caa Mon Sep 17 00:00:00 2001 From: Jeff Lloyd Date: Sun, 2 Aug 2026 19:18:47 -0400 Subject: [PATCH 7/8] Use authoritative availability for the buy-credits banner Drive the banner from the server-authoritative availability decision and refresh it when that decision changes. Rename the legacy request-limit helper to make its base-plan-only semantics explicit while preserving the few quota-specific consumers. Co-Authored-By: Oz --- app/src/ai/blocklist/controller.rs | 5 +- app/src/ai/blocklist/prompt/prompt_alert.rs | 7 +- app/src/ai/request_usage_model.rs | 44 +++-- app/src/ai/request_usage_model_tests.rs | 173 +++++++++++++++++++- app/src/settings_view/ai_page.rs | 3 +- app/src/terminal/buy_credits_banner.rs | 7 +- 6 files changed, 220 insertions(+), 19 deletions(-) diff --git a/app/src/ai/blocklist/controller.rs b/app/src/ai/blocklist/controller.rs index 10e8eafa465..a62ec799511 100644 --- a/app/src/ai/blocklist/controller.rs +++ b/app/src/ai/blocklist/controller.rs @@ -3217,11 +3217,12 @@ impl BlocklistAIController { // If a user is below their personal limits, then we know that they won't eat into overages, // so we don't need to refresh. - let has_no_requests_remaining = !AIRequestUsageModel::as_ref(ctx).has_requests_remaining(); + let has_no_base_plan_requests_remaining = + !AIRequestUsageModel::as_ref(ctx).has_base_plan_requests_remaining(); // If overages aren't enabled, we're not going to reap the benefit of refreshing at all anyway. let are_overages_enabled = workspace.are_overages_enabled(); - if are_overages_enabled && has_no_requests_remaining { + if are_overages_enabled && has_no_base_plan_requests_remaining { // Give a one second delay to ensure that Stripe has been charged and the database is completely updated, // before syncing new AI overages data. ctx.spawn( diff --git a/app/src/ai/blocklist/prompt/prompt_alert.rs b/app/src/ai/blocklist/prompt/prompt_alert.rs index 021555b90d5..08663f134b9 100644 --- a/app/src/ai/blocklist/prompt/prompt_alert.rs +++ b/app/src/ai/blocklist/prompt/prompt_alert.rs @@ -124,7 +124,10 @@ impl PromptAlertView { } let request_usage_model = AIRequestUsageModel::as_ref(app); - let has_requests_remaining = request_usage_model.has_requests_remaining(); + // Anonymous soft/hard gates are based on the base-plan request quota, + // not overall AI availability (bonus grants / BYO / etc.). + let has_base_plan_requests_remaining = + request_usage_model.has_base_plan_requests_remaining(); let auth_state = AuthStateProvider::as_ref(app).get(); // Next, if the user is anonymous, we check if they have reached a certain percentage of requests used. @@ -135,7 +138,7 @@ impl PromptAlertView { let percentage_used = request_usage_model.request_percentage_used(); if percentage_used >= ANONYMOUS_USER_REQUEST_LIMIT_SOFT_GATE_PERCENTAGE { - if has_requests_remaining { + if has_base_plan_requests_remaining { return PromptAlertState::AnonymousUserRequestLimitSoftGate; } else { return PromptAlertState::AnonymousUserRequestLimitHardGate; diff --git a/app/src/ai/request_usage_model.rs b/app/src/ai/request_usage_model.rs index 2fedca4db28..420159ed4e3 100644 --- a/app/src/ai/request_usage_model.rs +++ b/app/src/ai/request_usage_model.rs @@ -15,7 +15,7 @@ use warpui::{AppContext, Entity, ModelContext, SingletonEntity}; use crate::BlocklistAIHistoryModel; use crate::ai::agent::AIAgentExchangeId; use crate::ai::agent::conversation::AIConversationId; -use crate::ai::credit_availability::{AICreditAvailability, AICreditDenialReason}; +use crate::ai::credit_availability::{AICreditAvailability, AICreditDenialReason, AICreditSource}; use crate::auth::AuthStateProvider; use crate::pricing::PricingInfoModel; use crate::server::server_api::ai::AIClient; @@ -503,12 +503,14 @@ impl AIRequestUsageModel { } } - /// Returns `true` if the user has at least one request remaining before hitting the AI request - /// limit. + /// Returns `true` if the user still has unused **base-plan** AI request quota + /// (the monthly/weekly plan limit). /// - /// WARNING: This method doesn't account for add-on credits. Consider if you want - /// [`Self::has_any_ai_remaining`] instead. - pub fn has_requests_remaining(&self) -> bool { + /// This deliberately ignores add-on credits, bonus grants, overages, + /// auto-reload, and BYO credentials. It is **not** an AI-availability check — + /// use [`Self::has_any_ai_remaining`] when deciding whether the user can start + /// an interactive AI request. + pub(crate) fn has_base_plan_requests_remaining(&self) -> bool { self.requests_remaining() > 0 } @@ -578,7 +580,7 @@ impl AIRequestUsageModel { fn has_any_ai_remaining_from_local_state(&self, ctx: &AppContext) -> bool { let current_workspace = UserWorkspaces::as_ref(ctx).current_workspace(); - let has_base_plan_ai_requests = self.has_requests_remaining(); + let has_base_plan_ai_requests = self.has_base_plan_requests_remaining(); let user_bonus_credits = self.total_user_interactive_bonus_credits_remaining() > 0; let workspace_bonus_credits = current_workspace @@ -739,6 +741,12 @@ impl AIRequestUsageModel { /// Computes the current banner state based on live conditions. /// This is called on-demand and always returns fresh state. + /// + /// Once a server-authoritative availability decision is known, the banner + /// hides whenever interactive AI is permitted (including `OutOfCredits` + /// refined by a usable local BYO path). Ambient-only credit sources are an + /// intentional exception: they fund cloud/ambient agents, not interactive + /// usage, so they must not suppress this banner. pub fn compute_buy_addon_credits_banner_display_state( &self, ctx: &AppContext, @@ -753,6 +761,10 @@ impl AIRequestUsageModel { .purchase_policy() .is_some_and(|policy| policy.allows_purchases()); + if !policy_allows_purchasing { + return BuyCreditsBannerDisplayState::Hidden; + } + // TODO: we might want to suggest credits purchase if request_remain/bonus credits is below certain threshold // something to consider after launch // Ambient-only credits are usable for cloud agents and should not suppress this banner. @@ -769,10 +781,20 @@ impl AIRequestUsageModel { current_workspace.is_some_and(|workspace| workspace.uid == uid) } }); - if !policy_allows_purchasing - || self.has_requests_remaining() - || has_non_ambient_bonus_credits - { + + if let Some(availability) = self.server_availability.latest { + let only_ambient_server_source = availability.available + && matches!( + availability.credit_source, + Some(AICreditSource::AmbientBonusGrant) + ); + // Hide when interactive AI is permitted, except ambient-only sources + // which do not fund interactive requests. + if self.has_any_ai_remaining(ctx) && !only_ambient_server_source { + return BuyCreditsBannerDisplayState::Hidden; + } + } else if self.has_base_plan_requests_remaining() || has_non_ambient_bonus_credits { + // Legacy pre-fetch path: local base quota / non-ambient bonus only. return BuyCreditsBannerDisplayState::Hidden; } diff --git a/app/src/ai/request_usage_model_tests.rs b/app/src/ai/request_usage_model_tests.rs index 4d3a0a0e715..bb36ee59ed8 100644 --- a/app/src/ai/request_usage_model_tests.rs +++ b/app/src/ai/request_usage_model_tests.rs @@ -10,7 +10,7 @@ use warp_graphql::billing::{AddonCreditsOption, OveragesPricing, PricingInfo}; use warpui::{App, ModelHandle}; use super::*; -use crate::ai::credit_availability::{AICreditDenialReason, AICreditSource}; +use crate::ai::credit_availability::{AICreditAvailability, AICreditDenialReason, AICreditSource}; use crate::auth::AuthStateProvider; use crate::pricing::PricingInfoModel; use crate::server::server_api::ServerApiProvider; @@ -424,6 +424,177 @@ fn test_buy_credits_banner_shows_when_non_ambient_bonus_credits_are_depleted() { }); } +#[test] +fn test_buy_credits_banner_hidden_when_server_reports_available() { + App::test((), |mut app| async move { + let (_uid, mut workspace) = create_test_workspace(); + workspace + .billing_metadata + .tier + .purchase_add_on_credits_policy = Some(standard_purchase_policy()); + + add_user_workspaces_with_workspace(&mut app, workspace); + let request_usage_model = add_request_usage_model(&mut app); + + request_usage_model.update(&mut app, |model, ctx| { + // Local base quota is exhausted; without server availability the + // banner would show. A non-ambient server source must hide it. + model.request_limit_info = RequestLimitInfo::new_for_test(10, 10); + model.bonus_grants.clear(); + model.apply_server_availability( + Ok(AICreditAvailability::available_with_source(Some( + AICreditSource::BonusGrant, + ))), + ctx, + ); + + assert_eq!( + model.compute_buy_addon_credits_banner_display_state(ctx), + BuyCreditsBannerDisplayState::Hidden, + ); + }); + }); +} + +#[test] +fn test_buy_credits_banner_shows_when_server_reports_out_of_credits() { + App::test((), |mut app| async move { + let (_uid, mut workspace) = create_test_workspace(); + workspace + .billing_metadata + .tier + .purchase_add_on_credits_policy = Some(standard_purchase_policy()); + + add_user_workspaces_with_workspace(&mut app, workspace); + let request_usage_model = add_request_usage_model(&mut app); + + request_usage_model.update(&mut app, |model, ctx| { + // Stale local base quota must not suppress the banner once the + // server has denied interactive AI. + model.request_limit_info = RequestLimitInfo::new_for_test(10, 5); + model.bonus_grants.clear(); + model.apply_server_availability( + Ok(AICreditAvailability::unavailable( + AICreditDenialReason::OutOfCredits, + )), + ctx, + ); + + assert_eq!( + model.compute_buy_addon_credits_banner_display_state(ctx), + BuyCreditsBannerDisplayState::OutOfCredits, + ); + }); + }); +} + +#[test] +fn test_buy_credits_banner_shows_when_server_source_is_ambient_only() { + App::test((), |mut app| async move { + let (_uid, mut workspace) = create_test_workspace(); + workspace + .billing_metadata + .tier + .purchase_add_on_credits_policy = Some(standard_purchase_policy()); + + add_user_workspaces_with_workspace(&mut app, workspace); + let request_usage_model = add_request_usage_model(&mut app); + + request_usage_model.update(&mut app, |model, ctx| { + model.request_limit_info = RequestLimitInfo::new_for_test(10, 10); + model.bonus_grants.clear(); + model.apply_server_availability( + Ok(AICreditAvailability::available_with_source(Some( + AICreditSource::AmbientBonusGrant, + ))), + ctx, + ); + + assert_eq!( + model.compute_buy_addon_credits_banner_display_state(ctx), + BuyCreditsBannerDisplayState::OutOfCredits, + ); + }); + }); +} + +#[test] +fn test_buy_credits_banner_hidden_when_out_of_credits_refined_by_local_byo() { + App::test((), |mut app| async move { + let (_uid, mut workspace) = create_test_workspace(); + workspace + .billing_metadata + .tier + .purchase_add_on_credits_policy = Some(standard_purchase_policy()); + workspace.billing_metadata.tier.byo_api_key_policy = + Some(ByoApiKeyPolicy { enabled: true }); + + add_user_workspaces_with_workspace(&mut app, workspace); + let request_usage_model = add_request_usage_model(&mut app); + + ApiKeyManager::handle(&app).update(&mut app, |manager, ctx| { + manager.set_provider_key(LLMProvider::OpenAI, Some("test-key".to_string()), ctx); + }); + + request_usage_model.update(&mut app, |model, ctx| { + model.request_limit_info = RequestLimitInfo::new_for_test(10, 10); + model.bonus_grants.clear(); + model.apply_server_availability( + Ok(AICreditAvailability::unavailable( + AICreditDenialReason::OutOfCredits, + )), + ctx, + ); + + assert!( + model.has_any_ai_remaining(ctx), + "local BYO should refine OutOfCredits into available AI" + ); + assert_eq!( + model.compute_buy_addon_credits_banner_display_state(ctx), + BuyCreditsBannerDisplayState::Hidden, + ); + }); + }); +} + +#[test] +fn test_buy_credits_banner_respects_monthly_limit_under_server_out_of_credits() { + App::test((), |mut app| async move { + let (_uid, mut workspace) = create_test_workspace(); + workspace + .billing_metadata + .tier + .purchase_add_on_credits_policy = Some(standard_purchase_policy()); + enable_auto_reload(&mut workspace); + // Zero monthly spend limit means any auto-reload is blocked. + workspace + .settings + .addon_credits_settings + .max_monthly_spend_cents = Some(0); + + add_user_workspaces_with_workspace(&mut app, workspace); + let request_usage_model = add_request_usage_model(&mut app); + set_addon_credits_pricing_info(&mut app); + + request_usage_model.update(&mut app, |model, ctx| { + model.request_limit_info = RequestLimitInfo::new_for_test(10, 10); + model.bonus_grants.clear(); + model.apply_server_availability( + Ok(AICreditAvailability::unavailable( + AICreditDenialReason::OutOfCredits, + )), + ctx, + ); + + assert_eq!( + model.compute_buy_addon_credits_banner_display_state(ctx), + BuyCreditsBannerDisplayState::MonthlyLimitReached, + ); + }); + }); +} + #[test] fn test_ambient_credits_banner_dismissal_is_persisted() { App::test((), |mut app| async move { diff --git a/app/src/settings_view/ai_page.rs b/app/src/settings_view/ai_page.rs index 545ccfd1ff7..442f75ec22d 100644 --- a/app/src/settings_view/ai_page.rs +++ b/app/src/settings_view/ai_page.rs @@ -2227,7 +2227,8 @@ impl AISettingsPageView { let on_paid_plan = UserWorkspaces::as_ref(ctx) .current_workspace() .is_some_and(|workspace| workspace.billing_metadata.is_user_on_paid_plan()); - let out_of_monthly_credits = !AIRequestUsageModel::as_ref(ctx).has_requests_remaining(); + let out_of_monthly_credits = + !AIRequestUsageModel::as_ref(ctx).has_base_plan_requests_remaining(); !on_paid_plan && out_of_monthly_credits && !Self::active_base_model_is_byo_covered(ctx) } diff --git a/app/src/terminal/buy_credits_banner.rs b/app/src/terminal/buy_credits_banner.rs index 6ed20d580b9..402a18a5bba 100644 --- a/app/src/terminal/buy_credits_banner.rs +++ b/app/src/terminal/buy_credits_banner.rs @@ -81,8 +81,9 @@ impl BuyCreditsBanner { ctx.subscribe_to_model( &AIRequestUsageModel::handle(ctx), - |me, _handle, event, ctx| { - if let AIRequestUsageModelEvent::RequestUsageUpdated = event { + |me, _handle, event, ctx| match event { + AIRequestUsageModelEvent::RequestUsageUpdated + | AIRequestUsageModelEvent::CreditAvailabilityUpdated => { if me.checkout_pending && matches!( AIRequestUsageModel::as_ref(ctx) @@ -98,6 +99,8 @@ impl BuyCreditsBanner { } ctx.notify(); } + AIRequestUsageModelEvent::AmbientCreditsBannerDismissed + | AIRequestUsageModelEvent::RequestBonusRefunded { .. } => {} }, ); From 0e6dbf7af05b9c86f815e11a8544dc4b72f9273a Mon Sep 17 00:00:00 2001 From: Jeff Lloyd Date: Sun, 2 Aug 2026 22:22:42 -0400 Subject: [PATCH 8/8] Address credit availability review feedback Remove overly specific comments, rename the cold-start availability fallback, and avoid the redundant targeted availability request after auth when the immediate workspace metadata poll already supplies it. Co-Authored-By: Oz --- app/src/ai/request_usage_model.rs | 80 +++------------------ app/src/ai/request_usage_model_tests.rs | 14 ++-- app/src/ai_assistant/transcript_tests.rs | 3 - app/src/auth/auth_manager.rs | 1 - app/src/lib.rs | 6 -- app/src/workspaces/user_workspaces_tests.rs | 4 -- 6 files changed, 17 insertions(+), 91 deletions(-) diff --git a/app/src/ai/request_usage_model.rs b/app/src/ai/request_usage_model.rs index 420159ed4e3..c57fa706f62 100644 --- a/app/src/ai/request_usage_model.rs +++ b/app/src/ai/request_usage_model.rs @@ -185,21 +185,12 @@ fn get_cached_ambient_credits_banner_dismissed(app_mut: &mut AppContext) -> bool .unwrap_or_default() } -/// The server-authoritative AI credit availability state shared by all AI -/// surfaces. It is fed from exactly two paths: the workspace metadata refresh -/// piggyback (primary cadence) and targeted fetches after meaningful state -/// changes (auth completion, plan/billing, workspace selection, credentials). #[derive(Default)] struct ServerAvailabilityState { - /// The last successfully fetched decision. Retained as last-known-good - /// when a refresh fails; never cleared by transient errors. + /// Last successful decision; retained across transient refresh failures. latest: Option, - /// When `latest` was last updated from a successful response. last_success_time: Option, - /// Whether a targeted availability fetch is currently in flight. Used to - /// coalesce coincident event-triggered fetches. refresh_in_flight: bool, - /// The most recent refresh failure, kept until the next success. last_error: Option, } @@ -341,13 +332,7 @@ impl AIRequestUsageModel { self.server_availability.latest } - /// Records the outcome of an availability fetch, whether piggybacked on a - /// workspace metadata refresh or from a targeted fetch. - /// - /// A failure keeps the last-known-good decision: a transport or resolver - /// error must never flip availability in either direction, and legacy - /// locally derived availability must not be re-enabled once a valid server - /// decision has been received. + /// Applies an availability fetch result, keeping last-known-good on failure. pub fn apply_server_availability( &mut self, result: Result, @@ -368,13 +353,6 @@ impl AIRequestUsageModel { } } - /// Fetches the server-authoritative availability decision in response to a - /// meaningful state change (auth completion, plan/billing change, - /// workspace selection change, or API-key/credential change). - /// - /// Coincident triggers are coalesced: if a fetch is already in flight this - /// is a no-op. There is intentionally no retry loop — the next qualifying - /// trigger or workspace metadata refresh serves as the retry. pub fn request_availability_refresh(&mut self, ctx: &mut ModelContext) { if !AuthStateProvider::as_ref(ctx).get().is_logged_in() { return; @@ -503,39 +481,21 @@ impl AIRequestUsageModel { } } - /// Returns `true` if the user still has unused **base-plan** AI request quota - /// (the monthly/weekly plan limit). - /// - /// This deliberately ignores add-on credits, bonus grants, overages, - /// auto-reload, and BYO credentials. It is **not** an AI-availability check — - /// use [`Self::has_any_ai_remaining`] when deciding whether the user can start - /// an interactive AI request. + /// Whether unused base-plan request quota remains. pub(crate) fn has_base_plan_requests_remaining(&self) -> bool { self.requests_remaining() > 0 } /// Returns `true` if the user can start an interactive AI request. - /// Use this method as the starting point for AI availability checking. - /// - /// Once a server-authoritative availability decision has been received - /// this session, that decision (last-known-good on transient refresh - /// failures) is the authority. The locally derived fallback below is used - /// solely before the first successful fetch, e.g. right after startup or - /// against servers that don't support the availability field yet. + /// Prefers the server decision when present; otherwise uses the pre-fetch fallback. pub fn has_any_ai_remaining(&self, ctx: &AppContext) -> bool { if let Some(availability) = self.server_availability.latest { return Self::server_availability_permits_ai(availability, ctx); } - self.has_any_ai_remaining_from_local_state(ctx) + self.has_any_ai_remaining_before_server_decision(ctx) } - /// Interprets the server-authoritative decision. `available` means the - /// server knows for a fact that requests can run (a Warp credit source or - /// a server-managed BYO path) and is trusted outright. An `OutOfCredits` - /// denial means the server found no path *it can see* — locally stored - /// API keys are request-level parameters invisible to it — so the client - /// contributes that one fact. Every other denial is a hard rejection that - /// local keys cannot bypass. + /// Trusts `available`; only `OutOfCredits` may be refined by local BYO credentials. fn server_availability_permits_ai( availability: AICreditAvailability, ctx: &AppContext, @@ -549,11 +509,8 @@ impl AIRequestUsageModel { ) && Self::has_usable_byo_inference_path(ctx) } - /// Whether a BYO inference path is usable with credentials held on this - /// machine: a stored API key, custom endpoint, or Grok subscription (when - /// the BYOK policy allows it), or loaded local-chain AWS credentials for - /// an enabled Bedrock host. Server-managed paths are already reflected in - /// the server's availability decision. + /// Whether a local BYO path is usable: stored API key/endpoint/Grok when + /// BYOK is allowed, or loaded AWS credentials for an enabled Bedrock host. fn has_usable_byo_inference_path(ctx: &AppContext) -> bool { let user_workspaces = UserWorkspaces::as_ref(ctx); let api_keys = ApiKeyManager::as_ref(ctx); @@ -567,17 +524,8 @@ impl AIRequestUsageModel { ) } - /// Legacy locally derived availability check. Returns `true` if the user - /// meets one of the following conditions: - /// 1. user has ai credits from the plan base limit - /// 2. user has overage enabled - /// 3. user has bonus grants (either team grants or user grants) - /// 4. user's team plan has pay-as-you-go enabled (enterprise only) - /// 5. user's team has enterprise bonus grants auto-reload enabled (enterprise only) - /// 6. user's team has self-serve auto-reload enabled within its monthly spend limit - /// 7. user has BYOK enabled and has either provided at least one API key or - /// connected a Grok subscription - fn has_any_ai_remaining_from_local_state(&self, ctx: &AppContext) -> bool { + /// Prefetch fallback used only before any successful server availability decision this session. + fn has_any_ai_remaining_before_server_decision(&self, ctx: &AppContext) -> bool { let current_workspace = UserWorkspaces::as_ref(ctx).current_workspace(); let has_base_plan_ai_requests = self.has_base_plan_requests_remaining(); @@ -740,13 +688,6 @@ impl AIRequestUsageModel { } /// Computes the current banner state based on live conditions. - /// This is called on-demand and always returns fresh state. - /// - /// Once a server-authoritative availability decision is known, the banner - /// hides whenever interactive AI is permitted (including `OutOfCredits` - /// refined by a usable local BYO path). Ambient-only credit sources are an - /// intentional exception: they fund cloud/ambient agents, not interactive - /// usage, so they must not suppress this banner. pub fn compute_buy_addon_credits_banner_display_state( &self, ctx: &AppContext, @@ -794,7 +735,6 @@ impl AIRequestUsageModel { return BuyCreditsBannerDisplayState::Hidden; } } else if self.has_base_plan_requests_remaining() || has_non_ambient_bonus_credits { - // Legacy pre-fetch path: local base quota / non-ambient bonus only. return BuyCreditsBannerDisplayState::Hidden; } diff --git a/app/src/ai/request_usage_model_tests.rs b/app/src/ai/request_usage_model_tests.rs index bb36ee59ed8..c458b319a8d 100644 --- a/app/src/ai/request_usage_model_tests.rs +++ b/app/src/ai/request_usage_model_tests.rs @@ -1268,7 +1268,7 @@ fn test_availability_refresh_failure_keeps_last_known_good() { let request_usage_model = add_request_usage_model(&mut app); request_usage_model.update(&mut app, |model, ctx| { - // Local state says AI is available, so any legacy fallback would + // Local state says AI is available, so the pre-server-decision fallback would // report `true`. model.request_limit_info = RequestLimitInfo::new_for_test(10, 5); @@ -1278,7 +1278,7 @@ fn test_availability_refresh_failure_keeps_last_known_good() { // The last-known-good server decision is retained: the error is // recorded but neither flips availability nor re-enables the - // legacy locally derived decision. + // pre-server-decision fallback. assert_eq!(model.server_availability(), Some(denied)); assert!(!model.has_any_ai_remaining(ctx)); assert_eq!( @@ -1290,7 +1290,7 @@ fn test_availability_refresh_failure_keeps_last_known_good() { } #[test] -fn test_availability_refresh_failure_before_first_success_uses_legacy_fallback() { +fn test_availability_refresh_failure_before_first_success_uses_prefetch_fallback() { App::test((), |mut app| async move { app.add_singleton_model(UserWorkspaces::default_mock); let request_usage_model = add_request_usage_model(&mut app); @@ -1300,7 +1300,7 @@ fn test_availability_refresh_failure_before_first_success_uses_legacy_fallback() model.apply_server_availability(Err(anyhow::anyhow!("unsupported operation")), ctx); // Without any successful fetch (e.g. server doesn't support the - // field yet), the legacy locally derived decision still applies. + // field yet), the pre-server-decision fallback still applies. assert_eq!(model.server_availability(), None); assert!(model.has_any_ai_remaining(ctx)); }); @@ -1308,7 +1308,7 @@ fn test_availability_refresh_failure_before_first_success_uses_legacy_fallback() } #[test] -fn test_reset_server_availability_restores_legacy_fallback() { +fn test_reset_server_availability_restores_prefetch_fallback() { App::test((), |mut app| async move { app.add_singleton_model(UserWorkspaces::default_mock); let request_usage_model = add_request_usage_model(&mut app); @@ -1323,8 +1323,8 @@ fn test_reset_server_availability_restores_legacy_fallback() { ); assert!(!model.has_any_ai_remaining(ctx)); - // On logout the server decision is cleared and pre-fetch behavior - // is restored for the next principal. + // On logout the server decision is cleared and the pre-server-decision + // fallback is restored for the next principal. model.reset_server_availability(ctx); assert_eq!(model.server_availability(), None); assert!(model.has_any_ai_remaining(ctx)); diff --git a/app/src/ai_assistant/transcript_tests.rs b/app/src/ai_assistant/transcript_tests.rs index d5af8386d58..9bac6b0289c 100644 --- a/app/src/ai_assistant/transcript_tests.rs +++ b/app/src/ai_assistant/transcript_tests.rs @@ -48,9 +48,6 @@ fn initialize_app(app: &mut App) { initialize_settings_for_tests(app); appearance::register(app); app.add_singleton_model(UserWorkspaces::default_mock); - // The transcript render path reads the shared AI usage model to decide - // whether to show prepared responses, so it must be registered (as it - // always is in the real app). app.add_singleton_model(|_| AuthStateProvider::new_for_test()); app.add_singleton_model(|_| ServerApiProvider::new_for_test()); app.add_singleton_model(|ctx| { diff --git a/app/src/auth/auth_manager.rs b/app/src/auth/auth_manager.rs index 3efc22b0802..b415adfdf95 100644 --- a/app/src/auth/auth_manager.rs +++ b/app/src/auth/auth_manager.rs @@ -455,7 +455,6 @@ impl AuthManager { AIRequestUsageModel::handle(ctx).update(ctx, |usage_model, ctx| { usage_model.refresh_request_usage_async(ctx); - usage_model.request_availability_refresh(ctx); }); LLMPreferences::handle(ctx).update(ctx, |prefs, ctx| { diff --git a/app/src/lib.rs b/app/src/lib.rs index d30f4b6dd57..72b5d9aed17 100644 --- a/app/src/lib.rs +++ b/app/src/lib.rs @@ -1722,12 +1722,6 @@ pub(crate) fn initialize_app( manager }); - // Keep the server-authoritative AI credit availability fresh across - // meaningful state changes. The primary cadence is the workspace metadata - // refresh piggyback; these targeted triggers cover changes that don't - // immediately produce a metadata response. `TeamsChanged` is deliberately - // not a trigger: it fires on every metadata poll, whose response already - // carries the piggybacked availability. ctx.subscribe_to_model(&UserWorkspaces::handle(ctx), |_, event, ctx| { if matches!( event, diff --git a/app/src/workspaces/user_workspaces_tests.rs b/app/src/workspaces/user_workspaces_tests.rs index cee7310e917..2a241bb1008 100644 --- a/app/src/workspaces/user_workspaces_tests.rs +++ b/app/src/workspaces/user_workspaces_tests.rs @@ -135,10 +135,6 @@ fn initialize_window_team_test_app(app: &mut App, workspaces: Vec) { }); } -// Registers the shared AI usage model (and its dependencies) so that the -// `on_workspaces_updated` metadata-apply path — which now piggybacks the -// server-authoritative AI credit availability onto every refresh — can update -// it without panicking on a missing singleton. fn register_ai_usage_model(app: &mut App) { app.add_singleton_model(|_| ServerApiProvider::new_for_test()); if app.models_of_type::().is_empty() {