1use std::sync::Arc;
2
3use anyhow::{Context as _, anyhow};
4use chrono::Utc;
5use collections::HashMap;
6use stripe::SubscriptionStatus;
7use tokio::sync::RwLock;
8use uuid::Uuid;
9
10use crate::Result;
11use crate::db::billing_subscription::SubscriptionKind;
12use crate::llm::AGENT_EXTENDED_TRIAL_FEATURE_FLAG;
13use crate::stripe_client::{
14 RealStripeClient, StripeAutomaticTax, StripeBillingAddressCollection,
15 StripeCheckoutSessionMode, StripeCheckoutSessionPaymentMethodCollection, StripeClient,
16 StripeCreateCheckoutSessionLineItems, StripeCreateCheckoutSessionParams,
17 StripeCreateCheckoutSessionSubscriptionData, StripeCreateMeterEventParams,
18 StripeCreateMeterEventPayload, StripeCreateSubscriptionItems, StripeCreateSubscriptionParams,
19 StripeCustomerId, StripeCustomerUpdate, StripeCustomerUpdateAddress, StripeCustomerUpdateName,
20 StripePrice, StripePriceId, StripeSubscription, StripeSubscriptionId,
21 StripeSubscriptionTrialSettings, StripeSubscriptionTrialSettingsEndBehavior,
22 StripeSubscriptionTrialSettingsEndBehaviorMissingPaymentMethod, StripeTaxIdCollection,
23 UpdateSubscriptionItems, UpdateSubscriptionParams,
24};
25
26pub struct StripeBilling {
27 state: RwLock<StripeBillingState>,
28 client: Arc<dyn StripeClient>,
29}
30
31#[derive(Default)]
32struct StripeBillingState {
33 price_ids_by_meter_id: HashMap<String, StripePriceId>,
34 prices_by_lookup_key: HashMap<String, StripePrice>,
35}
36
37impl StripeBilling {
38 pub fn new(client: Arc<stripe::Client>) -> Self {
39 Self {
40 client: Arc::new(RealStripeClient::new(client.clone())),
41 state: RwLock::default(),
42 }
43 }
44
45 #[cfg(test)]
46 pub fn test(client: Arc<crate::stripe_client::FakeStripeClient>) -> Self {
47 Self {
48 client,
49 state: RwLock::default(),
50 }
51 }
52
53 pub fn client(&self) -> &Arc<dyn StripeClient> {
54 &self.client
55 }
56
57 pub async fn initialize(&self) -> Result<()> {
58 log::info!("StripeBilling: initializing");
59
60 let mut state = self.state.write().await;
61
62 let prices = self.client.list_prices().await?;
63
64 for price in prices {
65 if let Some(lookup_key) = price.lookup_key.clone() {
66 state.prices_by_lookup_key.insert(lookup_key, price.clone());
67 }
68
69 if let Some(recurring) = price.recurring {
70 if let Some(meter) = recurring.meter {
71 state.price_ids_by_meter_id.insert(meter, price.id);
72 }
73 }
74 }
75
76 log::info!("StripeBilling: initialized");
77
78 Ok(())
79 }
80
81 pub async fn zed_pro_price_id(&self) -> Result<StripePriceId> {
82 self.find_price_id_by_lookup_key("zed-pro").await
83 }
84
85 pub async fn zed_free_price_id(&self) -> Result<StripePriceId> {
86 self.find_price_id_by_lookup_key("zed-free").await
87 }
88
89 pub async fn find_price_id_by_lookup_key(&self, lookup_key: &str) -> Result<StripePriceId> {
90 self.state
91 .read()
92 .await
93 .prices_by_lookup_key
94 .get(lookup_key)
95 .map(|price| price.id.clone())
96 .ok_or_else(|| crate::Error::Internal(anyhow!("no price ID found for {lookup_key:?}")))
97 }
98
99 pub async fn find_price_by_lookup_key(&self, lookup_key: &str) -> Result<StripePrice> {
100 self.state
101 .read()
102 .await
103 .prices_by_lookup_key
104 .get(lookup_key)
105 .cloned()
106 .ok_or_else(|| crate::Error::Internal(anyhow!("no price found for {lookup_key:?}")))
107 }
108
109 pub async fn determine_subscription_kind(
110 &self,
111 subscription: &StripeSubscription,
112 ) -> Option<SubscriptionKind> {
113 let zed_pro_price_id = self.zed_pro_price_id().await.ok()?;
114 let zed_free_price_id = self.zed_free_price_id().await.ok()?;
115
116 subscription.items.iter().find_map(|item| {
117 let price = item.price.as_ref()?;
118
119 if price.id == zed_pro_price_id {
120 Some(if subscription.status == SubscriptionStatus::Trialing {
121 SubscriptionKind::ZedProTrial
122 } else {
123 SubscriptionKind::ZedPro
124 })
125 } else if price.id == zed_free_price_id {
126 Some(SubscriptionKind::ZedFree)
127 } else {
128 None
129 }
130 })
131 }
132
133 /// Returns the Stripe customer associated with the provided email address, or creates a new customer, if one does
134 /// not already exist.
135 ///
136 /// Always returns a new Stripe customer if the email address is `None`.
137 pub async fn find_or_create_customer_by_email(
138 &self,
139 email_address: Option<&str>,
140 ) -> Result<StripeCustomerId> {
141 let existing_customer = if let Some(email) = email_address {
142 let customers = self.client.list_customers_by_email(email).await?;
143
144 customers.first().cloned()
145 } else {
146 None
147 };
148
149 let customer_id = if let Some(existing_customer) = existing_customer {
150 existing_customer.id
151 } else {
152 let customer = self
153 .client
154 .create_customer(crate::stripe_client::CreateCustomerParams {
155 email: email_address,
156 })
157 .await?;
158
159 customer.id
160 };
161
162 Ok(customer_id)
163 }
164
165 pub async fn subscribe_to_price(
166 &self,
167 subscription_id: &StripeSubscriptionId,
168 price: &StripePrice,
169 ) -> Result<()> {
170 let subscription = self.client.get_subscription(subscription_id).await?;
171
172 if subscription_contains_price(&subscription, &price.id) {
173 return Ok(());
174 }
175
176 const BILLING_THRESHOLD_IN_CENTS: i64 = 20 * 100;
177
178 let price_per_unit = price.unit_amount.unwrap_or_default();
179 let _units_for_billing_threshold = BILLING_THRESHOLD_IN_CENTS / price_per_unit;
180
181 self.client
182 .update_subscription(
183 subscription_id,
184 UpdateSubscriptionParams {
185 items: Some(vec![UpdateSubscriptionItems {
186 price: Some(price.id.clone()),
187 }]),
188 trial_settings: Some(StripeSubscriptionTrialSettings {
189 end_behavior: StripeSubscriptionTrialSettingsEndBehavior {
190 missing_payment_method: StripeSubscriptionTrialSettingsEndBehaviorMissingPaymentMethod::Cancel
191 },
192 }),
193 },
194 )
195 .await?;
196
197 Ok(())
198 }
199
200 pub async fn bill_model_request_usage(
201 &self,
202 customer_id: &StripeCustomerId,
203 event_name: &str,
204 requests: i32,
205 ) -> Result<()> {
206 let timestamp = Utc::now().timestamp();
207 let idempotency_key = Uuid::new_v4();
208
209 self.client
210 .create_meter_event(StripeCreateMeterEventParams {
211 identifier: &format!("model_requests/{}", idempotency_key),
212 event_name,
213 payload: StripeCreateMeterEventPayload {
214 value: requests as u64,
215 stripe_customer_id: customer_id,
216 },
217 timestamp: Some(timestamp),
218 })
219 .await?;
220
221 Ok(())
222 }
223
224 pub async fn checkout_with_zed_pro(
225 &self,
226 customer_id: &StripeCustomerId,
227 github_login: &str,
228 success_url: &str,
229 ) -> Result<String> {
230 let zed_pro_price_id = self.zed_pro_price_id().await?;
231
232 let mut params = StripeCreateCheckoutSessionParams::default();
233 params.mode = Some(StripeCheckoutSessionMode::Subscription);
234 params.customer = Some(customer_id);
235 params.client_reference_id = Some(github_login);
236 params.line_items = Some(vec![StripeCreateCheckoutSessionLineItems {
237 price: Some(zed_pro_price_id.to_string()),
238 quantity: Some(1),
239 }]);
240 params.success_url = Some(success_url);
241 params.billing_address_collection = Some(StripeBillingAddressCollection::Required);
242 params.customer_update = Some(StripeCustomerUpdate {
243 address: Some(StripeCustomerUpdateAddress::Auto),
244 name: Some(StripeCustomerUpdateName::Auto),
245 shipping: None,
246 });
247 params.tax_id_collection = Some(StripeTaxIdCollection { enabled: true });
248
249 let session = self.client.create_checkout_session(params).await?;
250 Ok(session.url.context("no checkout session URL")?)
251 }
252
253 pub async fn checkout_with_zed_pro_trial(
254 &self,
255 customer_id: &StripeCustomerId,
256 github_login: &str,
257 feature_flags: Vec<String>,
258 success_url: &str,
259 ) -> Result<String> {
260 let zed_pro_price_id = self.zed_pro_price_id().await?;
261
262 let eligible_for_extended_trial = feature_flags
263 .iter()
264 .any(|flag| flag == AGENT_EXTENDED_TRIAL_FEATURE_FLAG);
265
266 let trial_period_days = if eligible_for_extended_trial { 60 } else { 14 };
267
268 let mut subscription_metadata = std::collections::HashMap::new();
269 if eligible_for_extended_trial {
270 subscription_metadata.insert(
271 "promo_feature_flag".to_string(),
272 AGENT_EXTENDED_TRIAL_FEATURE_FLAG.to_string(),
273 );
274 }
275
276 let mut params = StripeCreateCheckoutSessionParams::default();
277 params.subscription_data = Some(StripeCreateCheckoutSessionSubscriptionData {
278 trial_period_days: Some(trial_period_days),
279 trial_settings: Some(StripeSubscriptionTrialSettings {
280 end_behavior: StripeSubscriptionTrialSettingsEndBehavior {
281 missing_payment_method:
282 StripeSubscriptionTrialSettingsEndBehaviorMissingPaymentMethod::Cancel,
283 },
284 }),
285 metadata: if !subscription_metadata.is_empty() {
286 Some(subscription_metadata)
287 } else {
288 None
289 },
290 });
291 params.mode = Some(StripeCheckoutSessionMode::Subscription);
292 params.payment_method_collection =
293 Some(StripeCheckoutSessionPaymentMethodCollection::IfRequired);
294 params.customer = Some(customer_id);
295 params.client_reference_id = Some(github_login);
296 params.line_items = Some(vec![StripeCreateCheckoutSessionLineItems {
297 price: Some(zed_pro_price_id.to_string()),
298 quantity: Some(1),
299 }]);
300 params.success_url = Some(success_url);
301 params.billing_address_collection = Some(StripeBillingAddressCollection::Required);
302 params.customer_update = Some(StripeCustomerUpdate {
303 address: Some(StripeCustomerUpdateAddress::Auto),
304 name: Some(StripeCustomerUpdateName::Auto),
305 shipping: None,
306 });
307 params.tax_id_collection = Some(StripeTaxIdCollection { enabled: true });
308
309 let session = self.client.create_checkout_session(params).await?;
310 Ok(session.url.context("no checkout session URL")?)
311 }
312
313 pub async fn subscribe_to_zed_free(
314 &self,
315 customer_id: StripeCustomerId,
316 ) -> Result<StripeSubscription> {
317 let zed_free_price_id = self.zed_free_price_id().await?;
318
319 let existing_subscriptions = self
320 .client
321 .list_subscriptions_for_customer(&customer_id)
322 .await?;
323
324 let existing_active_subscription =
325 existing_subscriptions.into_iter().find(|subscription| {
326 subscription.status == SubscriptionStatus::Active
327 || subscription.status == SubscriptionStatus::Trialing
328 });
329 if let Some(subscription) = existing_active_subscription {
330 return Ok(subscription);
331 }
332
333 let params = StripeCreateSubscriptionParams {
334 customer: customer_id,
335 items: vec![StripeCreateSubscriptionItems {
336 price: Some(zed_free_price_id),
337 quantity: Some(1),
338 }],
339 automatic_tax: Some(StripeAutomaticTax { enabled: true }),
340 };
341
342 let subscription = self.client.create_subscription(params).await?;
343
344 Ok(subscription)
345 }
346}
347
348fn subscription_contains_price(
349 subscription: &StripeSubscription,
350 price_id: &StripePriceId,
351) -> bool {
352 subscription.items.iter().any(|item| {
353 item.price
354 .as_ref()
355 .map_or(false, |price| price.id == *price_id)
356 })
357}