Issue each account its own LiteLLM virtual key instead of the shared master key

register() now generates a per-account LiteLLM key (best-effort, non-blocking)
and stores it on the user. New internal endpoint GET /internal/users/{id}/llm-key
lets other services fetch (or lazily backfill) an account's key, authenticated
via the existing X-AI-Service-Key shared secret.

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>

apps/users/src/ai/litellm.rs

@@ -181,6 +181,36 @@ impl LiteLlmClient {
         Ok((content, usage, request_id))
     }
 
+    /// Issues a per-account LiteLLM virtual key via `/key/generate`, scoped to `user_id`.
+    /// Requires `self.api_key` to be the LiteLLM master key.
+    pub async fn generate_key(&self, user_id: uuid::Uuid) -> Result<String, LiteLlmError> {
+        let url = format!("{}/key/generate", self.base_url);
+
+        let mut req = self.client.post(&url).json(&serde_json::json!({
+            "user_id": user_id.to_string(),
+            "key_alias": format!("user-{}", user_id),
+        }));
+        if let Some(key) = &self.api_key {
+            req = req.header("Authorization", format!("Bearer {}", key));
+        }
+
+        let response = req.send().await?;
+        let status = response.status();
+        if !status.is_success() {
+            let body = response.text().await.unwrap_or_default();
+            return Err(LiteLlmError::Api {
+                status: status.as_u16(),
+                body,
+            });
+        }
+
+        let body: Value = response.json().await?;
+        body.get("key")
+            .and_then(|v| v.as_str())
+            .map(|s| s.to_string())
+            .ok_or(LiteLlmError::NoCompletion)
+    }
+
     /// Direct pass-through for callers that want the raw JSON response.
     #[allow(dead_code)]
     pub async fn raw_chat_completion(&self, body: Value) -> Result<Value, LiteLlmError> {

apps/users/src/handlers/auth.rs

@@ -308,6 +308,21 @@ async fn register(
         }
     })?;
 
+    // Issue this account its own LiteLLM virtual key (best-effort — a LiteLLM
+    // outage must not block registration; the internal /llm-key endpoint will
+    // lazily generate one on first use if this fails).
+    match crate::ai::litellm::LiteLlmClient::new() {
+        Ok(client) => match client.generate_key(user.id).await {
+            Ok(key) => {
+                if let Err(e) = db::models::user::UserRepository::set_litellm_key(&state.pool, user.id, &key).await {
+                    tracing::error!(user_id = %user.id, error = %e, "Failed to store LiteLLM key after registration");
+                }
+            }
+            Err(e) => tracing::warn!(user_id = %user.id, error = %e, "Failed to generate LiteLLM key at registration"),
+        },
+        Err(e) => tracing::warn!(error = %e, "LiteLLM client unavailable at registration"),
+    }
+
     // Check if this is a demo account (payment gateway integration)
     let is_demo_account = is_dummy_account_email(&email);
 

apps/users/src/handlers/internal.rs

@@ -0,0 +1,90 @@
+use crate::AppState;
+use axum::{
+    extract::{Path, State},
+    http::{HeaderMap, StatusCode},
+    response::IntoResponse,
+    routing::get,
+    Json, Router,
+};
+use db::models::user::UserRepository;
+use uuid::Uuid;
+
+/// Internal, service-to-service routes (not exposed on the public API gateway).
+/// Callers authenticate with the shared `X-AI-Service-Key` header, same as the
+/// existing `/api/support/tickets/ai/create` endpoint.
+pub fn router() -> Router<AppState> {
+    Router::new().route("/{id}/llm-key", get(get_llm_key))
+}
+
+fn check_service_key(headers: &HeaderMap) -> bool {
+    let expected_key = std::env::var("AI_SERVICE_KEY").unwrap_or_default();
+    let provided_key = headers
+        .get("X-AI-Service-Key")
+        .and_then(|v| v.to_str().ok())
+        .unwrap_or_default();
+    !expected_key.is_empty() && provided_key == expected_key
+}
+
+/// GET /internal/users/{id}/llm-key
+///
+/// Returns this account's per-user LiteLLM virtual key, generating and
+/// persisting one on the fly if it doesn't have one yet (covers accounts
+/// created before this feature existed, or where generation failed at
+/// registration time because LiteLLM was unreachable).
+async fn get_llm_key(
+    State(state): State<AppState>,
+    headers: HeaderMap,
+    Path(id): Path<Uuid>,
+) -> impl IntoResponse {
+    if !check_service_key(&headers) {
+        return (
+            StatusCode::UNAUTHORIZED,
+            Json(serde_json::json!({ "error": "Invalid or missing AI service key" })),
+        )
+            .into_response();
+    }
+
+    match UserRepository::get_litellm_key(&state.pool, id).await {
+        Ok(Some(key)) => (StatusCode::OK, Json(serde_json::json!({ "key": key }))).into_response(),
+        Ok(None) => generate_and_store(&state, id).await,
+        Err(e) => {
+            tracing::error!(user_id = %id, error = %e, "Failed to look up LiteLLM key");
+            (
+                StatusCode::INTERNAL_SERVER_ERROR,
+                Json(serde_json::json!({ "error": "Failed to look up LiteLLM key" })),
+            )
+                .into_response()
+        }
+    }
+}
+
+async fn generate_and_store(state: &AppState, user_id: Uuid) -> axum::response::Response {
+    let client = match crate::ai::litellm::LiteLlmClient::new() {
+        Ok(c) => c,
+        Err(e) => {
+            tracing::error!(error = %e, "LiteLLM client unavailable");
+            return (
+                StatusCode::BAD_GATEWAY,
+                Json(serde_json::json!({ "error": "LiteLLM unavailable" })),
+            )
+                .into_response();
+        }
+    };
+
+    match client.generate_key(user_id).await {
+        Ok(key) => {
+            if let Err(e) = UserRepository::set_litellm_key(&state.pool, user_id, &key).await {
+                tracing::error!(user_id = %user_id, error = %e, "Failed to store LiteLLM key");
+            }
+            (StatusCode::OK, Json(serde_json::json!({ "key": key }))).into_response()
+        }
+        Err(e) => {
+            tracing::error!(user_id = %user_id, error = %e, "Failed to generate LiteLLM key");
+            (
+                StatusCode::BAD_GATEWAY,
+                Json(serde_json::json!({ "error": "Failed to generate LiteLLM key" })),
+            )
+                .into_response()
+        }
+    }
+}

