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

collab: Rename symbols for existing Stripe synchronization (#29570)

Browse source 
This PR renames the symbols for the existing Stripe synchronization.

This will make things clearer once the new synchronization job for the
new billing is added.

Release Notes:

- N/A
This commit is contained in:
Marshall Bowers 2025-04-28 18:37:18 -04:00 • committed by GitHub
parent 1d7c86bf0d
commit 9abeedf0c6
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
3 changed files with 11 additions and 9 deletions
Download patch file Download diff file Expand all files Collapse all files

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

@ -1243,9 +1243,9 @@ async fn find_or_create_billing_customer(
Ok(Some(billing_customer)) Ok(Some(billing_customer))
} }
const SYNC_LLM_USAGE_WITH_STRIPE_INTERVAL: Duration = Duration::from_secs(60); const SYNC_LLM_TOKEN_USAGE_WITH_STRIPE_INTERVAL: Duration = Duration::from_secs(60);
pub fn sync_llm_usage_with_stripe_periodically(app: Arc<AppState>) { pub fn sync_llm_token_usage_with_stripe_periodically(app: Arc<AppState>) {
let Some(stripe_billing) = app.stripe_billing.clone() else { let Some(stripe_billing) = app.stripe_billing.clone() else {
log::warn!("failed to retrieve Stripe billing object"); log::warn!("failed to retrieve Stripe billing object");
return; return;
@ -1260,17 +1260,19 @@ pub fn sync_llm_usage_with_stripe_periodically(app: Arc<AppState>) {
let executor = executor.clone(); let executor = executor.clone();
async move { async move {
loop { loop {
sync_with_stripe(&app, &llm_db, &stripe_billing) sync_token_usage_with_stripe(&app, &llm_db, &stripe_billing)
.await .await
.context("failed to sync LLM usage to Stripe") .context("failed to sync LLM usage to Stripe")
.trace_err(); .trace_err();
executor.sleep(SYNC_LLM_USAGE_WITH_STRIPE_INTERVAL).await; executor
.sleep(SYNC_LLM_TOKEN_USAGE_WITH_STRIPE_INTERVAL)
.await;
} }
} }
}); });
} }
async fn sync_with_stripe( async fn sync_token_usage_with_stripe(
app: &Arc<AppState>, app: &Arc<AppState>,
llm_db: &Arc<LlmDatabase>, llm_db: &Arc<LlmDatabase>,
stripe_billing: &Arc<StripeBilling>, stripe_billing: &Arc<StripeBilling>,
@ -1306,7 +1308,7 @@ async fn sync_with_stripe(
.subscribe_to_model(&stripe_subscription_id, &stripe_model) .subscribe_to_model(&stripe_subscription_id, &stripe_model)
.await?; .await?;
stripe_billing stripe_billing
.bill_model_usage(&stripe_customer_id, &stripe_model, &event) .bill_model_token_usage(&stripe_customer_id, &stripe_model, &event)
.await?; .await?;
llm_db.consume_billing_event(event.id).await?; llm_db.consume_billing_event(event.id).await?;
} }

4
crates/collab/src/main.rs
View file

@ -8,7 +8,7 @@ use axum::{
}; };
use collab::api::CloudflareIpCountryHeader; use collab::api::CloudflareIpCountryHeader;
use collab::api::billing::sync_llm_usage_with_stripe_periodically; use collab::api::billing::sync_llm_token_usage_with_stripe_periodically;
use collab::llm::db::LlmDatabase; use collab::llm::db::LlmDatabase;
use collab::migrations::run_database_migrations; use collab::migrations::run_database_migrations;
use collab::user_backfiller::spawn_user_backfiller; use collab::user_backfiller::spawn_user_backfiller;
@ -152,7 +152,7 @@ async fn main() -> Result<()> {
if let Some(mut llm_db) = llm_db { if let Some(mut llm_db) = llm_db {
llm_db.initialize().await?; llm_db.initialize().await?;
sync_llm_usage_with_stripe_periodically(state.clone()); sync_llm_token_usage_with_stripe_periodically(state.clone());
} }
app = app app = app

2
crates/collab/src/stripe_billing.rs
View file

@ -268,7 +268,7 @@ impl StripeBilling {
Ok(()) Ok(())
} }
pub async fn bill_model_usage( pub async fn bill_model_token_usage(
&self, &self,
customer_id: &stripe::CustomerId, customer_id: &stripe::CustomerId,
model: &StripeModel, model: &StripeModel,