collab: Add lifetime spending limit for LLM usage (#16780)

Marshall Bowers created 1 year ago

This PR adds a lifetime spending limit on LLM usage.

Exceeding this limit will prevent further use of the Zed LLM provider.

Currently the cap is $1,000.

Release Notes:

- N/A

Change summary

crates/collab/src/llm.rs | 12 ++++++++++++
1 file changed, 12 insertions(+)

Detailed changes

crates/collab/src/llm.rs

@@ -411,6 +411,11 @@ fn normalize_model_name(known_models: Vec<String>, name: String) -> String {
     }
 }
 
+/// The maximum lifetime spending an individual user can reach before being cut off.
+///
+/// Represented in cents.
+const LIFETIME_SPENDING_LIMIT_IN_CENTS: usize = 1_000 * 100;
+
 async fn check_usage_limit(
     state: &Arc<LlmState>,
     provider: LanguageModelProvider,
@@ -428,6 +433,13 @@ async fn check_usage_limit(
         )
         .await?;
 
+    if usage.lifetime_spending >= LIFETIME_SPENDING_LIMIT_IN_CENTS {
+        return Err(Error::http(
+            StatusCode::FORBIDDEN,
+            "Maximum spending limit reached.".to_string(),
+        ));
+    }
+
     let active_users = state.get_active_user_count(provider, model_name).await?;
 
     let users_in_recent_minutes = active_users.users_in_recent_minutes.max(1);