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 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(StripeSubscriptionTrialSettings {
194 end_behavior: StripeSubscriptionTrialSettingsEndBehavior {
195 missing_payment_method: StripeSubscriptionTrialSettingsEndBehaviorMissingPaymentMethod::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: &StripeCustomerId,
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 = StripeCreateCheckoutSessionParams::default();
238 params.mode = Some(StripeCheckoutSessionMode::Subscription);
239 params.customer = Some(customer_id);
240 params.client_reference_id = Some(github_login);
241 params.line_items = Some(vec![StripeCreateCheckoutSessionLineItems {
242 price: Some(zed_pro_price_id.to_string()),
243 quantity: Some(1),
244 }]);
245 params.success_url = Some(success_url);
246
247 let session = self.client.create_checkout_session(params).await?;
248 Ok(session.url.context("no checkout session URL")?)
249 }
250
251 pub async fn checkout_with_zed_pro_trial(
252 &self,
253 customer_id: &StripeCustomerId,
254 github_login: &str,
255 feature_flags: Vec<String>,
256 success_url: &str,
257 ) -> Result<String> {
258 let zed_pro_price_id = self.zed_pro_price_id().await?;
259
260 let eligible_for_extended_trial = feature_flags
261 .iter()
262 .any(|flag| flag == AGENT_EXTENDED_TRIAL_FEATURE_FLAG);
263
264 let trial_period_days = if eligible_for_extended_trial { 60 } else { 14 };
265
266 let mut subscription_metadata = std::collections::HashMap::new();
267 if eligible_for_extended_trial {
268 subscription_metadata.insert(
269 "promo_feature_flag".to_string(),
270 AGENT_EXTENDED_TRIAL_FEATURE_FLAG.to_string(),
271 );
272 }
273
274 let mut params = StripeCreateCheckoutSessionParams::default();
275 params.subscription_data = Some(StripeCreateCheckoutSessionSubscriptionData {
276 trial_period_days: Some(trial_period_days),
277 trial_settings: Some(StripeSubscriptionTrialSettings {
278 end_behavior: StripeSubscriptionTrialSettingsEndBehavior {
279 missing_payment_method:
280 StripeSubscriptionTrialSettingsEndBehaviorMissingPaymentMethod::Cancel,
281 },
282 }),
283 metadata: if !subscription_metadata.is_empty() {
284 Some(subscription_metadata)
285 } else {
286 None
287 },
288 });
289 params.mode = Some(StripeCheckoutSessionMode::Subscription);
290 params.payment_method_collection =
291 Some(StripeCheckoutSessionPaymentMethodCollection::IfRequired);
292 params.customer = Some(customer_id);
293 params.client_reference_id = Some(github_login);
294 params.line_items = Some(vec![StripeCreateCheckoutSessionLineItems {
295 price: Some(zed_pro_price_id.to_string()),
296 quantity: Some(1),
297 }]);
298 params.success_url = Some(success_url);
299
300 let session = self.client.create_checkout_session(params).await?;
301 Ok(session.url.context("no checkout session URL")?)
302 }
303
304 pub async fn subscribe_to_zed_free(
305 &self,
306 customer_id: StripeCustomerId,
307 ) -> Result<StripeSubscription> {
308 let zed_free_price_id = self.zed_free_price_id().await?;
309
310 let existing_subscriptions = self
311 .client
312 .list_subscriptions_for_customer(&customer_id)
313 .await?;
314
315 let existing_active_subscription =
316 existing_subscriptions.into_iter().find(|subscription| {
317 subscription.status == SubscriptionStatus::Active
318 || subscription.status == SubscriptionStatus::Trialing
319 });
320 if let Some(subscription) = existing_active_subscription {
321 return Ok(subscription);
322 }
323
324 let params = StripeCreateSubscriptionParams {
325 customer: customer_id,
326 items: vec![StripeCreateSubscriptionItems {
327 price: Some(zed_free_price_id),
328 quantity: Some(1),
329 }],
330 };
331
332 let subscription = self.client.create_subscription(params).await?;
333
334 Ok(subscription)
335 }
336}
337
338fn subscription_contains_price(
339 subscription: &StripeSubscription,
340 price_id: &StripePriceId,
341) -> bool {
342 subscription.items.iter().any(|item| {
343 item.price
344 .as_ref()
345 .map_or(false, |price| price.id == *price_id)
346 })
347}