diff --git a/apps/users/src/ai/litellm.rs b/apps/users/src/ai/litellm.rs index 7b7d350..676da08 100644 --- a/apps/users/src/ai/litellm.rs +++ b/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 { + 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 { diff --git a/apps/users/src/handlers/auth.rs b/apps/users/src/handlers/auth.rs index 23702f6..1d0b711 100644 --- a/apps/users/src/handlers/auth.rs +++ b/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); diff --git a/apps/users/src/handlers/internal.rs b/apps/users/src/handlers/internal.rs new file mode 100644 index 0000000..1175e4e --- /dev/null +++ b/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 { + 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, + headers: HeaderMap, + Path(id): Path, +) -> 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() + } + } +} diff --git a/apps/users/src/handlers/mod.rs b/apps/users/src/handlers/mod.rs index 4f6ff91..04b709a 100644 --- a/apps/users/src/handlers/mod.rs +++ b/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; diff --git a/apps/users/src/main.rs b/apps/users/src/main.rs index 8bbc50c..249e948 100644 --- a/apps/users/src/main.rs +++ b/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); diff --git a/crates/db/migrations/20260716500000_users_litellm_key.down.sql b/crates/db/migrations/20260716500000_users_litellm_key.down.sql new file mode 100644 index 0000000..d09a3ca --- /dev/null +++ b/crates/db/migrations/20260716500000_users_litellm_key.down.sql @@ -0,0 +1,5 @@ +BEGIN; + +ALTER TABLE users DROP COLUMN litellm_key; + +COMMIT; diff --git a/crates/db/migrations/20260716500000_users_litellm_key.up.sql b/crates/db/migrations/20260716500000_users_litellm_key.up.sql new file mode 100644 index 0000000..192496d --- /dev/null +++ b/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; diff --git a/crates/db/src/models/user.rs b/crates/db/src/models/user.rs index 74d4214..8c65cdd 100644 --- a/crates/db/src/models/user.rs +++ b/crates/db/src/models/user.rs @@ -22,6 +22,7 @@ pub struct User { pub created_at: DateTime, pub updated_at: DateTime, pub deleted_at: Option>, + pub litellm_key: Option, } @@ -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, sqlx::Error> { + sqlx::query_scalar::<_, Option>( + "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,