warpdotdev
diff --git a/‎app/src/ai/blocklist/prompt/prompt_alert.rs‎
Lines changed: 12 additions & 9 deletions b/‎app/src/ai/blocklist/prompt/prompt_alert.rs‎
Lines changed: 12 additions & 9 deletions
diff --git a/‎app/src/ai/request_usage_model.rs‎
Lines changed: 5 additions & 9 deletions b/‎app/src/ai/request_usage_model.rs‎
Lines changed: 5 additions & 9 deletions
diff --git a/‎app/src/ai/request_usage_model_tests.rs‎
Lines changed: 142 additions & 6 deletions b/‎app/src/ai/request_usage_model_tests.rs‎
Lines changed: 142 additions & 6 deletions
diff --git a/‎app/src/server/server_api/workspace.rs‎
Lines changed: 23 additions & 6 deletions b/‎app/src/server/server_api/workspace.rs‎
Lines changed: 23 additions & 6 deletions
@@ -374,15 +374,18 @@ impl View for PromptAlertView {
         self.primary_text(&state, &mut text_fragments);
 
         let auth_state = AuthStateProvider::as_ref(app).get();
-        let current_team = UserWorkspaces::as_ref(app).team_for_view_handle(&self.view_handle, app);
-        let has_admin_permissions = auth_state
-            .user_email()
-            .zip(current_team)
-            .is_some_and(|(email, team)| team.has_admin_permissions(&email));
-
-        let can_purchase_addon_credits = current_team
-            .and_then(|team| team.billing_metadata.tier.purchase_add_on_credits_policy)
-            .is_some_and(|policy| policy.enabled);
+        let workspaces = UserWorkspaces::as_ref(app);
+        let current_team = workspaces.team_for_view_handle(&self.view_handle, app);
+        // A teamless user can be considered the admin of their non-existent team.
+        let has_admin_permissions = current_team.is_none_or(|team| {
+            auth_state
+                .user_email()
+                .is_some_and(|email| team.has_admin_permissions(&email))
+        });
+
+        let can_purchase_addon_credits = workspaces
+            .purchase_policy_for_team(current_team)
+            .is_some_and(|policy| policy.allows_purchases());
 
         let suggest_buy_credits = can_purchase_addon_credits
             && has_admin_permissions
 
@@ -603,15 +603,11 @@ impl AIRequestUsageModel {
         if self.buy_addon_credits_banner_dismissed {
             return BuyCreditsBannerDisplayState::Hidden;
         }
-        let current_workspace = UserWorkspaces::as_ref(ctx).current_workspace();
-        let policy_allows_purchasing = current_workspace
-            .map(|w| {
-                w.billing_metadata
-                    .tier
-                    .purchase_add_on_credits_policy
-                    .is_some_and(|p| p.enabled)
-            })
-            .unwrap_or(false);
+        let user_workspaces = UserWorkspaces::as_ref(ctx);
+        let current_workspace = user_workspaces.current_workspace();
+        let policy_allows_purchasing = user_workspaces
+            .purchase_policy()
+            .is_some_and(|policy| policy.allows_purchases());
 
         // TODO: we might want to suggest credits purchase if request_remain/bonus credits is below certain threshold
         // something to consider after launch
 
@@ -93,6 +93,22 @@ fn set_addon_credits_pricing_info(app: &mut App) {
     });
 }
 
+fn standard_purchase_policy() -> PurchaseAddOnCreditsPolicy {
+    PurchaseAddOnCreditsPolicy {
+        enabled: true,
+        premium_enabled: false,
+        price_premium_bps: 0,
+    }
+}
+
+fn premium_purchase_policy() -> PurchaseAddOnCreditsPolicy {
+    PurchaseAddOnCreditsPolicy {
+        enabled: false,
+        premium_enabled: true,
+        price_premium_bps: 1000,
+    }
+}
+
 fn enable_auto_reload(workspace: &mut Workspace) {
     workspace
         .settings
@@ -239,7 +255,7 @@ fn test_buy_credits_banner_shows_with_only_ambient_bonus_credits() {
         workspace
             .billing_metadata
             .tier
-            .purchase_add_on_credits_policy = Some(PurchaseAddOnCreditsPolicy { enabled: true });
+            .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);
@@ -266,14 +282,75 @@ fn test_buy_credits_banner_shows_with_only_ambient_bonus_credits() {
     });
 }
 
