Issue each account its own LiteLLM virtual key instead of the shared master key
All checks were successful
build-and-release / build (cron) (push) Successful in 5m1s
build-and-release / build (catering-services) (push) Successful in 9m1s
build-and-release / build (customers) (push) Successful in 9m3s
build-and-release / build (developers) (push) Successful in 9m23s
build-and-release / build (employees) (push) Successful in 10m36s
build-and-release / build (companies) (push) Successful in 10m59s
build-and-release / build (gateway) (push) Successful in 3m26s
build-and-release / build (fitness-trainers) (push) Successful in 8m52s
build-and-release / build (jobs) (push) Successful in 4m46s
build-and-release / build (graphic-designers) (push) Successful in 8m44s
build-and-release / build (job-seekers) (push) Successful in 9m22s
build-and-release / build (makeup-artists) (push) Successful in 8m33s
build-and-release / build (leads) (push) Successful in 10m17s
build-and-release / build (payments) (push) Successful in 8m40s
build-and-release / build (photographers) (push) Successful in 9m36s
build-and-release / build (social-media-managers) (push) Successful in 8m38s
build-and-release / build (tutors) (push) Successful in 8m42s
build-and-release / build (ugc-content-creators) (push) Successful in 7m27s
build-and-release / build (video-editors) (push) Successful in 7m46s
build-and-release / build (users) (push) Successful in 10m0s

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>
This commit is contained in:
Ashwin Kumar Sivakumar 2026-07-16 20:29:00 +05:30
parent 74664176da
commit f5201965d8
8 changed files with 177 additions and 5 deletions

View file

@ -181,6 +181,36 @@ impl LiteLlmClient {
Ok((content, usage, request_id)) 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. /// Direct pass-through for callers that want the raw JSON response.
#[allow(dead_code)] #[allow(dead_code)]
pub async fn raw_chat_completion(&self, body: Value) -> Result<Value, LiteLlmError> { pub async fn raw_chat_completion(&self, body: Value) -> Result<Value, LiteLlmError> {

View file

@ -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) // Check if this is a demo account (payment gateway integration)
let is_demo_account = is_dummy_account_email(&email); let is_demo_account = is_dummy_account_email(&email);

View file

@ -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()
}
}
}

View file

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

View file

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

View file

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

View file

@ -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;

View file

@ -22,6 +22,7 @@ pub struct User {
pub created_at: DateTime<Utc>, pub created_at: DateTime<Utc>,
pub updated_at: DateTime<Utc>, pub updated_at: DateTime<Utc>,
pub deleted_at: Option<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_verified, phone_verified, status,
email_verification_token, email_verification_expires_at, email_verification_token, email_verification_expires_at,
reset_password_token, reset_password_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) .bind(&payload.first_name)
@ -78,7 +79,7 @@ impl UserRepository {
email_verified, phone_verified, status, email_verified, phone_verified, status,
email_verification_token, email_verification_expires_at, email_verification_token, email_verification_expires_at,
reset_password_token, reset_password_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 FROM users
WHERE email = $1 AND deleted_at IS NULL WHERE email = $1 AND deleted_at IS NULL
"#, "#,
@ -95,7 +96,7 @@ impl UserRepository {
email_verified, phone_verified, status, email_verified, phone_verified, status,
email_verification_token, email_verification_expires_at, email_verification_token, email_verification_expires_at,
reset_password_token, reset_password_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 FROM users
WHERE id = $1 AND deleted_at IS NULL WHERE id = $1 AND deleted_at IS NULL
"#, "#,
@ -152,7 +153,7 @@ impl UserRepository {
email_verified, phone_verified, status, email_verified, phone_verified, status,
email_verification_token, email_verification_expires_at, email_verification_token, email_verification_expires_at,
reset_password_token, reset_password_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 FROM users
WHERE email_verification_token = $1 AND deleted_at IS NULL WHERE email_verification_token = $1 AND deleted_at IS NULL
"#, "#,
@ -200,7 +201,7 @@ impl UserRepository {
email_verified, phone_verified, status, email_verified, phone_verified, status,
email_verification_token, email_verification_expires_at, email_verification_token, email_verification_expires_at,
reset_password_token, reset_password_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 FROM users
WHERE reset_password_token = $1 AND deleted_at IS NULL WHERE reset_password_token = $1 AND deleted_at IS NULL
"#, "#,
@ -242,6 +243,26 @@ impl UserRepository {
Ok(()) 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( pub async fn store_refresh_token(
pool: &PgPool, pool: &PgPool,
user_id: Uuid, user_id: Uuid,