billing.rs

  1use std::str::FromStr;
  2use std::sync::Arc;
  3use std::time::Duration;
  4
  5use anyhow::{anyhow, bail, Context};
  6use axum::{
  7    extract::{self, Query},
  8    routing::{get, post},
  9    Extension, Json, Router,
 10};
 11use chrono::{DateTime, SecondsFormat};
 12use reqwest::StatusCode;
 13use sea_orm::ActiveValue;
 14use serde::{Deserialize, Serialize};
 15use stripe::{
 16    BillingPortalSession, CheckoutSession, CreateBillingPortalSession,
 17    CreateBillingPortalSessionFlowData, CreateBillingPortalSessionFlowDataAfterCompletion,
 18    CreateBillingPortalSessionFlowDataAfterCompletionRedirect,
 19    CreateBillingPortalSessionFlowDataType, CreateCheckoutSession, CreateCheckoutSessionLineItems,
 20    CreateCustomer, Customer, CustomerId, EventObject, EventType, Expandable, ListEvents,
 21    Subscription, SubscriptionId, SubscriptionStatus,
 22};
 23use util::ResultExt;
 24
 25use crate::db::billing_subscription::StripeSubscriptionStatus;
 26use crate::db::{
 27    billing_customer, BillingSubscriptionId, CreateBillingCustomerParams,
 28    CreateBillingSubscriptionParams, CreateProcessedStripeEventParams, UpdateBillingCustomerParams,
 29    UpdateBillingSubscriptionParams,
 30};
 31use crate::{AppState, Error, Result};
 32
 33pub fn router() -> Router {
 34    Router::new()
 35        .route(
 36            "/billing/subscriptions",
 37            get(list_billing_subscriptions).post(create_billing_subscription),
 38        )
 39        .route(
 40            "/billing/subscriptions/manage",
 41            post(manage_billing_subscription),
 42        )
 43}
 44
 45#[derive(Debug, Deserialize)]
 46struct ListBillingSubscriptionsParams {
 47    github_user_id: i32,
 48}
 49
 50#[derive(Debug, Serialize)]
 51struct BillingSubscriptionJson {
 52    id: BillingSubscriptionId,
 53    name: String,
 54    status: StripeSubscriptionStatus,
 55    cancel_at: Option<String>,
 56    /// Whether this subscription can be canceled.
 57    is_cancelable: bool,
 58}
 59
 60#[derive(Debug, Serialize)]
 61struct ListBillingSubscriptionsResponse {
 62    subscriptions: Vec<BillingSubscriptionJson>,
 63}
 64
 65async fn list_billing_subscriptions(
 66    Extension(app): Extension<Arc<AppState>>,
 67    Query(params): Query<ListBillingSubscriptionsParams>,
 68) -> Result<Json<ListBillingSubscriptionsResponse>> {
 69    let user = app
 70        .db
 71        .get_user_by_github_user_id(params.github_user_id)
 72        .await?
 73        .ok_or_else(|| anyhow!("user not found"))?;
 74
 75    let subscriptions = app.db.get_billing_subscriptions(user.id).await?;
 76
 77    Ok(Json(ListBillingSubscriptionsResponse {
 78        subscriptions: subscriptions
 79            .into_iter()
 80            .map(|subscription| BillingSubscriptionJson {
 81                id: subscription.id,
 82                name: "Zed Pro".to_string(),
 83                status: subscription.stripe_subscription_status,
 84                cancel_at: subscription.stripe_cancel_at.map(|cancel_at| {
 85                    cancel_at
 86                        .and_utc()
 87                        .to_rfc3339_opts(SecondsFormat::Millis, true)
 88                }),
 89                is_cancelable: subscription.stripe_subscription_status.is_cancelable()
 90                    && subscription.stripe_cancel_at.is_none(),
 91            })
 92            .collect(),
 93    }))
 94}
 95
 96#[derive(Debug, Deserialize)]
 97struct CreateBillingSubscriptionBody {
 98    github_user_id: i32,
 99}
