Yehowshua/ZIm
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

collab: Don't transfer existing usage when upgrading to Zed Pro (#30045)

Browse source 
This PR makes it so we don't transfer existing usage over when upgrading
from a trial to Zed Pro.

Release Notes:

- N/A
This commit is contained in:
Marshall Bowers 2025-05-06 17:29:06 -04:00 • committed by GitHub
parent 3cdf5ce947
commit cec1d2584b
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
4 changed files with 1 additions and 252 deletions
Download patch file Download diff file Expand all files Collapse all files

23
crates/collab/src/api/billing.rs
View file

@ -1023,29 +1023,6 @@ async fn handle_customer_subscription_event(
.get_billing_subscription_by_stripe_subscription_id(&subscription.id)
.await?
{
let llm_db = app
.llm_db
.clone()
.ok_or_else(|| anyhow!("LLM DB not initialized"))?;
let new_period_start_at =
chrono::DateTime::from_timestamp(subscription.current_period_start, 0)
.ok_or_else(|| anyhow!("No subscription period start"))?;
let new_period_end_at =
chrono::DateTime::from_timestamp(subscription.current_period_end, 0)
.ok_or_else(|| anyhow!("No subscription period end"))?;
llm_db
.transfer_existing_subscription_usage(
billing_customer.user_id,
&existing_subscription,
subscription_kind,
subscription.status.into(),
new_period_start_at,
new_period_end_at,
)
.await?;
app.db
.update_billing_subscription(
existing_subscription.id,

110
crates/collab/src/llm/db/queries/subscription_usages.rs
View file

@ -1,8 +1,6 @@
use chrono::Timelike;
use time::PrimitiveDateTime;
use crate::db::billing_subscription::{StripeSubscriptionStatus, SubscriptionKind};
use crate::db::{UserId, billing_subscription};
use crate::db::UserId;
use super::*;
@ -26,62 +24,6 @@ pub fn convert_chrono_to_time(datetime: DateTimeUtc) -> anyhow::Result<Primitive
}
impl LlmDatabase {
pub async fn create_subscription_usage(
&self,
user_id: UserId,
period_start_at: DateTimeUtc,
period_end_at: DateTimeUtc,
plan: SubscriptionKind,
model_requests: i32,
edit_predictions: i32,
) -> Result<subscription_usage::Model> {
self.transaction(|tx| async move {
self.create_subscription_usage_in_tx(
user_id,
period_start_at,
period_end_at,
plan,
model_requests,
edit_predictions,
&tx,
)
.await
})
.await
}
async fn create_subscription_usage_in_tx(
&self,
user_id: UserId,
period_start_at: DateTimeUtc,
period_end_at: DateTimeUtc,
plan: SubscriptionKind,
model_requests: i32,
edit_predictions: i32,
tx: &DatabaseTransaction,
) -> Result<subscription_usage::Model> {
// Clear out the nanoseconds so that these timestamps are comparable with Unix timestamps.
let period_start_at = period_start_at.with_nanosecond(0).unwrap();
let period_end_at = period_end_at.with_nanosecond(0).unwrap();
let period_start_at = convert_chrono_to_time(period_start_at)?;
let period_end_at = convert_chrono_to_time(period_end_at)?;
Ok(
subscription_usage::Entity::insert(subscription_usage::ActiveModel {
id: ActiveValue::set(Uuid::now_v7()),
user_id: ActiveValue::set(user_id),
period_start_at: ActiveValue::set(period_start_at),
period_end_at: ActiveValue::set(period_end_at),
plan: ActiveValue::set(plan),
model_requests: ActiveValue::set(model_requests),
edit_predictions: ActiveValue::set(edit_predictions),
})
.exec_with_returning(tx)
.await?,
)
}
pub async fn get_subscription_usage_for_period(
&self,
user_id: UserId,
@ -114,54 +56,4 @@ impl LlmDatabase {
.one(tx)
.await?)
}
pub async fn transfer_existing_subscription_usage(
&self,
user_id: UserId,
existing_subscription: &billing_subscription::Model,
new_subscription_kind: Option<SubscriptionKind>,
new_subscription_status: StripeSubscriptionStatus,
new_period_start_at: DateTimeUtc,
new_period_end_at: DateTimeUtc,
) -> Result<Option<subscription_usage::Model>> {
self.transaction(|tx| async move {
match (existing_subscription.kind, new_subscription_status) {
(Some(SubscriptionKind::ZedProTrial), StripeSubscriptionStatus::Active) => {
let trial_period_start_at = existing_subscription
.current_period_start_at()
.ok_or_else(|| anyhow!("No trial subscription period start"))?;
let trial_period_end_at = existing_subscription
.current_period_end_at()
.ok_or_else(|| anyhow!("No trial subscription period end"))?;
let existing_usage = self
.get_subscription_usage_for_period_in_tx(
user_id,
trial_period_start_at,
trial_period_end_at,
&tx,
)
.await?;
if let Some(existing_usage) = existing_usage {
return Ok(Some(
self.create_subscription_usage_in_tx(
user_id,
new_period_start_at,
new_period_end_at,
new_subscription_kind.unwrap_or(existing_usage.plan),
existing_usage.model_requests,
existing_usage.edit_predictions,
&tx,
)
.await?,
));
}
}
_ => {}
}
Ok(None)
})
.await
}
}

1
crates/collab/src/llm/db/tests.rs
View file

@ -1,5 +1,4 @@
mod provider_tests;
mod subscription_usage_tests;
use gpui::BackgroundExecutor;
use parking_lot::Mutex;

119
crates/collab/src/llm/db/tests/subscription_usage_tests.rs
View file

@ -1,119 +0,0 @@
use chrono::{Duration, Utc};
use pretty_assertions::assert_eq;
use crate::db::billing_subscription::{StripeSubscriptionStatus, SubscriptionKind};
use crate::db::{UserId, billing_subscription};
use crate::llm::db::LlmDatabase;
use crate::test_llm_db;
test_llm_db!(
test_transfer_existing_subscription_usage,
test_transfer_existing_subscription_usage_postgres
);
async fn test_transfer_existing_subscription_usage(db: &mut LlmDatabase) {
// Test when an existing Zed Pro trial subscription is upgraded to Zed Pro.
{
let user_id = UserId(1);
let now = Utc::now();
let trial_period_start_at = now - Duration::days(14);
let trial_period_end_at = now;
let new_period_start_at = now;
let new_period_end_at = now + Duration::days(30);
let existing_subscription = billing_subscription::Model {
kind: Some(SubscriptionKind::ZedProTrial),
stripe_current_period_start: Some(trial_period_start_at.timestamp()),
stripe_current_period_end: Some(trial_period_end_at.timestamp()),
..Default::default()
};
let existing_usage = db
.create_subscription_usage(
user_id,
trial_period_start_at,
trial_period_end_at,
SubscriptionKind::ZedProTrial,
25,
1_000,
)
.await
.unwrap();
let transferred_usage = db
.transfer_existing_subscription_usage(
user_id,
&existing_subscription,
Some(SubscriptionKind::ZedPro),
StripeSubscriptionStatus::Active,
new_period_start_at,
new_period_end_at,
)
.await
.unwrap();
assert!(
transferred_usage.is_some(),
"subscription usage not transferred successfully"
);
let transferred_usage = transferred_usage.unwrap();
assert_eq!(
transferred_usage.model_requests,
existing_usage.model_requests
);
assert_eq!(
transferred_usage.edit_predictions,
existing_usage.edit_predictions
);
}
// Test when an existing Zed Pro trial subscription is canceled.
{
let user_id = UserId(2);
let now = Utc::now();
let trial_period_start_at = now - Duration::days(14);
let trial_period_end_at = now;
let existing_subscription = billing_subscription::Model {
kind: Some(SubscriptionKind::ZedProTrial),
stripe_current_period_start: Some(trial_period_start_at.timestamp()),
stripe_current_period_end: Some(trial_period_end_at.timestamp()),
..Default::default()
};
let _existing_usage = db
.create_subscription_usage(
user_id,
trial_period_start_at,
trial_period_end_at,
SubscriptionKind::ZedProTrial,
25,
1_000,
)
.await
.unwrap();
let transferred_usage = db
.transfer_existing_subscription_usage(
user_id,
&existing_subscription,
Some(SubscriptionKind::ZedPro),
StripeSubscriptionStatus::Canceled,
trial_period_start_at,
trial_period_end_at,
)
.await
.unwrap();
assert!(
transferred_usage.is_none(),
"subscription usage was transferred when it should not have been"
);
}
}