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