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 6b73ff1e963..fcde3963ae5 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/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 101bff4a315..08663f134b9 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; @@ -123,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. @@ -134,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; @@ -142,6 +146,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 +167,46 @@ 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 => { + // 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) + } + } + } + + /// 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(); @@ -481,3 +533,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..4fe7931f2ce --- /dev/null +++ b/app/src/ai/blocklist/prompt/prompt_alert_tests.rs @@ -0,0 +1,165 @@ +use std::sync::Arc; + +use ai::LLMProvider; +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; +use crate::server::telemetry::context_provider::AppTelemetryContextProvider; +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()); + app.add_singleton_model(AppTelemetryContextProvider::new_context_provider); + app.add_singleton_model(|ctx| { + UserWorkspaces::mock( + Arc::new(MockTeamClient::new()), + Arc::new(MockWorkspaceClient::new()), + workspaces, + 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); + }); +} + +#[test] +fn test_server_managed_availability_maps_to_no_alert() { + App::test((), |mut app| async move { + initialize_app(&mut app); + // `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::NoAlert); + }); +} + +#[test] +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); + 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_provider_key(LLMProvider::OpenAI, Some("test-key".to_string()), ctx); + }); + + // 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/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 cda81dbdfaf..c57fa706f62 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}; @@ -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, AICreditDenialReason, AICreditSource}; use crate::auth::AuthStateProvider; use crate::pricing::PricingInfoModel; use crate::server::server_api::ai::AIClient; @@ -184,6 +185,15 @@ fn get_cached_ambient_credits_banner_dismissed(app_mut: &mut AppContext) -> bool .unwrap_or_default() } +#[derive(Default)] +struct ServerAvailabilityState { + /// Last successful decision; retained across transient refresh failures. + latest: Option, + last_success_time: Option, + refresh_in_flight: bool, + last_error: Option, +} + pub struct AIRequestUsageModel { ai_client: Arc, @@ -194,6 +204,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 +219,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 +242,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 +255,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 +326,59 @@ 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 + } + + /// Applies an availability fetch result, keeping last-known-good on failure. + 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:#}")); + } + } + } + + 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, @@ -412,29 +481,54 @@ impl AIRequestUsageModel { } } - /// Returns `true` if the user has at least one request remaining before hitting the AI request - /// 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 { + /// 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 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 - /// Use this method as the starting point for AI availability checking. + /// Returns `true` if the user can start an interactive AI request. + /// 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_before_server_decision(ctx) + } + + /// Trusts `available`; only `OutOfCredits` may be refined by local BYO credentials. + fn server_availability_permits_ai( + availability: AICreditAvailability, + ctx: &AppContext, + ) -> bool { + if availability.available { + return true; + } + matches!( + availability.denial_reason, + AICreditDenialReason::OutOfCredits + ) && Self::has_usable_byo_inference_path(ctx) + } + + /// 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); + 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 { .. } + ) + } + + /// 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_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 @@ -594,7 +688,6 @@ impl AIRequestUsageModel { } /// Computes the current banner state based on live conditions. - /// This is called on-demand and always returns fresh state. pub fn compute_buy_addon_credits_banner_display_state( &self, ctx: &AppContext, @@ -609,6 +702,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. @@ -625,10 +722,19 @@ 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 { return BuyCreditsBannerDisplayState::Hidden; } diff --git a/app/src/ai/request_usage_model_tests.rs b/app/src/ai/request_usage_model_tests.rs index 40e27b49145..c458b319a8d 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_core::telemetry::testing::MockTelemetryContextProvider; @@ -9,9 +10,11 @@ use warp_graphql::billing::{AddonCreditsOption, OveragesPricing, PricingInfo}; use warpui::{App, ModelHandle}; use super::*; +use crate::ai::credit_availability::{AICreditAvailability, 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; @@ -75,6 +78,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( @@ -411,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 { @@ -1042,3 +1226,304 @@ 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 the pre-server-decision 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 + // pre-server-decision fallback. + 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_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); + + 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 pre-server-decision fallback still applies. + assert_eq!(model.server_availability(), None); + assert!(model.has_any_ai_remaining(ctx)); + }); + }); +} + +#[test] +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); + + 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 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)); + }); + }); +} + +#[test] +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(); + 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); + + // 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::unavailable( + AICreditDenialReason::OutOfCredits, + )), + ctx, + ); + assert!( + !model.has_any_ai_remaining(ctx), + "out of credits without a stored key should gate AI", + ); + }); + + // Storing a key supplies the one fact the server cannot know. + ApiKeyManager::handle(&app).update(&mut app, |manager, ctx| { + manager.set_provider_key(LLMProvider::OpenAI, Some("test-key".to_string()), ctx); + }); + request_usage_model.read(&app, |model, ctx| { + assert!( + model.has_any_ai_remaining(ctx), + "out of credits with a stored key should permit AI", + ); + }); + }); +} + +#[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 { + // `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| { + 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 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 }); + 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.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 { + 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/ai_assistant/transcript_tests.rs b/app/src/ai_assistant/transcript_tests.rs index bcb6a316287..9bac6b0289c 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,11 @@ fn initialize_app(app: &mut App) { initialize_settings_for_tests(app); appearance::register(app); app.add_singleton_model(UserWorkspaces::default_mock); + 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] diff --git a/app/src/auth/mod.rs b/app/src/auth/mod.rs index 75f6c5f374e..823b30b994f 100644 --- a/app/src/auth/mod.rs +++ b/app/src/auth/mod.rs @@ -35,6 +35,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}; @@ -288,6 +289,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 d03a71cb417..72b5d9aed17 100644 --- a/app/src/lib.rs +++ b/app/src/lib.rs @@ -1722,6 +1722,28 @@ pub(crate) fn initialize_app( manager }); + 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 9d9cf31b293..fe8cef06d17 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::{ @@ -142,6 +145,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; @@ -1142,6 +1146,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. @@ -1795,6 +1803,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 3fbceba6980..442f75ec22d 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 => {} } @@ -2226,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/settings_view/teams_page.rs b/app/src/settings_view/teams_page.rs index fdc9b0e205f..e89f6a98496 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/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 { .. } => {} }, ); diff --git a/app/src/workspaces/gql_convert.rs b/app/src/workspaces/gql_convert.rs index 5ad5500f450..86f306a4552 100644 --- a/app/src/workspaces/gql_convert.rs +++ b/app/src/workspaces/gql_convert.rs @@ -1179,6 +1179,7 @@ impl From for WorkspacesMetadataResponse { joinable_teams, experiments, feature_model_choices, + ai_credit_availability: Some(gql_user.ai_credit_availability.into()), user_purchase_policy, } } diff --git a/app/src/workspaces/update_manager.rs b/app/src/workspaces/update_manager.rs index 6964a2ec3df..066b065dc00 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, user_purchase_policy: None, }, pricing_info: None, @@ -348,6 +350,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; let user_purchase_policy = response.metadata.user_purchase_policy; @@ -473,6 +481,12 @@ impl TeamUpdateManager { let experiments = user_workspaces_access.experiments; let user_purchase_policy = user_workspaces_access.user_purchase_policy; + 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.set_user_purchase_policy(user_purchase_policy); user_workspaces.update_workspaces(workspaces.clone(), ctx); diff --git a/app/src/workspaces/update_manager_tests.rs b/app/src/workspaces/update_manager_tests.rs index b90a95888ac..9f0acfcdff4 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, user_purchase_policy: None, }, pricing_info: None, @@ -166,6 +168,7 @@ fn test_leaving_team_removes_objects() { joinable_teams: vec![], experiments: None, feature_model_choices: None, + ai_credit_availability: None, user_purchase_policy: None, }, pricing_info: None, @@ -210,6 +213,49 @@ 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), + user_purchase_policy: None, + }), + ctx, + ); + }); + + AIRequestUsageModel::handle(&app).read(&app, |model, _| { + assert_eq!(model.server_availability(), Some(availability)); + }); + }); +} + #[test] fn test_poll_path_apply_refreshes_user_purchase_policy() { App::test((), |mut app| async move { @@ -228,6 +274,7 @@ fn test_poll_path_apply_refreshes_user_purchase_policy() { joinable_teams: vec![], experiments: None, feature_model_choices: None, + ai_credit_availability: None, user_purchase_policy: Some(PurchaseAddOnCreditsPolicy { enabled: false, premium_enabled: true, @@ -253,6 +300,7 @@ fn test_poll_path_apply_refreshes_user_purchase_policy() { joinable_teams: vec![], experiments: None, feature_model_choices: None, + ai_credit_availability: None, user_purchase_policy: None, }; team_update_manager.update(&mut app, |manager, ctx| { diff --git a/app/src/workspaces/user_workspaces.rs b/app/src/workspaces/user_workspaces.rs index 3a91b620c4c..cd062d2ba8e 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; @@ -84,6 +86,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, @@ -122,6 +126,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, /// The user-level add-on credits purchase policy; the teamless-purchase /// fallback (see [`UserWorkspaces::purchase_policy`]). pub user_purchase_policy: Option, @@ -520,9 +527,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. @@ -1041,6 +1052,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 22fa0c97ef2..2a241bb1008 100644 --- a/app/src/workspaces/user_workspaces_tests.rs +++ b/app/src/workspaces/user_workspaces_tests.rs @@ -135,6 +135,16 @@ fn initialize_window_team_test_app(app: &mut App, workspaces: Vec) { }); } +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); @@ -192,6 +202,7 @@ fn test_loading_all_spaces_after_switching_from_offline() { joinable_teams: vec![], experiments: None, feature_model_choices: None, + ai_credit_availability: None, user_purchase_policy: None, }, pricing_info: None, @@ -210,6 +221,7 @@ fn test_loading_all_spaces_after_switching_from_offline() { joinable_teams: vec![], experiments: None, feature_model_choices: None, + ai_credit_availability: None, user_purchase_policy: None, }, pricing_info: None, @@ -348,6 +360,7 @@ fn test_aws_bedrock_credentials_respect_user_setting() { joinable_teams: vec![], experiments: None, feature_model_choices: None, + ai_credit_availability: None, user_purchase_policy: None, }, pricing_info: None, @@ -405,6 +418,7 @@ fn test_aws_bedrock_credentials_enforced_by_admin() { joinable_teams: vec![], experiments: None, feature_model_choices: None, + ai_credit_availability: None, user_purchase_policy: None, }, pricing_info: None, @@ -1720,6 +1734,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), @@ -1735,6 +1754,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 @@ -1790,6 +1810,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, diff --git a/crates/graphql/src/api/ai.rs b/crates/graphql/src/api/ai.rs index 75c4a669a36..38488ac8099 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 6991b60be65..366d9e96403 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, PurchaseAddOnCreditsPolicy}; use crate::experiment::Experiment; use crate::request_context::RequestContext; @@ -13,6 +14,11 @@ query GetWorkspacesMetadataForUser($requestContext: RequestContext!) { profile { uid } + aiCreditAvailability { + available + denialReason + creditSource + } billingMetadata { tier { purchaseAddOnCreditsPolicy { @@ -235,6 +241,7 @@ pub enum PricingInfoResult { #[derive(cynic::QueryFragment, Debug)] pub struct User { pub profile: UserProfile, + pub ai_credit_availability: AICreditAvailability, pub billing_metadata: Option, pub workspaces: Vec, pub experiments: Option>, diff --git a/crates/graphql/src/api/queries/mod.rs b/crates/graphql/src/api/queries/mod.rs index 3818fee1d5b..6eb17ba0c24 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 c8658c87097..3cfe92bea54 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 @@ -4345,6 +4368,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