apps/users/src/handlers/mod.rs

@@ -12,6 +12,7 @@ pub mod auth;
 pub mod config;
 pub mod coupons;
 pub mod dashboard;
+pub mod internal;
 pub mod kb;
 pub mod modules;
 pub mod notifications;

apps/users/src/main.rs

@@ -128,6 +128,8 @@ async fn main() {
         .nest("/api/ai", handlers::ai::ai_router())
         // ── Ask Ash AI Credits (Phase 1) ────────────────────────────────────
         .nest("/api/ai-credits", handlers::ai_credits::router())
+
+        .nest("/internal/users", handlers::internal::router())
         .nest("/api/admin/ai-credits", handlers::ai_credits::admin_router())
         .route("/health", get(|| async { "Users OK" }))
         .with_state(state);

crates/db/migrations/20260716500000_users_litellm_key.down.sql

@@ -0,0 +1,5 @@
+BEGIN;
+
+ALTER TABLE users DROP COLUMN litellm_key;
+
+COMMIT;

crates/db/migrations/20260716500000_users_litellm_key.up.sql

@@ -0,0 +1,8 @@
+-- Per-account LiteLLM virtual key, so AI features are billed/tracked per user
+-- instead of every account sharing the LiteLLM master key.
+
+BEGIN;
+
+ALTER TABLE users ADD COLUMN litellm_key TEXT;
+
+COMMIT;

crates/db/src/models/user.rs

@@ -22,6 +22,7 @@ pub struct User {
     pub created_at: DateTime<Utc>,
     pub updated_at: DateTime<Utc>,
     pub deleted_at: Option<DateTime<Utc>>,
+    pub litellm_key: Option<String>,
 }
 
 
@@ -58,7 +59,7 @@ impl UserRepository {
                 email_verified, phone_verified, status,
                 email_verification_token, email_verification_expires_at,
                 reset_password_token, reset_password_expires_at,
-                created_at, updated_at, deleted_at
+                created_at, updated_at, deleted_at, litellm_key
             "#,
         )
         .bind(&payload.first_name)
@@ -78,7 +79,7 @@ impl UserRepository {
                    email_verified, phone_verified, status,
                    email_verification_token, email_verification_expires_at,
                    reset_password_token, reset_password_expires_at,
-                   created_at, updated_at, deleted_at
+                   created_at, updated_at, deleted_at, litellm_key
             FROM users
             WHERE email = $1 AND deleted_at IS NULL
             "#,
@@ -95,7 +96,7 @@ impl UserRepository {
                    email_verified, phone_verified, status,
                    email_verification_token, email_verification_expires_at,
                    reset_password_token, reset_password_expires_at,
-                   created_at, updated_at, deleted_at
+                   created_at, updated_at, deleted_at, litellm_key
             FROM users
             WHERE id = $1 AND deleted_at IS NULL
             "#,
@@ -152,7 +153,7 @@ impl UserRepository {
                    email_verified, phone_verified, status,
                    email_verification_token, email_verification_expires_at,
                    reset_password_token, reset_password_expires_at,
-                   created_at, updated_at, deleted_at
+                   created_at, updated_at, deleted_at, litellm_key
             FROM users
             WHERE email_verification_token = $1 AND deleted_at IS NULL
             "#,
@@ -200,7 +201,7 @@ impl UserRepository {
                    email_verified, phone_verified, status,
                    email_verification_token, email_verification_expires_at,
                    reset_password_token, reset_password_expires_at,
-                   created_at, updated_at, deleted_at
+                   created_at, updated_at, deleted_at, litellm_key
             FROM users
             WHERE reset_password_token = $1 AND deleted_at IS NULL
             "#,
@@ -242,6 +243,26 @@ impl UserRepository {
         Ok(())
     }
 
+    pub async fn set_litellm_key(pool: &PgPool, user_id: Uuid, key: &str) -> Result<(), sqlx::Error> {
+        sqlx::query(
+            "UPDATE users SET litellm_key = $1, updated_at = NOW() WHERE id = $2",
+        )
+        .bind(key)
+        .bind(user_id)
+        .execute(pool)
+        .await?;
+        Ok(())
+    }
+
+    pub async fn get_litellm_key(pool: &PgPool, user_id: Uuid) -> Result<Option<String>, sqlx::Error> {
+        sqlx::query_scalar::<_, Option<String>>(
+            "SELECT litellm_key FROM users WHERE id = $1 AND deleted_at IS NULL",
+        )
+        .bind(user_id)
+        .fetch_one(pool)
+        .await
+    }
+
     pub async fn store_refresh_token(
         pool: &PgPool,
         user_id: Uuid,