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