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: &StripeSubscription,
115 ) -> Option<SubscriptionKind> {
116 let zed_pro_price_id = self.zed_pro_price_id().await.ok()?;
117 let zed_free_price_id = self.zed_free_price_id().await.ok()?;
118
119 subscription.items.iter().find_map(|item| {
120 let price = item.price.as_ref()?;
121
122 if price.id == zed_pro_price_id {
123 Some(if subscription.status == SubscriptionStatus::Trialing {
124 SubscriptionKind::ZedProTrial
125 } else {
126 SubscriptionKind::ZedPro
127 })
128 } else if price.id == zed_free_price_id {
129 Some(SubscriptionKind::ZedFree)
130 } else {
131 None
132 }
133 })
134 }
135
136 /// Returns the Stripe customer associated with the provided email address, or creates a new customer, if one does
137 /// not already exist.
138 ///
139 /// Always returns a new Stripe customer if the email address is `None`.
140 pub async fn find_or_create_customer_by_email(
141 &self,
142 email_address: Option<&str>,
143 ) -> Result<StripeCustomerId> {
144 let existing_customer = if let Some(email) = email_address {
145 let customers = self.client.list_customers_by_email(email).await?;
146
147 customers.first().cloned()
148 } else {
149 None
150 };
151
152 let customer_id = if let Some(existing_customer) = existing_customer {
153 existing_customer.id
154 } else {
155 let customer = self
156 .client
157 .create_customer(crate::stripe_client::CreateCustomerParams {
158 email: email_address,
159 })
160 .await?;
161
162 customer.id
163 };
164
165 Ok(customer_id)
166 }
167
168 pub async fn subscribe_to_price(
169 &self,
170 subscription_id: &StripeSubscriptionId,
171 price: &StripePrice,
172 ) -> Result<()> {
173 let subscription = self.client.get_subscription(subscription_id).await?;
174
175 if subscription_contains_price(&subscription, &price.id) {
176 return Ok(());
177 }
178
179 const BILLING_THRESHOLD_IN_CENTS: i64 = 20 * 100;
180
181 let price_per_unit = price.unit_amount.unwrap_or_default();
182 let _units_for_billing_threshold = BILLING_THRESHOLD_IN_CENTS / price_per_unit;
183
184 self.client
185 .update_subscription(
186 subscription_id,
187 UpdateSubscriptionParams {
188 items: Some(vec![UpdateSubscriptionItems {
189 price: Some(price.id.clone()),
190 }]),
191 trial_settings: Some(StripeSubscriptionTrialSettings {
192 end_behavior: StripeSubscriptionTrialSettingsEndBehavior {
193 missing_payment_method: StripeSubscriptionTrialSettingsEndBehaviorMissingPaymentMethod::Cancel
194 },
195 }),
196 },
197 )
198 .await?;
199
200 Ok(())
201 }
202
203 pub async fn bill_model_request_usage(
204 &self,
205 customer_id: &StripeCustomerId,
206 event_name: &str,
207 requests: i32,
208 ) -> Result<()> {
209 let timestamp = Utc::now().timestamp();
210 let idempotency_key = Uuid::new_v4();
211
212 self.client
213 .create_meter_event(StripeCreateMeterEventParams {
214 identifier: &format!("model_requests/{}", idempotency_key),
215 event_name,
216 payload: StripeCreateMeterEventPayload {
217 value: requests as u64,
218 stripe_customer_id: customer_id,
219 },
220 timestamp: Some(timestamp),
221 })
222 .await?;
223
224 Ok(())
225 }
226
227 pub async fn checkout_with_zed_pro(
228 &self,
229 customer_id: &StripeCustomerId,
230 github_login: &str,
231 success_url: &str,
232 ) -> Result<String> {
233 let zed_pro_price_id = self.zed_pro_price_id().await?;
234
235 let mut params = StripeCreateCheckoutSessionParams::default();
236 params.mode = Some(StripeCheckoutSessionMode::Subscription);
237 params.customer = Some(customer_id);
238 params.client_reference_id = Some(github_login);
239 params.line_items = Some(vec![StripeCreateCheckoutSessionLineItems {
240 price: Some(zed_pro_price_id.to_string()),
241 quantity: Some(1),
242 }]);
243 params.success_url = Some(success_url);
244
245 let session = self.client.create_checkout_session(params).await?;
246 Ok(session.url.context("no checkout session URL")?)
247 }
248
249 pub async fn checkout_with_zed_pro_trial(
250 &self,
251 customer_id: &StripeCustomerId,
252 github_login: &str,
253 feature_flags: Vec<String>,
254 success_url: &str,
255 ) -> Result<String> {
256 let zed_pro_price_id = self.zed_pro_price_id().await?;
257
258 let eligible_for_extended_trial = feature_flags
259 .iter()
260 .any(|flag| flag == AGENT_EXTENDED_TRIAL_FEATURE_FLAG);
261
262 let trial_period_days = if eligible_for_extended_trial { 60 } else { 14 };
263
264 let mut subscription_metadata = std::collections::HashMap::new();
265 if eligible_for_extended_trial {
266 subscription_metadata.insert(
267 "promo_feature_flag".to_string(),
268 AGENT_EXTENDED_TRIAL_FEATURE_FLAG.to_string(),
269 );
270 }
271
272 let mut params = StripeCreateCheckoutSessionParams::default();
273 params.subscription_data = Some(StripeCreateCheckoutSessionSubscriptionData {
274 trial_period_days: Some(trial_period_days),
275 trial_settings: Some(StripeSubscriptionTrialSettings {
276 end_behavior: StripeSubscriptionTrialSettingsEndBehavior {
277 missing_payment_method:
278 StripeSubscriptionTrialSettingsEndBehaviorMissingPaymentMethod::Cancel,
279 },
280 }),
281 metadata: if !subscription_metadata.is_empty() {
282 Some(subscription_metadata)
283 } else {
284 None
285 },
286 });
287 params.mode = Some(StripeCheckoutSessionMode::Subscription);
288 params.payment_method_collection =
289 Some(StripeCheckoutSessionPaymentMethodCollection::IfRequired);
290 params.customer = Some(customer_id);
291 params.client_reference_id = Some(github_login);
292 params.line_items = Some(vec![StripeCreateCheckoutSessionLineItems {
293 price: Some(zed_pro_price_id.to_string()),
294 quantity: Some(1),
295 }]);
296 params.success_url = Some(success_url);
297
298 let session = self.client.create_checkout_session(params).await?;
299 Ok(session.url.context("no checkout session URL")?)
300 }
301
302 pub async fn subscribe_to_zed_free(
303 &self,
304 customer_id: StripeCustomerId,
305 ) -> Result<StripeSubscription> {
306 let zed_free_price_id = self.zed_free_price_id().await?;
307
308 let existing_subscriptions = self
309 .client
310 .list_subscriptions_for_customer(&customer_id)
311 .await?;
312
313 let existing_active_subscription =
314 existing_subscriptions.into_iter().find(|subscription| {
315 subscription.status == SubscriptionStatus::Active
316 || subscription.status == SubscriptionStatus::Trialing
317 });
318 if let Some(subscription) = existing_active_subscription {
319 return Ok(subscription);
320 }
321
322 let params = StripeCreateSubscriptionParams {
323 customer: customer_id,
324 items: vec![StripeCreateSubscriptionItems {
325 price: Some(zed_free_price_id),
326 quantity: Some(1),
327 }],
328 };
329
330 let subscription = self.client.create_subscription(params).await?;
331
332 Ok(subscription)
333 }
334}
335
336fn subscription_contains_price(
337 subscription: &StripeSubscription,
338 price_id: &StripePriceId,
339) -> bool {
340 subscription.items.iter().any(|item| {
341 item.price
342 .as_ref()
343 .map_or(false, |price| price.id == *price_id)
344 })
345}