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 let llm_db = app
1027 .llm_db
1028 .clone()
1029 .ok_or_else(|| anyhow!("LLM DB not initialized"))?;
1030
1031 let new_period_start_at =
1032 chrono::DateTime::from_timestamp(subscription.current_period_start, 0)
1033 .ok_or_else(|| anyhow!("No subscription period start"))?;
1034 let new_period_end_at =
1035 chrono::DateTime::from_timestamp(subscription.current_period_end, 0)
1036 .ok_or_else(|| anyhow!("No subscription period end"))?;
1037
1038 llm_db
1039 .transfer_existing_subscription_usage(
1040 billing_customer.user_id,
1041 &existing_subscription,
1042 subscription_kind,
1043 subscription.status.into(),
1044 new_period_start_at,
1045 new_period_end_at,
1046 )
1047 .await?;
1048
1049 app.db
1050 .update_billing_subscription(
1051 existing_subscription.id,
1052 &UpdateBillingSubscriptionParams {
1053 billing_customer_id: ActiveValue::set(billing_customer.id),
1054 kind: ActiveValue::set(subscription_kind),
1055 stripe_subscription_id: ActiveValue::set(subscription.id.to_string()),
1056 stripe_subscription_status: ActiveValue::set(subscription.status.into()),
1057 stripe_cancel_at: ActiveValue::set(
1058 subscription
1059 .cancel_at
1060 .and_then(|cancel_at| DateTime::from_timestamp(cancel_at, 0))
1061 .map(|time| time.naive_utc()),
1062 ),
1063 stripe_cancellation_reason: ActiveValue::set(
1064 subscription
1065 .cancellation_details
1066 .and_then(|details| details.reason)
1067 .map(|reason| reason.into()),
1068 ),
1069 stripe_current_period_start: ActiveValue::set(Some(
1070 subscription.current_period_start,
1071 )),
1072 stripe_current_period_end: ActiveValue::set(Some(
1073 subscription.current_period_end,
1074 )),
1075 },
1076 )
1077 .await?;
1078 } else {
1079 // If the user already has an active billing subscription, ignore the
1080 // event and return an `Ok` to signal that it was processed
1081 // successfully.
1082 //
1083 // There is the possibility that this could cause us to not create a
1084 // subscription in the following scenario:
1085 //
1086 // 1. User has an active subscription A
1087 // 2. User cancels subscription A
1088 // 3. User creates a new subscription B
1089 // 4. We process the new subscription B before the cancellation of subscription A
1090 // 5. User ends up with no subscriptions
1091 //
1092 // In theory this situation shouldn't arise as we try to process the events in the order they occur.
1093 if app
1094 .db
1095 .has_active_billing_subscription(billing_customer.user_id)
1096 .await?
1097 {
1098 log::info!(
1099 "user {user_id} already has an active subscription, skipping creation of subscription {subscription_id}",
1100 user_id = billing_customer.user_id,
1101 subscription_id = subscription.id
1102 );
1103 return Ok(());
1104 }
1105
1106 app.db
1107 .create_billing_subscription(&CreateBillingSubscriptionParams {
1108 billing_customer_id: billing_customer.id,
1109 kind: subscription_kind,
1110 stripe_subscription_id: subscription.id.to_string(),
1111 stripe_subscription_status: subscription.status.into(),
1112 stripe_cancellation_reason: subscription
1113 .cancellation_details
1114 .and_then(|details| details.reason)
1115 .map(|reason| reason.into()),
1116 stripe_current_period_start: Some(subscription.current_period_start),
1117 stripe_current_period_end: Some(subscription.current_period_end),
1118 })
1119 .await?;
1120 }
1121
1122 // When the user's subscription changes, we want to refresh their LLM tokens
1123 // to either grant/revoke access.
1124 rpc_server
1125 .refresh_llm_tokens_for_user(billing_customer.user_id)
1126 .await;
1127
1128 Ok(())
1129}
1130
1131#[derive(Debug, Deserialize)]
1132struct GetMonthlySpendParams {
1133 github_user_id: i32,
1134}
1135
1136#[derive(Debug, Serialize)]
1137struct GetMonthlySpendResponse {
1138 monthly_free_tier_spend_in_cents: u32,
1139 monthly_free_tier_allowance_in_cents: u32,
1140 monthly_spend_in_cents: u32,
1141}
1142
1143async fn get_monthly_spend(
1144 Extension(app): Extension<Arc<AppState>>,
1145 Query(params): Query<GetMonthlySpendParams>,
1146) -> Result<Json<GetMonthlySpendResponse>> {
1147 let user = app
1148 .db
1149 .get_user_by_github_user_id(params.github_user_id)
1150 .await?
1151 .ok_or_else(|| anyhow!("user not found"))?;
1152
1153 let Some(llm_db) = app.llm_db.clone() else {
1154 return Err(Error::http(
1155 StatusCode::NOT_IMPLEMENTED,
1156 "LLM database not available".into(),
1157 ));
1158 };
1159
1160 let free_tier = user
1161 .custom_llm_monthly_allowance_in_cents
1162 .map(|allowance| Cents(allowance as u32))
1163 .unwrap_or(FREE_TIER_MONTHLY_SPENDING_LIMIT);
1164
1165 let spending_for_month = llm_db
1166 .get_user_spending_for_month(user.id, Utc::now())
1167 .await?;
1168
1169 let free_tier_spend = Cents::min(spending_for_month, free_tier);
1170 let monthly_spend = spending_for_month.saturating_sub(free_tier);
1171
1172 Ok(Json(GetMonthlySpendResponse {
1173 monthly_free_tier_spend_in_cents: free_tier_spend.0,
1174 monthly_free_tier_allowance_in_cents: free_tier.0,
1175 monthly_spend_in_cents: monthly_spend.0,
1176 }))
1177}
1178
1179#[derive(Debug, Deserialize)]
1180struct GetCurrentUsageParams {
1181 github_user_id: i32,
1182}
1183
1184#[derive(Debug, Serialize)]
1185struct UsageCounts {
1186 pub used: i32,
1187 pub limit: Option<i32>,
1188 pub remaining: Option<i32>,
1189}
1190
1191#[derive(Debug, Serialize)]
1192struct ModelRequestUsage {
1193 pub model: String,
1194 pub mode: CompletionMode,
1195 pub requests: i32,
1196}
1197
1198#[derive(Debug, Serialize)]
1199struct CurrentUsage {
1200 pub model_requests: UsageCounts,
1201 pub model_request_usage: Vec<ModelRequestUsage>,
1202 pub edit_predictions: UsageCounts,
1203}
1204
1205#[derive(Debug, Default, Serialize)]
1206struct GetCurrentUsageResponse {
1207 pub plan: String,
1208 pub current_usage: Option<CurrentUsage>,
1209}
1210
1211async fn get_current_usage(
1212 Extension(app): Extension<Arc<AppState>>,
1213 Query(params): Query<GetCurrentUsageParams>,
1214) -> Result<Json<GetCurrentUsageResponse>> {
1215 let user = app
1216 .db
1217 .get_user_by_github_user_id(params.github_user_id)
1218 .await?
1219 .ok_or_else(|| anyhow!("user not found"))?;
1220
1221 let feature_flags = app.db.get_user_flags(user.id).await?;
1222 let has_extended_trial = feature_flags
1223 .iter()
1224 .any(|flag| flag == AGENT_EXTENDED_TRIAL_FEATURE_FLAG);
1225
1226 let Some(llm_db) = app.llm_db.clone() else {
1227 return Err(Error::http(
1228 StatusCode::NOT_IMPLEMENTED,
1229 "LLM database not available".into(),
1230 ));
1231 };
1232
1233 let Some(subscription) = app.db.get_active_billing_subscription(user.id).await? else {
1234 return Ok(Json(GetCurrentUsageResponse::default()));
1235 };
1236
1237 let subscription_period = maybe!({
1238 let period_start_at = subscription.current_period_start_at()?;
1239 let period_end_at = subscription.current_period_end_at()?;
1240
1241 Some((period_start_at, period_end_at))
1242 });
1243
1244 let Some((period_start_at, period_end_at)) = subscription_period else {
1245 return Ok(Json(GetCurrentUsageResponse::default()));
1246 };
1247
1248 let usage = llm_db
1249 .get_subscription_usage_for_period(user.id, period_start_at, period_end_at)
1250 .await?;
1251
1252 let plan = usage
1253 .as_ref()
1254 .map(|usage| usage.plan.into())
1255 .unwrap_or_else(|| {
1256 subscription
1257 .kind
1258 .map(Into::into)
1259 .unwrap_or(zed_llm_client::Plan::Free)
1260 });
1261
1262 let model_requests_limit = match plan.model_requests_limit() {
1263 zed_llm_client::UsageLimit::Limited(limit) => {
1264 let limit = if plan == zed_llm_client::Plan::ZedProTrial && has_extended_trial {
1265 1_000
1266 } else {
1267 limit
1268 };
1269
1270 Some(limit)
1271 }
1272 zed_llm_client::UsageLimit::Unlimited => None,
1273 };
1274
1275 let edit_predictions_limit = match plan.edit_predictions_limit() {
1276 zed_llm_client::UsageLimit::Limited(limit) => Some(limit),
1277 zed_llm_client::UsageLimit::Unlimited => None,
1278 };
1279
1280 let Some(usage) = usage else {
1281 return Ok(Json(GetCurrentUsageResponse {
1282 plan: plan.as_str().to_string(),
1283 current_usage: Some(CurrentUsage {
1284 model_requests: UsageCounts {
1285 used: 0,
1286 limit: model_requests_limit,
1287 remaining: model_requests_limit,
1288 },
1289 model_request_usage: Vec::new(),
1290 edit_predictions: UsageCounts {
1291 used: 0,
1292 limit: edit_predictions_limit,
1293 remaining: edit_predictions_limit,
1294 },
1295 }),
1296 }));
1297 };
1298
1299 let subscription_usage_meters = llm_db
1300 .get_current_subscription_usage_meters_for_user(user.id, Utc::now())
1301 .await?;
1302
1303 let model_request_usage = subscription_usage_meters
1304 .into_iter()
1305 .filter_map(|(usage_meter, _usage)| {
1306 let model = llm_db.model_by_id(usage_meter.model_id).ok()?;
1307
1308 Some(ModelRequestUsage {
1309 model: model.name.clone(),
1310 mode: usage_meter.mode,
1311 requests: usage_meter.requests,
1312 })
1313 })
1314 .collect::<Vec<_>>();
1315
1316 Ok(Json(GetCurrentUsageResponse {
1317 plan: plan.as_str().to_string(),
1318 current_usage: Some(CurrentUsage {
1319 model_requests: UsageCounts {
1320 used: usage.model_requests,
1321 limit: model_requests_limit,
1322 remaining: model_requests_limit.map(|limit| (limit - usage.model_requests).max(0)),
1323 },
1324 model_request_usage,
1325 edit_predictions: UsageCounts {
1326 used: usage.edit_predictions,
1327 limit: edit_predictions_limit,
1328 remaining: edit_predictions_limit
1329 .map(|limit| (limit - usage.edit_predictions).max(0)),
1330 },
1331 }),
1332 }))
1333}
1334
1335impl From<SubscriptionStatus> for StripeSubscriptionStatus {
1336 fn from(value: SubscriptionStatus) -> Self {
1337 match value {
1338 SubscriptionStatus::Incomplete => Self::Incomplete,
1339 SubscriptionStatus::IncompleteExpired => Self::IncompleteExpired,
1340 SubscriptionStatus::Trialing => Self::Trialing,
1341 SubscriptionStatus::Active => Self::Active,
1342 SubscriptionStatus::PastDue => Self::PastDue,
1343 SubscriptionStatus::Canceled => Self::Canceled,
1344 SubscriptionStatus::Unpaid => Self::Unpaid,
1345 SubscriptionStatus::Paused => Self::Paused,
1346 }
1347 }
1348}
1349
1350impl From<CancellationDetailsReason> for StripeCancellationReason {
1351 fn from(value: CancellationDetailsReason) -> Self {
1352 match value {
1353 CancellationDetailsReason::CancellationRequested => Self::CancellationRequested,
1354 CancellationDetailsReason::PaymentDisputed => Self::PaymentDisputed,
1355 CancellationDetailsReason::PaymentFailed => Self::PaymentFailed,
1356 }
1357 }
1358}
1359
1360/// Finds or creates a billing customer using the provided customer.
1361async fn find_or_create_billing_customer(
1362 app: &Arc<AppState>,
1363 stripe_client: &stripe::Client,
1364 customer_or_id: Expandable<Customer>,
1365) -> anyhow::Result<Option<billing_customer::Model>> {
1366 let customer_id = match &customer_or_id {
1367 Expandable::Id(id) => id,
1368 Expandable::Object(customer) => customer.id.as_ref(),
1369 };
1370
1371 // If we already have a billing customer record associated with the Stripe customer,
1372 // there's nothing more we need to do.
1373 if let Some(billing_customer) = app
1374 .db
1375 .get_billing_customer_by_stripe_customer_id(customer_id)
1376 .await?
1377 {
1378 return Ok(Some(billing_customer));
1379 }
1380
1381 // If all we have is a customer ID, resolve it to a full customer record by
1382 // hitting the Stripe API.
1383 let customer = match customer_or_id {
1384 Expandable::Id(id) => Customer::retrieve(stripe_client, &id, &[]).await?,
1385 Expandable::Object(customer) => *customer,
1386 };
1387
1388 let Some(email) = customer.email else {
1389 return Ok(None);
1390 };
1391
1392 let Some(user) = app.db.get_user_by_email(&email).await? else {
1393 return Ok(None);
1394 };
1395
1396 let billing_customer = app
1397 .db
1398 .create_billing_customer(&CreateBillingCustomerParams {
1399 user_id: user.id,
1400 stripe_customer_id: customer.id.to_string(),
1401 })
1402 .await?;
1403
1404 Ok(Some(billing_customer))
1405}
1406
1407const SYNC_LLM_TOKEN_USAGE_WITH_STRIPE_INTERVAL: Duration = Duration::from_secs(60);
1408
1409pub fn sync_llm_token_usage_with_stripe_periodically(app: Arc<AppState>) {
1410 let Some(stripe_billing) = app.stripe_billing.clone() else {
1411 log::warn!("failed to retrieve Stripe billing object");
1412 return;
1413 };
1414 let Some(llm_db) = app.llm_db.clone() else {
1415 log::warn!("failed to retrieve LLM database");
1416 return;
1417 };
1418
1419 let executor = app.executor.clone();
1420 executor.spawn_detached({
1421 let executor = executor.clone();
1422 async move {
1423 loop {
1424 sync_token_usage_with_stripe(&app, &llm_db, &stripe_billing)
1425 .await
1426 .context("failed to sync LLM usage to Stripe")
1427 .trace_err();
1428 executor
1429 .sleep(SYNC_LLM_TOKEN_USAGE_WITH_STRIPE_INTERVAL)
1430 .await;
1431 }
1432 }
1433 });
1434}
1435
1436async fn sync_token_usage_with_stripe(
1437 app: &Arc<AppState>,
1438 llm_db: &Arc<LlmDatabase>,
1439 stripe_billing: &Arc<StripeBilling>,
1440) -> anyhow::Result<()> {
1441 let events = llm_db.get_billing_events().await?;
1442 let user_ids = events
1443 .iter()
1444 .map(|(event, _)| event.user_id)
1445 .collect::<HashSet<UserId>>();
1446 let stripe_subscriptions = app.db.get_active_billing_subscriptions(user_ids).await?;
1447
1448 for (event, model) in events {
1449 let Some((stripe_db_customer, stripe_db_subscription)) =
1450 stripe_subscriptions.get(&event.user_id)
1451 else {
1452 tracing::warn!(
1453 user_id = event.user_id.0,
1454 "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."
1455 );
1456 continue;
1457 };
1458 let stripe_subscription_id: stripe::SubscriptionId = stripe_db_subscription
1459 .stripe_subscription_id
1460 .parse()
1461 .context("failed to parse stripe subscription id from db")?;
1462 let stripe_customer_id: stripe::CustomerId = stripe_db_customer
1463 .stripe_customer_id
1464 .parse()
1465 .context("failed to parse stripe customer id from db")?;
1466
1467 let stripe_model = stripe_billing
1468 .register_model_for_token_based_usage(&model)
1469 .await?;
1470 stripe_billing
1471 .subscribe_to_model(&stripe_subscription_id, &stripe_model)
1472 .await?;
1473 stripe_billing
1474 .bill_model_token_usage(&stripe_customer_id, &stripe_model, &event)
1475 .await?;
1476 llm_db.consume_billing_event(event.id).await?;
1477 }
1478
1479 Ok(())
1480}
1481
1482const SYNC_LLM_REQUEST_USAGE_WITH_STRIPE_INTERVAL: Duration = Duration::from_secs(60);
1483
1484pub fn sync_llm_request_usage_with_stripe_periodically(app: Arc<AppState>) {
1485 let Some(stripe_billing) = app.stripe_billing.clone() else {
1486 log::warn!("failed to retrieve Stripe billing object");
1487 return;
1488 };
1489 let Some(llm_db) = app.llm_db.clone() else {
1490 log::warn!("failed to retrieve LLM database");
1491 return;
1492 };
1493
1494 let executor = app.executor.clone();
1495 executor.spawn_detached({
1496 let executor = executor.clone();
1497 async move {
1498 loop {
1499 sync_model_request_usage_with_stripe(&app, &llm_db, &stripe_billing)
1500 .await
1501 .context("failed to sync LLM request usage to Stripe")
1502 .trace_err();
1503 executor
1504 .sleep(SYNC_LLM_REQUEST_USAGE_WITH_STRIPE_INTERVAL)
1505 .await;
1506 }
1507 }
1508 });
1509}
1510
1511async fn sync_model_request_usage_with_stripe(
1512 app: &Arc<AppState>,
1513 llm_db: &Arc<LlmDatabase>,
1514 stripe_billing: &Arc<StripeBilling>,
1515) -> anyhow::Result<()> {
1516 let staff_users = app.db.get_staff_users().await?;
1517 let staff_user_ids = staff_users
1518 .iter()
1519 .map(|user| user.id)
1520 .collect::<HashSet<UserId>>();
1521
1522 let usage_meters = llm_db
1523 .get_current_subscription_usage_meters(Utc::now())
1524 .await?;
1525 let usage_meters = usage_meters
1526 .into_iter()
1527 .filter(|(_, usage)| !staff_user_ids.contains(&usage.user_id))
1528 .collect::<Vec<_>>();
1529 let user_ids = usage_meters
1530 .iter()
1531 .map(|(_, usage)| usage.user_id)
1532 .collect::<HashSet<UserId>>();
1533 let billing_subscriptions = app
1534 .db
1535 .get_active_zed_pro_billing_subscriptions(user_ids)
1536 .await?;
1537
1538 let claude_3_5_sonnet = stripe_billing
1539 .find_price_by_lookup_key("claude-3-5-sonnet-requests")
1540 .await?;
1541 let claude_3_7_sonnet = stripe_billing
1542 .find_price_by_lookup_key("claude-3-7-sonnet-requests")
1543 .await?;
1544 let claude_3_7_sonnet_max = stripe_billing
1545 .find_price_by_lookup_key("claude-3-7-sonnet-requests-max")
1546 .await?;
1547
1548 for (usage_meter, usage) in usage_meters {
1549 maybe!(async {
1550 let Some((billing_customer, billing_subscription)) =
1551 billing_subscriptions.get(&usage.user_id)
1552 else {
1553 bail!(
1554 "Attempted to sync usage meter for user who is not a Stripe customer: {}",
1555 usage.user_id
1556 );
1557 };
1558
1559 let stripe_customer_id = billing_customer
1560 .stripe_customer_id
1561 .parse::<stripe::CustomerId>()
1562 .context("failed to parse Stripe customer ID from database")?;
1563 let stripe_subscription_id = billing_subscription
1564 .stripe_subscription_id
1565 .parse::<stripe::SubscriptionId>()
1566 .context("failed to parse Stripe subscription ID from database")?;
1567
1568 let model = llm_db.model_by_id(usage_meter.model_id)?;
1569
1570 let (price, meter_event_name) = match model.name.as_str() {
1571 "claude-3-5-sonnet" => (&claude_3_5_sonnet, "claude_3_5_sonnet/requests"),
1572 "claude-3-7-sonnet" => match usage_meter.mode {
1573 CompletionMode::Normal => (&claude_3_7_sonnet, "claude_3_7_sonnet/requests"),
1574 CompletionMode::Max => {
1575 (&claude_3_7_sonnet_max, "claude_3_7_sonnet/requests/max")
1576 }
1577 },
1578 model_name => {
1579 bail!("Attempted to sync usage meter for unsupported model: {model_name:?}")
1580 }
1581 };
1582
1583 stripe_billing
1584 .subscribe_to_price(&stripe_subscription_id, price)
1585 .await?;
1586 stripe_billing
1587 .bill_model_request_usage(
1588 &stripe_customer_id,
1589 meter_event_name,
1590 usage_meter.requests,
1591 )
1592 .await?;
1593
1594 Ok(())
1595 })
1596 .await
1597 .log_err();
1598 }
1599
1600 Ok(())
1601}