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

collab: Return increased limit for extended trials from GET /billing/usage (#29771)

Browse source 
This PR updates the `GET /billing/usage` endpoint to return the
increased usage limit for users in the extended trial.

Release Notes:

- N/A
This commit is contained in:
Marshall Bowers 2025-05-01 22:31:30 -04:00 • committed by GitHub
parent 51f1998107
commit 0c0a4ed866
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
1 changed files with 17 additions and 2 deletions
Download patch file Download diff file Expand all files Collapse all files

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

@ -27,7 +27,9 @@ use crate::db::billing_subscription::{
StripeCancellationReason, StripeSubscriptionStatus, SubscriptionKind,
};
use crate::llm::db::subscription_usage_meter::CompletionMode;
use crate::llm::{DEFAULT_MAX_MONTHLY_SPEND, FREE_TIER_MONTHLY_SPENDING_LIMIT};
use crate::llm::{
AGENT_EXTENDED_TRIAL_FEATURE_FLAG, DEFAULT_MAX_MONTHLY_SPEND, FREE_TIER_MONTHLY_SPENDING_LIMIT,
};
use crate::rpc::{ResultExt as _, Server};
use crate::{AppState, Cents, Error, Result};
use crate::{db::UserId, llm::db::LlmDatabase};
@ -1257,8 +1259,21 @@ async fn get_current_usage(
SubscriptionKind::ZedFree => zed_llm_client::Plan::Free,
};
let feature_flags = app.db.get_user_flags(user.id).await?;
let has_extended_trial = feature_flags
.iter()
.any(|flag| flag == AGENT_EXTENDED_TRIAL_FEATURE_FLAG);
let model_requests_limit = match plan.model_requests_limit() {
zed_llm_client::UsageLimit::Limited(limit) => Some(limit),
zed_llm_client::UsageLimit::Limited(limit) => {
let limit = if plan == zed_llm_client::Plan::ZedProTrial && has_extended_trial {
1_000
} else {
limit
};
Some(limit)
}
zed_llm_client::UsageLimit::Unlimited => None,
};
let edit_prediction_limit = match plan.edit_predictions_limit() {