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