collab: Look up Stripe prices with lookup keys (#29715)

This PR makes it so we look up Stripe prices via lookup keys instead of
passing in the price IDs as environment variables.

Release Notes:

- N/A
This commit is contained in:
Marshall Bowers 2025-04-30 20:26:31 -04:00 committed by GitHub
parent afeb3d4fd9
commit f046d70625
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
4 changed files with 40 additions and 50 deletions

View file

@ -362,12 +362,7 @@ async fn create_billing_subscription(
let checkout_session_url = match body.product { let checkout_session_url = match body.product {
Some(ProductCode::ZedPro) => { Some(ProductCode::ZedPro) => {
stripe_billing stripe_billing
.checkout_with_price( .checkout_with_zed_pro(customer_id, &user.github_login, &success_url)
app.config.zed_pro_price_id()?,
customer_id,
&user.github_login,
&success_url,
)
.await? .await?
} }
Some(ProductCode::ZedProTrial) => { Some(ProductCode::ZedProTrial) => {
@ -384,7 +379,6 @@ async fn create_billing_subscription(
stripe_billing stripe_billing
.checkout_with_zed_pro_trial( .checkout_with_zed_pro_trial(
app.config.zed_pro_price_id()?,
customer_id, customer_id,
&user.github_login, &user.github_login,
feature_flags, feature_flags,
@ -458,6 +452,14 @@ async fn manage_billing_subscription(
))? ))?
}; };
let Some(stripe_billing) = app.stripe_billing.clone() else {
log::error!("failed to retrieve Stripe billing object");
Err(Error::http(
StatusCode::NOT_IMPLEMENTED,
"not supported".into(),
))?
};
let customer = app let customer = app
.db .db
.get_billing_customer_by_user_id(user.id) .get_billing_customer_by_user_id(user.id)
@ -508,8 +510,8 @@ async fn manage_billing_subscription(
let flow = match body.intent { let flow = match body.intent {
ManageSubscriptionIntent::ManageSubscription => None, ManageSubscriptionIntent::ManageSubscription => None,
ManageSubscriptionIntent::UpgradeToPro => { ManageSubscriptionIntent::UpgradeToPro => {
let zed_pro_price_id = app.config.zed_pro_price_id()?; let zed_pro_price_id = stripe_billing.zed_pro_price_id().await?;
let zed_free_price_id = app.config.zed_free_price_id()?; let zed_free_price_id = stripe_billing.zed_free_price_id().await?;
let stripe_subscription = let stripe_subscription =
Subscription::retrieve(&stripe_client, &subscription_id, &[]).await?; Subscription::retrieve(&stripe_client, &subscription_id, &[]).await?;
@ -856,9 +858,11 @@ async fn handle_customer_subscription_event(
log::info!("handling Stripe {} event: {}", event.type_, event.id); log::info!("handling Stripe {} event: {}", event.type_, event.id);
let subscription_kind = maybe!({ let subscription_kind = maybe!(async {
let zed_pro_price_id = app.config.zed_pro_price_id().ok()?; let stripe_billing = app.stripe_billing.clone()?;
let zed_free_price_id = app.config.zed_free_price_id().ok()?;
let zed_pro_price_id = stripe_billing.zed_pro_price_id().await.ok()?;
let zed_free_price_id = stripe_billing.zed_free_price_id().await.ok()?;
subscription.items.data.iter().find_map(|item| { subscription.items.data.iter().find_map(|item| {
let price = item.price.as_ref()?; let price = item.price.as_ref()?;
@ -875,7 +879,8 @@ async fn handle_customer_subscription_event(
None None
} }
}) })
}); })
.await;
let billing_customer = let billing_customer =
find_or_create_billing_customer(app, stripe_client, subscription.customer) find_or_create_billing_customer(app, stripe_client, subscription.customer)
@ -1398,13 +1403,13 @@ async fn sync_model_request_usage_with_stripe(
.await?; .await?;
let claude_3_5_sonnet = stripe_billing let claude_3_5_sonnet = stripe_billing
.find_price_by_lookup_key("claude-3-5-sonnet-requests") .find_price_id_by_lookup_key("claude-3-5-sonnet-requests")
.await?; .await?;
let claude_3_7_sonnet = stripe_billing let claude_3_7_sonnet = stripe_billing
.find_price_by_lookup_key("claude-3-7-sonnet-requests") .find_price_id_by_lookup_key("claude-3-7-sonnet-requests")
.await?; .await?;
let claude_3_7_sonnet_max = stripe_billing let claude_3_7_sonnet_max = stripe_billing
.find_price_by_lookup_key("claude-3-7-sonnet-requests-max") .find_price_id_by_lookup_key("claude-3-7-sonnet-requests-max")
.await?; .await?;
for (usage_meter, usage) in usage_meters { for (usage_meter, usage) in usage_meters {
@ -1430,11 +1435,11 @@ async fn sync_model_request_usage_with_stripe(
let model = llm_db.model_by_id(usage_meter.model_id)?; let model = llm_db.model_by_id(usage_meter.model_id)?;
let (price_id, meter_event_name) = match model.name.as_str() { let (price_id, meter_event_name) = match model.name.as_str() {
"claude-3-5-sonnet" => (&claude_3_5_sonnet.id, "claude_3_5_sonnet/requests"), "claude-3-5-sonnet" => (&claude_3_5_sonnet, "claude_3_5_sonnet/requests"),
"claude-3-7-sonnet" => match usage_meter.mode { "claude-3-7-sonnet" => match usage_meter.mode {
CompletionMode::Normal => (&claude_3_7_sonnet.id, "claude_3_7_sonnet/requests"), CompletionMode::Normal => (&claude_3_7_sonnet, "claude_3_7_sonnet/requests"),
CompletionMode::Max => { CompletionMode::Max => {
(&claude_3_7_sonnet_max.id, "claude_3_7_sonnet/requests/max") (&claude_3_7_sonnet_max, "claude_3_7_sonnet/requests/max")
} }
}, },
model_name => { model_name => {

View file

@ -180,9 +180,6 @@ pub struct Config {
pub slack_panics_webhook: Option<String>, pub slack_panics_webhook: Option<String>,
pub auto_join_channel_id: Option<ChannelId>, pub auto_join_channel_id: Option<ChannelId>,
pub stripe_api_key: Option<String>, pub stripe_api_key: Option<String>,
pub stripe_zed_pro_price_id: Option<String>,
pub stripe_zed_pro_trial_price_id: Option<String>,
pub stripe_zed_free_price_id: Option<String>,
pub supermaven_admin_api_key: Option<Arc<str>>, pub supermaven_admin_api_key: Option<Arc<str>>,
pub user_backfiller_github_access_token: Option<Arc<str>>, pub user_backfiller_github_access_token: Option<Arc<str>>,
} }
@ -201,22 +198,6 @@ impl Config {
} }
} }
pub fn zed_pro_price_id(&self) -> anyhow::Result<stripe::PriceId> {
Self::parse_stripe_price_id("Zed Pro", self.stripe_zed_pro_price_id.as_deref())
}
pub fn zed_free_price_id(&self) -> anyhow::Result<stripe::PriceId> {
Self::parse_stripe_price_id("Zed Free", self.stripe_zed_pro_price_id.as_deref())
}
fn parse_stripe_price_id(name: &str, value: Option<&str>) -> anyhow::Result<stripe::PriceId> {
use std::str::FromStr as _;
let price_id = value.ok_or_else(|| anyhow!("{name} price ID not set"))?;
Ok(stripe::PriceId::from_str(price_id)?)
}
#[cfg(test)] #[cfg(test)]
pub fn test() -> Self { pub fn test() -> Self {
Self { Self {
@ -254,9 +235,6 @@ impl Config {
migrations_path: None, migrations_path: None,
seed_path: None, seed_path: None,
stripe_api_key: None, stripe_api_key: None,
stripe_zed_pro_price_id: None,
stripe_zed_pro_trial_price_id: None,
stripe_zed_free_price_id: None,
supermaven_admin_api_key: None, supermaven_admin_api_key: None,
user_backfiller_github_access_token: None, user_backfiller_github_access_token: None,
kinesis_region: None, kinesis_region: None,

View file

@ -81,13 +81,21 @@ impl StripeBilling {
Ok(()) Ok(())
} }
pub async fn find_price_by_lookup_key(&self, lookup_key: &str) -> Result<stripe::Price> { pub async fn zed_pro_price_id(&self) -> Result<PriceId> {
self.find_price_id_by_lookup_key("zed-pro").await
}
pub async fn zed_free_price_id(&self) -> Result<PriceId> {
self.find_price_id_by_lookup_key("zed-free").await
}
pub async fn find_price_id_by_lookup_key(&self, lookup_key: &str) -> Result<PriceId> {
self.state self.state
.read() .read()
.await .await
.prices_by_lookup_key .prices_by_lookup_key
.get(lookup_key) .get(lookup_key)
.cloned() .map(|price| price.id.clone())
.ok_or_else(|| crate::Error::Internal(anyhow!("no price ID found for {lookup_key:?}"))) .ok_or_else(|| crate::Error::Internal(anyhow!("no price ID found for {lookup_key:?}")))
} }
@ -463,19 +471,20 @@ impl StripeBilling {
Ok(session.url.context("no checkout session URL")?) Ok(session.url.context("no checkout session URL")?)
} }
pub async fn checkout_with_price( pub async fn checkout_with_zed_pro(
&self, &self,
price_id: PriceId,
customer_id: stripe::CustomerId, customer_id: stripe::CustomerId,
github_login: &str, github_login: &str,
success_url: &str, success_url: &str,
) -> Result<String> { ) -> Result<String> {
let zed_pro_price_id = self.zed_pro_price_id().await?;
let mut params = stripe::CreateCheckoutSession::new(); let mut params = stripe::CreateCheckoutSession::new();
params.mode = Some(stripe::CheckoutSessionMode::Subscription); params.mode = Some(stripe::CheckoutSessionMode::Subscription);
params.customer = Some(customer_id); params.customer = Some(customer_id);
params.client_reference_id = Some(github_login); params.client_reference_id = Some(github_login);
params.line_items = Some(vec![stripe::CreateCheckoutSessionLineItems { params.line_items = Some(vec![stripe::CreateCheckoutSessionLineItems {
price: Some(price_id.to_string()), price: Some(zed_pro_price_id.to_string()),
quantity: Some(1), quantity: Some(1),
..Default::default() ..Default::default()
}]); }]);
@ -487,12 +496,13 @@ impl StripeBilling {
pub async fn checkout_with_zed_pro_trial( pub async fn checkout_with_zed_pro_trial(
&self, &self,
zed_pro_price_id: PriceId,
customer_id: stripe::CustomerId, customer_id: stripe::CustomerId,
github_login: &str, github_login: &str,
feature_flags: Vec<String>, feature_flags: Vec<String>,
success_url: &str, success_url: &str,
) -> Result<String> { ) -> Result<String> {
let zed_pro_price_id = self.zed_pro_price_id().await?;
let eligible_for_extended_trial = feature_flags let eligible_for_extended_trial = feature_flags
.iter() .iter()
.any(|flag| flag == AGENT_EXTENDED_TRIAL_FEATURE_FLAG); .any(|flag| flag == AGENT_EXTENDED_TRIAL_FEATURE_FLAG);

View file

@ -554,9 +554,6 @@ impl TestServer {
migrations_path: None, migrations_path: None,
seed_path: None, seed_path: None,
stripe_api_key: None, stripe_api_key: None,
stripe_zed_pro_price_id: None,
stripe_zed_pro_trial_price_id: None,
stripe_zed_free_price_id: None,
supermaven_admin_api_key: None, supermaven_admin_api_key: None,
user_backfiller_github_access_token: None, user_backfiller_github_access_token: None,
kinesis_region: None, kinesis_region: None,