+#[test]
+fn test_buy_credits_banner_shows_for_premium_enabled_plan_out_of_credits() {
+    App::test((), |mut app| async move {
+        // The test workspace has no teams: this covers the teamless fresh
+        // free user, whose purchase policy lives on the workspace billing
+        // metadata until their first purchase creates a team server-side.
+        let (_uid, mut workspace) = create_test_workspace();
+        workspace
+            .billing_metadata
+            .tier
+            .purchase_add_on_credits_policy = Some(premium_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| {
+            assert!(
+                !UserWorkspaces::as_ref(ctx).has_teams(),
+                "this test covers the teamless case"
+            );
+            model.request_limit_info = RequestLimitInfo::new_for_test(10, 10);
+            model.bonus_grants.clear();
+
+            assert_eq!(
+                model.compute_buy_addon_credits_banner_display_state(ctx),
+                BuyCreditsBannerDisplayState::OutOfCredits,
+            );
+        });
+    });
+}
+
+#[test]
+fn test_buy_credits_banner_hidden_when_policy_fully_disabled() {
+    App::test((), |mut app| async move {
+        // Also a teamless workspace: without premiumEnabled the purchase
+        // surfaces must stay hidden for fresh free users.
+        let (_uid, mut workspace) = create_test_workspace();
+        workspace
+            .billing_metadata
+            .tier
+            .purchase_add_on_credits_policy = Some(PurchaseAddOnCreditsPolicy {
+            enabled: false,
+            premium_enabled: false,
+            price_premium_bps: 0,
+        });
+
+        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();
+
+            assert_eq!(
+                model.compute_buy_addon_credits_banner_display_state(ctx),
+                BuyCreditsBannerDisplayState::Hidden,
+            );
+        });
+    });
+}
+
 #[test]
 fn test_buy_credits_banner_hidden_with_non_ambient_bonus_credits() {
     App::test((), |mut app| async move {
         let (_uid, mut workspace) = create_test_workspace();
         workspace
             .billing_metadata
             .tier
-            .purchase_add_on_credits_policy = Some(PurchaseAddOnCreditsPolicy { enabled: true });
+            .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);
@@ -307,7 +384,7 @@ fn test_buy_credits_banner_shows_when_non_ambient_bonus_credits_are_depleted() {
         workspace
             .billing_metadata
             .tier
-            .purchase_add_on_credits_policy = Some(PurchaseAddOnCreditsPolicy { enabled: true });
+            .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);
@@ -568,7 +645,7 @@ fn test_has_any_ai_remaining_true_with_self_serve_auto_reload() {
         workspace
             .billing_metadata
             .tier
-            .purchase_add_on_credits_policy = Some(PurchaseAddOnCreditsPolicy { enabled: true });
+            .purchase_add_on_credits_policy = Some(standard_purchase_policy());
         enable_auto_reload(&mut workspace);
 
         add_user_workspaces_with_workspace(&mut app, workspace);
@@ -587,6 +664,65 @@ fn test_has_any_ai_remaining_true_with_self_serve_auto_reload() {
     });
 }
 
+#[test]
+fn test_has_any_ai_remaining_true_with_premium_auto_reload() {
+    App::test((), |mut app| async move {
+        let (_uid, mut workspace) = create_test_workspace();
+        workspace
+            .billing_metadata
+            .tier
+            .purchase_add_on_credits_policy = Some(premium_purchase_policy());
+        enable_auto_reload(&mut workspace);
+
+        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();
+
+            assert!(
+                model.has_any_ai_remaining(ctx),
+                "expected has_any_ai_remaining to be true when premium-plan auto-reload is enabled",
+            );
+        });
+    });
+}
+
+#[test]
+fn test_has_any_ai_remaining_false_when_premium_auto_reload_would_exceed_limit() {
+    App::test((), |mut app| async move {
+        let (_uid, mut workspace) = create_test_workspace();
+        workspace
+            .billing_metadata
+            .tier
+            .purchase_add_on_credits_policy = Some(premium_purchase_policy());
+        enable_auto_reload(&mut workspace);
+        // The reload's list price is $10.00 but the premium price is $11.00;
+        // a $10.50 monthly limit only blocks the reload when the premium
+        // surcharge is included in the check.
+        workspace
+            .settings
+            .addon_credits_settings
+            .max_monthly_spend_cents = Some(1050);
+
+        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();
+
+            assert!(
+                !model.has_any_ai_remaining(ctx),
+                "expected has_any_ai_remaining to be false when the premium-priced reload would exceed the monthly spend limit",
+            );
+        });
+    });
+}
+
 #[test]
 fn test_has_any_ai_remaining_true_with_self_serve_auto_reload_and_billing_v2_disabled() {
     App::test((), |mut app| async move {
@@ -596,7 +732,7 @@ fn test_has_any_ai_remaining_true_with_self_serve_auto_reload_and_billing_v2_dis
         workspace
             .billing_metadata
             .tier
-            .purchase_add_on_credits_policy = Some(PurchaseAddOnCreditsPolicy { enabled: true });
+            .purchase_add_on_credits_policy = Some(standard_purchase_policy());
         enable_auto_reload(&mut workspace);
 
         add_user_workspaces_with_workspace(&mut app, workspace);
@@ -622,7 +758,7 @@ fn test_has_any_ai_remaining_false_with_add_on_credits_policy_when_purchase_woul
         workspace
             .billing_metadata
             .tier
-            .purchase_add_on_credits_policy = Some(PurchaseAddOnCreditsPolicy { enabled: true });
+            .purchase_add_on_credits_policy = Some(standard_purchase_policy());
         enable_auto_reload(&mut workspace);
         workspace
             .settings
 
@@ -28,6 +28,18 @@ use crate::server::ids::ServerId;
 use crate::workspaces::user_workspaces::WorkspacesMetadataResponse;
 use crate::workspaces::workspace::AiOverages;
 
+/// Outcome of a successful `purchaseAddonCredits` mutation. Mirrors the
+/// server's `PurchaseAddonCreditsResult` union members one-to-one.
+pub enum PurchaseAddonCreditsOutcome {
+    /// The saved payment method was charged synchronously and credits were
+    /// granted immediately. Carries refreshed workspace metadata.
+    Completed(Box<WorkspacesMetadataResponse>),
+    /// There was no saved payment method to charge. The user must complete
+    /// the purchase in the browser at `checkout_url`; credits are granted
+    /// via webhook shortly after checkout completes.
+    CheckoutRequired { checkout_url: String },
+}
+
 #[cfg_attr(test, automock)]
 #[cfg_attr(not(target_family = "wasm"), async_trait)]
 #[cfg_attr(target_family = "wasm", async_trait(?Send))]
@@ -45,9 +57,9 @@ pub trait WorkspaceClient: 'static + Send + Sync {
 
     async fn purchase_addon_credits(
         &self,
-        team_uid: ServerId,
+        team_uid: Option<ServerId>,
         credits: i32,
-    ) -> Result<WorkspacesMetadataResponse>;
+    ) -> Result<PurchaseAddonCreditsOutcome>;
 
     async fn update_addon_credits_settings(
         &self,
@@ -151,12 +163,12 @@ impl WorkspaceClient for ServerApi {
 
     async fn purchase_addon_credits(
         &self,
-        team_uid: ServerId,
+        team_uid: Option<ServerId>,
         credits: i32,
-    ) -> Result<WorkspacesMetadataResponse> {
+    ) -> Result<PurchaseAddonCreditsOutcome> {
         let variables = PurchaseAddonCreditsVariables {
             input: PurchaseAddonCreditsInput {
-                team_uid: team_uid.into(),
+                team_uid: team_uid.map(Into::into),
                 credits,
             },
             request_context: get_request_context(),
@@ -170,7 +182,12 @@ impl WorkspaceClient for ServerApi {
                 PurchaseAddonCreditsResult::PurchaseAddonCreditsOutput(_) => {
                     TeamClient::workspaces_metadata(self)
                         .await
-                        .map(|w| w.metadata)
+                        .map(|w| PurchaseAddonCreditsOutcome::Completed(Box::new(w.metadata)))
+                }
+                PurchaseAddonCreditsResult::PurchaseAddonCreditsCheckoutOutput(output) => {
+                    Ok(PurchaseAddonCreditsOutcome::CheckoutRequired {
+                        checkout_url: output.checkout_url,
+                    })
                 }
                 PurchaseAddonCreditsResult::UserFacingError(error) => match error.error {
                     UserFacingErrorInterface::BudgetExceededError(budget_error) => {