collab: Require billing address in all Stripe checkouts (#32980)

morgankrey and Marshall Bowers created 5 months ago

Summary

I've successfully implemented the required billing address collection
feature for Stripe Checkout sessions. Here's what was done:

### 1. **Added New Data Structures** (`stripe_client.rs`):
- Added `StripeBillingAddressCollection` enum with `Auto` and `Required`
variants
- Added `billing_address_collection` field to
`StripeCreateCheckoutSessionParams`

### 2. **Updated Stripe Client Implementation**
(`real_stripe_client.rs`):
- Added conversion from `StripeBillingAddressCollection` to Stripe's
`CheckoutSessionBillingAddressCollection`
- Updated the `TryFrom` implementation to map the billing address
collection field when creating checkout sessions
- Added the necessary import

### 3. **Updated Billing Service** (`stripe_billing.rs`):
- Set `billing_address_collection` to `Required` in both
`checkout_with_zed_pro()` and `checkout_with_zed_pro_trial()` methods
- Added the necessary import

### 4. **Updated Test Infrastructure** (`fake_stripe_client.rs`):
- Added `billing_address_collection` field to
`StripeCreateCheckoutSessionCall`
- Updated the `create_checkout_session` implementation to capture the
new field
- Added the necessary import

### 5. **Updated Tests** (`stripe_billing_tests.rs`):
- Added assertions to verify that `billing_address_collection` is set to
`Required` in all three test cases:
  - `test_checkout_with_zed_pro`
  - `test_checkout_with_zed_pro_trial` (regular trial)
  - `test_checkout_with_zed_pro_trial` (extended trial)
- Added the necessary import

The implementation follows the pattern established in the codebase and
ensures that whenever a Stripe Checkout session is created for Zed Pro
subscriptions (both regular and trial), the billing address will be
required from customers. This aligns with the Stripe documentation you
provided, which shows that setting `billing_address_collection=required`
will ensure the billing address is always collected during checkout.

Release Notes:

- N/A

Co-authored-by: Marshall Bowers <git@maxdeviant.com>

Change summary

crates/collab/src/stripe_billing.rs                   |  7 ++-
crates/collab/src/stripe_client.rs                    |  7 +++
crates/collab/src/stripe_client/fake_stripe_client.rs |  6 ++-
crates/collab/src/stripe_client/real_stripe_client.rs | 21 ++++++++++-
crates/collab/src/tests/stripe_billing_tests.rs       | 23 ++++++++++--
5 files changed, 52 insertions(+), 12 deletions(-)

Detailed changes

crates/collab/src/stripe_billing.rs 🔗

@@ -11,8 +11,9 @@ use crate::Result;
 use crate::db::billing_subscription::SubscriptionKind;
 use crate::llm::AGENT_EXTENDED_TRIAL_FEATURE_FLAG;
 use crate::stripe_client::{
-    RealStripeClient, StripeCheckoutSessionMode, StripeCheckoutSessionPaymentMethodCollection,
-    StripeClient, StripeCreateCheckoutSessionLineItems, StripeCreateCheckoutSessionParams,
+    RealStripeClient, StripeBillingAddressCollection, StripeCheckoutSessionMode,
+    StripeCheckoutSessionPaymentMethodCollection, StripeClient,
+    StripeCreateCheckoutSessionLineItems, StripeCreateCheckoutSessionParams,
     StripeCreateCheckoutSessionSubscriptionData, StripeCreateMeterEventParams,
     StripeCreateMeterEventPayload, StripeCreateSubscriptionItems, StripeCreateSubscriptionParams,
     StripeCustomerId, StripeMeter, StripePrice, StripePriceId, StripeSubscription,
@@ -245,6 +246,7 @@ impl StripeBilling {
             quantity: Some(1),
         }]);
         params.success_url = Some(success_url);
+        params.billing_address_collection = Some(StripeBillingAddressCollection::Required);
 
         let session = self.client.create_checkout_session(params).await?;
         Ok(session.url.context("no checkout session URL")?)
