diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index 5da517b88b3..20bb696c46c 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -3820,6 +3820,7 @@ dependencies = [ "futures-util", "http 1.4.2", "litellm-cost", + "litellm-gateway-auth", "litellm-model-catalog", "memchr", "rstest", diff --git a/litellm-rust/crates/db/.sqlx/query-aebecf84ec08b203635d59ae2a616bd622e7ed0102166fed8b09b452741c107d.json b/litellm-rust/crates/db/.sqlx/query-08aa4e832176f266c2200c676ece399bc28aaeab7ee4da8e53f2ce87a52a3da9.json similarity index 56% rename from litellm-rust/crates/db/.sqlx/query-aebecf84ec08b203635d59ae2a616bd622e7ed0102166fed8b09b452741c107d.json rename to litellm-rust/crates/db/.sqlx/query-08aa4e832176f266c2200c676ece399bc28aaeab7ee4da8e53f2ce87a52a3da9.json index 0283a5854a0..1a9ad366cfa 100644 --- a/litellm-rust/crates/db/.sqlx/query-aebecf84ec08b203635d59ae2a616bd622e7ed0102166fed8b09b452741c107d.json +++ b/litellm-rust/crates/db/.sqlx/query-08aa4e832176f266c2200c676ece399bc28aaeab7ee4da8e53f2ce87a52a3da9.json @@ -1,6 +1,6 @@ { "db_name": "PostgreSQL", - "query": "SELECT v.token,\n (EXTRACT(EPOCH FROM v.expires) * 1000)::bigint AS \"expires_ms?\",\n v.blocked,\n v.models,\n v.spend,\n v.max_budget,\n v.allowed_routes,\n v.key_alias,\n v.key_name,\n v.metadata,\n (v.rpm_limit IS NOT NULL OR v.tpm_limit IS NOT NULL OR v.tpd_limit IS NOT NULL\n OR v.max_parallel_requests IS NOT NULL\n OR kb.rpm_limit IS NOT NULL OR kb.tpm_limit IS NOT NULL\n OR kb.tpd_limit IS NOT NULL OR kb.max_parallel_requests IS NOT NULL\n OR COALESCE(jsonb_path_exists(v.model_max_budget,\n '$.* ? (@.rpm_limit != null || @.tpm_limit != null)'), false)\n OR COALESCE(jsonb_path_exists(kb.model_max_budget,\n '$.* ? (@.rpm_limit != null || @.tpm_limit != null)'), false))\n AS \"key_rate_limits!\",\n COALESCE(cardinality(v.policies) > 0, false) AS \"key_policies!\",\n (v.aliases NOT IN ('{}', 'null')\n OR COALESCE(v.router_settings NOT IN ('{}', 'null'), false)) AS \"key_model_aliases!\",\n (v.object_permission_id IS NOT NULL) AS \"key_object_permissions!\",\n COALESCE(cardinality(v.allowed_cache_controls) > 0, false) AS \"key_cache_controls!\",\n (v.project_id IS NOT NULL) AS \"key_project!\",\n (v.organization_id IS NOT NULL) AS \"key_organization!\",\n (v.agent_id IS NOT NULL) AS \"key_agent!\",\n v.team_id,\n t.team_id AS \"team_row_id?\",\n t.team_alias AS \"team_alias?\",\n t.blocked AS \"team_blocked?\",\n t.models AS \"team_models?\",\n t.spend AS \"team_spend?\",\n t.max_budget AS \"team_max_budget?\",\n t.metadata AS \"team_metadata?\",\n (t.rpm_limit IS NOT NULL OR t.tpm_limit IS NOT NULL OR t.tpd_limit IS NOT NULL\n OR t.max_parallel_requests IS NOT NULL\n OR COALESCE(jsonb_path_exists(t.model_max_budget,\n '$.* ? (@.rpm_limit != null || @.tpm_limit != null)'), false))\n AS \"team_rate_limits!\",\n COALESCE(cardinality(t.policies) > 0, false) AS \"team_policies!\",\n (t.model_id IS NOT NULL\n OR COALESCE(t.router_settings NOT IN ('{}', 'null'), false)) AS \"team_model_aliases!\",\n (t.object_permission_id IS NOT NULL) AS \"team_object_permissions!\",\n (t.organization_id IS NOT NULL) AS \"team_organization!\",\n (mb.rpm_limit IS NOT NULL OR mb.tpm_limit IS NOT NULL OR mb.tpd_limit IS NOT NULL\n OR mb.max_parallel_requests IS NOT NULL\n OR COALESCE(jsonb_path_exists(mb.model_max_budget,\n '$.* ? (@.rpm_limit != null || @.tpm_limit != null)'), false))\n AS \"member_rate_limits!\",\n mb.allowed_models AS \"member_allowed_models?\",\n v.user_id,\n u.user_id AS \"user_row_id?\",\n u.models AS \"user_models?\",\n (u.rpm_limit IS NOT NULL OR u.tpm_limit IS NOT NULL\n OR u.max_parallel_requests IS NOT NULL\n OR COALESCE(jsonb_path_exists(u.model_max_budget,\n '$.* ? (@.rpm_limit != null || @.tpm_limit != null)'), false))\n AS \"user_rate_limits!\",\n COALESCE(cardinality(u.policies) > 0, false) AS \"user_policies!\",\n (u.object_permission_id IS NOT NULL) AS \"user_object_permissions!\"\n FROM \"LiteLLM_VerificationToken\" AS v\n LEFT JOIN \"LiteLLM_BudgetTable\" AS kb ON v.budget_id = kb.budget_id\n LEFT JOIN \"LiteLLM_TeamTable\" AS t ON v.team_id = t.team_id\n LEFT JOIN \"LiteLLM_TeamMembership\" AS tm\n ON v.team_id = tm.team_id AND v.user_id = tm.user_id\n LEFT JOIN \"LiteLLM_BudgetTable\" AS mb ON tm.budget_id = mb.budget_id\n LEFT JOIN \"LiteLLM_UserTable\" AS u ON v.user_id = u.user_id\n WHERE v.token = $1", + "query": "SELECT v.token,\n (EXTRACT(EPOCH FROM v.expires) * 1000)::bigint AS \"expires_ms?\",\n v.blocked,\n v.models,\n v.spend,\n COALESCE(v.max_budget, kb.max_budget) AS max_budget,\n v.allowed_routes,\n v.key_alias,\n v.key_name,\n v.metadata,\n (v.rpm_limit IS NOT NULL OR v.tpm_limit IS NOT NULL OR v.tpd_limit IS NOT NULL\n OR v.max_parallel_requests IS NOT NULL\n OR kb.rpm_limit IS NOT NULL OR kb.tpm_limit IS NOT NULL\n OR kb.tpd_limit IS NOT NULL OR kb.max_parallel_requests IS NOT NULL\n OR COALESCE(jsonb_path_exists(v.model_max_budget,\n '$.* ? (@.rpm_limit != null || @.tpm_limit != null)'), false)\n OR COALESCE(jsonb_path_exists(kb.model_max_budget,\n '$.* ? (@.rpm_limit != null || @.tpm_limit != null)'), false))\n AS \"key_rate_limits!\",\n COALESCE(cardinality(v.policies) > 0, false) AS \"key_policies!\",\n (v.aliases NOT IN ('{}', 'null')\n OR COALESCE(v.router_settings NOT IN ('{}', 'null'), false)) AS \"key_model_aliases!\",\n (v.object_permission_id IS NOT NULL) AS \"key_object_permissions!\",\n COALESCE(cardinality(v.allowed_cache_controls) > 0, false) AS \"key_cache_controls!\",\n (v.project_id IS NOT NULL) AS \"key_project!\",\n (v.organization_id IS NOT NULL) AS \"key_organization!\",\n (v.agent_id IS NOT NULL) AS \"key_agent!\",\n v.team_id,\n t.team_id AS \"team_row_id?\",\n t.team_alias AS \"team_alias?\",\n t.blocked AS \"team_blocked?\",\n t.models AS \"team_models?\",\n t.spend AS \"team_spend?\",\n t.max_budget AS \"team_max_budget?\",\n t.metadata AS \"team_metadata?\",\n (t.rpm_limit IS NOT NULL OR t.tpm_limit IS NOT NULL OR t.tpd_limit IS NOT NULL\n OR t.max_parallel_requests IS NOT NULL\n OR COALESCE(jsonb_path_exists(t.model_max_budget,\n '$.* ? (@.rpm_limit != null || @.tpm_limit != null)'), false))\n AS \"team_rate_limits!\",\n COALESCE(cardinality(t.policies) > 0, false) AS \"team_policies!\",\n (t.model_id IS NOT NULL\n OR COALESCE(t.router_settings NOT IN ('{}', 'null'), false)) AS \"team_model_aliases!\",\n (t.object_permission_id IS NOT NULL) AS \"team_object_permissions!\",\n (t.organization_id IS NOT NULL) AS \"team_organization!\",\n (mb.rpm_limit IS NOT NULL OR mb.tpm_limit IS NOT NULL OR mb.tpd_limit IS NOT NULL\n OR mb.max_parallel_requests IS NOT NULL\n OR COALESCE(jsonb_path_exists(mb.model_max_budget,\n '$.* ? (@.rpm_limit != null || @.tpm_limit != null)'), false))\n AS \"member_rate_limits!\",\n mb.allowed_models AS \"member_allowed_models?\",\n tm.spend AS \"member_spend?\",\n mb.max_budget AS \"member_max_budget?\",\n CASE WHEN mb.temp_budget_expiry > (now() AT TIME ZONE 'UTC')\n THEN mb.temp_budget_increase END AS \"member_budget_increase?\",\n tdb.max_budget AS \"member_default_budget?\",\n v.user_id,\n u.user_id AS \"user_row_id?\",\n u.models AS \"user_models?\",\n u.spend AS \"user_spend?\",\n u.max_budget AS \"user_max_budget?\",\n (u.rpm_limit IS NOT NULL OR u.tpm_limit IS NOT NULL\n OR u.max_parallel_requests IS NOT NULL\n OR COALESCE(jsonb_path_exists(u.model_max_budget,\n '$.* ? (@.rpm_limit != null || @.tpm_limit != null)'), false))\n AS \"user_rate_limits!\",\n COALESCE(cardinality(u.policies) > 0, false) AS \"user_policies!\",\n (u.object_permission_id IS NOT NULL) AS \"user_object_permissions!\"\n FROM \"LiteLLM_VerificationToken\" AS v\n LEFT JOIN \"LiteLLM_BudgetTable\" AS kb ON v.budget_id = kb.budget_id\n LEFT JOIN \"LiteLLM_TeamTable\" AS t ON v.team_id = t.team_id\n LEFT JOIN \"LiteLLM_TeamMembership\" AS tm\n ON v.team_id = tm.team_id AND v.user_id = tm.user_id\n LEFT JOIN \"LiteLLM_BudgetTable\" AS mb ON tm.budget_id = mb.budget_id\n LEFT JOIN \"LiteLLM_BudgetTable\" AS tdb\n ON tdb.budget_id = t.metadata ->> 'team_member_budget_id'\n LEFT JOIN \"LiteLLM_UserTable\" AS u ON v.user_id = u.user_id\n WHERE v.token = $1", "describe": { "columns": [ { @@ -57,12 +57,7 @@ "ordinal": 5, "name": "max_budget", "type_info": "Float8", - "origin": { - "Table": { - "table": "\"LiteLLM_VerificationToken\"", - "name": "max_budget" - } - } + "origin": "Expression" }, { "ordinal": 6, @@ -293,6 +288,45 @@ }, { "ordinal": 33, + "name": "member_spend?", + "type_info": "Float8", + "origin": { + "Table": { + "table": "\"LiteLLM_TeamMembership\"", + "name": "spend" + } + } + }, + { + "ordinal": 34, + "name": "member_max_budget?", + "type_info": "Float8", + "origin": { + "Table": { + "table": "\"LiteLLM_BudgetTable\"", + "name": "max_budget" + } + } + }, + { + "ordinal": 35, + "name": "member_budget_increase?", + "type_info": "Float8", + "origin": "Expression" + }, + { + "ordinal": 36, + "name": "member_default_budget?", + "type_info": "Float8", + "origin": { + "Table": { + "table": "\"LiteLLM_BudgetTable\"", + "name": "max_budget" + } + } + }, + { + "ordinal": 37, "name": "user_id", "type_info": "Text", "origin": { @@ -303,7 +337,7 @@ } }, { - "ordinal": 34, + "ordinal": 38, "name": "user_row_id?", "type_info": "Text", "origin": { @@ -314,7 +348,7 @@ } }, { - "ordinal": 35, + "ordinal": 39, "name": "user_models?", "type_info": "TextArray", "origin": { @@ -325,19 +359,41 @@ } }, { - "ordinal": 36, + "ordinal": 40, + "name": "user_spend?", + "type_info": "Float8", + "origin": { + "Table": { + "table": "\"LiteLLM_UserTable\"", + "name": "spend" + } + } + }, + { + "ordinal": 41, + "name": "user_max_budget?", + "type_info": "Float8", + "origin": { + "Table": { + "table": "\"LiteLLM_UserTable\"", + "name": "max_budget" + } + } + }, + { + "ordinal": 42, "name": "user_rate_limits!", "type_info": "Bool", "origin": "Expression" }, { - "ordinal": 37, + "ordinal": 43, "name": "user_policies!", "type_info": "Bool", "origin": "Expression" }, { - "ordinal": 38, + "ordinal": 44, "name": "user_object_permissions!", "type_info": "Bool", "origin": "Expression" @@ -354,7 +410,7 @@ true, true, false, - true, + null, true, true, true, @@ -384,11 +440,17 @@ true, true, true, + null, + true, + true, + true, + true, + true, true, null, null, null ] }, - "hash": "aebecf84ec08b203635d59ae2a616bd622e7ed0102166fed8b09b452741c107d" + "hash": "08aa4e832176f266c2200c676ece399bc28aaeab7ee4da8e53f2ce87a52a3da9" } diff --git a/litellm-rust/crates/db/.sqlx/query-a907e2963ee0cf25b34af9ed477f92fc7030de81c903adb593a5eca47237cfbe.json b/litellm-rust/crates/db/.sqlx/query-a907e2963ee0cf25b34af9ed477f92fc7030de81c903adb593a5eca47237cfbe.json new file mode 100644 index 00000000000..c9805198db5 --- /dev/null +++ b/litellm-rust/crates/db/.sqlx/query-a907e2963ee0cf25b34af9ed477f92fc7030de81c903adb593a5eca47237cfbe.json @@ -0,0 +1,60 @@ +{ + "db_name": "PostgreSQL", + "query": "SELECT e.spend AS \"spend?\",\n e.blocked AS \"blocked?\",\n b.max_budget AS \"max_budget?\",\n (SELECT d.max_budget FROM \"LiteLLM_BudgetTable\" AS d\n WHERE d.budget_id = $2) AS \"default_max_budget?\"\n FROM (SELECT 1) AS one\n LEFT JOIN \"LiteLLM_EndUserTable\" AS e ON e.user_id = $1\n LEFT JOIN \"LiteLLM_BudgetTable\" AS b ON b.budget_id = e.budget_id", + "describe": { + "columns": [ + { + "ordinal": 0, + "name": "spend?", + "type_info": "Float8", + "origin": { + "Table": { + "table": "\"LiteLLM_EndUserTable\"", + "name": "spend" + } + } + }, + { + "ordinal": 1, + "name": "blocked?", + "type_info": "Bool", + "origin": { + "Table": { + "table": "\"LiteLLM_EndUserTable\"", + "name": "blocked" + } + } + }, + { + "ordinal": 2, + "name": "max_budget?", + "type_info": "Float8", + "origin": { + "Table": { + "table": "\"LiteLLM_BudgetTable\"", + "name": "max_budget" + } + } + }, + { + "ordinal": 3, + "name": "default_max_budget?", + "type_info": "Float8", + "origin": "Expression" + } + ], + "parameters": { + "Left": [ + "Text", + "Text" + ] + }, + "nullable": [ + true, + true, + true, + null + ] + }, + "hash": "a907e2963ee0cf25b34af9ed477f92fc7030de81c903adb593a5eca47237cfbe" +} diff --git a/litellm-rust/crates/db/src/end_user.rs b/litellm-rust/crates/db/src/end_user.rs new file mode 100644 index 00000000000..a4d98eef2aa --- /dev/null +++ b/litellm-rust/crates/db/src/end_user.rs @@ -0,0 +1,40 @@ +use sqlx::PgPool; + +use crate::Error; + +/// An end user's spend and budget, as the proxy's end-user budget checks read them. +#[derive(Clone, Debug, Default, PartialEq)] +pub struct EndUser { + /// `LiteLLM_EndUserTable.spend`, when the end user has a row. + pub spend: Option, + pub blocked: bool, + /// The row's budget, else the default budget's. + pub max_budget: Option, +} + +/// Reads `end_user_id`, with `default_budget_id` as the budget of an end user without its own. +pub async fn find_end_user( + pool: &PgPool, + end_user_id: &str, + default_budget_id: Option<&str>, +) -> Result { + let row = sqlx::query!( + r#"SELECT e.spend AS "spend?", + e.blocked AS "blocked?", + b.max_budget AS "max_budget?", + (SELECT d.max_budget FROM "LiteLLM_BudgetTable" AS d + WHERE d.budget_id = $2) AS "default_max_budget?" + FROM (SELECT 1) AS one + LEFT JOIN "LiteLLM_EndUserTable" AS e ON e.user_id = $1 + LEFT JOIN "LiteLLM_BudgetTable" AS b ON b.budget_id = e.budget_id"#, + end_user_id, + default_budget_id, + ) + .fetch_one(pool) + .await?; + Ok(EndUser { + spend: row.spend, + blocked: row.blocked.unwrap_or(false), + max_budget: row.max_budget.or(row.default_max_budget), + }) +} diff --git a/litellm-rust/crates/db/src/lib.rs b/litellm-rust/crates/db/src/lib.rs index 39ef58241b2..b41ee5ace78 100644 --- a/litellm-rust/crates/db/src/lib.rs +++ b/litellm-rust/crates/db/src/lib.rs @@ -1,9 +1,11 @@ +mod end_user; mod error; mod failure; mod pool; mod spend; mod verification_token; +pub use end_user::{EndUser, find_end_user}; pub use error::Error; pub use failure::DbFailure; pub use pool::{PoolSettings, connect_lazy}; diff --git a/litellm-rust/crates/db/src/verification_token.rs b/litellm-rust/crates/db/src/verification_token.rs index aa48a881767..3ff008b0cff 100644 --- a/litellm-rust/crates/db/src/verification_token.rs +++ b/litellm-rust/crates/db/src/verification_token.rs @@ -41,9 +41,15 @@ pub struct Team { pub restrictions: Restrictions, } -/// What a team membership's budget table sets beyond its dollar budget. -#[derive(Clone, Debug, Default, Eq, PartialEq)] +/// The key's team membership: its spend and budget, and what its budget table sets beyond them. +#[derive(Clone, Debug, Default, PartialEq)] pub struct TeamMembership { + /// The membership row's spend, when the key's user has a row for its team. + pub spend: Option, + /// The membership's own budget, with any temporary increase still in effect. + pub max_budget: Option, + /// The team's default member budget, `metadata.team_member_budget_id`. + pub default_max_budget: Option, /// `rpm_limit`, `tpm_limit`, `tpd_limit` or `max_parallel_requests`. pub rate_limits: bool, /// The member's own model scope; empty inherits the team's. @@ -52,6 +58,8 @@ pub struct TeamMembership { #[derive(Clone, Debug, PartialEq)] pub struct User { + pub spend: f64, + pub max_budget: Option, pub models: Vec, pub restrictions: Restrictions, } @@ -85,7 +93,7 @@ pub async fn find_verification_token( v.blocked, v.models, v.spend, - v.max_budget, + COALESCE(v.max_budget, kb.max_budget) AS max_budget, v.allowed_routes, v.key_alias, v.key_name, @@ -131,9 +139,16 @@ pub async fn find_verification_token( '$.* ? (@.rpm_limit != null || @.tpm_limit != null)'), false)) AS "member_rate_limits!", mb.allowed_models AS "member_allowed_models?", + tm.spend AS "member_spend?", + mb.max_budget AS "member_max_budget?", + CASE WHEN mb.temp_budget_expiry > (now() AT TIME ZONE 'UTC') + THEN mb.temp_budget_increase END AS "member_budget_increase?", + tdb.max_budget AS "member_default_budget?", v.user_id, u.user_id AS "user_row_id?", u.models AS "user_models?", + u.spend AS "user_spend?", + u.max_budget AS "user_max_budget?", (u.rpm_limit IS NOT NULL OR u.tpm_limit IS NOT NULL OR u.max_parallel_requests IS NOT NULL OR COALESCE(jsonb_path_exists(u.model_max_budget, @@ -147,6 +162,8 @@ pub async fn find_verification_token( LEFT JOIN "LiteLLM_TeamMembership" AS tm ON v.team_id = tm.team_id AND v.user_id = tm.user_id LEFT JOIN "LiteLLM_BudgetTable" AS mb ON tm.budget_id = mb.budget_id + LEFT JOIN "LiteLLM_BudgetTable" AS tdb + ON tdb.budget_id = t.metadata ->> 'team_member_budget_id' LEFT JOIN "LiteLLM_UserTable" AS u ON v.user_id = u.user_id WHERE v.token = $1"#, token @@ -192,11 +209,18 @@ pub async fn find_verification_token( }, }), team_membership: TeamMembership { + spend: row.member_spend, + max_budget: row + .member_max_budget + .map(|budget| budget + row.member_budget_increase.unwrap_or(0.0)), + default_max_budget: row.member_default_budget, rate_limits: row.member_rate_limits, allowed_models: row.member_allowed_models.unwrap_or_default(), }, user_id: row.user_id, user: row.user_row_id.map(|_| User { + spend: row.user_spend.unwrap_or(0.0), + max_budget: row.user_max_budget, models: row.user_models.unwrap_or_default(), restrictions: Restrictions { rate_limits: row.user_rate_limits, diff --git a/litellm-rust/crates/db/tests/verification_token.rs b/litellm-rust/crates/db/tests/verification_token.rs index 169429f4839..ffe0b4d66f6 100644 --- a/litellm-rust/crates/db/tests/verification_token.rs +++ b/litellm-rust/crates/db/tests/verification_token.rs @@ -1,7 +1,8 @@ use std::time::{Duration, SystemTime}; use litellm_db::{ - Restrictions, Team, TeamMembership, User, VerificationToken, find_verification_token, + EndUser, Restrictions, Team, TeamMembership, User, VerificationToken, find_end_user, + find_verification_token, }; use litellm_db_testing::MigratedPostgres; use rstest::rstest; @@ -21,9 +22,12 @@ INSERT INTO "LiteLLM_BudgetTable" (budget_id, max_budget, rpm_limit, model_max_b VALUES ('b-1', 5.0, NULL, NULL, ARRAY['gpt-a'], 'x', 'x'), ('b-2', NULL, 3, NULL, ARRAY[]::text[], 'x', 'x'), ('b-3', NULL, NULL, '{"gpt-a": {"tpm_limit": 7}}', ARRAY[]::text[], 'x', 'x'); -INSERT INTO "LiteLLM_TeamMembership" (user_id, team_id, budget_id) - VALUES ('user-2', 'team-2', 'b-1'), - ('user-1', 'team-1', 'b-2'); +INSERT INTO "LiteLLM_TeamMembership" (user_id, team_id, budget_id, spend) + VALUES ('user-2', 'team-2', 'b-1', 0.75), + ('user-1', 'team-1', 'b-2', 0); +INSERT INTO "LiteLLM_EndUserTable" (user_id, spend, budget_id, blocked) + VALUES ('end-1', 1.5, 'b-1', false), + ('end-2', 0.25, NULL, true); INSERT INTO "LiteLLM_OrganizationTable" (organization_id, organization_alias, budget_id, created_by, updated_by) VALUES ('org-1', 'org', 'b-1', 'x', 'x'); INSERT INTO "LiteLLM_ProjectTable" (project_id, created_by, updated_by) VALUES ('project-1', 'x', 'x'); @@ -58,19 +62,20 @@ INSERT INTO "LiteLLM_VerificationToken" '{"gpt-a": {"budget_limit": 1.0, "rpm_limit": 2}}'); "#; +#[expect( + clippy::disallowed_methods, + reason = "test fixture rows; the lookups under test use the checked macro" +)] +async fn setup(pool: &sqlx::PgPool) { + sqlx::raw_sql(FIXTURES).execute(pool).await.unwrap(); +} + #[rstest] #[tokio::test] async fn reads_the_gating_columns_of_a_verification_token_and_its_owners() { let postgres = MigratedPostgres::start().await.unwrap(); - #[expect( - clippy::disallowed_methods, - reason = "test fixture rows; the lookup under test uses the checked macro" - )] - sqlx::raw_sql(FIXTURES) - .execute(postgres.pool()) - .await - .unwrap(); let pool = postgres.pool(); + setup(pool).await; let find = |token: &'static str| async move { find_verification_token(pool, token).await.unwrap().unwrap() }; @@ -125,11 +130,16 @@ async fn reads_the_gating_columns_of_a_verification_token_and_its_owners() { }, }), team_membership: TeamMembership { + spend: Some(0.0), + max_budget: None, + default_max_budget: None, rate_limits: true, allowed_models: vec![], }, user_id: Some("user-1".into()), user: Some(User { + spend: 0.0, + max_budget: Some(1.0), models: vec!["gpt-b".into()], restrictions: Restrictions { rate_limits: true, @@ -175,10 +185,16 @@ async fn reads_the_gating_columns_of_a_verification_token_and_its_owners() { assert_eq!( member.team_membership, TeamMembership { + spend: Some(0.75), + max_budget: Some(5.0), + default_max_budget: None, rate_limits: false, allowed_models: vec!["gpt-a".into()], } ); + // A key without its own budget takes its budget row's. + assert_eq!(find("restricted").await.max_budget, Some(5.0)); + assert_eq!(find("budgeted").await.max_budget, Some(4.0)); assert_eq!(member.team.unwrap().restrictions, Restrictions::default()); assert_eq!(member.user.unwrap().restrictions, Restrictions::default()); let orphan = find("orphan").await; @@ -193,3 +209,39 @@ async fn reads_the_gating_columns_of_a_verification_token_and_its_owners() { None ); } + +#[rstest] +#[tokio::test] +async fn reads_end_user_budgets_with_their_default() { + let postgres = MigratedPostgres::start().await.unwrap(); + let pool = postgres.pool(); + setup(pool).await; + let find = + async |id: &str, default: Option<&str>| find_end_user(pool, id, default).await.unwrap(); + + assert_eq!( + find("end-1", Some("b-2")).await, + EndUser { + spend: Some(1.5), + blocked: false, + max_budget: Some(5.0) + } + ); + assert_eq!( + find("end-2", Some("b-1")).await, + EndUser { + spend: Some(0.25), + blocked: true, + max_budget: Some(5.0) + } + ); + assert_eq!(find("end-2", None).await.max_budget, None); + assert_eq!( + find("unknown", Some("b-1")).await, + EndUser { + spend: None, + blocked: false, + max_budget: Some(5.0) + } + ); +} diff --git a/litellm-rust/crates/gateway-auth/src/keys.rs b/litellm-rust/crates/gateway-auth/src/keys.rs index 2ce5338a674..ff99572a1ee 100644 --- a/litellm-rust/crates/gateway-auth/src/keys.rs +++ b/litellm-rust/crates/gateway-auth/src/keys.rs @@ -297,8 +297,46 @@ pub struct KeyProfile { pub team_alias: Option, pub metadata: Arc>, pub team_metadata: Arc>, + pub budgets: KeyBudgets, } +/// Stored spend against an optional budget, as the key store last read it. +#[derive(Clone, Copy, Debug, Default, PartialEq)] +pub struct Spend { + pub spend: f64, + pub max_budget: Option, +} + +impl Eq for Spend {} + +/// The spend and budgets of a key and of who it belongs to, read with the key and refreshed +/// with it. +#[derive(Clone, Debug, Default, Eq, PartialEq)] +pub struct KeyBudgets { + /// The key's `max_budget`, else its budget table's. + pub key: Spend, + /// `key_name`, the key's masked form, which budget errors show while it still looks masked. + pub key_name: Option, + /// The key's team, when its row exists. + pub team: Option, + /// The key's user, when its row exists. + pub user: Option, + /// The key's user within its team, when a membership row exists. + pub team_member: Option, + /// `metadata.end_user_budget_id`: the budget of end users without their own. + pub end_user_budget_id: Option, +} + +/// A team member's spend, with the membership's own budget and the team's default one. +#[derive(Clone, Copy, Debug, Default, PartialEq)] +pub struct MemberSpend { + pub spend: f64, + pub max_budget: Option, + pub default_max_budget: Option, +} + +impl Eq for MemberSpend {} + #[derive(Clone, Debug, Eq, PartialEq)] pub enum KeyStatus { Active { diff --git a/litellm-rust/crates/gateway-hooks/Cargo.toml b/litellm-rust/crates/gateway-hooks/Cargo.toml index b0d6c67f15f..e072cd72c7a 100644 --- a/litellm-rust/crates/gateway-hooks/Cargo.toml +++ b/litellm-rust/crates/gateway-hooks/Cargo.toml @@ -10,6 +10,7 @@ bytes.workspace = true futures-util.workspace = true http.workspace = true litellm-cost.workspace = true +litellm-gateway-auth.workspace = true litellm-model-catalog.workspace = true memchr = "2" serde.workspace = true diff --git a/litellm-rust/crates/gateway-hooks/src/facts.rs b/litellm-rust/crates/gateway-hooks/src/facts.rs index 3e7e2033e13..437768707d0 100644 --- a/litellm-rust/crates/gateway-hooks/src/facts.rs +++ b/litellm-rust/crates/gateway-hooks/src/facts.rs @@ -8,6 +8,8 @@ use bytes::Bytes; use http::HeaderMap; use serde_json::{Map, Value}; +use litellm_gateway_auth::keys::KeyBudgets; + use crate::usage::Usage; /// The facts a hook reads or an observer asks for. The gateway builds only what some hook or @@ -101,6 +103,7 @@ pub struct KeyFacts { pub org_id: Option, pub metadata: Arc>, pub team_metadata: Arc>, + pub budgets: KeyBudgets, } /// What the gateway knows about a request before it is sent, shared by every hook and observer. diff --git a/litellm-rust/crates/gateway-hooks/src/hook.rs b/litellm-rust/crates/gateway-hooks/src/hook.rs index 0392cd0f88f..3e57e6b12c8 100644 --- a/litellm-rust/crates/gateway-hooks/src/hook.rs +++ b/litellm-rust/crates/gateway-hooks/src/hook.rs @@ -12,7 +12,10 @@ use futures_util::future::BoxFuture; use http::{HeaderName, HeaderValue}; use serde_json::{Map, Value}; -use crate::facts::{Annotations, Facts, Record, RequestFacts}; +use crate::{ + facts::{Annotations, Facts, Record, RequestFacts}, + pricing::Prices, +}; /// The parts of a request a hook may change. #[derive(Clone, Copy, Debug, Default, Eq, PartialEq)] @@ -93,6 +96,8 @@ pub struct Refusal { pub status: u16, /// The proxy's error `type`, such as `budget_exceeded`. pub kind: Cow<'static, str>, + /// The error's `param`: the string `"None"` for most proxy errors, `null` for some. + pub param: Option<&'static str>, pub message: String, pub headers: Vec<(HeaderName, HeaderValue)>, } @@ -106,11 +111,18 @@ impl Refusal { Self { status, kind: kind.into(), + param: Some("None"), message: message.into(), headers: Vec::new(), } } + /// Sends a `null` `param`, as the proxy's budget errors do. + pub fn without_param(mut self) -> Self { + self.param = None; + self + } + pub fn with_header(mut self, name: HeaderName, value: HeaderValue) -> Self { self.headers.push((name, value)); self @@ -229,6 +241,9 @@ pub trait Hook: Send + Sync + 'static { /// The parts of the request this hook may change. fn writes(&self) -> Targets; + /// Called once with the deployments' prices before the first request. + fn prepare(&self, _prices: &Arc) {} + /// Runs before the request is routed. Returning a [`Settle`] asks to be told how the request /// ended. fn run<'a>( diff --git a/litellm-rust/crates/gateway-hooks/src/pipeline.rs b/litellm-rust/crates/gateway-hooks/src/pipeline.rs index 384babaf8d2..04d7076a85e 100644 --- a/litellm-rust/crates/gateway-hooks/src/pipeline.rs +++ b/litellm-rust/crates/gateway-hooks/src/pipeline.rs @@ -24,6 +24,7 @@ impl Pipeline { let mut reads = observers.reads(); for hook in &hooks { reads |= hook.reads(); + hook.prepare(&prices); } // Cost is computed from usage and the deployment that served the request. if reads.contains(Facts::COST) { diff --git a/litellm-rust/crates/gateway-hooks/src/pricing.rs b/litellm-rust/crates/gateway-hooks/src/pricing.rs index 813676c6f5b..be1bca921ba 100644 --- a/litellm-rust/crates/gateway-hooks/src/pricing.rs +++ b/litellm-rust/crates/gateway-hooks/src/pricing.rs @@ -17,12 +17,16 @@ use crate::usage::{ServiceTier, Usage}; #[derive(Debug, Default)] pub struct Prices { plans: HashMap, PricingPlan>, + max_output_tokens: HashMap, u64>, } impl Prices { /// Adds a deployment's prices from its cost map fields. A deployment whose fields hold no /// usable prices is left out, and its cost is unknown. pub fn insert(&mut self, id: Arc, fields: &Map) { + if let Some(limit) = fields.get("max_output_tokens").and_then(Value::as_u64) { + self.max_output_tokens.insert(id.clone(), limit); + } let parsed = ParsedRates::parse(fields); if parsed.standard == Rates::EMPTY && parsed.thresholds.is_empty() { return; @@ -43,6 +47,11 @@ impl Prices { self.plans.is_empty() } + /// The most output tokens the deployment's model can return, when the cost map says. + pub fn max_output_tokens(&self, deployment: &str) -> Option { + self.max_output_tokens.get(deployment).copied() + } + /// The cost in dollars of `usage` on a deployment, or `None` when its prices are unknown. pub fn cost(&self, deployment: &str, usage: &Usage) -> Option { let plan = self.plans.get(deployment)?; diff --git a/litellm-rust/crates/gateway-inference/src/error.rs b/litellm-rust/crates/gateway-inference/src/error.rs index ce4b5aed0ac..05009604fba 100644 --- a/litellm-rust/crates/gateway-inference/src/error.rs +++ b/litellm-rust/crates/gateway-inference/src/error.rs @@ -180,7 +180,7 @@ fn refused(refusal: Refusal) -> Response { let body = json!({"error": { "message": refusal.message, "type": refusal.kind, - "param": "None", + "param": refusal.param, "code": status.as_u16().to_string(), }}); let mut response = (status, Json(body)).into_response(); diff --git a/litellm-rust/crates/gateway-inference/src/hooks.rs b/litellm-rust/crates/gateway-inference/src/hooks.rs index 408810bf095..e60bcd140bc 100644 --- a/litellm-rust/crates/gateway-inference/src/hooks.rs +++ b/litellm-rust/crates/gateway-inference/src/hooks.rs @@ -120,6 +120,7 @@ fn key_facts(hooks: &Hooks, identity: &AuthenticatedRequest) -> Option Some(Arc::new(KeyFacts { diff --git a/litellm-rust/crates/gateway/src/budgets.rs b/litellm-rust/crates/gateway/src/budgets.rs new file mode 100644 index 00000000000..be04ae028be --- /dev/null +++ b/litellm-rust/crates/gateway/src/budgets.rs @@ -0,0 +1,1209 @@ +//! The budgets hook: refuses a request when its key, team, user, team member or end user is over +//! budget, with the proxy's checks, order and messages, and keeps the proxy's `spend:*` counters +//! in Redis so gateway replicas and Python proxies sharing that Redis see each other's spend. +//! +//! Each counter is read with the stored spend as its floor: a missing counter is seeded from it +//! and one below it is raised to it, as the proxy reseeds and repairs counters. A request then +//! reserves its most expensive possible cost on every counter with a budget, and settles the +//! reservation to its real cost when it ends, so concurrent requests cannot all slip under one +//! budget. Reading, checking and reserving is one Redis script, atomic per request. Without +//! Redis, or while it is unreachable, the counters live in this replica's memory. + +use std::{ + collections::HashMap, + sync::{Arc, Mutex, OnceLock}, + time::{Duration, Instant}, +}; + +use futures_util::future::BoxFuture; +use litellm_config::Config; +use litellm_db::{EndUser, PgPool}; +use litellm_gateway_hooks::{ + Facts, Hook, Prices, Record, Refusal, RequestEdit, RequestFacts, Settle, Targets, Usage, +}; +use redis::aio::ConnectionManager; + +use crate::coordination::Coordination; + +/// The output tokens the proxy assumes when neither the request nor the model caps them. +const DEFAULT_OUTPUT_TOKENS: u64 = 16_384; +/// The proxy's slack when a reservation exactly meets a budget. +const EPSILON: f64 = 1e-12; +/// How long an end user's budget is reused before it is read again. +const END_USER_TTL: Duration = Duration::from_secs(60); +/// End users held in memory; the cache starts over past this. +const MAX_END_USERS: usize = 10_000; +const FAIL_CLOSED_READ: &str = "{'error': 'Budget enforcement unavailable: current spend could not be verified against Redis or the database, and fail_closed_budget_enforcement is enabled, so the request was rejected to avoid exceeding the configured budget. Retry shortly.'}"; + +/// Seeds or raises each counter to its floor, runs the budget checks in the order given and +/// reserves the estimate on every counter that takes a reservation, shrinking it to what is +/// left when a counter would go over, as the proxy's reservation does. +const ADMIT: &str = r#" +local ttl = tonumber(ARGV[1]) +local estimate = tonumber(ARGV[2]) +local n = #KEYS +local current, max, mode, reserve = {}, {}, {}, {} +for i = 1, n do + local base = 3 + (i - 1) * 4 + local floor = ARGV[base + 1] + max[i] = tonumber(ARGV[base + 2]) + mode[i] = ARGV[base + 3] + reserve[i] = ARGV[base + 4] == '1' + local raw = redis.call('GET', KEYS[i]) + local value = raw and tonumber(raw) or nil + if value == nil or value < tonumber(floor) then + redis.call('SET', KEYS[i], floor, 'EX', ttl) + raw = floor + end + current[i] = raw +end +for index in string.gmatch(ARGV[3], '%d+') do + local i = tonumber(index) + local value = tonumber(current[i]) + if (mode[i] == '1' and value >= max[i]) or (mode[i] == '2' and value > max[i]) then + return {'check', i, current[i]} + end +end +local charged = {} +local applied = {} +local cost = estimate +for i = 1, n do + charged[i] = 0 + if reserve[i] and cost > 0 then + local raw = redis.call('INCRBYFLOAT', KEYS[i], cost) + redis.call('EXPIRE', KEYS[i], ttl) + charged[i] = cost + applied[#applied + 1] = i + local post = tonumber(raw) + if post > max[i] then + local remaining = max[i] - (post - cost) + if remaining <= 1e-12 then + for _, j in ipairs(applied) do + redis.call('INCRBYFLOAT', KEYS[j], -charged[j]) + end + return {'reserve', i, raw} + end + for _, j in ipairs(applied) do + redis.call('INCRBYFLOAT', KEYS[j], -(charged[j] - remaining)) + charged[j] = remaining + end + cost = remaining + end + end +end +local out = {'ok'} +for i = 1, n do + out[#out + 1] = string.format('%.17g', charged[i]) +end +return out +"#; + +/// Adds each request's settled cost: seeds a missing counter from its floor (charging the whole +/// cost, since its reservation was lost), raises one below its floor, then applies the change. +const SETTLE: &str = r#" +local ttl = tonumber(ARGV[1]) +for i, key in ipairs(KEYS) do + local base = 1 + (i - 1) * 3 + local floor, delta, fresh = ARGV[base + 1], ARGV[base + 2], ARGV[base + 3] + local raw = redis.call('GET', key) + local value = raw and tonumber(raw) or nil + if value == nil then + redis.call('SET', key, floor) + delta = fresh + elseif value < tonumber(floor) then + redis.call('SET', key, floor) + end + if tonumber(delta) ~= 0 then + redis.call('INCRBYFLOAT', key, delta) + end + redis.call('EXPIRE', key, ttl) +end +return #KEYS +"#; + +#[derive(Clone, Debug, PartialEq)] +pub struct BudgetSettings { + /// `litellm_settings.budget_exceeded_status_code`, 422 by default. + pub status: u16, + /// Off with `general_settings.disable_budget_reservation`. + pub reservations: bool, + /// `general_settings.apply_user_budget_to_team_keys`: a user's budget also covers their + /// team keys, not only their personal ones. + pub user_budget_on_team_keys: bool, + /// `litellm_settings.max_end_user_budget_id`: the budget of end users without their own. + pub end_user_budget_id: Option, + /// `general_settings.fail_closed_budget_enforcement`: refuse rather than fall back to this + /// replica's counters while Redis is unreachable. + pub fail_closed: bool, + /// `default_redis_ttl` (in `litellm_settings` or `cache_params`), 60 s by default. + pub ttl: Duration, +} + +impl BudgetSettings { + pub fn from_config(config: &Config) -> Self { + let general = &config.general_settings.additional_fields; + let litellm = &config.litellm_settings; + let ttl = litellm + .additional_fields + .get("default_redis_ttl") + .or_else(|| { + litellm + .cache_params + .as_ref() + .and_then(|params| params.get("default_redis_ttl")) + }) + .and_then(number) + .filter(|seconds| *seconds >= 1.0) + .unwrap_or(60.0); + Self { + status: litellm + .additional_fields + .get("budget_exceeded_status_code") + .and_then(number) + .and_then(|status| u16::try_from(status as i64).ok()) + .filter(|status| (400..600).contains(status)) + .unwrap_or(422), + reservations: !flag(general.get("disable_budget_reservation")), + user_budget_on_team_keys: flag(general.get("apply_user_budget_to_team_keys")), + end_user_budget_id: litellm + .additional_fields + .get("max_end_user_budget_id") + .and_then(litellm_config::Value::as_str) + .filter(|id| !id.is_empty()) + .map(str::to_owned), + fail_closed: flag(general.get("fail_closed_budget_enforcement")), + ttl: Duration::from_secs(ttl as u64), + } + } +} + +fn number(value: &litellm_config::Value) -> Option { + match value { + litellm_config::Value::Number(number) => number.as_f64(), + litellm_config::Value::String(text) => text.trim().parse().ok(), + _ => None, + } +} + +fn flag(value: Option<&litellm_config::Value>) -> bool { + match value { + Some(litellm_config::Value::Bool(flag)) => *flag, + Some(litellm_config::Value::String(text)) => text.eq_ignore_ascii_case("true"), + _ => false, + } +} + +/// When a counter refuses at check time. +#[derive(Clone, Copy, Debug, PartialEq)] +enum Check { + None, + /// `spend >= max_budget`. + AtLeast, + /// `spend > max_budget`, for teams and end users. + Above, +} + +/// Whose counter it is, which names it in refusals. +#[derive(Clone, Debug, PartialEq)] +enum Owner { + Key { + hash: String, + descriptor: String, + }, + Team(String), + User(String), + /// The team member check the proxy runs while building the key: a positive budget. + Member { + user_id: String, + team_id: String, + }, + /// The proxy's later team member check, which also refuses on a zero budget. + ZeroBudgetMember { + user_id: String, + team_id: String, + }, + EndUser(String), +} + +impl Owner { + /// The counter, before the namespace, as the proxy names it. + fn counter(&self) -> String { + match self { + Self::Key { hash, .. } => format!("spend:key:{hash}"), + Self::Team(id) => format!("spend:team:{id}"), + Self::User(id) => format!("spend:user:{id}"), + Self::Member { user_id, team_id } | Self::ZeroBudgetMember { user_id, team_id } => { + format!("spend:team_member:{user_id}:{team_id}") + } + Self::EndUser(id) => format!("spend:end_user:{id}"), + } + } + + /// The auth-time refusal message. + fn check_message(&self, spend: f64, max_budget: f64) -> String { + let (spend, max_budget) = (python_float(spend), python_float(max_budget)); + match self { + Self::Key { descriptor, .. } => format!( + "Budget has been exceeded! Key={descriptor} Current cost: {spend}, Max budget: {max_budget}" + ), + Self::Team(id) => format!( + "Budget has been exceeded! Team={id} Current cost: {spend}, Max budget: {max_budget}" + ), + Self::User(id) => { + format!("ExceededBudget: User={id} over budget. Spend={spend}, Budget={max_budget}") + } + Self::Member { user_id, team_id } => format!( + "Budget has been exceeded! TeamMember={user_id}:{team_id} Current cost: {spend}, Max budget: {max_budget}" + ), + Self::ZeroBudgetMember { user_id, team_id } => format!( + "Budget has been exceeded! User={user_id} in Team={team_id} Current cost: {spend}, Max budget: {max_budget}" + ), + Self::EndUser(id) => format!( + "ExceededBudget: End User={id} over budget. Spend={spend}, Budget={max_budget}" + ), + } + } + + /// The reservation's refusal message, which names keys by their hash. + fn reservation_message(&self, spend: f64, max_budget: f64) -> String { + let owner = match self { + Self::Key { hash, .. } => format!("Key={hash}"), + Self::Team(id) => format!("Team={id}"), + Self::User(id) => format!("User={id}"), + Self::Member { user_id, team_id } | Self::ZeroBudgetMember { user_id, team_id } => { + format!("TeamMember={user_id}:{team_id}") + } + Self::EndUser(id) => format!("EndUser={id}"), + }; + format!( + "Budget has been exceeded! {owner} Current cost: {}, Max budget: {}", + python_float(spend), + python_float(max_budget) + ) + } +} + +/// One spend counter a request touches. +#[derive(Clone, Debug, PartialEq)] +struct Counter { + owner: Owner, + /// The stored spend, below which the counter is never read. + floor: f64, + max_budget: Option, + check: Check, + reserve: bool, +} + +/// The counters a request touches, in the proxy's reservation order, and the order its checks +/// run in. +#[derive(Debug, Default, PartialEq)] +struct Plan { + counters: Vec, + checks: Vec, +} + +/// How admission went. +#[derive(Debug, PartialEq)] +enum Admitted { + /// What was reserved on each counter. + Charged(Vec), + /// A check refused, with the counter's spend. + Over(usize, f64), + /// The reservation could not fit, with the counter's spend after it. + Reserved(usize, f64), +} + +/// Counters in this replica's memory: the proxy's in-memory spend counters, used without Redis +/// and while it is unreachable. +#[derive(Default)] +struct Local(Mutex>); + +impl Local { + fn admit(&self, keys: &[String], plan: &Plan, estimate: f64, ttl: Duration) -> Admitted { + let mut counters = self.0.lock().unwrap_or_else(|e| e.into_inner()); + let now = Instant::now(); + if counters.len() > 100_000 { + counters.retain(|_, (_, expires)| *expires > now); + } + let mut current = Vec::with_capacity(keys.len()); + for (key, counter) in keys.iter().zip(&plan.counters) { + let entry = counters + .entry(key.clone()) + .or_insert((counter.floor, now + ttl)); + if entry.1 <= now || entry.0 < counter.floor { + *entry = (counter.floor, now + ttl); + } + current.push(entry.0); + } + for &index in &plan.checks { + let counter = &plan.counters[index]; + if over(counter.check, current[index], counter.max_budget) { + return Admitted::Over(index, current[index]); + } + } + let mut charged = vec![0.0; keys.len()]; + let mut applied = Vec::new(); + let mut cost = estimate; + for (index, counter) in plan.counters.iter().enumerate() { + if !counter.reserve || cost <= 0.0 { + continue; + } + let Some(max_budget) = counter.max_budget else { + continue; + }; + let entry = counters.get_mut(&keys[index]).expect("seeded above"); + entry.0 += cost; + entry.1 = now + ttl; + let post = entry.0; + charged[index] = cost; + applied.push(index); + if post > max_budget { + let remaining = max_budget - (post - cost); + let refund = + |counters: &mut HashMap, index: usize, amount| { + if let Some(entry) = counters.get_mut(&keys[index]) { + entry.0 -= amount; + } + }; + if remaining <= EPSILON { + for &applied in &applied { + refund(&mut counters, applied, charged[applied]); + } + return Admitted::Reserved(index, post); + } + for &applied in &applied { + refund(&mut counters, applied, charged[applied] - remaining); + charged[applied] = remaining; + } + cost = remaining; + } + } + Admitted::Charged(charged) + } + + fn settle(&self, keys: &[String], changes: &[Change], ttl: Duration) { + let mut counters = self.0.lock().unwrap_or_else(|e| e.into_inner()); + let now = Instant::now(); + for (key, change) in keys.iter().zip(changes) { + let entry = counters.entry(key.clone()).or_insert((f64::NAN, now)); + let delta = if entry.0.is_nan() || entry.1 <= now { + entry.0 = change.floor; + change.fresh + } else { + entry.0 = entry.0.max(change.floor); + change.delta + }; + entry.0 += delta; + entry.1 = now + ttl; + } + } +} + +fn over(check: Check, spend: f64, max_budget: Option) -> bool { + match (check, max_budget) { + (Check::AtLeast, Some(max_budget)) => spend >= max_budget, + (Check::Above, Some(max_budget)) => spend > max_budget, + _ => false, + } +} + +/// How a counter changes when a request ends. +#[derive(Clone, Copy, Debug, PartialEq)] +struct Change { + floor: f64, + /// Added to a counter that still exists. + delta: f64, + /// Added to a counter that expired while the request ran. + fresh: f64, +} + +/// Where the counters live. +struct Counters { + redis: Option>, + local: Local, + ttl: Duration, + fail_closed: bool, +} + +/// Redis could not be used and the settings say to refuse rather than fall back. +struct Unavailable; + +impl Counters { + fn connection(&self) -> Option<(ConnectionManager, &Coordination)> { + let redis = self.redis.as_deref()?; + Some((redis.connection()?, redis)) + } + + async fn admit(&self, plan: &Plan, estimate: f64) -> Result<(Admitted, bool), Unavailable> { + let keys: Vec = plan + .counters + .iter() + .map(|counter| counter.owner.counter()) + .collect(); + if self.redis.is_some() { + match self.connection() { + Some((mut connection, redis)) => { + let keys: Vec = keys.iter().map(|key| redis.namespaced(key)).collect(); + match admit_redis(&mut connection, &keys, plan, estimate, self.ttl).await { + Ok(admitted) => return Ok((admitted, true)), + Err(error) => { + tracing::warn!(%error, "reading spend counters failed; using this replica's counters") + } + } + } + None => tracing::debug!("Redis not connected; using this replica's spend counters"), + } + if self.fail_closed { + return Err(Unavailable); + } + } + Ok((self.local.admit(&keys, plan, estimate, self.ttl), false)) + } + + async fn settle(&self, keys: Vec, changes: Vec, shared: bool) { + if shared && let Some((mut connection, redis)) = self.connection() { + let keys: Vec = keys.iter().map(|key| redis.namespaced(key)).collect(); + let mut command = redis::cmd("EVAL"); + command + .arg(SETTLE) + .arg(keys.len()) + .arg(&keys) + .arg(self.ttl.as_secs()); + for change in &changes { + command + .arg(change.floor.to_string()) + .arg(change.delta.to_string()) + .arg(change.fresh.to_string()); + } + if let Err(error) = command.query_async::(&mut connection).await { + // A counter left half-written would mislead every replica; drop it so the next + // read reseeds it from the database, as the proxy does. + tracing::warn!(%error, "updating spend counters failed; dropping them"); + let _: Result = redis::cmd("DEL") + .arg(&keys) + .query_async(&mut connection) + .await; + } + return; + } + self.local.settle(&keys, &changes, self.ttl); + } + + /// Keeps reserved counters alive while a request runs, renewing them every half TTL. + fn lease(&self, keys: Vec) -> Option> { + let (mut connection, redis) = self.connection()?; + let keys: Vec = keys.iter().map(|key| redis.namespaced(key)).collect(); + let ttl = self.ttl; + let (done, mut finished) = tokio::sync::oneshot::channel::<()>(); + tokio::spawn(async move { + let mut renew = tokio::time::interval(ttl / 2); + renew.tick().await; + loop { + tokio::select! { + _ = &mut finished => return, + _ = renew.tick() => { + let mut pipeline = redis::pipe(); + for key in &keys { + pipeline.cmd("EXPIRE").arg(key).arg(ttl.as_secs()).ignore(); + } + let _: Result<(), _> = pipeline.query_async(&mut connection).await; + } + } + } + }); + Some(done) + } +} + +async fn admit_redis( + connection: &mut ConnectionManager, + keys: &[String], + plan: &Plan, + estimate: f64, + ttl: Duration, +) -> redis::RedisResult { + let mut command = redis::cmd("EVAL"); + command + .arg(ADMIT) + .arg(keys.len()) + .arg(keys) + .arg(ttl.as_secs()) + .arg(estimate.to_string()) + .arg( + plan.checks + .iter() + .map(|index| (index + 1).to_string()) + .collect::>() + .join(","), + ); + for counter in &plan.counters { + let mode = match (counter.check, counter.max_budget) { + (_, None) | (Check::None, _) => "0", + (Check::AtLeast, Some(_)) => "1", + (Check::Above, Some(_)) => "2", + }; + command + .arg(counter.floor.to_string()) + .arg( + counter + .max_budget + .map(|max| max.to_string()) + .unwrap_or_default(), + ) + .arg(mode) + .arg(if counter.reserve && counter.max_budget.is_some() { + "1" + } else { + "0" + }); + } + let reply: Vec = command.query_async(connection).await?; + let text = |value: &redis::Value| -> Option { + match value { + redis::Value::BulkString(bytes) => String::from_utf8(bytes.clone()).ok(), + redis::Value::SimpleString(text) => Some(text.clone()), + redis::Value::Int(number) => Some(number.to_string()), + _ => None, + } + }; + let malformed = || { + redis::RedisError::from(( + redis::ErrorKind::UnexpectedReturnType, + "unexpected budget script reply", + )) + }; + let status = reply.first().and_then(text).ok_or_else(malformed)?; + let float = |value: Option<&redis::Value>| { + value + .and_then(text) + .and_then(|value| value.parse::().ok()) + .ok_or_else(malformed) + }; + let index = || { + reply + .get(1) + .and_then(text) + .and_then(|value| value.parse::().ok()) + .and_then(|index| index.checked_sub(1)) + .filter(|index| *index < keys.len()) + .ok_or_else(malformed) + }; + match status.as_str() { + "ok" => reply[1..] + .iter() + .map(|value| float(Some(value))) + .collect::>() + .map(Admitted::Charged), + "check" => Ok(Admitted::Over(index()?, float(reply.get(2))?)), + "reserve" => Ok(Admitted::Reserved(index()?, float(reply.get(2))?)), + _ => Err(malformed()), + } +} + +/// An end user and the default budget it was read with. +type EndUserKey = (String, Option); + +/// End users' budgets, read from the proxy's database and reused for a minute. +struct EndUsers { + pool: PgPool, + cache: Mutex>, +} + +impl EndUsers { + async fn get(&self, id: &str, default_budget_id: Option<&str>) -> Option { + let key = (id.to_owned(), default_budget_id.map(str::to_owned)); + let now = Instant::now(); + if let Some((end_user, read_at)) = self + .cache + .lock() + .unwrap_or_else(|e| e.into_inner()) + .get(&key) + && now.duration_since(*read_at) < END_USER_TTL + { + return Some(end_user.clone()); + } + match litellm_db::find_end_user(&self.pool, id, default_budget_id).await { + Ok(end_user) => { + let mut cache = self.cache.lock().unwrap_or_else(|e| e.into_inner()); + if cache.len() >= MAX_END_USERS { + cache.clear(); + } + cache.insert(key, (end_user.clone(), now)); + Some(end_user) + } + Err(error) => { + tracing::warn!(%error, "reading an end user's budget failed; not checking it"); + None + } + } + } +} + +/// The `budgets` hook. +pub struct Budgets { + settings: BudgetSettings, + counters: Arc, + end_users: Option, + /// Each model group's deployment ids, for pricing a request's most expensive outcome. + groups: HashMap>>, + prices: OnceLock>, +} + +impl Budgets { + /// `deployments` pairs each deployment's model group with its id. `redis`, when given, holds + /// the shared counters; `database` gives end users' budgets. + pub fn new( + settings: BudgetSettings, + deployments: impl IntoIterator)>, + redis: Option>, + database: Option, + ) -> Self { + let mut groups: HashMap>> = HashMap::new(); + for (group, id) in deployments { + groups.entry(group).or_default().push(id); + } + Self { + counters: Arc::new(Counters { + redis, + local: Local::default(), + ttl: settings.ttl, + fail_closed: settings.fail_closed, + }), + settings, + end_users: database.map(|pool| EndUsers { + pool, + cache: Mutex::default(), + }), + groups, + prices: OnceLock::new(), + } + } + + /// The hook for a gateway: its deployments, the coordination Redis and the proxy database. + pub fn for_gateway( + config: &Config, + inference: &litellm_gateway_inference::Gateway, + redis: Option>, + database: Option, + ) -> Self { + let deployments: Vec<(String, Arc)> = config + .model_list + .iter() + .map(|model| model.model_name.clone()) + .zip(inference.models.deployment_ids().map(Arc::from)) + .collect(); + Self::new( + BudgetSettings::from_config(config), + deployments, + redis, + database, + ) + } + + /// Every deployment of the requested group has explicit zero prices: the proxy skips budget + /// checks for free models. + fn free(&self, request: &RequestFacts) -> bool { + let (Some(prices), Some(ids)) = ( + self.prices.get(), + self.groups.get(&*request.requested_model), + ) else { + return false; + }; + let usage = Usage { + input_tokens: 1_000_000, + output_tokens: 1_000_000, + ..Usage::default() + }; + !ids.is_empty() && ids.iter().all(|id| prices.cost(id, &usage) == Some(0.0)) + } + + /// The proxy's estimate of a request's most expensive cost: its input, and as much output + /// as it may ask for, priced on the group's most expensive deployment. `None` when a + /// deployment's prices are unknown. + fn estimate(&self, request: &RequestFacts) -> Option { + let prices = self.prices.get()?; + let ids = self.groups.get(&*request.requested_model)?; + let input_tokens = request.input_tokens(); + let mut most = None::; + for id in ids { + let output_tokens = request.max_tokens.unwrap_or(DEFAULT_OUTPUT_TOKENS).min( + prices + .max_output_tokens(id) + .unwrap_or(DEFAULT_OUTPUT_TOKENS), + ); + let cost = prices.cost( + id, + &Usage { + input_tokens, + output_tokens, + ..Usage::default() + }, + )?; + most = Some(most.map_or(cost, |most| most.max(cost))); + } + most + } + + async fn plan(&self, request: &RequestFacts) -> Plan { + let mut plan = Plan::default(); + let Some(key) = request.key.as_deref() else { + return plan; + }; + let budgets = &key.budgets; + let reserve = |max_budget: Option| { + self.settings.reservations && max_budget.is_some_and(|max| max > 0.0 && max.is_finite()) + }; + let finite = |max_budget: Option| max_budget.filter(|max| max.is_finite()); + let push = |plan: &mut Plan, counter: Counter| { + plan.counters.push(counter); + plan.counters.len() - 1 + }; + // Reservation order: key, team, user, end user, team member. + let mut key_index = None; + if !key.hash.is_empty() { + let max_budget = finite(budgets.key.max_budget); + key_index = Some(push( + &mut plan, + Counter { + owner: Owner::Key { + hash: key.hash.clone(), + descriptor: key_descriptor( + key.alias.as_deref(), + budgets.key_name.as_deref(), + ), + }, + floor: budgets.key.spend, + max_budget, + check: Check::AtLeast, + reserve: reserve(max_budget), + }, + )); + } + let mut team_index = None; + if let Some(team_id) = &key.team_id { + let team = budgets.team.unwrap_or_default(); + let max_budget = finite(team.max_budget); + team_index = Some(push( + &mut plan, + Counter { + owner: Owner::Team(team_id.clone()), + floor: team.spend, + max_budget, + check: Check::Above, + reserve: reserve(max_budget), + }, + )); + } + let mut user_index = None; + if let Some(user_id) = &key.user_id { + let user = budgets.user.unwrap_or_default(); + let applies = key.team_id.is_none() || self.settings.user_budget_on_team_keys; + let max_budget = finite(user.max_budget).filter(|_| applies); + user_index = Some(push( + &mut plan, + Counter { + owner: Owner::User(user_id.clone()), + floor: user.spend, + max_budget, + check: Check::AtLeast, + reserve: reserve(max_budget), + }, + )); + } + let mut end_user_index = None; + if let Some(end_user_id) = request.end_user.as_deref() { + let default_budget_id = budgets + .end_user_budget_id + .as_deref() + .or(self.settings.end_user_budget_id.as_deref()); + let end_user = match &self.end_users { + Some(end_users) => end_users.get(end_user_id, default_budget_id).await, + None => None, + } + .unwrap_or_default(); + let max_budget = finite(end_user.max_budget); + end_user_index = Some(push( + &mut plan, + Counter { + owner: Owner::EndUser(end_user_id.to_owned()), + floor: end_user.spend.unwrap_or(0.0), + max_budget, + // The proxy's check needs an end-user row; its reservation does not. + check: match end_user.spend { + Some(_) => Check::Above, + None => Check::None, + }, + reserve: reserve(max_budget), + }, + )); + } + let mut member_index = None; + if let (Some(user_id), Some(team_id)) = (&key.user_id, &key.team_id) { + let member = budgets.team_member; + let max_budget = member.and_then(|member| { + finite(member.max_budget) + .or_else(|| finite(member.default_max_budget).filter(|max| *max > 0.0)) + }); + let (owner, check) = match max_budget { + Some(max) if max <= 0.0 => ( + Owner::ZeroBudgetMember { + user_id: user_id.clone(), + team_id: team_id.clone(), + }, + Check::AtLeast, + ), + _ => ( + Owner::Member { + user_id: user_id.clone(), + team_id: team_id.clone(), + }, + Check::AtLeast, + ), + }; + member_index = Some(push( + &mut plan, + Counter { + owner, + floor: member.map_or(0.0, |member| member.spend), + max_budget, + check, + reserve: reserve(max_budget), + }, + )); + } + // Check order: the positive member budget and the key while the proxy builds the key, + // then team, user, a zero member budget and the end user in its common checks. + let positive_member = member_index + .filter(|index| matches!(plan.counters[*index].owner, Owner::Member { .. })); + let zero_member = member_index + .filter(|index| matches!(plan.counters[*index].owner, Owner::ZeroBudgetMember { .. })); + plan.checks = [ + positive_member, + key_index, + team_index, + user_index, + zero_member, + end_user_index, + ] + .into_iter() + .flatten() + .filter(|index| plan.counters[*index].max_budget.is_some()) + .collect(); + plan + } + + fn refusal(&self, message: String) -> Refusal { + Refusal::new(self.settings.status, "budget_exceeded", message).without_param() + } +} + +impl Hook for Budgets { + fn name(&self) -> &'static str { + "budgets" + } + + fn reads(&self) -> Facts { + Facts::KEY | Facts::END_USER | Facts::INPUT_TOKENS | Facts::COST + } + + fn writes(&self) -> Targets { + Targets::ADMISSION + } + + fn prepare(&self, prices: &Arc) { + let _ = self.prices.set(prices.clone()); + } + + fn run<'a>( + &'a self, + request: &'a Arc, + _edit: &'a mut RequestEdit, + ) -> BoxFuture<'a, Result>, Refusal>> { + Box::pin(async move { + let plan = self.plan(request).await; + if plan.counters.is_empty() { + return Ok(None); + } + let checked = !self.free(request); + let reserving = checked && plan.counters.iter().any(|counter| counter.reserve); + let estimate = match reserving { + true => self.estimate(request).filter(|cost| *cost > 0.0), + false => None, + }; + let charged = if checked && (!plan.checks.is_empty() || estimate.is_some()) { + let (admitted, shared) = self + .counters + .admit(&plan, estimate.unwrap_or(0.0)) + .await + .map_err(|Unavailable| Refusal::new(503, "auth_error", FAIL_CLOSED_READ))?; + match admitted { + Admitted::Charged(charged) => Some((charged, shared)), + Admitted::Over(index, spend) => { + let counter = &plan.counters[index]; + let max_budget = counter.max_budget.unwrap_or_default(); + return Err(self.refusal(counter.owner.check_message(spend, max_budget))); + } + Admitted::Reserved(index, spend) => { + let counter = &plan.counters[index]; + let max_budget = counter.max_budget.unwrap_or_default(); + return Err( + self.refusal(counter.owner.reservation_message(spend, max_budget)) + ); + } + } + } else { + None + }; + let shared = charged + .as_ref() + .map_or(self.counters.redis.is_some(), |(_, shared)| *shared); + let charged = + charged.map_or_else(|| vec![0.0; plan.counters.len()], |(charged, _)| charged); + let reserved: Vec = plan + .counters + .iter() + .zip(&charged) + .filter(|(_, charged)| **charged != 0.0) + .map(|(counter, _)| counter.owner.counter()) + .collect(); + let lease = match shared && !reserved.is_empty() { + true => self.counters.lease(reserved), + false => None, + }; + Ok(Some(Box::new(Reservation { + counters: self.counters.clone(), + plan, + charged, + shared, + _lease: lease, + }) as Box)) + }) + } +} + +/// What a request reserved, settled to its real cost when it ends. +struct Reservation { + counters: Arc, + plan: Plan, + charged: Vec, + shared: bool, + /// Dropping it stops the lease. + _lease: Option>, +} + +impl Settle for Reservation { + fn settle(self: Box, record: &Record) { + let success = (200..300).contains(&record.outcome.status); + let cost = match success { + true => record.outcome.cost.unwrap_or(0.0), + false => 0.0, + }; + let mut keys = Vec::new(); + let mut changes = Vec::new(); + for (counter, charged) in self.plan.counters.iter().zip(&self.charged) { + let delta = cost - charged; + if delta == 0.0 && *charged == 0.0 { + continue; + } + keys.push(counter.owner.counter()); + changes.push(Change { + floor: counter.floor, + delta, + fresh: cost, + }); + } + if keys.is_empty() { + return; + } + let Reservation { + counters, + shared, + _lease, + .. + } = *self; + tokio::spawn(async move { + counters.settle(keys, changes, shared).await; + drop(_lease); + }); + } +} + +/// The proxy names a key by its alias, adding `key_name` only while it still looks masked. +fn key_descriptor(alias: Option<&str>, key_name: Option<&str>) -> String { + let alias = alias.filter(|alias| !alias.is_empty()).unwrap_or("key"); + match key_name.filter(|name| is_masked(name)) { + Some(name) => format!("{alias} ({name})"), + None => alias.to_owned(), + } +} + +/// The proxy's `_MASKED_KEY_NAME_RE`, `sk-\.\.\.(?:[^\s\x00-\x1f\x7f-\x9f]{4})?`. +fn is_masked(name: &str) -> bool { + name.strip_prefix("sk-...").is_some_and(|rest| { + rest.is_empty() + || (rest.chars().count() == 4 + && rest + .chars() + .all(|character| !character.is_whitespace() && !character.is_control())) + }) +} + +/// Python's `repr` of a float. +fn python_float(value: f64) -> String { + if value.is_nan() { + return "nan".into(); + } + if value.is_infinite() { + return if value > 0.0 { "inf" } else { "-inf" }.into(); + } + let magnitude = value.abs(); + if magnitude != 0.0 && !(1e-4..1e16).contains(&magnitude) { + let scientific = format!("{value:e}"); + let (mantissa, exponent) = scientific.split_once('e').unwrap_or((&scientific, "0")); + let exponent: i32 = exponent.parse().unwrap_or(0); + let sign = if exponent < 0 { '-' } else { '+' }; + return format!("{mantissa}e{sign}{:02}", exponent.abs()); + } + format!("{value:?}") +} + +#[cfg(test)] +mod tests { + use rstest::rstest; + + use super::*; + + #[rstest] + #[case::no_alias(None, None, "key")] + #[case::empty_alias(Some(""), None, "key")] + #[case::alias(Some("prod"), None, "prod")] + #[case::masked_name(Some("prod"), Some("sk-...abcd"), "prod (sk-...abcd)")] + #[case::bare_masked_name(None, Some("sk-..."), "key (sk-...)")] + #[case::unmasked_name(Some("prod"), Some("sk-unmasked90"), "prod")] + #[case::name_with_space(Some("prod"), Some("sk-...ab d"), "prod")] + #[case::name_too_long(Some("prod"), Some("sk-...abcde"), "prod")] + fn describes_keys_like_the_proxy( + #[case] alias: Option<&str>, + #[case] name: Option<&str>, + #[case] expected: &str, + ) { + assert_eq!(key_descriptor(alias, name), expected); + } + + #[rstest] + #[case::whole(10.0, "10.0")] + #[case::fraction(0.1, "0.1")] + #[case::zero(0.0, "0.0")] + #[case::small(0.0001, "0.0001")] + #[case::tiny(0.00001, "1e-05")] + #[case::large(1e15, "1000000000000000.0")] + #[case::huge(1e16, "1e+16")] + #[case::huge_fraction(1.5e17, "1.5e+17")] + #[case::negative(-2.5, "-2.5")] + #[case::infinite(f64::INFINITY, "inf")] + fn prints_floats_like_python(#[case] value: f64, #[case] expected: &str) { + assert_eq!(python_float(value), expected); + } + + fn counter(owner: Owner, floor: f64, max_budget: Option, check: Check) -> Counter { + Counter { + owner, + floor, + max_budget, + check, + reserve: max_budget.is_some_and(|max| max > 0.0), + } + } + + fn plan(counters: Vec) -> (Vec, Plan) { + let keys = counters.iter().map(|c| c.owner.counter()).collect(); + let checks = (0..counters.len()).collect(); + (keys, Plan { counters, checks }) + } + + #[rstest] + fn local_counters_check_reserve_shrink_and_refuse() { + let local = Local::default(); + let ttl = Duration::from_secs(60); + let (keys, plan) = plan(vec![ + counter(Owner::Team("t".into()), 4.0, Some(10.0), Check::Above), + counter(Owner::User("u".into()), 0.0, Some(100.0), Check::AtLeast), + ]); + // Fits: both counters take the whole estimate. + assert_eq!( + local.admit(&keys, &plan, 5.0, ttl), + Admitted::Charged(vec![5.0, 5.0]) + ); + // Team has 1.0 left: the reservation shrinks to it on every counter. + assert_eq!( + local.admit(&keys, &plan, 5.0, ttl), + Admitted::Charged(vec![1.0, 1.0]) + ); + // Exactly at budget passes the team's strict check, but nothing is left to reserve. + assert_eq!( + local.admit(&keys, &plan, 5.0, ttl), + Admitted::Reserved(0, 15.0) + ); + // Settling the first request to 0.5 frees 4.5. + local.settle( + &keys, + &[ + Change { + floor: 4.0, + delta: -4.5, + fresh: 0.5, + }, + Change { + floor: 0.0, + delta: -4.5, + fresh: 0.5, + }, + ], + ttl, + ); + assert_eq!( + local.admit(&keys, &plan, 1.0, ttl), + Admitted::Charged(vec![1.0, 1.0]) + ); + // A stored spend above the counter raises it, and the check refuses. + let (keys, plan) = plan_with_floor(30.0); + assert_eq!(local.admit(&keys, &plan, 1.0, ttl), Admitted::Over(0, 30.0)); + } + + fn plan_with_floor(floor: f64) -> (Vec, Plan) { + plan(vec![counter( + Owner::Team("t".into()), + floor, + Some(10.0), + Check::Above, + )]) + } + + #[rstest] + fn refusals_read_like_the_proxy() { + let key = Owner::Key { + hash: "abc".into(), + descriptor: "prod (sk-...abcd)".into(), + }; + assert_eq!( + key.check_message(10.0, 10.0), + "Budget has been exceeded! Key=prod (sk-...abcd) Current cost: 10.0, Max budget: 10.0" + ); + assert_eq!( + key.reservation_message(12.5, 10.0), + "Budget has been exceeded! Key=abc Current cost: 12.5, Max budget: 10.0" + ); + assert_eq!( + Owner::EndUser("e".into()).check_message(2.0, 1.0), + "ExceededBudget: End User=e over budget. Spend=2.0, Budget=1.0" + ); + assert_eq!( + Owner::ZeroBudgetMember { + user_id: "u".into(), + team_id: "t".into() + } + .check_message(0.0, 0.0), + "Budget has been exceeded! User=u in Team=t Current cost: 0.0, Max budget: 0.0" + ); + } +} diff --git a/litellm-rust/crates/gateway/src/coordination.rs b/litellm-rust/crates/gateway/src/coordination.rs index eeb01158e59..48cd6ac14cf 100644 --- a/litellm-rust/crates/gateway/src/coordination.rs +++ b/litellm-rust/crates/gateway/src/coordination.rs @@ -267,10 +267,16 @@ impl Coordination { })) } - fn connection(&self) -> Option { + /// The shared connection, once Redis has answered. + pub(crate) fn connection(&self) -> Option { self.connection.get().cloned() } + /// `key` under the configured namespace, as the proxy prefixes its keys. + pub(crate) fn namespaced(&self, key: &str) -> String { + self.settings.key(key) + } + /// Whether the connection shared state is written and read through is up yet. pub fn connected(&self) -> bool { self.connection.get().is_some() diff --git a/litellm-rust/crates/gateway/src/key_status.rs b/litellm-rust/crates/gateway/src/key_status.rs index f4f54e9ac3d..50d8bf65567 100644 --- a/litellm-rust/crates/gateway/src/key_status.rs +++ b/litellm-rust/crates/gateway/src/key_status.rs @@ -2,7 +2,8 @@ use litellm_db::{Restrictions, VerificationToken}; use std::sync::Arc; use litellm_gateway_auth::keys::{ - AllowedRoutes, KeyModels, KeyProfile, KeyStatus, ModelAccess, Owner, Refusal, Restriction, + AllowedRoutes, KeyBudgets, KeyModels, KeyProfile, KeyStatus, MemberSpend, ModelAccess, Owner, + Refusal, Restriction, Spend, }; use serde_json::Value; @@ -37,6 +38,7 @@ pub(crate) fn status(row: VerificationToken, personal_keys_use_teams: bool) -> K Value::Object(fields) => Arc::new(fields), _ => Arc::default(), }; + let budgets = budgets(&row); let (team_models, team_alias, team_metadata) = match row.team { Some(team) => ( Some(ModelAccess::new(team.models)), @@ -59,10 +61,41 @@ pub(crate) fn status(row: VerificationToken, personal_keys_use_teams: bool) -> K team_alias, metadata: object(row.metadata), team_metadata, + budgets, }), } } +/// The spend and budgets the budgets hook checks for this key. +fn budgets(row: &VerificationToken) -> KeyBudgets { + let membership = &row.team_membership; + KeyBudgets { + key: Spend { + spend: row.spend, + max_budget: row.max_budget, + }, + key_name: row.key_name.clone(), + team: row.team.as_ref().map(|team| Spend { + spend: team.spend, + max_budget: team.max_budget, + }), + user: row.user.as_ref().map(|user| Spend { + spend: user.spend, + max_budget: user.max_budget, + }), + team_member: membership.spend.map(|spend| MemberSpend { + spend, + max_budget: membership.max_budget, + default_max_budget: membership.default_max_budget, + }), + end_user_budget_id: row + .metadata + .get("end_user_budget_id") + .and_then(Value::as_str) + .map(str::to_owned), + } +} + fn refusal(row: &VerificationToken, personal_keys_use_teams: bool) -> Option { let team = match (&row.team_id, &row.team) { (Some(team_id), None) => { @@ -80,29 +113,6 @@ fn refusal(row: &VerificationToken, personal_keys_use_teams: bool) -> Option= max_budget - { - return Some(budget_exceeded( - format!("Key={}", key_label(row)), - row.spend, - max_budget, - )); - } - if let Some((team_id, team)) = team - && let Some(max_budget) = team.max_budget - && max_budget.is_finite() - && team.spend > max_budget - { - return Some(budget_exceeded( - format!("Team={team_id}"), - team.spend, - max_budget, - )); - } unsupported(row, personal_keys_use_teams) .map(|(restriction, owner)| Refusal::Unsupported { restriction, owner }) } @@ -197,57 +207,6 @@ fn is_set(value: &Value) -> bool { } } -fn budget_exceeded(label: String, spend: f64, max_budget: f64) -> Refusal { - Refusal::BudgetExceeded { - owner: label, - spend: python_float(spend), - max_budget: python_float(max_budget), - } -} - -/// The proxy names a key by its alias, adding `key_name` only while it still looks masked. -fn key_label(row: &VerificationToken) -> String { - let alias = row - .key_alias - .as_deref() - .filter(|alias| !alias.is_empty()) - .unwrap_or("key"); - match row.key_name.as_deref().filter(|name| is_masked(name)) { - Some(name) => format!("{alias} ({name})"), - None => alias.to_owned(), - } -} - -/// The proxy's `_MASKED_KEY_NAME_RE`, `sk-\.\.\.(?:[^\s\x00-\x1f\x7f-\x9f]{4})?`. -fn is_masked(name: &str) -> bool { - name.strip_prefix("sk-...").is_some_and(|rest| { - rest.is_empty() - || (rest.chars().count() == 4 - && rest - .chars() - .all(|character| !character.is_whitespace() && !character.is_control())) - }) -} - -/// Python's `repr` of a float. -fn python_float(value: f64) -> String { - if value.is_nan() { - return "nan".into(); - } - if value.is_infinite() { - return if value > 0.0 { "inf" } else { "-inf" }.into(); - } - let magnitude = value.abs(); - if magnitude != 0.0 && !(1e-4..1e16).contains(&magnitude) { - let scientific = format!("{value:e}"); - let (mantissa, exponent) = scientific.split_once('e').unwrap_or((&scientific, "0")); - let exponent: i32 = exponent.parse().unwrap_or(0); - let sign = if exponent < 0 { '-' } else { '+' }; - return format!("{mantissa}e{sign}{:02}", exponent.abs()); - } - format!("{value:?}") -} - #[cfg(test)] mod tests { use std::time::SystemTime; @@ -307,7 +266,13 @@ mod tests { }, allowed_routes: AllowedRoutes::new(strings(&["llm_api_routes"])), team_id: team.then(|| "team-1".to_owned()), - profile: Default::default(), + profile: Arc::new(KeyProfile { + budgets: KeyBudgets { + team: team.then(Spend::default), + ..KeyBudgets::default() + }, + ..KeyProfile::default() + }), } } @@ -318,14 +283,6 @@ mod tests { } } - fn exceeded(owner: &str, spend: &str, max_budget: &str) -> KeyStatus { - refused(Refusal::BudgetExceeded { - owner: owner.into(), - spend: spend.into(), - max_budget: max_budget.into(), - }) - } - fn unsupported(restriction: Restriction, owner: Owner) -> KeyStatus { refused(Refusal::Unsupported { restriction, owner }) } @@ -358,57 +315,6 @@ mod tests { 0.0, refused(Refusal::TeamBlocked { team_id: "team-1".into() }) )] - #[case::key_at_budget( - false, - None, - None, - Some(10.0), - 10.0, - exceeded("Key=key", "10.0", "10.0") - )] - #[case::key_under_budget(false, None, None, Some(10.0), 9.5, active(false))] - #[case::key_zero_budget(false, None, None, Some(0.0), 0.0, exceeded("Key=key", "0.0", "0.0"))] - #[case::key_infinite_budget(false, None, None, Some(f64::INFINITY), 1e20, active(false))] - #[case::team_over_budget( - false, - Some("team-1"), - team(false, Some(5.0), 5.25), - None, - 0.0, - exceeded("Team=team-1", "5.25", "5.0") - )] - #[case::team_at_budget( - false, - Some("team-1"), - team(false, Some(5.0), 5.0), - None, - 0.0, - active(true) - )] - #[case::team_over_budget_under_key_budget( - false, - Some("team-1"), - team(false, Some(5.0), 9.0), - Some(10.0), - 2.0, - exceeded("Team=team-1", "9.0", "5.0") - )] - #[case::key_and_team_under_budget( - false, - Some("team-1"), - team(false, Some(5.0), 1.0), - Some(10.0), - 2.0, - active(true) - )] - #[case::key_budget_checked_before_team( - false, - Some("team-1"), - team(false, Some(5.0), 9.0), - Some(1.0), - 2.0, - exceeded("Key=key", "2.0", "1.0") - )] fn refuses_like_the_proxy( row: VerificationToken, #[case] blocked: bool, @@ -467,6 +373,8 @@ mod tests { team: team_id.and_then(|_| team(false, None, 0.0)), user_id: Some("user-1".into()), user: Some(User { + spend: 0.0, + max_budget: None, models: strings(models), restrictions, }), @@ -483,6 +391,7 @@ mod tests { team_membership: TeamMembership { rate_limits, allowed_models: strings(allowed_models), + ..TeamMembership::default() }, ..on_team(NONE, json!({}), row) } @@ -691,7 +600,7 @@ mod tests { models, allowed_routes, team_id, - .. + profile, } => KeyStatus::Active { expires_at, models, @@ -699,7 +608,7 @@ mod tests { team_id, profile: Arc::new(KeyProfile { user_id: user_id.map(str::to_owned), - ..KeyProfile::default() + ..(*profile).clone() }), }, other => other, @@ -714,42 +623,4 @@ mod tests { } ); } - - #[rstest] - #[case::no_alias(None, None, "key")] - #[case::empty_alias(Some(""), None, "key")] - #[case::alias(Some("prod"), None, "prod")] - #[case::masked_name(Some("prod"), Some("sk-...abcd"), "prod (sk-...abcd)")] - #[case::bare_masked_name(None, Some("sk-..."), "key (sk-...)")] - #[case::unmasked_name(Some("prod"), Some("sk-unmasked90"), "prod")] - #[case::name_with_space(Some("prod"), Some("sk-...ab d"), "prod")] - #[case::name_too_long(Some("prod"), Some("sk-...abcde"), "prod")] - fn labels_keys_like_the_proxy( - row: VerificationToken, - #[case] alias: Option<&str>, - #[case] name: Option<&str>, - #[case] expected: &str, - ) { - let row = VerificationToken { - key_alias: alias.map(str::to_owned), - key_name: name.map(str::to_owned), - ..row - }; - assert_eq!(key_label(&row), expected); - } - - #[rstest] - #[case::whole(10.0, "10.0")] - #[case::fraction(0.1, "0.1")] - #[case::zero(0.0, "0.0")] - #[case::small(0.0001, "0.0001")] - #[case::tiny(0.00001, "1e-05")] - #[case::large(1e15, "1000000000000000.0")] - #[case::huge(1e16, "1e+16")] - #[case::huge_fraction(1.5e17, "1.5e+17")] - #[case::negative(-2.5, "-2.5")] - #[case::infinite(f64::INFINITY, "inf")] - fn prints_floats_like_python(#[case] value: f64, #[case] expected: &str) { - assert_eq!(python_float(value), expected); - } } diff --git a/litellm-rust/crates/gateway/src/lib.rs b/litellm-rust/crates/gateway/src/lib.rs index 7e5e50fcab1..dcd3aa994a3 100644 --- a/litellm-rust/crates/gateway/src/lib.rs +++ b/litellm-rust/crates/gateway/src/lib.rs @@ -1,4 +1,5 @@ mod auth; +pub mod budgets; pub mod coordination; pub mod drain; mod error; diff --git a/litellm-rust/crates/gateway/src/main.rs b/litellm-rust/crates/gateway/src/main.rs index fccfa6cfc66..0954fd7634b 100644 --- a/litellm-rust/crates/gateway/src/main.rs +++ b/litellm-rust/crates/gateway/src/main.rs @@ -112,13 +112,22 @@ async fn main() -> Result<(), Box> { coordination.start(inference.clone(), keys.clone(), shutdown.clone()); } // Built-in hooks and observers are added here as the gateway gains them. + let mut hooks: Vec> = Vec::new(); + if database.is_some() { + hooks.push(Arc::new(litellm_gateway::budgets::Budgets::for_gateway( + &config, + &inference, + coordination.clone(), + database.clone(), + ))); + } let mut observers: Vec> = Vec::new(); let spend = database.zip(litellm_gateway::spend::SpendSettings::from_config(&config)); let spend_tracking = spend.is_some(); if let Some((pool, settings)) = spend { observers.push(litellm_gateway::spend::SpendWriter::start(pool, settings)); } - litellm_gateway::hooks::install(&inference, &config, Vec::new(), observers).await?; + litellm_gateway::hooks::install(&inference, &config, hooks, observers).await?; let listener = tokio::net::TcpListener::bind((settings.host.as_str(), settings.port)).await?; let address = listener.local_addr()?; // Without TCP_NODELAY, Nagle holds small SSE chunks until the client's delayed ACK (~40 ms per stream). diff --git a/litellm-rust/crates/gateway/tests/budgets.rs b/litellm-rust/crates/gateway/tests/budgets.rs new file mode 100644 index 00000000000..83a7778c732 --- /dev/null +++ b/litellm-rust/crates/gateway/tests/budgets.rs @@ -0,0 +1,495 @@ +//! The budgets hook on real requests: refusals as the proxy words them, reservations that keep +//! concurrent spend under a budget, and counters shared through Redis. Set +//! `LITELLM_TEST_REDIS_URL` to also run the shared cases against a live Redis. + +use std::{ + collections::HashMap, + sync::{ + Arc, + atomic::{AtomicUsize, Ordering}, + }, + time::Duration, +}; + +use axum::{ + Router, + body::{Body, to_bytes}, + http::{Request, header::CONTENT_TYPE}, + response::IntoResponse, +}; +use litellm_auth_types::SecretValue; +use litellm_config::Config; +use litellm_gateway::{ + budgets::{BudgetSettings, Budgets}, + coordination::{Coordination, RedisSettings}, +}; +use litellm_gateway_auth::keys::{ + AllowedRoutes, KeyBudgets, KeyHash, KeyLookup, KeyModels, KeyProfile, KeyStatus, LookupFuture, + MemberSpend, Spend, +}; +use rstest::rstest; +use serde_json::{Value, json}; +use tokio::net::TcpListener; +use tower::ServiceExt; + +/// Costs 3 * 0.001 + 4 * 0.002 = 0.011 at the test prices. +const MESSAGE: &str = r#"{"id":"msg_1","type":"message","role":"assistant","model":"claude-x","content":[{"type":"text","text":"hi"}],"stop_reason":"end_turn","stop_sequence":null,"usage":{"input_tokens":3,"output_tokens":4}}"#; +const COST: f64 = 0.011; + +async fn upstream() -> (String, Arc) { + let calls = Arc::new(AtomicUsize::new(0)); + let counted = calls.clone(); + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let address = listener.local_addr().unwrap(); + let app = Router::new().fallback(move |_: Request| { + let counted = counted.clone(); + async move { + counted.fetch_add(1, Ordering::SeqCst); + ([(CONTENT_TYPE, "application/json")], MESSAGE).into_response() + } + }); + tokio::spawn(async move { axum::serve(listener, app).await }); + (format!("http://{address}"), calls) +} + +struct Keys(HashMap); + +impl KeyLookup for Keys { + fn lookup<'a>(&'a self, hash: &'a KeyHash) -> LookupFuture<'a> { + Box::pin(async move { Ok(self.0.get(hash).cloned()) }) + } +} + +fn hash(token: &str) -> KeyHash { + KeyHash::from_token(&SecretValue::new(token)) +} + +fn key(team_id: Option<&str>, budgets: KeyBudgets) -> KeyStatus { + KeyStatus::Active { + expires_at: None, + models: KeyModels::default(), + allowed_routes: AllowedRoutes::default(), + team_id: team_id.map(str::to_owned), + profile: Arc::new(KeyProfile { + alias: Some("prod".into()), + user_id: Some("user-1".into()), + budgets, + ..KeyProfile::default() + }), + } +} + +fn budget(spend: f64, max_budget: f64) -> Spend { + Spend { + spend, + max_budget: Some(max_budget), + } +} + +struct Gateway { + app: Router, + calls: Arc, +} + +async fn gateway( + keys: Vec<(&str, KeyStatus)>, + redis: Option>, + settings: &str, +) -> Gateway { + gateway_with_database(keys, redis, None, settings).await +} + +async fn gateway_with_database( + keys: Vec<(&str, KeyStatus)>, + redis: Option>, + database: Option, + settings: &str, +) -> Gateway { + let (upstream, calls) = upstream().await; + let config = Config::from_yaml(&format!( + "model_list: + - model_name: claude + litellm_params: + model: anthropic/claude-x + api_base: {upstream} + api_key: sk-upstream + input_cost_per_token: 0.001 + output_cost_per_token: 0.002 + model_info: + id: claude-1 + - model_name: free + litellm_params: + model: anthropic/claude-x + api_base: {upstream} + api_key: sk-upstream + input_cost_per_token: 0 + output_cost_per_token: 0 + model_info: + id: free-1 +router_settings: + num_retries: 0 +general_settings: + master_key: sk-master +{settings}" + )) + .unwrap(); + let inference = litellm_gateway::build_inference(&config).unwrap(); + let budgets = Budgets::for_gateway(&config, &inference, redis, database); + litellm_gateway::hooks::install(&inference, &config, vec![Arc::new(budgets)], Vec::new()) + .await + .unwrap(); + let keys: Arc = Arc::new(Keys( + keys.into_iter() + .map(|(token, status)| (hash(token), status)) + .collect(), + )); + let app = litellm_gateway::router_with_keys(inference, &config, None, None, Some(keys)); + Gateway { app, calls } +} + +async fn send(app: &Router, token: &str, model: &str) -> (u16, Value) { + send_as(app, token, model, None).await +} + +async fn send_as(app: &Router, token: &str, model: &str, end_user: Option<&str>) -> (u16, Value) { + let mut body = json!({ + "model": model, + "max_tokens": 8, + "messages": [{"role": "user", "content": "hi"}], + }); + if let Some(end_user) = end_user { + body["metadata"] = json!({"user_id": end_user}); + } + let response = app + .clone() + .oneshot( + Request::post("/v1/messages") + .header("content-type", "application/json") + .header("authorization", format!("Bearer {token}")) + .body(Body::from(body.to_string())) + .unwrap(), + ) + .await + .unwrap(); + let status = response.status().as_u16(); + let body = to_bytes(response.into_body(), 1 << 20).await.unwrap(); + (status, serde_json::from_slice(&body).unwrap_or(Value::Null)) +} + +/// Settling runs after the response, off the request path. +async fn settled() { + tokio::time::sleep(Duration::from_millis(100)).await; +} + +#[rstest] +#[case::key_at_budget( + None, + KeyBudgets { key: budget(10.0, 10.0), ..KeyBudgets::default() }, + "Budget has been exceeded! Key=prod Current cost: 10.0, Max budget: 10.0" +)] +#[case::zero_key_budget( + None, + KeyBudgets { key: budget(0.0, 0.0), ..KeyBudgets::default() }, + "Budget has been exceeded! Key=prod Current cost: 0.0, Max budget: 0.0" +)] +#[case::team_over_budget( + Some("team-1"), + KeyBudgets { team: Some(budget(5.25, 5.0)), ..KeyBudgets::default() }, + "Budget has been exceeded! Team=team-1 Current cost: 5.25, Max budget: 5.0" +)] +#[case::personal_user_over_budget( + None, + KeyBudgets { user: Some(budget(3.0, 3.0)), ..KeyBudgets::default() }, + "ExceededBudget: User=user-1 over budget. Spend=3.0, Budget=3.0" +)] +#[case::member_before_key( + Some("team-1"), + KeyBudgets { + key: budget(10.0, 10.0), + team_member: Some(MemberSpend { spend: 2.0, max_budget: Some(1.0), default_max_budget: None }), + ..KeyBudgets::default() + }, + "Budget has been exceeded! TeamMember=user-1:team-1 Current cost: 2.0, Max budget: 1.0" +)] +#[case::zero_member_budget( + Some("team-1"), + KeyBudgets { + team_member: Some(MemberSpend { spend: 0.0, max_budget: Some(0.0), default_max_budget: None }), + ..KeyBudgets::default() + }, + "Budget has been exceeded! User=user-1 in Team=team-1 Current cost: 0.0, Max budget: 0.0" +)] +#[case::team_default_member_budget( + Some("team-1"), + KeyBudgets { + team_member: Some(MemberSpend { spend: 4.0, max_budget: None, default_max_budget: Some(4.0) }), + ..KeyBudgets::default() + }, + "Budget has been exceeded! TeamMember=user-1:team-1 Current cost: 4.0, Max budget: 4.0" +)] +#[tokio::test] +async fn over_budget_requests_are_refused_like_the_proxy( + #[case] team_id: Option<&str>, + #[case] budgets: KeyBudgets, + #[case] message: &str, +) { + let gateway = gateway(vec![("sk-key", key(team_id, budgets))], None, "").await; + let (status, body) = send(&gateway.app, "sk-key", "claude").await; + assert_eq!(status, 422); + assert_eq!( + body, + json!({"error": {"message": message, "type": "budget_exceeded", "param": null, "code": "422"}}) + ); + assert_eq!(gateway.calls.load(Ordering::SeqCst), 0); +} + +#[rstest] +#[case::team_at_budget( + Some("team-1"), + KeyBudgets { team: Some(budget(5.0, 5.0)), ..KeyBudgets::default() } +)] +#[case::user_budget_on_team_key( + Some("team-1"), + KeyBudgets { user: Some(budget(3.0, 3.0)), ..KeyBudgets::default() } +)] +#[tokio::test] +async fn budgets_the_proxy_does_not_apply_pass( + #[case] team_id: Option<&str>, + #[case] budgets: KeyBudgets, +) { + // A team exactly at its budget passes the proxy's check; only a reservation would refuse it. + let gateway = gateway( + vec![("sk-key", key(team_id, budgets))], + None, + " disable_budget_reservation: true\n", + ) + .await; + assert_eq!(send(&gateway.app, "sk-key", "claude").await.0, 200); +} + +#[rstest] +#[tokio::test] +async fn free_models_skip_budgets_and_the_status_code_is_configurable() { + let budgets = KeyBudgets { + key: budget(10.0, 10.0), + ..KeyBudgets::default() + }; + let gateway = gateway( + vec![("sk-key", key(None, budgets))], + None, + "litellm_settings:\n budget_exceeded_status_code: 429\n", + ) + .await; + assert_eq!(send(&gateway.app, "sk-key", "free").await.0, 200); + assert_eq!(send(&gateway.app, "sk-key", "claude").await.0, 429); +} + +/// Spends a key with a 0.02 budget down: the first request fits, the second's reservation +/// shrinks to what is left, and the third finds the budget spent. +async fn spend_down(app: &Router) { + assert_eq!(send(app, "sk-key", "claude").await.0, 200); + settled().await; + assert_eq!(send(app, "sk-key", "claude").await.0, 200); + settled().await; + let (status, body) = send(app, "sk-key", "claude").await; + assert_eq!(status, 422, "{body}"); + let message = body["error"]["message"].as_str().unwrap(); + assert!( + message.starts_with("Budget has been exceeded! Key=prod Current cost: 0.02"), + "{message}" + ); +} + +fn budgeted_key() -> KeyStatus { + key( + None, + KeyBudgets { + key: budget(0.0, 0.02), + ..KeyBudgets::default() + }, + ) +} + +#[rstest] +#[tokio::test] +async fn spend_is_counted_on_the_replica_without_redis() { + let gateway = gateway(vec![("sk-key", budgeted_key())], None, "").await; + spend_down(&gateway.app).await; + assert_eq!(gateway.calls.load(Ordering::SeqCst), 2); +} + +#[rstest] +#[tokio::test] +async fn reservations_hold_back_concurrent_requests() { + // The first request's reservation takes the whole budget, so the second, sent before the + // first settles, finds it spent. + let key = key( + None, + KeyBudgets { + key: budget(0.0, 0.0001), + ..KeyBudgets::default() + }, + ); + let gateway = gateway(vec![("sk-key", key)], None, "").await; + let (first, second) = tokio::join!( + send(&gateway.app, "sk-key", "claude"), + send(&gateway.app, "sk-key", "claude") + ); + let mut statuses = [first.0, second.0]; + statuses.sort_unstable(); + assert_eq!(statuses, [200, 422]); +} + +fn redis_url() -> Option { + std::env::var("LITELLM_TEST_REDIS_URL").ok() +} + +async fn redis(url: &str, namespace: &str) -> Arc { + let coordination = + Coordination::new(RedisSettings::new(url, Some(namespace.to_owned()))).unwrap(); + let router: Arc = Arc::new(NoRouter( + litellm_gateway_inference::ModelRouter::from_config( + &Config::from_yaml("model_list: []").unwrap(), + ) + .unwrap(), + )); + coordination.start(router, None, tokio_util::sync::CancellationToken::new()); + for _ in 0..100 { + if coordination.connected() { + return coordination; + } + tokio::time::sleep(Duration::from_millis(20)).await; + } + panic!("Redis did not connect"); +} + +struct NoRouter(litellm_gateway_inference::ModelRouter); + +impl litellm_gateway::coordination::RouterHandle for NoRouter { + fn router(&self) -> &litellm_gateway_inference::ModelRouter { + &self.0 + } +} + +#[rstest] +#[tokio::test] +async fn replicas_share_spend_through_the_proxy_counters() { + let Some(url) = redis_url() else { return }; + let namespace = format!("budgets-{}", uuid::Uuid::new_v4()); + let first = gateway( + vec![("sk-key", budgeted_key())], + Some(redis(&url, &namespace).await), + "", + ) + .await; + let second = gateway( + vec![("sk-key", budgeted_key())], + Some(redis(&url, &namespace).await), + "", + ) + .await; + assert_eq!(send(&first.app, "sk-key", "claude").await.0, 200); + settled().await; + assert_eq!(send(&second.app, "sk-key", "claude").await.0, 200); + settled().await; + assert_eq!(send(&first.app, "sk-key", "claude").await.0, 422); + + let mut connection = redis::Client::open(url.as_str()) + .unwrap() + .get_multiplexed_async_connection() + .await + .unwrap(); + let counter = format!("{namespace}:spend:key:{}", hash("sk-key").as_str()); + let spend: String = redis::cmd("GET") + .arg(&counter) + .query_async(&mut connection) + .await + .unwrap(); + assert!( + (spend.parse::().unwrap() - 2.0 * COST).abs() < 1e-9, + "{spend}" + ); + let ttl: i64 = redis::cmd("TTL") + .arg(&counter) + .query_async(&mut connection) + .await + .unwrap(); + assert!((1..=60).contains(&ttl), "{ttl}"); + let user: String = redis::cmd("GET") + .arg(format!("{namespace}:spend:user:user-1")) + .query_async(&mut connection) + .await + .unwrap(); + assert!( + (user.parse::().unwrap() - 2.0 * COST).abs() < 1e-9, + "{user}" + ); +} + +#[expect( + clippy::disallowed_methods, + reason = "test fixture rows; the lookups under test use the checked macro" +)] +async fn end_users(pool: &litellm_db::PgPool) { + sqlx::raw_sql( + r#" +INSERT INTO "LiteLLM_BudgetTable" (budget_id, max_budget, created_by, updated_by) + VALUES ('small', 1.0, 'x', 'x'), ('default', 0.0001, 'x', 'x'); +INSERT INTO "LiteLLM_EndUserTable" (user_id, spend, budget_id) + VALUES ('over', 1.5, 'small'), ('at', 1.0, 'small'), ('unbudgeted', 9.0, NULL); +"#, + ) + .execute(pool) + .await + .unwrap(); +} + +#[rstest] +#[tokio::test] +async fn end_users_are_held_to_their_budgets() { + let postgres = litellm_db_testing::MigratedPostgres::start().await.unwrap(); + end_users(postgres.pool()).await; + let key = key(None, KeyBudgets::default()); + let gateway = gateway_with_database( + vec![("sk-key", key)], + None, + Some(postgres.pool().clone()), + " disable_budget_reservation: true\nlitellm_settings:\n max_end_user_budget_id: default\n", + ) + .await; + let (status, body) = send_as(&gateway.app, "sk-key", "claude", Some("over")).await; + assert_eq!(status, 422); + assert_eq!( + body["error"]["message"], + "ExceededBudget: End User=over over budget. Spend=1.5, Budget=1.0" + ); + // Strictly over: an end user exactly at budget passes the check. + assert_eq!( + send_as(&gateway.app, "sk-key", "claude", Some("at")) + .await + .0, + 200 + ); + // Without a row, only a reservation applies the default budget, and reservations are off. + assert_eq!( + send_as(&gateway.app, "sk-key", "claude", Some("new")) + .await + .0, + 200 + ); + // A row without a budget takes the default one. + let (status, body) = send_as(&gateway.app, "sk-key", "claude", Some("unbudgeted")).await; + assert_eq!(status, 422); + assert_eq!( + body["error"]["message"], + "ExceededBudget: End User=unbudgeted over budget. Spend=9.0, Budget=0.0001" + ); +} + +#[rstest] +fn settings_follow_the_proxy_defaults() { + let settings = BudgetSettings::from_config(&Config::from_yaml("model_list: []").unwrap()); + assert_eq!(settings.status, 422); + assert!(settings.reservations); + assert_eq!(settings.ttl, Duration::from_secs(60)); +} diff --git a/litellm-rust/docs/customer-proxy/budgets.svg b/litellm-rust/docs/customer-proxy/budgets.svg new file mode 100644 index 00000000000..3ff96dd5c81 --- /dev/null +++ b/litellm-rust/docs/customer-proxy/budgets.svg @@ -0,0 +1,51 @@ + + + + + + + Budgets hook: proxy checks on the proxy's shared spend counters + + + Key store row (cached with the key) + key spend + max_budget (or budget row's) + team, user spend + budgets + team member spend, member budget + (+ active temp increase, or team default) + end users: own lookup, cached 60 s + + + One Redis script per request (atomic) + 1. read spend:key / team / user / end_user / team_member counters; missing or below + the stored spend: seed or raise to it (the proxy's reseed + set-max repair) + 2. checks in the proxy's order: member, key, team (>), user, zero member, end user (>) + 3. reserve the max-cost estimate on every counter with a positive budget; + shrink to what is left, refuse when nothing is (1e-12), as the proxy + + + + Refusal: 422 (budget_exceeded_status_code) + {"error": {"message", "type": "budget_exceeded", + "param": null, "code": "422"}} on both endpoints + proxy messages, Python float repr; no provider call + + + + Admitted: request runs + estimate = input tokens x input price + + min(max_tokens or 16384, model max) x output price + lease renews reserved counters every TTL/2 + + + + Settle (off the request path) + success: INCRBYFLOAT (actual - reserved) on reserved counters, actual on the others; failure: release the reservation + a counter that expired mid-request is reseeded and charged the whole cost; EXPIRE ttl (default_redis_ttl, 60 s) refreshed + script failure: the counters are deleted so the next read reseeds them, as the proxy does + + + + Same keys, values and TTLs as the Python proxy, under its namespace + gateway replicas and Python proxies on one Redis enforce one budget; the proxy's reset job deletes counters and they reseed. + No Redis, or Redis unreachable: the same logic on this replica's counters (fail_closed_budget_enforcement: 503 instead) +