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