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: 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 Some((_billing_customer, billing_subscription)) =
654 old_billing_subscriptions_by_user.get(&user.id)
655 else {
656 return Err(Error::http(
657 StatusCode::NOT_FOUND,
658 "No active billing subscriptions to migrate".into(),
659 ));
660 };
661
662 let stripe_subscription_id = billing_subscription
663 .stripe_subscription_id
664 .parse::<stripe::SubscriptionId>()
665 .context("failed to parse Stripe subscription ID from database")?;
666
667 Subscription::cancel(
668 &stripe_client,
669 &stripe_subscription_id,
670 stripe::CancelSubscription {
671 invoice_now: Some(true),
672 ..Default::default()
673 },
674 )
675 .await?;
676
677 let feature_flags = app.db.list_feature_flags().await?;
678
679 for feature_flag in ["new-billing", "assistant2"] {
680 let already_in_feature_flag = feature_flags.iter().any(|flag| flag.flag == feature_flag);
681 if already_in_feature_flag {
682 continue;
683 }
684
685 let feature_flag = 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: stripe_subscription_id.to_string(),
695 }))
696}
697
698/// The amount of time we wait in between each poll of Stripe events.
699///
700/// This value should strike a balance between:
701/// 1. Being short enough that we update quickly when something in Stripe changes
702/// 2. Being long enough that we don't eat into our rate limits.
703///
704/// As a point of reference, the Sequin folks say they have this at **500ms**:
705///
706/// > We poll the Stripe /events endpoint every 500ms per account
707/// >
708/// > — https://blog.sequinstream.com/events-not-webhooks/
709const POLL_EVENTS_INTERVAL: Duration = Duration::from_secs(5);
710
711/// The maximum number of events to return per page.
712///
713/// We set this to 100 (the max) so we have to make fewer requests to Stripe.
714///
715/// > Limit can range between 1 and 100, and the default is 10.
716const EVENTS_LIMIT_PER_PAGE: u64 = 100;
717
718/// The number of pages consisting entirely of already-processed events that we
719/// will see before we stop retrieving events.
720///
721/// This is used to prevent over-fetching the Stripe events API for events we've
722/// already seen and processed.
723const NUMBER_OF_ALREADY_PROCESSED_PAGES_BEFORE_WE_STOP: usize = 4;
724
725/// Polls the Stripe events API periodically to reconcile the records in our
726/// database with the data in Stripe.
727pub fn poll_stripe_events_periodically(app: Arc<AppState>, rpc_server: Arc<Server>) {
728 let Some(stripe_client) = app.stripe_client.clone() else {
729 log::warn!("failed to retrieve Stripe client");
730 return;
731 };
732
733 let executor = app.executor.clone();
734 executor.spawn_detached({
735 let executor = executor.clone();
736 async move {
737 loop {
738 poll_stripe_events(&app, &rpc_server, &stripe_client)
739 .await
740 .log_err();
741
742 executor.sleep(POLL_EVENTS_INTERVAL).await;
743 }
744 }
745 });
746}
747
748async fn poll_stripe_events(
749 app: &Arc<AppState>,
750 rpc_server: &Arc<Server>,
751 stripe_client: &stripe::Client,
752) -> anyhow::Result<()> {
753 fn event_type_to_string(event_type: EventType) -> String {
754 // Calling `to_string` on `stripe::EventType` members gives us a quoted string,
755 // so we need to unquote it.
756 event_type.to_string().trim_matches('"').to_string()
757 }
758
759 let event_types = [
760 EventType::CustomerCreated,
761 EventType::CustomerUpdated,
762 EventType::CustomerSubscriptionCreated,
763 EventType::CustomerSubscriptionUpdated,
764 EventType::CustomerSubscriptionPaused,
765 EventType::CustomerSubscriptionResumed,
766 EventType::CustomerSubscriptionDeleted,
767 ]
768 .into_iter()
769 .map(event_type_to_string)
770 .collect::<Vec<_>>();
771
772 let mut pages_of_already_processed_events = 0;
773 let mut unprocessed_events = Vec::new();
774
775 log::info!(
776 "Stripe events: starting retrieval for {}",
777 event_types.join(", ")
778 );
779 let mut params = ListEvents::new();
780 params.types = Some(event_types.clone());
781 params.limit = Some(EVENTS_LIMIT_PER_PAGE);
782
783 let mut event_pages = stripe::Event::list(&stripe_client, ¶ms)
784 .await?
785 .paginate(params);
786
787 loop {
788 let processed_event_ids = {
789 let event_ids = event_pages
790 .page
791 .data
792 .iter()
793 .map(|event| event.id.as_str())
794 .collect::<Vec<_>>();
795 app.db
796 .get_processed_stripe_events_by_event_ids(&event_ids)
797 .await?
798 .into_iter()
799 .map(|event| event.stripe_event_id)
800 .collect::<Vec<_>>()
801 };
802
803 let mut processed_events_in_page = 0;
804 let events_in_page = event_pages.page.data.len();
805 for event in &event_pages.page.data {
806 if processed_event_ids.contains(&event.id.to_string()) {
807 processed_events_in_page += 1;
808 log::debug!("Stripe events: already processed '{}', skipping", event.id);
809 } else {
810 unprocessed_events.push(event.clone());
811 }
812 }
813
814 if processed_events_in_page == events_in_page {
815 pages_of_already_processed_events += 1;
816 }
817
818 if event_pages.page.has_more {
819 if pages_of_already_processed_events >= NUMBER_OF_ALREADY_PROCESSED_PAGES_BEFORE_WE_STOP
820 {
821 log::info!(
822 "Stripe events: stopping, saw {pages_of_already_processed_events} pages of already-processed events"
823 );
824 break;
825 } else {
826 log::info!("Stripe events: retrieving next page");
827 event_pages = event_pages.next(&stripe_client).await?;
828 }
829 } else {
830 break;
831 }
832 }
833
834 log::info!("Stripe events: unprocessed {}", unprocessed_events.len());
835
836 // Sort all of the unprocessed events in ascending order, so we can handle them in the order they occurred.
837 unprocessed_events.sort_by(|a, b| a.created.cmp(&b.created).then_with(|| a.id.cmp(&b.id)));
838
839 for event in unprocessed_events {
840 let event_id = event.id.clone();
841 let processed_event_params = CreateProcessedStripeEventParams {
842 stripe_event_id: event.id.to_string(),
843 stripe_event_type: event_type_to_string(event.type_),
844 stripe_event_created_timestamp: event.created,
845 };
846
847 // If the event has happened too far in the past, we don't want to
848 // process it and risk overwriting other more-recent updates.
849 //
850 // 1 day was chosen arbitrarily. This could be made longer or shorter.
851 let one_day = Duration::from_secs(24 * 60 * 60);
852 let a_day_ago = Utc::now() - one_day;
853 if a_day_ago.timestamp() > event.created {
854 log::info!(
855 "Stripe events: event '{}' is more than {one_day:?} old, marking as processed",
856 event_id
857 );
858 app.db
859 .create_processed_stripe_event(&processed_event_params)
860 .await?;
861
862 return Ok(());
863 }
864
865 let process_result = match event.type_ {
866 EventType::CustomerCreated | EventType::CustomerUpdated => {
867 handle_customer_event(app, stripe_client, event).await
868 }
869 EventType::CustomerSubscriptionCreated
870 | EventType::CustomerSubscriptionUpdated
871 | EventType::CustomerSubscriptionPaused
872 | EventType::CustomerSubscriptionResumed
873 | EventType::CustomerSubscriptionDeleted => {
874 handle_customer_subscription_event(app, rpc_server, stripe_client, event).await
875 }
876 _ => Ok(()),
877 };
878
879 if let Some(()) = process_result
880 .with_context(|| format!("failed to process event {event_id} successfully"))
881 .log_err()
882 {
883 app.db
884 .create_processed_stripe_event(&processed_event_params)
885 .await?;
886 }
887 }
888
889 Ok(())
890}
891
892async fn handle_customer_event(
893 app: &Arc<AppState>,
894 _stripe_client: &stripe::Client,
895 event: stripe::Event,
896) -> anyhow::Result<()> {
897 let EventObject::Customer(customer) = event.data.object else {
898 bail!("unexpected event payload for {}", event.id);
899 };
900
901 log::info!("handling Stripe {} event: {}", event.type_, event.id);
902
903 let Some(email) = customer.email else {
904 log::info!("Stripe customer has no email: skipping");
905 return Ok(());
906 };
907
908 let Some(user) = app.db.get_user_by_email(&email).await? else {
909 log::info!("no user found for email: skipping");
910 return Ok(());
911 };
912
913 if let Some(existing_customer) = app
914 .db
915 .get_billing_customer_by_stripe_customer_id(&customer.id)
916 .await?
917 {
918 app.db
919 .update_billing_customer(
920 existing_customer.id,
921 &UpdateBillingCustomerParams {
922 // For now we just leave the information as-is, as it is not
923 // likely to change.
924 ..Default::default()
925 },
926 )
927 .await?;
928 } else {
929 app.db
930 .create_billing_customer(&CreateBillingCustomerParams {
931 user_id: user.id,
932 stripe_customer_id: customer.id.to_string(),
933 })
934 .await?;
935 }
936
937 Ok(())
938}
939
940async fn handle_customer_subscription_event(
941 app: &Arc<AppState>,
942 rpc_server: &Arc<Server>,
943 stripe_client: &stripe::Client,
944 event: stripe::Event,
945) -> anyhow::Result<()> {
946 let EventObject::Subscription(subscription) = event.data.object else {
947 bail!("unexpected event payload for {}", event.id);
948 };
949
950 log::info!("handling Stripe {} event: {}", event.type_, event.id);
951
952 let subscription_kind = maybe!(async {
953 let stripe_billing = app.stripe_billing.clone()?;
954
955 let zed_pro_price_id = stripe_billing.zed_pro_price_id().await.ok()?;
956 let zed_free_price_id = stripe_billing.zed_free_price_id().await.ok()?;
957
958 subscription.items.data.iter().find_map(|item| {
959 let price = item.price.as_ref()?;
960
961 if price.id == zed_pro_price_id {
962 Some(if subscription.status == SubscriptionStatus::Trialing {
963 SubscriptionKind::ZedProTrial
964 } else {
965 SubscriptionKind::ZedPro
966 })
967 } else if price.id == zed_free_price_id {
968 Some(SubscriptionKind::ZedFree)
969 } else {
970 None
971 }
972 })
973 })
974 .await;
975
976 let billing_customer =
977 find_or_create_billing_customer(app, stripe_client, subscription.customer)
978 .await?
979 .ok_or_else(|| anyhow!("billing customer not found"))?;
980
981 if let Some(SubscriptionKind::ZedProTrial) = subscription_kind {
982 if subscription.status == SubscriptionStatus::Trialing {
983 let current_period_start =
984 DateTime::from_timestamp(subscription.current_period_start, 0)
985 .ok_or_else(|| anyhow!("No trial subscription period start"))?;
986
987 app.db
988 .update_billing_customer(
989 billing_customer.id,
990 &UpdateBillingCustomerParams {
991 trial_started_at: ActiveValue::set(Some(current_period_start.naive_utc())),
992 ..Default::default()
993 },
994 )
995 .await?;
996 }
997 }
998
999 let was_canceled_due_to_payment_failure = subscription.status == SubscriptionStatus::Canceled
1000 && subscription
1001 .cancellation_details
1002 .as_ref()
1003 .and_then(|details| details.reason)
1004 .map_or(false, |reason| {
1005 reason == CancellationDetailsReason::PaymentFailed
1006 });
1007
1008 if was_canceled_due_to_payment_failure {
1009 app.db
1010 .update_billing_customer(
1011 billing_customer.id,
1012 &UpdateBillingCustomerParams {
1013 has_overdue_invoices: ActiveValue::set(true),
1014 ..Default::default()
1015 },
1016 )
1017 .await?;
1018 }
1019
1020 if let Some(existing_subscription) = app
1021 .db
1022 .get_billing_subscription_by_stripe_subscription_id(&subscription.id)
1023 .await?
1024 {
1025 let llm_db = app
1026 .llm_db
1027 .clone()
1028 .ok_or_else(|| anyhow!("LLM DB not initialized"))?;
1029
1030 let new_period_start_at =
1031 chrono::DateTime::from_timestamp(subscription.current_period_start, 0)
1032 .ok_or_else(|| anyhow!("No subscription period start"))?;
1033 let new_period_end_at =
1034 chrono::DateTime::from_timestamp(subscription.current_period_end, 0)
1035 .ok_or_else(|| anyhow!("No subscription period end"))?;
1036
1037 llm_db
1038 .transfer_existing_subscription_usage(
1039 billing_customer.user_id,
1040 &existing_subscription,
1041 subscription_kind,
1042 new_period_start_at,
1043 new_period_end_at,
1044 )
1045 .await?;
1046
1047 app.db
1048 .update_billing_subscription(
1049 existing_subscription.id,
1050 &UpdateBillingSubscriptionParams {
1051 billing_customer_id: ActiveValue::set(billing_customer.id),
1052 kind: ActiveValue::set(subscription_kind),
1053 stripe_subscription_id: ActiveValue::set(subscription.id.to_string()),
1054 stripe_subscription_status: ActiveValue::set(subscription.status.into()),
1055 stripe_cancel_at: ActiveValue::set(
1056 subscription
1057 .cancel_at
1058 .and_then(|cancel_at| DateTime::from_timestamp(cancel_at, 0))
1059 .map(|time| time.naive_utc()),
1060 ),
1061 stripe_cancellation_reason: ActiveValue::set(
1062 subscription
1063 .cancellation_details
1064 .and_then(|details| details.reason)
1065 .map(|reason| reason.into()),
1066 ),
1067 stripe_current_period_start: ActiveValue::set(Some(
1068 subscription.current_period_start,
1069 )),
1070 stripe_current_period_end: ActiveValue::set(Some(
1071 subscription.current_period_end,
1072 )),
1073 },
1074 )
1075 .await?;
1076 } else {
1077 // If the user already has an active billing subscription, ignore the
1078 // event and return an `Ok` to signal that it was processed
1079 // successfully.
1080 //
1081 // There is the possibility that this could cause us to not create a
1082 // subscription in the following scenario:
1083 //
1084 // 1. User has an active subscription A
1085 // 2. User cancels subscription A
1086 // 3. User creates a new subscription B
1087 // 4. We process the new subscription B before the cancellation of subscription A
1088 // 5. User ends up with no subscriptions
1089 //
1090 // In theory this situation shouldn't arise as we try to process the events in the order they occur.
1091 if app
1092 .db
1093 .has_active_billing_subscription(billing_customer.user_id)
1094 .await?
1095 {
1096 log::info!(
1097 "user {user_id} already has an active subscription, skipping creation of subscription {subscription_id}",
1098 user_id = billing_customer.user_id,
1099 subscription_id = subscription.id
1100 );
1101 return Ok(());
1102 }
1103
1104 app.db
1105 .create_billing_subscription(&CreateBillingSubscriptionParams {
1106 billing_customer_id: billing_customer.id,
1107 kind: subscription_kind,
1108 stripe_subscription_id: subscription.id.to_string(),
1109 stripe_subscription_status: subscription.status.into(),
1110 stripe_cancellation_reason: subscription
1111 .cancellation_details
1112 .and_then(|details| details.reason)
1113 .map(|reason| reason.into()),
1114 stripe_current_period_start: Some(subscription.current_period_start),
1115 stripe_current_period_end: Some(subscription.current_period_end),
1116 })
1117 .await?;
1118 }
1119
1120 // When the user's subscription changes, we want to refresh their LLM tokens
1121 // to either grant/revoke access.
1122 rpc_server
1123 .refresh_llm_tokens_for_user(billing_customer.user_id)
1124 .await;
1125
1126 Ok(())
1127}
1128
1129#[derive(Debug, Deserialize)]
1130struct GetMonthlySpendParams {
1131 github_user_id: i32,
1132}
1133
1134#[derive(Debug, Serialize)]
1135struct GetMonthlySpendResponse {
1136 monthly_free_tier_spend_in_cents: u32,
1137 monthly_free_tier_allowance_in_cents: u32,
1138 monthly_spend_in_cents: u32,
1139}
1140
1141async fn get_monthly_spend(
1142 Extension(app): Extension<Arc<AppState>>,
1143 Query(params): Query<GetMonthlySpendParams>,
1144) -> Result<Json<GetMonthlySpendResponse>> {
1145 let user = app
1146 .db
1147 .get_user_by_github_user_id(params.github_user_id)
1148 .await?
1149 .ok_or_else(|| anyhow!("user not found"))?;
1150
1151 let Some(llm_db) = app.llm_db.clone() else {
1152 return Err(Error::http(
1153 StatusCode::NOT_IMPLEMENTED,
1154 "LLM database not available".into(),
1155 ));
1156 };
1157
1158 let free_tier = user
1159 .custom_llm_monthly_allowance_in_cents
1160 .map(|allowance| Cents(allowance as u32))
1161 .unwrap_or(FREE_TIER_MONTHLY_SPENDING_LIMIT);
1162
1163 let spending_for_month = llm_db
1164 .get_user_spending_for_month(user.id, Utc::now())
1165 .await?;
1166
1167 let free_tier_spend = Cents::min(spending_for_month, free_tier);
1168 let monthly_spend = spending_for_month.saturating_sub(free_tier);
1169
1170 Ok(Json(GetMonthlySpendResponse {
1171 monthly_free_tier_spend_in_cents: free_tier_spend.0,
1172 monthly_free_tier_allowance_in_cents: free_tier.0,
1173 monthly_spend_in_cents: monthly_spend.0,
1174 }))
1175}
1176
1177#[derive(Debug, Deserialize)]
1178struct GetCurrentUsageParams {
1179 github_user_id: i32,
1180}
1181
1182#[derive(Debug, Serialize)]
1183struct UsageCounts {
1184 pub used: i32,
1185 pub limit: Option<i32>,
1186 pub remaining: Option<i32>,
1187}
1188
1189#[derive(Debug, Serialize)]
1190struct ModelRequestUsage {
1191 pub model: String,
1192 pub mode: CompletionMode,
1193 pub requests: i32,
1194}
1195
1196#[derive(Debug, Serialize)]
1197struct GetCurrentUsageResponse {
1198 pub model_requests: UsageCounts,
1199 pub model_request_usage: Vec<ModelRequestUsage>,
1200 pub edit_predictions: UsageCounts,
1201}
1202
1203async fn get_current_usage(
1204 Extension(app): Extension<Arc<AppState>>,
1205 Query(params): Query<GetCurrentUsageParams>,
1206) -> Result<Json<GetCurrentUsageResponse>> {
1207 let user = app
1208 .db
1209 .get_user_by_github_user_id(params.github_user_id)
1210 .await?
1211 .ok_or_else(|| anyhow!("user not found"))?;
1212
1213 let Some(llm_db) = app.llm_db.clone() else {
1214 return Err(Error::http(
1215 StatusCode::NOT_IMPLEMENTED,
1216 "LLM database not available".into(),
1217 ));
1218 };
1219
1220 let empty_usage = GetCurrentUsageResponse {
1221 model_requests: UsageCounts {
1222 used: 0,
1223 limit: Some(0),
1224 remaining: Some(0),
1225 },
1226 model_request_usage: Vec::new(),
1227 edit_predictions: UsageCounts {
1228 used: 0,
1229 limit: Some(0),
1230 remaining: Some(0),
1231 },
1232 };
1233
1234 let Some(subscription) = app.db.get_active_billing_subscription(user.id).await? else {
1235 return Ok(Json(empty_usage));
1236 };
1237
1238 let subscription_period = maybe!({
1239 let period_start_at = subscription.current_period_start_at()?;
1240 let period_end_at = subscription.current_period_end_at()?;
1241
1242 Some((period_start_at, period_end_at))
1243 });
1244
1245 let Some((period_start_at, period_end_at)) = subscription_period else {
1246 return Ok(Json(empty_usage));
1247 };
1248
1249 let usage = llm_db
1250 .get_subscription_usage_for_period(user.id, period_start_at, period_end_at)
1251 .await?;
1252 let Some(usage) = usage else {
1253 return Ok(Json(empty_usage));
1254 };
1255
1256 let plan = match usage.plan {
1257 SubscriptionKind::ZedPro => zed_llm_client::Plan::ZedPro,
1258 SubscriptionKind::ZedProTrial => zed_llm_client::Plan::ZedProTrial,
1259 SubscriptionKind::ZedFree => zed_llm_client::Plan::Free,
1260 };
1261
1262 let feature_flags = app.db.get_user_flags(user.id).await?;
1263 let has_extended_trial = feature_flags
1264 .iter()
1265 .any(|flag| flag == AGENT_EXTENDED_TRIAL_FEATURE_FLAG);
1266
1267 let model_requests_limit = match plan.model_requests_limit() {
1268 zed_llm_client::UsageLimit::Limited(limit) => {
1269 let limit = if plan == zed_llm_client::Plan::ZedProTrial && has_extended_trial {
1270 1_000
1271 } else {
1272 limit
1273 };
1274
1275 Some(limit)
1276 }
1277 zed_llm_client::UsageLimit::Unlimited => None,
1278 };
1279 let edit_prediction_limit = match plan.edit_predictions_limit() {
1280 zed_llm_client::UsageLimit::Limited(limit) => Some(limit),
1281 zed_llm_client::UsageLimit::Unlimited => None,
1282 };
1283
1284 let subscription_usage_meters = llm_db
1285 .get_current_subscription_usage_meters_for_user(user.id, Utc::now())
1286 .await?;
1287
1288 let model_request_usage = subscription_usage_meters
1289 .into_iter()
1290 .filter_map(|(usage_meter, _usage)| {
1291 let model = llm_db.model_by_id(usage_meter.model_id).ok()?;
1292
1293 Some(ModelRequestUsage {
1294 model: model.name.clone(),
1295 mode: usage_meter.mode,
1296 requests: usage_meter.requests,
1297 })
1298 })
1299 .collect::<Vec<_>>();
1300
1301 Ok(Json(GetCurrentUsageResponse {
1302 model_requests: UsageCounts {
1303 used: usage.model_requests,
1304 limit: model_requests_limit,
1305 remaining: model_requests_limit.map(|limit| (limit - usage.model_requests).max(0)),
1306 },
1307 model_request_usage,
1308 edit_predictions: UsageCounts {
1309 used: usage.edit_predictions,
1310 limit: edit_prediction_limit,
1311 remaining: edit_prediction_limit.map(|limit| (limit - usage.edit_predictions).max(0)),
1312 },
1313 }))
1314}
1315
1316impl From<SubscriptionStatus> for StripeSubscriptionStatus {
1317 fn from(value: SubscriptionStatus) -> Self {
1318 match value {
1319 SubscriptionStatus::Incomplete => Self::Incomplete,
1320 SubscriptionStatus::IncompleteExpired => Self::IncompleteExpired,
1321 SubscriptionStatus::Trialing => Self::Trialing,
1322 SubscriptionStatus::Active => Self::Active,
1323 SubscriptionStatus::PastDue => Self::PastDue,
1324 SubscriptionStatus::Canceled => Self::Canceled,
1325 SubscriptionStatus::Unpaid => Self::Unpaid,
1326 SubscriptionStatus::Paused => Self::Paused,
1327 }
1328 }
1329}
1330
1331impl From<CancellationDetailsReason> for StripeCancellationReason {
1332 fn from(value: CancellationDetailsReason) -> Self {
1333 match value {
1334 CancellationDetailsReason::CancellationRequested => Self::CancellationRequested,
1335 CancellationDetailsReason::PaymentDisputed => Self::PaymentDisputed,
1336 CancellationDetailsReason::PaymentFailed => Self::PaymentFailed,
1337 }
1338 }
1339}
1340
1341/// Finds or creates a billing customer using the provided customer.
1342async fn find_or_create_billing_customer(
1343 app: &Arc<AppState>,
1344 stripe_client: &stripe::Client,
1345 customer_or_id: Expandable<Customer>,
1346) -> anyhow::Result<Option<billing_customer::Model>> {
1347 let customer_id = match &customer_or_id {
1348 Expandable::Id(id) => id,
1349 Expandable::Object(customer) => customer.id.as_ref(),
1350 };
1351
1352 // If we already have a billing customer record associated with the Stripe customer,
1353 // there's nothing more we need to do.
1354 if let Some(billing_customer) = app
1355 .db
1356 .get_billing_customer_by_stripe_customer_id(customer_id)
1357 .await?
1358 {
1359 return Ok(Some(billing_customer));
1360 }
1361
1362 // If all we have is a customer ID, resolve it to a full customer record by
1363 // hitting the Stripe API.
1364 let customer = match customer_or_id {
1365 Expandable::Id(id) => Customer::retrieve(stripe_client, &id, &[]).await?,
1366 Expandable::Object(customer) => *customer,
1367 };
1368
1369 let Some(email) = customer.email else {
1370 return Ok(None);
1371 };
1372
1373 let Some(user) = app.db.get_user_by_email(&email).await? else {
1374 return Ok(None);
1375 };
1376
1377 let billing_customer = app
1378 .db
1379 .create_billing_customer(&CreateBillingCustomerParams {
1380 user_id: user.id,
1381 stripe_customer_id: customer.id.to_string(),
1382 })
1383 .await?;
1384
1385 Ok(Some(billing_customer))
1386}
1387
1388const SYNC_LLM_TOKEN_USAGE_WITH_STRIPE_INTERVAL: Duration = Duration::from_secs(60);
1389
1390pub fn sync_llm_token_usage_with_stripe_periodically(app: Arc<AppState>) {
1391 let Some(stripe_billing) = app.stripe_billing.clone() else {
1392 log::warn!("failed to retrieve Stripe billing object");
1393 return;
1394 };
1395 let Some(llm_db) = app.llm_db.clone() else {
1396 log::warn!("failed to retrieve LLM database");
1397 return;
1398 };
1399
1400 let executor = app.executor.clone();
1401 executor.spawn_detached({
1402 let executor = executor.clone();
1403 async move {
1404 loop {
1405 sync_token_usage_with_stripe(&app, &llm_db, &stripe_billing)
1406 .await
1407 .context("failed to sync LLM usage to Stripe")
1408 .trace_err();
1409 executor
1410 .sleep(SYNC_LLM_TOKEN_USAGE_WITH_STRIPE_INTERVAL)
1411 .await;
1412 }
1413 }
1414 });
1415}
1416
1417async fn sync_token_usage_with_stripe(
1418 app: &Arc<AppState>,
1419 llm_db: &Arc<LlmDatabase>,
1420 stripe_billing: &Arc<StripeBilling>,
1421) -> anyhow::Result<()> {
1422 let events = llm_db.get_billing_events().await?;
1423 let user_ids = events
1424 .iter()
1425 .map(|(event, _)| event.user_id)
1426 .collect::<HashSet<UserId>>();
1427 let stripe_subscriptions = app.db.get_active_billing_subscriptions(user_ids).await?;
1428
1429 for (event, model) in events {
1430 let Some((stripe_db_customer, stripe_db_subscription)) =
1431 stripe_subscriptions.get(&event.user_id)
1432 else {
1433 tracing::warn!(
1434 user_id = event.user_id.0,
1435 "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."
1436 );
1437 continue;
1438 };
1439 let stripe_subscription_id: stripe::SubscriptionId = stripe_db_subscription
1440 .stripe_subscription_id
1441 .parse()
1442 .context("failed to parse stripe subscription id from db")?;
1443 let stripe_customer_id: stripe::CustomerId = stripe_db_customer
1444 .stripe_customer_id
1445 .parse()
1446 .context("failed to parse stripe customer id from db")?;
1447
1448 let stripe_model = stripe_billing
1449 .register_model_for_token_based_usage(&model)
1450 .await?;
1451 stripe_billing
1452 .subscribe_to_model(&stripe_subscription_id, &stripe_model)
1453 .await?;
1454 stripe_billing
1455 .bill_model_token_usage(&stripe_customer_id, &stripe_model, &event)
1456 .await?;
1457 llm_db.consume_billing_event(event.id).await?;
1458 }
1459
1460 Ok(())
1461}
1462
1463const SYNC_LLM_REQUEST_USAGE_WITH_STRIPE_INTERVAL: Duration = Duration::from_secs(60);
1464
1465pub fn sync_llm_request_usage_with_stripe_periodically(app: Arc<AppState>) {
1466 let Some(stripe_billing) = app.stripe_billing.clone() else {
1467 log::warn!("failed to retrieve Stripe billing object");
1468 return;
1469 };
1470 let Some(llm_db) = app.llm_db.clone() else {
1471 log::warn!("failed to retrieve LLM database");
1472 return;
1473 };
1474
1475 let executor = app.executor.clone();
1476 executor.spawn_detached({
1477 let executor = executor.clone();
1478 async move {
1479 loop {
1480 sync_model_request_usage_with_stripe(&app, &llm_db, &stripe_billing)
1481 .await
1482 .context("failed to sync LLM request usage to Stripe")
1483 .trace_err();
1484 executor
1485 .sleep(SYNC_LLM_REQUEST_USAGE_WITH_STRIPE_INTERVAL)
1486 .await;
1487 }
1488 }
1489 });
1490}
1491
1492async fn sync_model_request_usage_with_stripe(
1493 app: &Arc<AppState>,
1494 llm_db: &Arc<LlmDatabase>,
1495 stripe_billing: &Arc<StripeBilling>,
1496) -> anyhow::Result<()> {
1497 let usage_meters = llm_db
1498 .get_current_subscription_usage_meters(Utc::now())
1499 .await?;
1500 let user_ids = usage_meters
1501 .iter()
1502 .map(|(_, usage)| usage.user_id)
1503 .collect::<HashSet<UserId>>();
1504 let billing_subscriptions = app
1505 .db
1506 .get_active_zed_pro_billing_subscriptions(user_ids)
1507 .await?;
1508
1509 let claude_3_5_sonnet = stripe_billing
1510 .find_price_by_lookup_key("claude-3-5-sonnet-requests")
1511 .await?;
1512 let claude_3_7_sonnet = stripe_billing
1513 .find_price_by_lookup_key("claude-3-7-sonnet-requests")
1514 .await?;
1515 let claude_3_7_sonnet_max = stripe_billing
1516 .find_price_by_lookup_key("claude-3-7-sonnet-requests-max")
1517 .await?;
1518
1519 for (usage_meter, usage) in usage_meters {
1520 maybe!(async {
1521 let Some((billing_customer, billing_subscription)) =
1522 billing_subscriptions.get(&usage.user_id)
1523 else {
1524 bail!(
1525 "Attempted to sync usage meter for user who is not a Stripe customer: {}",
1526 usage.user_id
1527 );
1528 };
1529
1530 let stripe_customer_id = billing_customer
1531 .stripe_customer_id
1532 .parse::<stripe::CustomerId>()
1533 .context("failed to parse Stripe customer ID from database")?;
1534 let stripe_subscription_id = billing_subscription
1535 .stripe_subscription_id
1536 .parse::<stripe::SubscriptionId>()
1537 .context("failed to parse Stripe subscription ID from database")?;
1538
1539 let model = llm_db.model_by_id(usage_meter.model_id)?;
1540
1541 let (price, meter_event_name) = match model.name.as_str() {
1542 "claude-3-5-sonnet" => (&claude_3_5_sonnet, "claude_3_5_sonnet/requests"),
1543 "claude-3-7-sonnet" => match usage_meter.mode {
1544 CompletionMode::Normal => (&claude_3_7_sonnet, "claude_3_7_sonnet/requests"),
1545 CompletionMode::Max => {
1546 (&claude_3_7_sonnet_max, "claude_3_7_sonnet/requests/max")
1547 }
1548 },
1549 model_name => {
1550 bail!("Attempted to sync usage meter for unsupported model: {model_name:?}")
1551 }
1552 };
1553
1554 stripe_billing
1555 .subscribe_to_price(&stripe_subscription_id, price)
1556 .await?;
1557 stripe_billing
1558 .bill_model_request_usage(
1559 &stripe_customer_id,
1560 meter_event_name,
1561 usage_meter.requests,
1562 )
1563 .await?;
1564
1565 Ok(())
1566 })
1567 .await
1568 .log_err();
1569 }
1570
1571 Ok(())
1572}