100
101#[derive(Debug, Serialize)]
102struct CreateBillingSubscriptionResponse {
103    checkout_session_url: String,
104}
105
106/// Initiates a Stripe Checkout session for creating a billing subscription.
107async fn create_billing_subscription(
108    Extension(app): Extension<Arc<AppState>>,
109    extract::Json(body): extract::Json<CreateBillingSubscriptionBody>,
110) -> Result<Json<CreateBillingSubscriptionResponse>> {
111    let user = app
112        .db
113        .get_user_by_github_user_id(body.github_user_id)
114        .await?
115        .ok_or_else(|| anyhow!("user not found"))?;
116
117    let Some((stripe_client, stripe_price_id)) = app
118        .stripe_client
119        .clone()
120        .zip(app.config.stripe_price_id.clone())
121    else {
122        log::error!("failed to retrieve Stripe client or price ID");
123        Err(Error::Http(
124            StatusCode::NOT_IMPLEMENTED,
125            "not supported".into(),
126        ))?
127    };
128
129    let customer_id =
130        if let Some(existing_customer) = app.db.get_billing_customer_by_user_id(user.id).await? {
131            CustomerId::from_str(&existing_customer.stripe_customer_id)
132                .context("failed to parse customer ID")?
133        } else {
134            let customer = Customer::create(
135                &stripe_client,
136                CreateCustomer {
137                    email: user.email_address.as_deref(),
138                    ..Default::default()
139                },
140            )
141            .await?;
142
143            customer.id
144        };
145
146    let checkout_session = {
147        let mut params = CreateCheckoutSession::new();
148        params.mode = Some(stripe::CheckoutSessionMode::Subscription);
149        params.customer = Some(customer_id);
150        params.client_reference_id = Some(user.github_login.as_str());
151        params.line_items = Some(vec![CreateCheckoutSessionLineItems {
152            price: Some(stripe_price_id.to_string()),
153            quantity: Some(1),
154            ..Default::default()
155        }]);
156        params.success_url = Some("https://zed.dev/billing/success");
157
158        CheckoutSession::create(&stripe_client, params).await?
159    };
160
161    Ok(Json(CreateBillingSubscriptionResponse {
162        checkout_session_url: checkout_session
163            .url
164            .ok_or_else(|| anyhow!("no checkout session URL"))?,
165    }))
166}
167
168#[derive(Debug, PartialEq, Deserialize)]
169#[serde(rename_all = "snake_case")]
170enum ManageSubscriptionIntent {
171    /// The user intends to cancel their subscription.
172    Cancel,
173    /// The user intends to stop the cancelation of their subscription.
174    StopCancelation,
175}
176
177#[derive(Debug, Deserialize)]
178struct ManageBillingSubscriptionBody {
179    github_user_id: i32,
180    intent: ManageSubscriptionIntent,
181    /// The ID of the subscription to manage.
182    subscription_id: BillingSubscriptionId,
183}
184
185#[derive(Debug, Serialize)]
186struct ManageBillingSubscriptionResponse {
187    billing_portal_session_url: Option<String>,
188}
189
190/// Initiates a Stripe customer portal session for managing a billing subscription.
191async fn manage_billing_subscription(
192    Extension(app): Extension<Arc<AppState>>,
193    extract::Json(body): extract::Json<ManageBillingSubscriptionBody>,
194) -> Result<Json<ManageBillingSubscriptionResponse>> {
195    let user = app
196        .db
197        .get_user_by_github_user_id(body.github_user_id)
198        .await?
199        .ok_or_else(|| anyhow!("user not found"))?;
200
201    let Some(stripe_client) = app.stripe_client.clone() else {
202        log::error!("failed to retrieve Stripe client");
203        Err(Error::Http(
204            StatusCode::NOT_IMPLEMENTED,
205            "not supported".into(),
206        ))?
207    };
208
209    let customer = app
210        .db
211        .get_billing_customer_by_user_id(user.id)
212        .await?
213        .ok_or_else(|| anyhow!("billing customer not found"))?;
214    let customer_id = CustomerId::from_str(&customer.stripe_customer_id)
215        .context("failed to parse customer ID")?;
216
217    let subscription = app
218        .db
219        .get_billing_subscription_by_id(body.subscription_id)
220        .await?
221        .ok_or_else(|| anyhow!("subscription not found"))?;
222
223    if body.intent == ManageSubscriptionIntent::StopCancelation {
224        let subscription_id = SubscriptionId::from_str(&subscription.stripe_subscription_id)
225            .context("failed to parse subscription ID")?;
226
227        let updated_stripe_subscription = Subscription::update(
228            &stripe_client,
229            &subscription_id,
230            stripe::UpdateSubscription {
231                cancel_at_period_end: Some(false),
232                ..Default::default()
233            },
234        )
235        .await?;
236
237        app.db
238            .update_billing_subscription(
239                subscription.id,
240                &UpdateBillingSubscriptionParams {
241                    stripe_cancel_at: ActiveValue::set(
242                        updated_stripe_subscription
243                            .cancel_at
244                            .and_then(|cancel_at| DateTime::from_timestamp(cancel_at, 0))
245                            .map(|time| time.naive_utc()),
246                    ),
247                    ..Default::default()
248                },
249            )
250            .await?;
251
252        return Ok(Json(ManageBillingSubscriptionResponse {
253            billing_portal_session_url: None,
254        }));
255    }
256
257    let flow = match body.intent {
258        ManageSubscriptionIntent::Cancel => CreateBillingPortalSessionFlowData {
259            type_: CreateBillingPortalSessionFlowDataType::SubscriptionCancel,
260            after_completion: Some(CreateBillingPortalSessionFlowDataAfterCompletion {
261                type_: stripe::CreateBillingPortalSessionFlowDataAfterCompletionType::Redirect,
262                redirect: Some(CreateBillingPortalSessionFlowDataAfterCompletionRedirect {
263                    return_url: "https://zed.dev/settings".into(),
264                }),
265                ..Default::default()
266            }),
267            subscription_cancel: Some(
268                stripe::CreateBillingPortalSessionFlowDataSubscriptionCancel {
269                    subscription: subscription.stripe_subscription_id,
270                    retention: None,
271                },
272            ),
273            ..Default::default()
274        },
275        ManageSubscriptionIntent::StopCancelation => unreachable!(),
276    };
277
278    let mut params = CreateBillingPortalSession::new(customer_id);
279    params.flow_data = Some(flow);
280    params.return_url = Some("https://zed.dev/settings");
281
282    let session = BillingPortalSession::create(&stripe_client, params).await?;
283
284    Ok(Json(ManageBillingSubscriptionResponse {
285        billing_portal_session_url: Some(session.url),
286    }))
287}
288
289const POLL_EVENTS_INTERVAL: Duration = Duration::from_secs(5 * 60);
290
291/// Polls the Stripe events API periodically to reconcile the records in our
292/// database with the data in Stripe.
293pub fn poll_stripe_events_periodically(app: Arc<AppState>) {
294    let Some(stripe_client) = app.stripe_client.clone() else {
295        log::warn!("failed to retrieve Stripe client");
296        return;
297    };
298
299    let executor = app.executor.clone();
300    executor.spawn_detached({
301        let executor = executor.clone();
302        async move {
303            loop {
304                poll_stripe_events(&app, &stripe_client).await.log_err();
305
306                executor.sleep(POLL_EVENTS_INTERVAL).await;
307            }
308        }
309    });
310}
311
312async fn poll_stripe_events(
313    app: &Arc<AppState>,
314    stripe_client: &stripe::Client,
315) -> anyhow::Result<()> {
316    fn event_type_to_string(event_type: EventType) -> String {
317        // Calling `to_string` on `stripe::EventType` members gives us a quoted string,
318        // so we need to unquote it.
319        event_type.to_string().trim_matches('"').to_string()
320    }
321
322    let event_types = [
323        EventType::CustomerCreated,
324        EventType::CustomerUpdated,
325        EventType::CustomerSubscriptionCreated,
326        EventType::CustomerSubscriptionUpdated,
327        EventType::CustomerSubscriptionPaused,
328        EventType::CustomerSubscriptionResumed,
329        EventType::CustomerSubscriptionDeleted,
330    ]
331    .into_iter()
332    .map(event_type_to_string)
333    .collect::<Vec<_>>();
334
335    let mut unprocessed_events = Vec::new();
336
337    loop {
338        log::info!("retrieving events from Stripe: {}", event_types.join(", "));
339
340        let mut params = ListEvents::new();
341        params.types = Some(event_types.clone());
342        params.limit = Some(100);
343
344        let events = stripe::Event::list(stripe_client, &params).await?;
345
346        let processed_event_ids = {
347            let event_ids = &events
348                .data
349                .iter()
350                .map(|event| event.id.as_str())
351                .collect::<Vec<_>>();
352
353            app.db
354                .get_processed_stripe_events_by_event_ids(event_ids)
355                .await?
356                .into_iter()
357                .map(|event| event.stripe_event_id)
358                .collect::<Vec<_>>()
359        };
360
361        for event in events.data {
362            if processed_event_ids.contains(&event.id.to_string()) {
363                log::info!("Stripe event {} already processed: skipping", event.id);
364            } else {
365                unprocessed_events.push(event);
366            }
367        }
368
369        if !events.has_more {
370            break;
371        }
372    }
373
374    log::info!(
375        "unprocessed events from Stripe: {}",
376        unprocessed_events.len()
377    );
378
379    // Sort all of the unprocessed events in ascending order, so we can handle them in the order they occurred.
380    unprocessed_events.sort_by(|a, b| a.created.cmp(&b.created).then_with(|| a.id.cmp(&b.id)));
381
382    for event in unprocessed_events {
383        let processed_event_params = CreateProcessedStripeEventParams {
384            stripe_event_id: event.id.to_string(),
385            stripe_event_type: event_type_to_string(event.type_),
386            stripe_event_created_timestamp: event.created,
387        };
388
389        match event.type_ {
390            EventType::CustomerCreated | EventType::CustomerUpdated => {
391                handle_customer_event(app, stripe_client, event)
392                    .await
393                    .log_err();
394            }
395            EventType::CustomerSubscriptionCreated
396            | EventType::CustomerSubscriptionUpdated
397            | EventType::CustomerSubscriptionPaused
398            | EventType::CustomerSubscriptionResumed
399            | EventType::CustomerSubscriptionDeleted => {
400                handle_customer_subscription_event(app, stripe_client, event)
401                    .await
402                    .log_err();
403            }
404            _ => {}
405        }
406
407        app.db
408            .create_processed_stripe_event(&processed_event_params)
409            .await?;
410    }
411
412    Ok(())
413}
414
415async fn handle_customer_event(
416    app: &Arc<AppState>,
417    _stripe_client: &stripe::Client,
418    event: stripe::Event,
419) -> anyhow::Result<()> {
420    let EventObject::Customer(customer) = event.data.object else {
421        bail!("unexpected event payload for {}", event.id);
422    };
423
424    log::info!("handling Stripe {} event: {}", event.type_, event.id);
425
426    let Some(email) = customer.email else {
427        log::info!("Stripe customer has no email: skipping");
428        return Ok(());
429    };
430
431    let Some(user) = app.db.get_user_by_email(&email).await? else {
432        log::info!("no user found for email: skipping");
433        return Ok(());
434    };
435
436    if let Some(existing_customer) = app
437        .db
438        .get_billing_customer_by_stripe_customer_id(&customer.id)
439        .await?
440    {
441        app.db
442            .update_billing_customer(
443                existing_customer.id,
444                &UpdateBillingCustomerParams {
445                    // For now we just leave the information as-is, as it is not
446                    // likely to change.
447                    ..Default::default()
448                },
449            )
450            .await?;
451    } else {
452        app.db
453            .create_billing_customer(&CreateBillingCustomerParams {
454                user_id: user.id,
455                stripe_customer_id: customer.id.to_string(),
456            })
457            .await?;
458    }
459
460    Ok(())
461}
462
463async fn handle_customer_subscription_event(
464    app: &Arc<AppState>,
465    stripe_client: &stripe::Client,
466    event: stripe::Event,
467) -> anyhow::Result<()> {
468    let EventObject::Subscription(subscription) = event.data.object else {
469        bail!("unexpected event payload for {}", event.id);
470    };
471
472    log::info!("handling Stripe {} event: {}", event.type_, event.id);
473
474    let billing_customer =
475        find_or_create_billing_customer(app, stripe_client, subscription.customer)
476            .await?
477            .ok_or_else(|| anyhow!("billing customer not found"))?;
478
479    if let Some(existing_subscription) = app
480        .db
481        .get_billing_subscription_by_stripe_subscription_id(&subscription.id)
482        .await?
483    {
484        app.db
485            .update_billing_subscription(
486                existing_subscription.id,
487                &UpdateBillingSubscriptionParams {
488                    billing_customer_id: ActiveValue::set(billing_customer.id),
489                    stripe_subscription_id: ActiveValue::set(subscription.id.to_string()),
490                    stripe_subscription_status: ActiveValue::set(subscription.status.into()),
491                    stripe_cancel_at: ActiveValue::set(
492                        subscription
493                            .cancel_at
494                            .and_then(|cancel_at| DateTime::from_timestamp(cancel_at, 0))
495                            .map(|time| time.naive_utc()),
496                    ),
497                },
498            )
499            .await?;
500    } else {
501        app.db
502            .create_billing_subscription(&CreateBillingSubscriptionParams {
503                billing_customer_id: billing_customer.id,
504                stripe_subscription_id: subscription.id.to_string(),
505                stripe_subscription_status: subscription.status.into(),
506            })
507            .await?;
508    }
509
510    Ok(())
511}
512
513impl From<SubscriptionStatus> for StripeSubscriptionStatus {
514    fn from(value: SubscriptionStatus) -> Self {
515        match value {
516            SubscriptionStatus::Incomplete => Self::Incomplete,
517            SubscriptionStatus::IncompleteExpired => Self::IncompleteExpired,
518            SubscriptionStatus::Trialing => Self::Trialing,
519            SubscriptionStatus::Active => Self::Active,
520            SubscriptionStatus::PastDue => Self::PastDue,
521            SubscriptionStatus::Canceled => Self::Canceled,
522            SubscriptionStatus::Unpaid => Self::Unpaid,
523            SubscriptionStatus::Paused => Self::Paused,
524        }
525    }
526}
527
528/// Finds or creates a billing customer using the provided customer.
529async fn find_or_create_billing_customer(
530    app: &Arc<AppState>,
531    stripe_client: &stripe::Client,
532    customer_or_id: Expandable<Customer>,
533) -> anyhow::Result<Option<billing_customer::Model>> {
534    let customer_id = match &customer_or_id {
535        Expandable::Id(id) => id,
536        Expandable::Object(customer) => customer.id.as_ref(),
537    };
538
539    // If we already have a billing customer record associated with the Stripe customer,
540    // there's nothing more we need to do.
541    if let Some(billing_customer) = app
542        .db
543        .get_billing_customer_by_stripe_customer_id(&customer_id)
544        .await?
545    {
546        return Ok(Some(billing_customer));
547    }
548
549    // If all we have is a customer ID, resolve it to a full customer record by
550    // hitting the Stripe API.
551    let customer = match customer_or_id {
552        Expandable::Id(id) => Customer::retrieve(&stripe_client, &id, &[]).await?,
553        Expandable::Object(customer) => *customer,
554    };
555
556    let Some(email) = customer.email else {
557        return Ok(None);
558    };
559
560    let Some(user) = app.db.get_user_by_email(&email).await? else {
561        return Ok(None);
562    };
563
564    let billing_customer = app
565        .db
566        .create_billing_customer(&CreateBillingCustomerParams {
567            user_id: user.id,
568            stripe_customer_id: customer.id.to_string(),
569        })
570        .await?;
571
572    Ok(Some(billing_customer))
573}