@@ -298,6 +300,7 @@ impl StripeBilling {
             quantity: Some(1),
         }]);
         params.success_url = Some(success_url);
+        params.billing_address_collection = Some(StripeBillingAddressCollection::Required);
 
         let session = self.client.create_checkout_session(params).await?;
         Ok(session.url.context("no checkout session URL")?)

crates/collab/src/stripe_client.rs 🔗

@@ -148,6 +148,12 @@ pub struct StripeCreateMeterEventPayload<'a> {
     pub stripe_customer_id: &'a StripeCustomerId,
 }
 
+#[derive(Debug, PartialEq, Eq, Clone, Copy)]
+pub enum StripeBillingAddressCollection {
+    Auto,
+    Required,
+}
+
 #[derive(Debug, Default)]
 pub struct StripeCreateCheckoutSessionParams<'a> {
     pub customer: Option<&'a StripeCustomerId>,
@@ -157,6 +163,7 @@ pub struct StripeCreateCheckoutSessionParams<'a> {
     pub payment_method_collection: Option<StripeCheckoutSessionPaymentMethodCollection>,
     pub subscription_data: Option<StripeCreateCheckoutSessionSubscriptionData>,
     pub success_url: Option<&'a str>,
+    pub billing_address_collection: Option<StripeBillingAddressCollection>,
 }
 
 #[derive(Debug, PartialEq, Eq, Clone, Copy)]

crates/collab/src/stripe_client/fake_stripe_client.rs 🔗

@@ -8,8 +8,8 @@ use parking_lot::Mutex;
 use uuid::Uuid;
 
 use crate::stripe_client::{
-    CreateCustomerParams, StripeCheckoutSession, StripeCheckoutSessionMode,
-    StripeCheckoutSessionPaymentMethodCollection, StripeClient,
+    CreateCustomerParams, StripeBillingAddressCollection, StripeCheckoutSession,
+    StripeCheckoutSessionMode, StripeCheckoutSessionPaymentMethodCollection, StripeClient,
     StripeCreateCheckoutSessionLineItems, StripeCreateCheckoutSessionParams,
     StripeCreateCheckoutSessionSubscriptionData, StripeCreateMeterEventParams,
     StripeCreateSubscriptionParams, StripeCustomer, StripeCustomerId, StripeMeter, StripeMeterId,
@@ -35,6 +35,7 @@ pub struct StripeCreateCheckoutSessionCall {
     pub payment_method_collection: Option<StripeCheckoutSessionPaymentMethodCollection>,
     pub subscription_data: Option<StripeCreateCheckoutSessionSubscriptionData>,
     pub success_url: Option<String>,
+    pub billing_address_collection: Option<StripeBillingAddressCollection>,
 }
 
 pub struct FakeStripeClient {
@@ -231,6 +232,7 @@ impl StripeClient for FakeStripeClient {
                 payment_method_collection: params.payment_method_collection,
                 subscription_data: params.subscription_data,
                 success_url: params.success_url.map(|url| url.to_string()),
+                billing_address_collection: params.billing_address_collection,
             });
 
         Ok(StripeCheckoutSession {

crates/collab/src/stripe_client/real_stripe_client.rs 🔗

@@ -17,9 +17,10 @@ use stripe::{
 };
 
 use crate::stripe_client::{
-    CreateCustomerParams, StripeCancellationDetails, StripeCancellationDetailsReason,
-    StripeCheckoutSession, StripeCheckoutSessionMode, StripeCheckoutSessionPaymentMethodCollection,
-    StripeClient, StripeCreateCheckoutSessionLineItems, StripeCreateCheckoutSessionParams,
+    CreateCustomerParams, StripeBillingAddressCollection, StripeCancellationDetails,
+    StripeCancellationDetailsReason, StripeCheckoutSession, StripeCheckoutSessionMode,
+    StripeCheckoutSessionPaymentMethodCollection, StripeClient,
+    StripeCreateCheckoutSessionLineItems, StripeCreateCheckoutSessionParams,
     StripeCreateCheckoutSessionSubscriptionData, StripeCreateMeterEventParams,
     StripeCreateSubscriptionParams, StripeCustomer, StripeCustomerId, StripeMeter, StripePrice,
     StripePriceId, StripePriceRecurring, StripeSubscription, StripeSubscriptionId,
@@ -444,6 +445,7 @@ impl<'a> TryFrom<StripeCreateCheckoutSessionParams<'a>> for CreateCheckoutSessio
             payment_method_collection: value.payment_method_collection.map(Into::into),
             subscription_data: value.subscription_data.map(Into::into),
             success_url: value.success_url,
+            billing_address_collection: value.billing_address_collection.map(Into::into),
             ..Default::default()
         })
     }
@@ -526,3 +528,16 @@ impl From<CheckoutSession> for StripeCheckoutSession {
         Self { url: value.url }
     }
 }
