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, StripeCheckoutSessionMode, StripeCheckoutSessionPaymentMethodCollection,
15 StripeClient, StripeCreateCheckoutSessionLineItems, StripeCreateCheckoutSessionParams,
16 StripeCreateCheckoutSessionSubscriptionData, StripeCreateMeterEventParams,
17 StripeCreateMeterEventPayload, StripeCreateSubscriptionItems, StripeCreateSubscriptionParams,
18 StripeCustomerId, StripeMeter, StripePrice, StripePriceId, StripeSubscription,
19 StripeSubscriptionId, StripeSubscriptionTrialSettings,
20 StripeSubscriptionTrialSettingsEndBehavior,
21 StripeSubscriptionTrialSettingsEndBehaviorMissingPaymentMethod, UpdateSubscriptionItems,
22 UpdateSubscriptionParams,
23};
24
25pub struct StripeBilling {
26 state: RwLock<StripeBillingState>,
27 client: Arc<dyn StripeClient>,
28}
29
30#[derive(Default)]
31struct StripeBillingState {
32 meters_by_event_name: HashMap<String, StripeMeter>,
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 (meters, prices) =
63 futures::try_join!(self.client.list_meters(), self.client.list_prices())?;
64
65 for meter in meters {
66 state
67 .meters_by_event_name
68 .insert(meter.event_name.clone(), meter);
69 }
70
71 for price in prices {
72 if let Some(lookup_key) = price.lookup_key.clone() {
73 state.prices_by_lookup_key.insert(lookup_key, price.clone());
74 }
75
76 if let Some(recurring) = price.recurring {
77 if let Some(meter) = recurring.meter {
78 state.price_ids_by_meter_id.insert(meter, price.id);
79 }
80 }
81 }
82
83 log::info!("StripeBilling: initialized");
84
85 Ok(())
86 }
87
88 pub async fn zed_pro_price_id(&self) -> Result<StripePriceId> {
89 self.find_price_id_by_lookup_key("zed-pro").await
90 }
91
92 pub async fn zed_free_price_id(&self) -> Result<StripePriceId> {
93 self.find_price_id_by_lookup_key("zed-free").await
94 }
95
96 pub async fn find_price_id_by_lookup_key(&self, lookup_key: &str) -> Result<StripePriceId> {
97 self.state
98 .read()
99 .await
100 .prices_by_lookup_key
101 .get(lookup_key)
102 .map(|price| price.id.clone())
103 .ok_or_else(|| crate::Error::Internal(anyhow!("no price ID found for {lookup_key:?}")))
104 }
105
106 pub async fn find_price_by_lookup_key(&self, lookup_key: &str) -> Result<StripePrice> {
107 self.state
108 .read()
109 .await
110 .prices_by_lookup_key
111 .get(lookup_key)
112 .cloned()
113 .ok_or_else(|| crate::Error::Internal(anyhow!("no price found for {lookup_key:?}")))
114 }
115
116 pub async fn determine_subscription_kind(
117 &self,
118 subscription: &StripeSubscription,
119 ) -> Option<SubscriptionKind> {
120 let zed_pro_price_id = self.zed_pro_price_id().await.ok()?;
121 let zed_free_price_id = self.zed_free_price_id().await.ok()?;
122
123 subscription.items.iter().find_map(|item| {
124 let price = item.price.as_ref()?;
125
126 if price.id == zed_pro_price_id {
127 Some(if subscription.status == SubscriptionStatus::Trialing {
128 SubscriptionKind::ZedProTrial
129 } else {
130 SubscriptionKind::ZedPro
131 })
132 } else if price.id == zed_free_price_id {
133 Some(SubscriptionKind::ZedFree)
134 } else {
135 None
136 }
137 })
138 }
139
140 /// Returns the Stripe customer associated with the provided email address, or creates a new customer, if one does
141 /// not already exist.
142 ///
143 /// Always returns a new Stripe customer if the email address is `None`.
144 pub async fn find_or_create_customer_by_email(
145 &self,
146 email_address: Option<&str>,
147 ) -> Result<StripeCustomerId> {
148 let existing_customer = if let Some(email) = email_address {
149 let customers = self.client.list_customers_by_email(email).await?;
150
151 customers.first().cloned()
152 } else {
153 None
154 };
155
156 let customer_id = if let Some(existing_customer) = existing_customer {
157 existing_customer.id
158 } else {
159 let customer = self
160 .client
161 .create_customer(crate::stripe_client::CreateCustomerParams {
162 email: email_address,
163 })
164 .await?;
165
166 customer.id
167 };
168
169 Ok(customer_id)
170 }
171
172 pub async fn subscribe_to_price(
173 &self,
174 subscription_id: &StripeSubscriptionId,
175 price: &StripePrice,
176 ) -> Result<()> {
177 let subscription = self.client.get_subscription(subscription_id).await?;
178
179 if subscription_contains_price(&subscription, &price.id) {
180 return Ok(());
181 }
182
183 const BILLING_THRESHOLD_IN_CENTS: i64 = 20 * 100;
184
185 let price_per_unit = price.unit_amount.unwrap_or_default();
186 let _units_for_billing_threshold = BILLING_THRESHOLD_IN_CENTS / price_per_unit;
187
188 self.client
189 .update_subscription(
190 subscription_id,
191 UpdateSubscriptionParams {
192 items: Some(vec![UpdateSubscriptionItems {
193 price: Some(price.id.clone()),
194 }]),
195 trial_settings: Some(StripeSubscriptionTrialSettings {
196 end_behavior: StripeSubscriptionTrialSettingsEndBehavior {
197 missing_payment_method: StripeSubscriptionTrialSettingsEndBehaviorMissingPaymentMethod::Cancel
198 },
199 }),
200 },
201 )
202 .await?;
203
204 Ok(())
205 }
206
207 pub async fn bill_model_request_usage(
208 &self,
209 customer_id: &StripeCustomerId,
210 event_name: &str,
211 requests: i32,
212 ) -> Result<()> {
213 let timestamp = Utc::now().timestamp();
214 let idempotency_key = Uuid::new_v4();
215
216 self.client
217 .create_meter_event(StripeCreateMeterEventParams {
218 identifier: &format!("model_requests/{}", idempotency_key),
219 event_name,
220 payload: StripeCreateMeterEventPayload {
221 value: requests as u64,
222 stripe_customer_id: customer_id,
223 },
224 timestamp: Some(timestamp),
225 })
226 .await?;
227
228 Ok(())
229 }
230
231 pub async fn checkout_with_zed_pro(
232 &self,
233 customer_id: &StripeCustomerId,
234 github_login: &str,
235 success_url: &str,
236 ) -> Result<String> {
237 let zed_pro_price_id = self.zed_pro_price_id().await?;
238
239 let mut params = StripeCreateCheckoutSessionParams::default();
240 params.mode = Some(StripeCheckoutSessionMode::Subscription);
241 params.customer = Some(customer_id);
242 params.client_reference_id = Some(github_login);
243 params.line_items = Some(vec![StripeCreateCheckoutSessionLineItems {
244 price: Some(zed_pro_price_id.to_string()),
245 quantity: Some(1),
246 }]);
247 params.success_url = Some(success_url);
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
302 let session = self.client.create_checkout_session(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: StripeCustomerId,
309 ) -> Result<StripeSubscription> {
310 let zed_free_price_id = self.zed_free_price_id().await?;
311
312 let existing_subscriptions = self
313 .client
314 .list_subscriptions_for_customer(&customer_id)
315 .await?;
316
317 let existing_active_subscription =
318 existing_subscriptions.into_iter().find(|subscription| {
319 subscription.status == SubscriptionStatus::Active
320 || subscription.status == SubscriptionStatus::Trialing
321 });
322 if let Some(subscription) = existing_active_subscription {
323 return Ok(subscription);
324 }
325
326 let params = StripeCreateSubscriptionParams {
327 customer: customer_id,
328 items: vec![StripeCreateSubscriptionItems {
329 price: Some(zed_free_price_id),
330 quantity: Some(1),
331 }],
332 };
333
334 let subscription = self.client.create_subscription(params).await?;
335
336 Ok(subscription)
337 }
338}
339
340fn subscription_contains_price(
341 subscription: &StripeSubscription,
342 price_id: &StripePriceId,
343) -> bool {
344 subscription.items.iter().any(|item| {
345 item.price
346 .as_ref()
347 .map_or(false, |price| price.id == *price_id)
348 })
349}