From b83353906b07553e2b490a1db9da3e0120456d69 Mon Sep 17 00:00:00 2001 From: Nate Armstrong Date: Mon, 19 Oct 2026 09:58:44 -0700 Subject: [PATCH] feat(gateway): enforce key, team, user, member and end-user budgets on shared spend counters (#47044) Adds a built-in budgets hook that runs the proxy's budget checks, in its order and with its messages, on the proxy's Redis spend:* counters. Each request reads and seeds the counters, checks them and reserves its maximum cost in one Redis script, then settles the reservation to its real cost off the request path, so replicas and Python proxies sharing the Redis enforce one budget. Without Redis the counters are kept per replica. The key store now reads user, team member and end-user budgets, and auth no longer checks budgets against stored spend. Claude-Session: https://claude.ai/code/session_0138knkqKpkoSr6PApenZan5 Co-authored-by: Claude Opus 5.5 (1M context) --- litellm-rust/Cargo.lock | 1 + ...399bc28aaeab7ee4da8e53f2ce87a52a3da9.json} | 90 +- ...f92fc7030de81c903adb593a5eca47237cfbe.json | 60 + litellm-rust/crates/db/src/end_user.rs | 40 + litellm-rust/crates/db/src/lib.rs | 2 + .../crates/db/src/verification_token.rs | 30 +- .../crates/db/tests/verification_token.rs | 76 +- litellm-rust/crates/gateway-auth/src/keys.rs | 38 + litellm-rust/crates/gateway-hooks/Cargo.toml | 1 + .../crates/gateway-hooks/src/facts.rs | 3 + litellm-rust/crates/gateway-hooks/src/hook.rs | 17 +- .../crates/gateway-hooks/src/pipeline.rs | 1 + .../crates/gateway-hooks/src/pricing.rs | 9 + .../crates/gateway-inference/src/error.rs | 2 +- .../crates/gateway-inference/src/hooks.rs | 1 + litellm-rust/crates/gateway/src/budgets.rs | 1209 +++++++++++++++++ .../crates/gateway/src/coordination.rs | 8 +- litellm-rust/crates/gateway/src/key_status.rs | 221 +-- litellm-rust/crates/gateway/src/lib.rs | 1 + litellm-rust/crates/gateway/src/main.rs | 11 +- litellm-rust/crates/gateway/tests/budgets.rs | 495 +++++++ litellm-rust/docs/customer-proxy/budgets.svg | 51 + 22 files changed, 2159 insertions(+), 208 deletions(-) rename litellm-rust/crates/db/.sqlx/{query-aebecf84ec08b203635d59ae2a616bd622e7ed0102166fed8b09b452741c107d.json => query-08aa4e832176f266c2200c676ece399bc28aaeab7ee4da8e53f2ce87a52a3da9.json} (56%) create mode 100644 litellm-rust/crates/db/.sqlx/query-a907e2963ee0cf25b34af9ed477f92fc7030de81c903adb593a5eca47237cfbe.json create mode 100644 litellm-rust/crates/db/src/end_user.rs create mode 100644 litellm-rust/crates/gateway/src/budgets.rs create mode 100644 litellm-rust/crates/gateway/tests/budgets.rs create mode 100644 litellm-rust/docs/customer-proxy/budgets.svg 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) +