+
+impl From<StripeBillingAddressCollection> for stripe::CheckoutSessionBillingAddressCollection {
+    fn from(value: StripeBillingAddressCollection) -> Self {
+        match value {
+            StripeBillingAddressCollection::Auto => {
+                stripe::CheckoutSessionBillingAddressCollection::Auto
+            }
+            StripeBillingAddressCollection::Required => {
+                stripe::CheckoutSessionBillingAddressCollection::Required
+            }
+        }
+    }
+}

crates/collab/src/tests/stripe_billing_tests.rs 🔗

@@ -6,11 +6,12 @@ use pretty_assertions::assert_eq;
 use crate::llm::AGENT_EXTENDED_TRIAL_FEATURE_FLAG;
 use crate::stripe_billing::StripeBilling;
 use crate::stripe_client::{
-    FakeStripeClient, StripeCheckoutSessionMode, StripeCheckoutSessionPaymentMethodCollection,
-    StripeCreateCheckoutSessionLineItems, StripeCreateCheckoutSessionSubscriptionData,
-    StripeCustomerId, StripeMeter, StripeMeterId, StripePrice, StripePriceId, StripePriceRecurring,
-    StripeSubscription, StripeSubscriptionId, StripeSubscriptionItem, StripeSubscriptionItemId,
-    StripeSubscriptionTrialSettings, StripeSubscriptionTrialSettingsEndBehavior,
+    FakeStripeClient, StripeBillingAddressCollection, StripeCheckoutSessionMode,
+    StripeCheckoutSessionPaymentMethodCollection, StripeCreateCheckoutSessionLineItems,
+    StripeCreateCheckoutSessionSubscriptionData, StripeCustomerId, StripeMeter, StripeMeterId,
+    StripePrice, StripePriceId, StripePriceRecurring, StripeSubscription, StripeSubscriptionId,
+    StripeSubscriptionItem, StripeSubscriptionItemId, StripeSubscriptionTrialSettings,
+    StripeSubscriptionTrialSettingsEndBehavior,
     StripeSubscriptionTrialSettingsEndBehaviorMissingPaymentMethod, UpdateSubscriptionItems,
 };
 
@@ -426,6 +427,10 @@ async fn test_checkout_with_zed_pro() {
         assert_eq!(call.payment_method_collection, None);
         assert_eq!(call.subscription_data, None);
         assert_eq!(call.success_url.as_deref(), Some(success_url));
+        assert_eq!(
+            call.billing_address_collection,
+            Some(StripeBillingAddressCollection::Required)
+        );
     }
 }
 
@@ -507,6 +512,10 @@ async fn test_checkout_with_zed_pro_trial() {
             })
         );
         assert_eq!(call.success_url.as_deref(), Some(success_url));
+        assert_eq!(
+            call.billing_address_collection,
+            Some(StripeBillingAddressCollection::Required)
+        );
     }
 
     // Successful checkout with extended trial.
@@ -561,5 +570,9 @@ async fn test_checkout_with_zed_pro_trial() {
             })
         );
         assert_eq!(call.success_url.as_deref(), Some(success_url));
+        assert_eq!(
+            call.billing_address_collection,
+            Some(StripeBillingAddressCollection::Required)
+        );
     }
 }