1use anyhow::{Context, anyhow, bail};
2use axum::{
3 Extension, Json, Router,
4 extract::{self, Query},
5 routing::{get, post},
6};
7use chrono::{DateTime, SecondsFormat, Utc};
8use collections::HashSet;
9use reqwest::StatusCode;
10use sea_orm::ActiveValue;
11use serde::{Deserialize, Serialize};
12use serde_json::json;
13use std::{str::FromStr, sync::Arc, time::Duration};
14use stripe::{
15 BillingPortalSession, CancellationDetailsReason, CreateBillingPortalSession,
16 CreateBillingPortalSessionFlowData, CreateBillingPortalSessionFlowDataAfterCompletion,
17 CreateBillingPortalSessionFlowDataAfterCompletionRedirect,
18 CreateBillingPortalSessionFlowDataSubscriptionUpdateConfirm,
19 CreateBillingPortalSessionFlowDataSubscriptionUpdateConfirmItems,
20 CreateBillingPortalSessionFlowDataType, CreateCustomer, Customer, CustomerId, EventObject,
21 EventType, Expandable, ListEvents, Subscription, SubscriptionId, SubscriptionStatus,
22};
23use util::{ResultExt, maybe};
24
25use crate::api::events::SnowflakeRow;
26use crate::db::billing_subscription::{
27 StripeCancellationReason, StripeSubscriptionStatus, SubscriptionKind,
28};
29use crate::llm::db::subscription_usage_meter::CompletionMode;
30use crate::llm::{
31 AGENT_EXTENDED_TRIAL_FEATURE_FLAG, DEFAULT_MAX_MONTHLY_SPEND, FREE_TIER_MONTHLY_SPENDING_LIMIT,
32};
33use crate::rpc::{ResultExt as _, Server};
34use crate::{AppState, Cents, Error, Result};
35use crate::{db::UserId, llm::db::LlmDatabase};
36use crate::{
37 db::{
38 BillingSubscriptionId, CreateBillingCustomerParams, CreateBillingSubscriptionParams,
39 CreateProcessedStripeEventParams, UpdateBillingCustomerParams,
40 UpdateBillingPreferencesParams, UpdateBillingSubscriptionParams, billing_customer,
41 },
42 stripe_billing::StripeBilling,
43};
44
45pub fn router() -> Router {
46 Router::new()
47 .route(
48 "/billing/preferences",
49 get(get_billing_preferences).put(update_billing_preferences),
50 )
51 .route(
52 "/billing/subscriptions",
53 get(list_billing_subscriptions).post(create_billing_subscription),
54 )
55 .route(
56 "/billing/subscriptions/manage",
57 post(manage_billing_subscription),
58 )
59 .route(
60 "/billing/subscriptions/migrate",
61 post(migrate_to_new_billing),
62 )
63 .route("/billing/monthly_spend", get(get_monthly_spend))
64 .route("/billing/usage", get(get_current_usage))
65}
66
67#[derive(Debug, Deserialize)]
68struct GetBillingPreferencesParams {
69 github_user_id: i32,
70}
71
72#[derive(Debug, Serialize)]
73struct BillingPreferencesResponse {
74 max_monthly_llm_usage_spending_in_cents: i32,
75 model_request_overages_enabled: bool,
76 model_request_overages_spend_limit_in_cents: i32,
77}
78
79async fn get_billing_preferences(
80 Extension(app): Extension<Arc<AppState>>,
81 Query(params): Query<GetBillingPreferencesParams>,
82) -> Result<Json<BillingPreferencesResponse>> {
83 let user = app
84 .db
85 .get_user_by_github_user_id(params.github_user_id)
86 .await?
87 .ok_or_else(|| anyhow!("user not found"))?;
88
89 let preferences = app.db.get_billing_preferences(user.id).await?;
90
91 Ok(Json(BillingPreferencesResponse {
92 max_monthly_llm_usage_spending_in_cents: preferences
93 .as_ref()
94 .map_or(DEFAULT_MAX_MONTHLY_SPEND.0 as i32, |preferences| {
95 preferences.max_monthly_llm_usage_spending_in_cents
96 }),
97 model_request_overages_enabled: preferences.as_ref().map_or(false, |preferences| {
98 preferences.model_request_overages_enabled
99 }),
100 model_request_overages_spend_limit_in_cents: preferences
101 .as_ref()
102 .map_or(0, |preferences| {
103 preferences.model_request_overages_spend_limit_in_cents
104 }),
105 }))
106}
107
108#[derive(Debug, Deserialize)]
109struct UpdateBillingPreferencesBody {
110 github_user_id: i32,
111 #[serde(default)]
112 max_monthly_llm_usage_spending_in_cents: i32,
113 #[serde(default)]
114 model_request_overages_enabled: bool,
115 #[serde(default)]
116 model_request_overages_spend_limit_in_cents: i32,
117}
118
119async fn update_billing_preferences(
120 Extension(app): Extension<Arc<AppState>>,
121 Extension(rpc_server): Extension<Arc<crate::rpc::Server>>,
122 extract::Json(body): extract::Json<UpdateBillingPreferencesBody>,
123) -> Result<Json<BillingPreferencesResponse>> {
124 let user = app
125 .db
126 .get_user_by_github_user_id(body.github_user_id)
127 .await?
128 .ok_or_else(|| anyhow!("user not found"))?;
129
130 let max_monthly_llm_usage_spending_in_cents =
131 body.max_monthly_llm_usage_spending_in_cents.max(0);
132 let model_request_overages_spend_limit_in_cents =
133 body.model_request_overages_spend_limit_in_cents.max(0);
134
135 let billing_preferences =
136 if let Some(_billing_preferences) = app.db.get_billing_preferences(user.id).await? {
137 app.db
138 .update_billing_preferences(
139 user.id,
140 &UpdateBillingPreferencesParams {
141 max_monthly_llm_usage_spending_in_cents: ActiveValue::set(
142 max_monthly_llm_usage_spending_in_cents,
143 ),
144 model_request_overages_enabled: ActiveValue::set(
145 body.model_request_overages_enabled,
146 ),
147 model_request_overages_spend_limit_in_cents: ActiveValue::set(
148 model_request_overages_spend_limit_in_cents,
149 ),
150 },
151 )
152 .await?
153 } else {
154 app.db
155 .create_billing_preferences(
156 user.id,
157 &crate::db::CreateBillingPreferencesParams {
158 max_monthly_llm_usage_spending_in_cents,
159 model_request_overages_enabled: body.model_request_overages_enabled,
160 model_request_overages_spend_limit_in_cents,
161 },
162 )
163 .await?
164 };
165
166 SnowflakeRow::new(
167 "Billing Preferences Updated",
168 Some(user.metrics_id),
169 user.admin,
170 None,
171 json!({
172 "user_id": user.id,
173 "model_request_overages_enabled": billing_preferences.model_request_overages_enabled,
174 "model_request_overages_spend_limit_in_cents": billing_preferences.model_request_overages_spend_limit_in_cents,
175 "max_monthly_llm_usage_spending_in_cents": billing_preferences.max_monthly_llm_usage_spending_in_cents,
176 }),
177 )
178 .write(&app.kinesis_client, &app.config.kinesis_stream)
179 .await
180 .log_err();
181
182 rpc_server.refresh_llm_tokens_for_user(user.id).await;
183
184 Ok(Json(BillingPreferencesResponse {
185 max_monthly_llm_usage_spending_in_cents: billing_preferences
186 .max_monthly_llm_usage_spending_in_cents,
187 model_request_overages_enabled: billing_preferences.model_request_overages_enabled,
188 model_request_overages_spend_limit_in_cents: billing_preferences
189 .model_request_overages_spend_limit_in_cents,
190 }))
191}
192
193#[derive(Debug, Deserialize)]
194struct ListBillingSubscriptionsParams {
195 github_user_id: i32,
196}
197
198#[derive(Debug, Serialize)]
199struct BillingSubscriptionJson {
200 id: BillingSubscriptionId,
201 name: String,
202 status: StripeSubscriptionStatus,
203 trial_end_at: Option<String>,
204 cancel_at: Option<String>,
205 /// Whether this subscription can be canceled.
206 is_cancelable: bool,
207}
208
209#[derive(Debug, Serialize)]
210struct ListBillingSubscriptionsResponse {
211 subscriptions: Vec<BillingSubscriptionJson>,
212}
213
214async fn list_billing_subscriptions(
215 Extension(app): Extension<Arc<AppState>>,
216 Query(params): Query<ListBillingSubscriptionsParams>,
217) -> Result<Json<ListBillingSubscriptionsResponse>> {
218 let user = app
219 .db
220 .get_user_by_github_user_id(params.github_user_id)
221 .await?
222 .ok_or_else(|| anyhow!("user not found"))?;
223
224 let subscriptions = app.db.get_billing_subscriptions(user.id).await?;
225
226 Ok(Json(ListBillingSubscriptionsResponse {
227 subscriptions: subscriptions
228 .into_iter()
229 .map(|subscription| BillingSubscriptionJson {
230 id: subscription.id,
231 name: match subscription.kind {
232 Some(SubscriptionKind::ZedPro) => "Zed Pro".to_string(),
233 Some(SubscriptionKind::ZedProTrial) => "Zed Pro (Trial)".to_string(),
234 Some(SubscriptionKind::ZedFree) => "Zed Free".to_string(),
235 None => "Zed LLM Usage".to_string(),
236 },
237 status: subscription.stripe_subscription_status,
238 trial_end_at: if subscription.kind == Some(SubscriptionKind::ZedProTrial) {
239 maybe!({
240 let end_at = subscription.stripe_current_period_end?;
241 let end_at = DateTime::from_timestamp(end_at, 0)?;
242
243 Some(end_at.to_rfc3339_opts(SecondsFormat::Millis, true))
244 })
245 } else {
246 None
247 },
248 cancel_at: subscription.stripe_cancel_at.map(|cancel_at| {
249 cancel_at
250 .and_utc()
251 .to_rfc3339_opts(SecondsFormat::Millis, true)
252 }),
253 is_cancelable: subscription.stripe_subscription_status.is_cancelable()
254 && subscription.stripe_cancel_at.is_none(),
255 })
256 .collect(),
257 }))
258}
259
260#[derive(Debug, Clone, Copy, Deserialize)]
261#[serde(rename_all = "snake_case")]
262enum ProductCode {
263 ZedPro,
264 ZedProTrial,
265 ZedFree,
266}
267
268#[derive(Debug, Deserialize)]
269struct CreateBillingSubscriptionBody {
270 github_user_id: i32,
271 product: Option<ProductCode>,
272}
273
274#[derive(Debug, Serialize)]
275struct CreateBillingSubscriptionResponse {
276 checkout_session_url: String,
277}
278
279/// Initiates a Stripe Checkout session for creating a billing subscription.
280async fn create_billing_subscription(
281 Extension(app): Extension<Arc<AppState>>,
282 extract::Json(body): extract::Json<CreateBillingSubscriptionBody>,
283) -> Result<Json<CreateBillingSubscriptionResponse>> {
284 let user = app
285 .db
286 .get_user_by_github_user_id(body.github_user_id)
287 .await?
288 .ok_or_else(|| anyhow!("user not found"))?;
289
290 let Some(stripe_client) = app.stripe_client.clone() else {
291 log::error!("failed to retrieve Stripe client");
292 Err(Error::http(
293 StatusCode::NOT_IMPLEMENTED,
294 "not supported".into(),
295 ))?
296 };
297 let Some(stripe_billing) = app.stripe_billing.clone() else {
298 log::error!("failed to retrieve Stripe billing object");
299 Err(Error::http(
300 StatusCode::NOT_IMPLEMENTED,
301 "not supported".into(),
302 ))?
303 };
304 let Some(llm_db) = app.llm_db.clone() else {
305 log::error!("failed to retrieve LLM database");
306 Err(Error::http(
307 StatusCode::NOT_IMPLEMENTED,
308 "not supported".into(),
309 ))?
310 };
311
312 if app.db.has_active_billing_subscription(user.id).await? {
313 return Err(Error::http(
314 StatusCode::CONFLICT,
315 "user already has an active subscription".into(),
316 ));
317 }
318
319 let existing_billing_customer = app.db.get_billing_customer_by_user_id(user.id).await?;
320 if let Some(existing_billing_customer) = &existing_billing_customer {
321 if existing_billing_customer.has_overdue_invoices {
322 return Err(Error::http(
323 StatusCode::PAYMENT_REQUIRED,
324 "user has overdue invoices".into(),
325 ));
326 }
327 }
328
329 let customer_id = if let Some(existing_customer) = &existing_billing_customer {
330 CustomerId::from_str(&existing_customer.stripe_customer_id)
331 .context("failed to parse customer ID")?
332 } else {
333 let existing_customer = if let Some(email) = user.email_address.as_deref() {
334 let customers = Customer::list(
335 &stripe_client,
336 &stripe::ListCustomers {
337 email: Some(email),
338 ..Default::default()
339 },
340 )
341 .await?;
342
343 customers.data.first().cloned()
344 } else {
345 None
346 };
347
348 if let Some(existing_customer) = existing_customer {
349 existing_customer.id
350 } else {
351 let customer = Customer::create(
352 &stripe_client,
353 CreateCustomer {
354 email: user.email_address.as_deref(),
355 ..Default::default()
356 },
357 )
358 .await?;
359
360 customer.id
361 }
362 };
363
364 let success_url = format!(
365 "{}/account?checkout_complete=1",
366 app.config.zed_dot_dev_url()
367 );
368
369 let checkout_session_url = match body.product {
370 Some(ProductCode::ZedPro) => {
371 stripe_billing
372 .checkout_with_zed_pro(customer_id, &user.github_login, &success_url)
373 .await?
374 }
375 Some(ProductCode::ZedProTrial) => {
376 if let Some(existing_billing_customer) = &existing_billing_customer {
377 if existing_billing_customer.trial_started_at.is_some() {
378 return Err(Error::http(
379 StatusCode::FORBIDDEN,
380 "user already used free trial".into(),
381 ));
382 }
383 }
384
385 let feature_flags = app.db.get_user_flags(user.id).await?;
386
387 stripe_billing
388 .checkout_with_zed_pro_trial(
389 customer_id,
390 &user.github_login,
391 feature_flags,
392 &success_url,
393 )
394 .await?
395 }
396 Some(ProductCode::ZedFree) => {
397 stripe_billing
398 .checkout_with_zed_free(customer_id, &user.github_login, &success_url)
399 .await?
400 }
401 None => {
402 let default_model = llm_db.model(
403 zed_llm_client::LanguageModelProvider::Anthropic,
404 "claude-3-7-sonnet",
405 )?;
406 let stripe_model = stripe_billing
407 .register_model_for_token_based_usage(default_model)
408 .await?;
409 stripe_billing
410 .checkout(customer_id, &user.github_login, &stripe_model, &success_url)
411 .await?
412 }
413 };
414
415 Ok(Json(CreateBillingSubscriptionResponse {
416 checkout_session_url,
417 }))
418}
419
420#[derive(Debug, PartialEq, Deserialize)]
421#[serde(rename_all = "snake_case")]
422enum ManageSubscriptionIntent {
423 /// The user intends to manage their subscription.
424 ///
425 /// This will open the Stripe billing portal without putting the user in a specific flow.
426 ManageSubscription,
427 /// The user intends to upgrade to Zed Pro.
428 UpgradeToPro,
429 /// The user intends to cancel their subscription.
430 Cancel,
431 /// The user intends to stop the cancellation of their subscription.
432 StopCancellation,
433}
434
435#[derive(Debug, Deserialize)]
436struct ManageBillingSubscriptionBody {
437 github_user_id: i32,
438 intent: ManageSubscriptionIntent,
439 /// The ID of the subscription to manage.
440 subscription_id: BillingSubscriptionId,
441}
442
443#[derive(Debug, Serialize)]
444struct ManageBillingSubscriptionResponse {
445 billing_portal_session_url: Option<String>,
446}
447
448/// Initiates a Stripe customer portal session for managing a billing subscription.
449async fn manage_billing_subscription(
450 Extension(app): Extension<Arc<AppState>>,
451 extract::Json(body): extract::Json<ManageBillingSubscriptionBody>,
452) -> Result<Json<ManageBillingSubscriptionResponse>> {
453 let user = app
454 .db
455 .get_user_by_github_user_id(body.github_user_id)
456 .await?
457 .ok_or_else(|| anyhow!("user not found"))?;
458
459 let Some(stripe_client) = app.stripe_client.clone() else {
460 log::error!("failed to retrieve Stripe client");
461 Err(Error::http(
462 StatusCode::NOT_IMPLEMENTED,
463 "not supported".into(),
464 ))?
465 };
466
467 let Some(stripe_billing) = app.stripe_billing.clone() else {
468 log::error!("failed to retrieve Stripe billing object");
469 Err(Error::http(
470 StatusCode::NOT_IMPLEMENTED,
471 "not supported".into(),
472 ))?
473 };
474
475 let customer = app
476 .db
477 .get_billing_customer_by_user_id(user.id)
478 .await?
479 .ok_or_else(|| anyhow!("billing customer not found"))?;
480 let customer_id = CustomerId::from_str(&customer.stripe_customer_id)
481 .context("failed to parse customer ID")?;
482
483 let subscription = app
484 .db
485 .get_billing_subscription_by_id(body.subscription_id)
486 .await?
487 .ok_or_else(|| anyhow!("subscription not found"))?;
488 let subscription_id = SubscriptionId::from_str(&subscription.stripe_subscription_id)
489 .context("failed to parse subscription ID")?;
490
491 if body.intent == ManageSubscriptionIntent::StopCancellation {
492 let updated_stripe_subscription = Subscription::update(
493 &stripe_client,
494 &subscription_id,
495 stripe::UpdateSubscription {
496 cancel_at_period_end: Some(false),
497 ..Default::default()
498 },
499 )
500 .await?;
501
502 app.db
503 .update_billing_subscription(
504 subscription.id,
505 &UpdateBillingSubscriptionParams {
506 stripe_cancel_at: ActiveValue::set(
507 updated_stripe_subscription
508 .cancel_at
509 .and_then(|cancel_at| DateTime::from_timestamp(cancel_at, 0))
510 .map(|time| time.naive_utc()),
511 ),
512 ..Default::default()
513 },
514 )
515 .await?;
516
517 return Ok(Json(ManageBillingSubscriptionResponse {
518 billing_portal_session_url: None,
519 }));
520 }
521
522 let flow = match body.intent {
523 ManageSubscriptionIntent::ManageSubscription => None,
524 ManageSubscriptionIntent::UpgradeToPro => {
525 let zed_pro_price_id = stripe_billing.zed_pro_price_id().await?;
526 let zed_free_price_id = stripe_billing.zed_free_price_id().await?;
527
528 let stripe_subscription =
529 Subscription::retrieve(&stripe_client, &subscription_id, &[]).await?;
530
531 let is_on_zed_pro_trial = stripe_subscription.status == SubscriptionStatus::Trialing
532 && stripe_subscription.items.data.iter().any(|item| {
533 item.price
534 .as_ref()
535 .map_or(false, |price| price.id == zed_pro_price_id)
536 });
537 if is_on_zed_pro_trial {
538 // If the user is already on a Zed Pro trial and wants to upgrade to Pro, we just need to end their trial early.
539 Subscription::update(
540 &stripe_client,
541 &stripe_subscription.id,
542 stripe::UpdateSubscription {
543 trial_end: Some(stripe::Scheduled::now()),
544 ..Default::default()
545 },
546 )
547 .await?;
548
549 return Ok(Json(ManageBillingSubscriptionResponse {
550 billing_portal_session_url: None,
551 }));
552 }
553
554 let subscription_item_to_update = stripe_subscription
555 .items
556 .data
557 .iter()
558 .find_map(|item| {
559 let price = item.price.as_ref()?;
560
561 if price.id == zed_free_price_id {
562 Some(item.id.clone())
563 } else {
564 None
565 }
566 })
567 .ok_or_else(|| anyhow!("No subscription item to update"))?;
568
569 Some(CreateBillingPortalSessionFlowData {
570 type_: CreateBillingPortalSessionFlowDataType::SubscriptionUpdateConfirm,
571 subscription_update_confirm: Some(
572 CreateBillingPortalSessionFlowDataSubscriptionUpdateConfirm {
573 subscription: subscription.stripe_subscription_id,
574 items: vec![
575 CreateBillingPortalSessionFlowDataSubscriptionUpdateConfirmItems {
576 id: subscription_item_to_update.to_string(),
577 price: Some(zed_pro_price_id.to_string()),
578 quantity: Some(1),
579 },
580 ],
581 discounts: None,
582 },
583 ),
584 ..Default::default()
585 })
586 }
587 ManageSubscriptionIntent::Cancel => Some(CreateBillingPortalSessionFlowData {
588 type_: CreateBillingPortalSessionFlowDataType::SubscriptionCancel,
589 after_completion: Some(CreateBillingPortalSessionFlowDataAfterCompletion {
590 type_: stripe::CreateBillingPortalSessionFlowDataAfterCompletionType::Redirect,
591 redirect: Some(CreateBillingPortalSessionFlowDataAfterCompletionRedirect {
592 return_url: format!("{}/account", app.config.zed_dot_dev_url()),
593 }),
594 ..Default::default()
595 }),
596 subscription_cancel: Some(
597 stripe::CreateBillingPortalSessionFlowDataSubscriptionCancel {
598 subscription: subscription.stripe_subscription_id,
599 retention: None,
600 },
601 ),
602 ..Default::default()
603 }),
604 ManageSubscriptionIntent::StopCancellation => unreachable!(),
605 };
606
607 let mut params = CreateBillingPortalSession::new(customer_id);
608 params.flow_data = flow;
609 let return_url = format!("{}/account", app.config.zed_dot_dev_url());
610 params.return_url = Some(&return_url);
611
612 let session = BillingPortalSession::create(&stripe_client, params).await?;
613
614 Ok(Json(ManageBillingSubscriptionResponse {
615 billing_portal_session_url: Some(session.url),
616 }))
617}
618
619#[derive(Debug, Deserialize)]
620struct MigrateToNewBillingBody {
621 github_user_id: i32,
622}
623
624#[derive(Debug, Serialize)]
625struct MigrateToNewBillingResponse {
626 /// The ID of the subscription that was canceled.
627 canceled_subscription_id: Option<String>,
628}
629
630async fn migrate_to_new_billing(
631 Extension(app): Extension<Arc<AppState>>,
632 extract::Json(body): extract::Json<MigrateToNewBillingBody>,
633) -> Result<Json<MigrateToNewBillingResponse>> {
634 let Some(stripe_client) = app.stripe_client.clone() else {
635 log::error!("failed to retrieve Stripe client");
636 Err(Error::http(
637 StatusCode::NOT_IMPLEMENTED,
638 "not supported".into(),
639 ))?
640 };
641
642 let user = app
643 .db
644 .get_user_by_github_user_id(body.github_user_id)
645 .await?
646 .ok_or_else(|| anyhow!("user not found"))?;
647
648 let old_billing_subscriptions_by_user = app
649 .db
650 .get_active_billing_subscriptions(HashSet::from_iter([user.id]))
651 .await?;
652
653 let canceled_subscription_id = if let Some((_billing_customer, billing_subscription)) =
654 old_billing_subscriptions_by_user.get(&user.id)
655 {
656 let stripe_subscription_id = billing_subscription
657 .stripe_subscription_id
658 .parse::<stripe::SubscriptionId>()
659 .context("failed to parse Stripe subscription ID from database")?;
660
661 Subscription::cancel(
662 &stripe_client,
663 &stripe_subscription_id,
664 stripe::CancelSubscription {
665 invoice_now: Some(true),
666 ..Default::default()
667 },
668 )
669 .await?;
670
671 Some(stripe_subscription_id)
672 } else {
673 None
674 };
675
676 let all_feature_flags = app.db.list_feature_flags().await?;
677 let user_feature_flags = app.db.get_user_flags(user.id).await?;
678
679 for feature_flag in ["new-billing", "assistant2"] {
680 let already_in_feature_flag = user_feature_flags.iter().any(|flag| flag == feature_flag);
681 if already_in_feature_flag {
682 continue;
683 }
684
685 let feature_flag = all_feature_flags
686 .iter()
687 .find(|flag| flag.flag == feature_flag)
688 .context("failed to find feature flag: {feature_flag:?}")?;
689
690 app.db.add_user_flag(user.id, feature_flag.id).await?;
691 }
692
693 Ok(Json(MigrateToNewBillingResponse {
694 canceled_subscription_id: canceled_subscription_id
695 .map(|subscription_id| subscription_id.to_string()),
696 }))
697}
698
699/// The amount of time we wait in between each poll of Stripe events.
700///
701/// This value should strike a balance between:
702/// 1. Being short enough that we update quickly when something in Stripe changes
703/// 2. Being long enough that we don't eat into our rate limits.
704///
705/// As a point of reference, the Sequin folks say they have this at **500ms**:
706///
707/// > We poll the Stripe /events endpoint every 500ms per account
708/// >
709/// > — https://blog.sequinstream.com/events-not-webhooks/
710const POLL_EVENTS_INTERVAL: Duration = Duration::from_secs(5);
711
712/// The maximum number of events to return per page.
713///
714/// We set this to 100 (the max) so we have to make fewer requests to Stripe.
715///
716/// > Limit can range between 1 and 100, and the default is 10.
717const EVENTS_LIMIT_PER_PAGE: u64 = 100;
718
719/// The number of pages consisting entirely of already-processed events that we
720/// will see before we stop retrieving events.
721///
722/// This is used to prevent over-fetching the Stripe events API for events we've
723/// already seen and processed.
724const NUMBER_OF_ALREADY_PROCESSED_PAGES_BEFORE_WE_STOP: usize = 4;
725
726/// Polls the Stripe events API periodically to reconcile the records in our
727/// database with the data in Stripe.
728pub fn poll_stripe_events_periodically(app: Arc<AppState>, rpc_server: Arc<Server>) {
729 let Some(stripe_client) = app.stripe_client.clone() else {
730 log::warn!("failed to retrieve Stripe client");
731 return;
732 };
733
734 let executor = app.executor.clone();
735 executor.spawn_detached({
736 let executor = executor.clone();
737 async move {
738 loop {
739 poll_stripe_events(&app, &rpc_server, &stripe_client)
740 .await
741 .log_err();
742
743 executor.sleep(POLL_EVENTS_INTERVAL).await;
744 }
745 }
746 });
747}
748
749async fn poll_stripe_events(
750 app: &Arc<AppState>,
751 rpc_server: &Arc<Server>,
752 stripe_client: &stripe::Client,
753) -> anyhow::Result<()> {
754 fn event_type_to_string(event_type: EventType) -> String {
755 // Calling `to_string` on `stripe::EventType` members gives us a quoted string,
756 // so we need to unquote it.
757 event_type.to_string().trim_matches('"').to_string()
758 }
759
760 let event_types = [
761 EventType::CustomerCreated,
762 EventType::CustomerUpdated,
763 EventType::CustomerSubscriptionCreated,
764 EventType::CustomerSubscriptionUpdated,
765 EventType::CustomerSubscriptionPaused,
766 EventType::CustomerSubscriptionResumed,
767 EventType::CustomerSubscriptionDeleted,
768 ]
769 .into_iter()
770 .map(event_type_to_string)
771 .collect::<Vec<_>>();
772
773 let mut pages_of_already_processed_events = 0;
774 let mut unprocessed_events = Vec::new();
775
776 log::info!(
777 "Stripe events: starting retrieval for {}",
778 event_types.join(", ")
779 );
780 let mut params = ListEvents::new();
781 params.types = Some(event_types.clone());
782 params.limit = Some(EVENTS_LIMIT_PER_PAGE);
783
784 let mut event_pages = stripe::Event::list(&stripe_client, ¶ms)
785 .await?
786 .paginate(params);
787
788 loop {
789 let processed_event_ids = {
790 let event_ids = event_pages
791 .page
792 .data
793 .iter()
794 .map(|event| event.id.as_str())
795 .collect::<Vec<_>>();
796 app.db
797 .get_processed_stripe_events_by_event_ids(&event_ids)
798 .await?
799 .into_iter()
800 .map(|event| event.stripe_event_id)
801 .collect::<Vec<_>>()
802 };
803
804 let mut processed_events_in_page = 0;
805 let events_in_page = event_pages.page.data.len();
806 for event in &event_pages.page.data {
807 if processed_event_ids.contains(&event.id.to_string()) {
808 processed_events_in_page += 1;
809 log::debug!("Stripe events: already processed '{}', skipping", event.id);
810 } else {
811 unprocessed_events.push(event.clone());
812 }
813 }
814
815 if processed_events_in_page == events_in_page {
816 pages_of_already_processed_events += 1;
817 }
818
819 if event_pages.page.has_more {
820 if pages_of_already_processed_events >= NUMBER_OF_ALREADY_PROCESSED_PAGES_BEFORE_WE_STOP
821 {
822 log::info!(
823 "Stripe events: stopping, saw {pages_of_already_processed_events} pages of already-processed events"
824 );
825 break;
826 } else {
827 log::info!("Stripe events: retrieving next page");
828 event_pages = event_pages.next(&stripe_client).await?;
829 }
830 } else {
831 break;
832 }
833 }
834
835 log::info!("Stripe events: unprocessed {}", unprocessed_events.len());
836
837 // Sort all of the unprocessed events in ascending order, so we can handle them in the order they occurred.
838 unprocessed_events.sort_by(|a, b| a.created.cmp(&b.created).then_with(|| a.id.cmp(&b.id)));
839
840 for event in unprocessed_events {
841 let event_id = event.id.clone();
842 let processed_event_params = CreateProcessedStripeEventParams {
843 stripe_event_id: event.id.to_string(),
844 stripe_event_type: event_type_to_string(event.type_),
845 stripe_event_created_timestamp: event.created,
846 };
847
848 // If the event has happened too far in the past, we don't want to
849 // process it and risk overwriting other more-recent updates.
850 //
851 // 1 day was chosen arbitrarily. This could be made longer or shorter.
852 let one_day = Duration::from_secs(24 * 60 * 60);
853 let a_day_ago = Utc::now() - one_day;
854 if a_day_ago.timestamp() > event.created {
855 log::info!(
856 "Stripe events: event '{}' is more than {one_day:?} old, marking as processed",
857 event_id
858 );
859 app.db
860 .create_processed_stripe_event(&processed_event_params)
861 .await?;
862
863 return Ok(());
864 }
865
866 let process_result = match event.type_ {
867 EventType::CustomerCreated | EventType::CustomerUpdated => {
868 handle_customer_event(app, stripe_client, event).await
869 }
870 EventType::CustomerSubscriptionCreated
871 | EventType::CustomerSubscriptionUpdated
872 | EventType::CustomerSubscriptionPaused
873 | EventType::CustomerSubscriptionResumed
874 | EventType::CustomerSubscriptionDeleted => {
875 handle_customer_subscription_event(app, rpc_server, stripe_client, event).await
876 }
877 _ => Ok(()),
878 };
879
880 if let Some(()) = process_result
881 .with_context(|| format!("failed to process event {event_id} successfully"))
882 .log_err()
883 {
884 app.db
885 .create_processed_stripe_event(&processed_event_params)
886 .await?;
887 }
888 }
889
890 Ok(())
891}
892
893async fn handle_customer_event(
894 app: &Arc<AppState>,
895 _stripe_client: &stripe::Client,
896 event: stripe::Event,
897) -> anyhow::Result<()> {
898 let EventObject::Customer(customer) = event.data.object else {
899 bail!("unexpected event payload for {}", event.id);
900 };
901
902 log::info!("handling Stripe {} event: {}", event.type_, event.id);
903
904 let Some(email) = customer.email else {
905 log::info!("Stripe customer has no email: skipping");
906 return Ok(());
907 };
908
909 let Some(user) = app.db.get_user_by_email(&email).await? else {
910 log::info!("no user found for email: skipping");
911 return Ok(());
912 };
913
914 if let Some(existing_customer) = app
915 .db
916 .get_billing_customer_by_stripe_customer_id(&customer.id)
917 .await?
918 {
919 app.db
920 .update_billing_customer(
921 existing_customer.id,
922 &UpdateBillingCustomerParams {
923 // For now we just leave the information as-is, as it is not
924 // likely to change.
925 ..Default::default()
926 },
927 )
928 .await?;
929 } else {
930 app.db
931 .create_billing_customer(&CreateBillingCustomerParams {
932 user_id: user.id,
933 stripe_customer_id: customer.id.to_string(),
934 })
935 .await?;
936 }
937
938 Ok(())
939}
940
941async fn handle_customer_subscription_event(
942 app: &Arc<AppState>,
943 rpc_server: &Arc<Server>,
944 stripe_client: &stripe::Client,
945 event: stripe::Event,
946) -> anyhow::Result<()> {
947 let EventObject::Subscription(subscription) = event.data.object else {
948 bail!("unexpected event payload for {}", event.id);
949 };
950
951 log::info!("handling Stripe {} event: {}", event.type_, event.id);
952
953 let subscription_kind = maybe!(async {
954 let stripe_billing = app.stripe_billing.clone()?;
955
956 let zed_pro_price_id = stripe_billing.zed_pro_price_id().await.ok()?;
957 let zed_free_price_id = stripe_billing.zed_free_price_id().await.ok()?;
958
959 subscription.items.data.iter().find_map(|item| {
960 let price = item.price.as_ref()?;
961
962 if price.id == zed_pro_price_id {
963 Some(if subscription.status == SubscriptionStatus::Trialing {
964 SubscriptionKind::ZedProTrial
965 } else {
966 SubscriptionKind::ZedPro
967 })
968 } else if price.id == zed_free_price_id {
969 Some(SubscriptionKind::ZedFree)
970 } else {
971 None
972 }
973 })
974 })
975 .await;
976
977 let billing_customer =
978 find_or_create_billing_customer(app, stripe_client, subscription.customer)
979 .await?
980 .ok_or_else(|| anyhow!("billing customer not found"))?;
981
982 if let Some(SubscriptionKind::ZedProTrial) = subscription_kind {
983 if subscription.status == SubscriptionStatus::Trialing {
984 let current_period_start =
985 DateTime::from_timestamp(subscription.current_period_start, 0)
986 .ok_or_else(|| anyhow!("No trial subscription period start"))?;
987
988 app.db
989 .update_billing_customer(
990 billing_customer.id,
991 &UpdateBillingCustomerParams {
992 trial_started_at: ActiveValue::set(Some(current_period_start.naive_utc())),
993 ..Default::default()
994 },
995 )
996 .await?;
997 }
998 }
999
1000 let was_canceled_due_to_payment_failure = subscription.status == SubscriptionStatus::Canceled
1001 && subscription
1002 .cancellation_details
1003 .as_ref()
1004 .and_then(|details| details.reason)
1005 .map_or(false, |reason| {
1006 reason == CancellationDetailsReason::PaymentFailed
1007 });
1008
1009 if was_canceled_due_to_payment_failure {
1010 app.db
1011 .update_billing_customer(
1012 billing_customer.id,
1013 &UpdateBillingCustomerParams {
1014 has_overdue_invoices: ActiveValue::set(true),
1015 ..Default::default()
1016 },
1017 )
1018 .await?;
1019 }
1020
1021 if let Some(existing_subscription) = app
1022 .db
1023 .get_billing_subscription_by_stripe_subscription_id(&subscription.id)
1024 .await?
1025 {
1026 app.db
1027 .update_billing_subscription(
1028 existing_subscription.id,
1029 &UpdateBillingSubscriptionParams {
1030 billing_customer_id: ActiveValue::set(billing_customer.id),
1031 kind: ActiveValue::set(subscription_kind),
1032 stripe_subscription_id: ActiveValue::set(subscription.id.to_string()),
1033 stripe_subscription_status: ActiveValue::set(subscription.status.into()),
1034 stripe_cancel_at: ActiveValue::set(
1035 subscription
1036 .cancel_at
1037 .and_then(|cancel_at| DateTime::from_timestamp(cancel_at, 0))
1038 .map(|time| time.naive_utc()),
1039 ),
1040 stripe_cancellation_reason: ActiveValue::set(
1041 subscription
1042 .cancellation_details
1043 .and_then(|details| details.reason)
1044 .map(|reason| reason.into()),
1045 ),
1046 stripe_current_period_start: ActiveValue::set(Some(
1047 subscription.current_period_start,
1048 )),
1049 stripe_current_period_end: ActiveValue::set(Some(
1050 subscription.current_period_end,
1051 )),
1052 },
1053 )
1054 .await?;
1055 } else {
1056 // If the user already has an active billing subscription, ignore the
1057 // event and return an `Ok` to signal that it was processed
1058 // successfully.
1059 //
1060 // There is the possibility that this could cause us to not create a
1061 // subscription in the following scenario:
1062 //
1063 // 1. User has an active subscription A
1064 // 2. User cancels subscription A
1065 // 3. User creates a new subscription B
1066 // 4. We process the new subscription B before the cancellation of subscription A
1067 // 5. User ends up with no subscriptions
1068 //
1069 // In theory this situation shouldn't arise as we try to process the events in the order they occur.
1070 if app
1071 .db
1072 .has_active_billing_subscription(billing_customer.user_id)
1073 .await?
1074 {
1075 log::info!(
1076 "user {user_id} already has an active subscription, skipping creation of subscription {subscription_id}",
1077 user_id = billing_customer.user_id,
1078 subscription_id = subscription.id
1079 );
1080 return Ok(());
1081 }
1082
1083 app.db
1084 .create_billing_subscription(&CreateBillingSubscriptionParams {
1085 billing_customer_id: billing_customer.id,
1086 kind: subscription_kind,
1087 stripe_subscription_id: subscription.id.to_string(),
1088 stripe_subscription_status: subscription.status.into(),
1089 stripe_cancellation_reason: subscription
1090 .cancellation_details
1091 .and_then(|details| details.reason)
1092 .map(|reason| reason.into()),
1093 stripe_current_period_start: Some(subscription.current_period_start),
1094 stripe_current_period_end: Some(subscription.current_period_end),
1095 })
1096 .await?;
1097 }
1098
1099 // When the user's subscription changes, we want to refresh their LLM tokens
1100 // to either grant/revoke access.
1101 rpc_server
1102 .refresh_llm_tokens_for_user(billing_customer.user_id)
1103 .await;
1104
1105 Ok(())
1106}
1107
1108#[derive(Debug, Deserialize)]
1109struct GetMonthlySpendParams {
1110 github_user_id: i32,
1111}
1112
1113#[derive(Debug, Serialize)]
1114struct GetMonthlySpendResponse {
1115 monthly_free_tier_spend_in_cents: u32,
1116 monthly_free_tier_allowance_in_cents: u32,
1117 monthly_spend_in_cents: u32,
1118}
1119
1120async fn get_monthly_spend(
1121 Extension(app): Extension<Arc<AppState>>,
1122 Query(params): Query<GetMonthlySpendParams>,
1123) -> Result<Json<GetMonthlySpendResponse>> {
1124 let user = app
1125 .db
1126 .get_user_by_github_user_id(params.github_user_id)
1127 .await?
1128 .ok_or_else(|| anyhow!("user not found"))?;
1129
1130 let Some(llm_db) = app.llm_db.clone() else {
1131 return Err(Error::http(
1132 StatusCode::NOT_IMPLEMENTED,
1133 "LLM database not available".into(),
1134 ));
1135 };
1136
1137 let free_tier = user
1138 .custom_llm_monthly_allowance_in_cents
1139 .map(|allowance| Cents(allowance as u32))
1140 .unwrap_or(FREE_TIER_MONTHLY_SPENDING_LIMIT);
1141
1142 let spending_for_month = llm_db
1143 .get_user_spending_for_month(user.id, Utc::now())
1144 .await?;
1145
1146 let free_tier_spend = Cents::min(spending_for_month, free_tier);
1147 let monthly_spend = spending_for_month.saturating_sub(free_tier);
1148
1149 Ok(Json(GetMonthlySpendResponse {
1150 monthly_free_tier_spend_in_cents: free_tier_spend.0,
1151 monthly_free_tier_allowance_in_cents: free_tier.0,
1152 monthly_spend_in_cents: monthly_spend.0,
1153 }))
1154}
1155
1156#[derive(Debug, Deserialize)]
1157struct GetCurrentUsageParams {
1158 github_user_id: i32,
1159}
1160
1161#[derive(Debug, Serialize)]
1162struct UsageCounts {
1163 pub used: i32,
1164 pub limit: Option<i32>,
1165 pub remaining: Option<i32>,
1166}
1167
1168#[derive(Debug, Serialize)]
1169struct ModelRequestUsage {
1170 pub model: String,
1171 pub mode: CompletionMode,
1172 pub requests: i32,
1173}
1174
1175#[derive(Debug, Serialize)]
1176struct CurrentUsage {
1177 pub model_requests: UsageCounts,
1178 pub model_request_usage: Vec<ModelRequestUsage>,
1179 pub edit_predictions: UsageCounts,
1180}
1181
1182#[derive(Debug, Default, Serialize)]
1183struct GetCurrentUsageResponse {
1184 pub plan: String,
1185 pub current_usage: Option<CurrentUsage>,
1186}
1187
1188async fn get_current_usage(
1189 Extension(app): Extension<Arc<AppState>>,
1190 Query(params): Query<GetCurrentUsageParams>,
1191) -> Result<Json<GetCurrentUsageResponse>> {
1192 let user = app
1193 .db
1194 .get_user_by_github_user_id(params.github_user_id)
1195 .await?
1196 .ok_or_else(|| anyhow!("user not found"))?;
1197
1198 let feature_flags = app.db.get_user_flags(user.id).await?;
1199 let has_extended_trial = feature_flags
1200 .iter()
1201 .any(|flag| flag == AGENT_EXTENDED_TRIAL_FEATURE_FLAG);
1202
1203 let Some(llm_db) = app.llm_db.clone() else {
1204 return Err(Error::http(
1205 StatusCode::NOT_IMPLEMENTED,
1206 "LLM database not available".into(),
1207 ));
1208 };
1209
1210 let Some(subscription) = app.db.get_active_billing_subscription(user.id).await? else {
1211 return Ok(Json(GetCurrentUsageResponse::default()));
1212 };
1213
1214 let subscription_period = maybe!({
1215 let period_start_at = subscription.current_period_start_at()?;
1216 let period_end_at = subscription.current_period_end_at()?;
1217
1218 Some((period_start_at, period_end_at))
1219 });
1220
1221 let Some((period_start_at, period_end_at)) = subscription_period else {
1222 return Ok(Json(GetCurrentUsageResponse::default()));
1223 };
1224
1225 let usage = llm_db
1226 .get_subscription_usage_for_period(user.id, period_start_at, period_end_at)
1227 .await?;
1228
1229 let plan = usage
1230 .as_ref()
1231 .map(|usage| usage.plan.into())
1232 .unwrap_or_else(|| {
1233 subscription
1234 .kind
1235 .map(Into::into)
1236 .unwrap_or(zed_llm_client::Plan::Free)
1237 });
1238
1239 let model_requests_limit = match plan.model_requests_limit() {
1240 zed_llm_client::UsageLimit::Limited(limit) => {
1241 let limit = if plan == zed_llm_client::Plan::ZedProTrial && has_extended_trial {
1242 1_000
1243 } else {
1244 limit
1245 };
1246
1247 Some(limit)
1248 }
1249 zed_llm_client::UsageLimit::Unlimited => None,
1250 };
1251
1252 let edit_predictions_limit = match plan.edit_predictions_limit() {
1253 zed_llm_client::UsageLimit::Limited(limit) => Some(limit),
1254 zed_llm_client::UsageLimit::Unlimited => None,
1255 };
1256
1257 let Some(usage) = usage else {
1258 return Ok(Json(GetCurrentUsageResponse {
1259 plan: plan.as_str().to_string(),
1260 current_usage: Some(CurrentUsage {
1261 model_requests: UsageCounts {
1262 used: 0,
1263 limit: model_requests_limit,
1264 remaining: model_requests_limit,
1265 },
1266 model_request_usage: Vec::new(),
1267 edit_predictions: UsageCounts {
1268 used: 0,
1269 limit: edit_predictions_limit,
1270 remaining: edit_predictions_limit,
1271 },
1272 }),
1273 }));
1274 };
1275
1276 let subscription_usage_meters = llm_db
1277 .get_current_subscription_usage_meters_for_user(user.id, Utc::now())
1278 .await?;
1279
1280 let model_request_usage = subscription_usage_meters
1281 .into_iter()
1282 .filter_map(|(usage_meter, _usage)| {
1283 let model = llm_db.model_by_id(usage_meter.model_id).ok()?;
1284
1285 Some(ModelRequestUsage {
1286 model: model.name.clone(),
1287 mode: usage_meter.mode,
1288 requests: usage_meter.requests,
1289 })
1290 })
1291 .collect::<Vec<_>>();
1292
1293 Ok(Json(GetCurrentUsageResponse {
1294 plan: plan.as_str().to_string(),
1295 current_usage: Some(CurrentUsage {
1296 model_requests: UsageCounts {
1297 used: usage.model_requests,
1298 limit: model_requests_limit,
1299 remaining: model_requests_limit.map(|limit| (limit - usage.model_requests).max(0)),
1300 },
1301 model_request_usage,
1302 edit_predictions: UsageCounts {
1303 used: usage.edit_predictions,
1304 limit: edit_predictions_limit,
1305 remaining: edit_predictions_limit
1306 .map(|limit| (limit - usage.edit_predictions).max(0)),
1307 },
1308 }),
1309 }))
1310}
1311
1312impl From<SubscriptionStatus> for StripeSubscriptionStatus {
1313 fn from(value: SubscriptionStatus) -> Self {
1314 match value {
1315 SubscriptionStatus::Incomplete => Self::Incomplete,
1316 SubscriptionStatus::IncompleteExpired => Self::IncompleteExpired,
1317 SubscriptionStatus::Trialing => Self::Trialing,
1318 SubscriptionStatus::Active => Self::Active,
1319 SubscriptionStatus::PastDue => Self::PastDue,
1320 SubscriptionStatus::Canceled => Self::Canceled,
1321 SubscriptionStatus::Unpaid => Self::Unpaid,
1322 SubscriptionStatus::Paused => Self::Paused,
1323 }
1324 }
1325}
1326
1327impl From<CancellationDetailsReason> for StripeCancellationReason {
1328 fn from(value: CancellationDetailsReason) -> Self {
1329 match value {
1330 CancellationDetailsReason::CancellationRequested => Self::CancellationRequested,
1331 CancellationDetailsReason::PaymentDisputed => Self::PaymentDisputed,
1332 CancellationDetailsReason::PaymentFailed => Self::PaymentFailed,
1333 }
1334 }
1335}
1336
1337/// Finds or creates a billing customer using the provided customer.
1338async fn find_or_create_billing_customer(
1339 app: &Arc<AppState>,
1340 stripe_client: &stripe::Client,
1341 customer_or_id: Expandable<Customer>,
1342) -> anyhow::Result<Option<billing_customer::Model>> {
1343 let customer_id = match &customer_or_id {
1344 Expandable::Id(id) => id,
1345 Expandable::Object(customer) => customer.id.as_ref(),
1346 };
1347
1348 // If we already have a billing customer record associated with the Stripe customer,
1349 // there's nothing more we need to do.
1350 if let Some(billing_customer) = app
1351 .db
1352 .get_billing_customer_by_stripe_customer_id(customer_id)
1353 .await?
1354 {
1355 return Ok(Some(billing_customer));
1356 }
1357
1358 // If all we have is a customer ID, resolve it to a full customer record by
1359 // hitting the Stripe API.
1360 let customer = match customer_or_id {
1361 Expandable::Id(id) => Customer::retrieve(stripe_client, &id, &[]).await?,
1362 Expandable::Object(customer) => *customer,
1363 };
1364
1365 let Some(email) = customer.email else {
1366 return Ok(None);
1367 };
1368
1369 let Some(user) = app.db.get_user_by_email(&email).await? else {
1370 return Ok(None);
1371 };
1372
1373 let billing_customer = app
1374 .db
1375 .create_billing_customer(&CreateBillingCustomerParams {
1376 user_id: user.id,
1377 stripe_customer_id: customer.id.to_string(),
1378 })
1379 .await?;
1380
1381 Ok(Some(billing_customer))
1382}
1383
1384const SYNC_LLM_TOKEN_USAGE_WITH_STRIPE_INTERVAL: Duration = Duration::from_secs(60);
1385
1386pub fn sync_llm_token_usage_with_stripe_periodically(app: Arc<AppState>) {
1387 let Some(stripe_billing) = app.stripe_billing.clone() else {
1388 log::warn!("failed to retrieve Stripe billing object");
1389 return;
1390 };
1391 let Some(llm_db) = app.llm_db.clone() else {
1392 log::warn!("failed to retrieve LLM database");
1393 return;
1394 };
1395
1396 let executor = app.executor.clone();
1397 executor.spawn_detached({
1398 let executor = executor.clone();
1399 async move {
1400 loop {
1401 sync_token_usage_with_stripe(&app, &llm_db, &stripe_billing)
1402 .await
1403 .context("failed to sync LLM usage to Stripe")
1404 .trace_err();
1405 executor
1406 .sleep(SYNC_LLM_TOKEN_USAGE_WITH_STRIPE_INTERVAL)
1407 .await;
1408 }
1409 }
1410 });
1411}
1412
1413async fn sync_token_usage_with_stripe(
1414 app: &Arc<AppState>,
1415 llm_db: &Arc<LlmDatabase>,
1416 stripe_billing: &Arc<StripeBilling>,
1417) -> anyhow::Result<()> {
1418 let events = llm_db.get_billing_events().await?;
1419 let user_ids = events
1420 .iter()
1421 .map(|(event, _)| event.user_id)
1422 .collect::<HashSet<UserId>>();
1423 let stripe_subscriptions = app.db.get_active_billing_subscriptions(user_ids).await?;
1424
1425 for (event, model) in events {
1426 let Some((stripe_db_customer, stripe_db_subscription)) =
1427 stripe_subscriptions.get(&event.user_id)
1428 else {
1429 tracing::warn!(
1430 user_id = event.user_id.0,
1431 "Registered billing event for user who is not a Stripe customer. Billing events should only be created for users who are Stripe customers, so this is a mistake on our side."
1432 );
1433 continue;
1434 };
1435 let stripe_subscription_id: stripe::SubscriptionId = stripe_db_subscription
1436 .stripe_subscription_id
1437 .parse()
1438 .context("failed to parse stripe subscription id from db")?;
1439 let stripe_customer_id: stripe::CustomerId = stripe_db_customer
1440 .stripe_customer_id
1441 .parse()
1442 .context("failed to parse stripe customer id from db")?;
1443
1444 let stripe_model = stripe_billing
1445 .register_model_for_token_based_usage(&model)
1446 .await?;
1447 stripe_billing
1448 .subscribe_to_model(&stripe_subscription_id, &stripe_model)
1449 .await?;
1450 stripe_billing
1451 .bill_model_token_usage(&stripe_customer_id, &stripe_model, &event)
1452 .await?;
1453 llm_db.consume_billing_event(event.id).await?;
1454 }
1455
1456 Ok(())
1457}
1458
1459const SYNC_LLM_REQUEST_USAGE_WITH_STRIPE_INTERVAL: Duration = Duration::from_secs(60);
1460
1461pub fn sync_llm_request_usage_with_stripe_periodically(app: Arc<AppState>) {
1462 let Some(stripe_billing) = app.stripe_billing.clone() else {
1463 log::warn!("failed to retrieve Stripe billing object");
1464 return;
1465 };
1466 let Some(llm_db) = app.llm_db.clone() else {
1467 log::warn!("failed to retrieve LLM database");
1468 return;
1469 };
1470
1471 let executor = app.executor.clone();
1472 executor.spawn_detached({
1473 let executor = executor.clone();
1474 async move {
1475 loop {
1476 sync_model_request_usage_with_stripe(&app, &llm_db, &stripe_billing)
1477 .await
1478 .context("failed to sync LLM request usage to Stripe")
1479 .trace_err();
1480 executor
1481 .sleep(SYNC_LLM_REQUEST_USAGE_WITH_STRIPE_INTERVAL)
1482 .await;
1483 }
1484 }
1485 });
1486}
1487
1488async fn sync_model_request_usage_with_stripe(
1489 app: &Arc<AppState>,
1490 llm_db: &Arc<LlmDatabase>,
1491 stripe_billing: &Arc<StripeBilling>,
1492) -> anyhow::Result<()> {
1493 let staff_users = app.db.get_staff_users().await?;
1494 let staff_user_ids = staff_users
1495 .iter()
1496 .map(|user| user.id)
1497 .collect::<HashSet<UserId>>();
1498
1499 let usage_meters = llm_db
1500 .get_current_subscription_usage_meters(Utc::now())
1501 .await?;
1502 let usage_meters = usage_meters
1503 .into_iter()
1504 .filter(|(_, usage)| !staff_user_ids.contains(&usage.user_id))
1505 .collect::<Vec<_>>();
1506 let user_ids = usage_meters
1507 .iter()
1508 .map(|(_, usage)| usage.user_id)
1509 .collect::<HashSet<UserId>>();
1510 let billing_subscriptions = app
1511 .db
1512 .get_active_zed_pro_billing_subscriptions(user_ids)
1513 .await?;
1514
1515 let claude_3_5_sonnet = stripe_billing
1516 .find_price_by_lookup_key("claude-3-5-sonnet-requests")
1517 .await?;
1518 let claude_3_7_sonnet = stripe_billing
1519 .find_price_by_lookup_key("claude-3-7-sonnet-requests")
1520 .await?;
1521 let claude_3_7_sonnet_max = stripe_billing
1522 .find_price_by_lookup_key("claude-3-7-sonnet-requests-max")
1523 .await?;
1524
1525 for (usage_meter, usage) in usage_meters {
1526 maybe!(async {
1527 let Some((billing_customer, billing_subscription)) =
1528 billing_subscriptions.get(&usage.user_id)
1529 else {
1530 bail!(
1531 "Attempted to sync usage meter for user who is not a Stripe customer: {}",
1532 usage.user_id
1533 );
1534 };
1535
1536 let stripe_customer_id = billing_customer
1537 .stripe_customer_id
1538 .parse::<stripe::CustomerId>()
1539 .context("failed to parse Stripe customer ID from database")?;
1540 let stripe_subscription_id = billing_subscription
1541 .stripe_subscription_id
1542 .parse::<stripe::SubscriptionId>()
1543 .context("failed to parse Stripe subscription ID from database")?;
1544
1545 let model = llm_db.model_by_id(usage_meter.model_id)?;
1546
1547 let (price, meter_event_name) = match model.name.as_str() {
1548 "claude-3-5-sonnet" => (&claude_3_5_sonnet, "claude_3_5_sonnet/requests"),
1549 "claude-3-7-sonnet" => match usage_meter.mode {
1550 CompletionMode::Normal => (&claude_3_7_sonnet, "claude_3_7_sonnet/requests"),
1551 CompletionMode::Max => {
1552 (&claude_3_7_sonnet_max, "claude_3_7_sonnet/requests/max")
1553 }
1554 },
1555 model_name => {
1556 bail!("Attempted to sync usage meter for unsupported model: {model_name:?}")
1557 }
1558 };
1559
1560 stripe_billing
1561 .subscribe_to_price(&stripe_subscription_id, price)
1562 .await?;
1563 stripe_billing
1564 .bill_model_request_usage(
1565 &stripe_customer_id,
1566 meter_event_name,
1567 usage_meter.requests,
1568 )
1569 .await?;
1570
1571 Ok(())
1572 })
1573 .await
1574 .log_err();
1575 }
1576
1577 Ok(())
1578}