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