From 0919823385f46bf55baa74910c120abd67e2bf3c 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, user, team, member and end-user rate limits on shared counters (#47050) * feat(gateway): enforce key, user, team, member and end-user rate limits on shared counters Adds a built-in rate limits hook with the proxy's v3 limiter semantics: requests per window, tokens per window and in-flight requests per key, on the proxy's own {scope:id}:window|requests|tokens|max_parallel_requests Redis keys and with its scripts, so replicas and Python proxies sharing the Redis enforce one limit. Tokens are reserved from the proxy's estimate and settled to usage when the response ends. Refusals carry the proxy's message, body and headers on each endpoint, and complete responses carry its x-ratelimit-* headers. Without Redis the counters are kept per replica. The key store now reads key, user, team and member rpm/tpm limits and the key's max_parallel_requests, and end users' budget limits are read with their budgets; keys with plain rate limits are no longer refused. Co-Authored-By: Claude Opus 5.5 (1M context) Claude-Session: https://claude.ai/code/session_0138knkqKpkoSr6PApenZan5 * test(gateway-auth): expect the narrowed unsupported rate limit message Co-Authored-By: Claude Opus 5.5 (1M context) Claude-Session: https://claude.ai/code/session_0138knkqKpkoSr6PApenZan5 --------- Co-authored-by: Claude Opus 5.5 (1M context) --- ...ceebb7cf0688affae913dd9d25bddcd2a4d6.json} | 30 +- ...598589a6afcbd4fb2cf1c38c6f1fa46d7773.json} | 172 ++- litellm-rust/crates/db/src/end_user.rs | 23 +- litellm-rust/crates/db/src/lib.rs | 3 +- .../crates/db/src/verification_token.rs | 64 +- .../crates/db/tests/verification_token.rs | 91 +- litellm-rust/crates/gateway-auth/src/error.rs | 2 +- litellm-rust/crates/gateway-auth/src/keys.rs | 30 +- .../crates/gateway-hooks/src/facts.rs | 3 +- litellm-rust/crates/gateway-hooks/src/hook.rs | 10 + .../crates/gateway-inference/src/error.rs | 18 +- .../crates/gateway-inference/src/hooks.rs | 1 + litellm-rust/crates/gateway/src/budgets.rs | 54 +- litellm-rust/crates/gateway/src/end_users.rs | 70 + litellm-rust/crates/gateway/src/key_status.rs | 100 +- litellm-rust/crates/gateway/src/lib.rs | 2 + litellm-rust/crates/gateway/src/main.rs | 7 + .../crates/gateway/src/rate_limits.rs | 1193 +++++++++++++++++ litellm-rust/crates/gateway/src/spend.rs | 2 +- .../crates/gateway/tests/rate_limits.rs | 515 +++++++ .../docs/customer-proxy/rate-limits.svg | 55 + 21 files changed, 2293 insertions(+), 152 deletions(-) rename litellm-rust/crates/db/.sqlx/{query-a907e2963ee0cf25b34af9ed477f92fc7030de81c903adb593a5eca47237cfbe.json => query-623a70738ef371dca69ec8c04e63ceebb7cf0688affae913dd9d25bddcd2a4d6.json} (55%) rename litellm-rust/crates/db/.sqlx/{query-08aa4e832176f266c2200c676ece399bc28aaeab7ee4da8e53f2ce87a52a3da9.json => query-77f81c78f5498fc5485226ccad7c598589a6afcbd4fb2cf1c38c6f1fa46d7773.json} (64%) create mode 100644 litellm-rust/crates/gateway/src/end_users.rs create mode 100644 litellm-rust/crates/gateway/src/rate_limits.rs create mode 100644 litellm-rust/crates/gateway/tests/rate_limits.rs create mode 100644 litellm-rust/docs/customer-proxy/rate-limits.svg diff --git a/litellm-rust/crates/db/.sqlx/query-a907e2963ee0cf25b34af9ed477f92fc7030de81c903adb593a5eca47237cfbe.json b/litellm-rust/crates/db/.sqlx/query-623a70738ef371dca69ec8c04e63ceebb7cf0688affae913dd9d25bddcd2a4d6.json similarity index 55% rename from litellm-rust/crates/db/.sqlx/query-a907e2963ee0cf25b34af9ed477f92fc7030de81c903adb593a5eca47237cfbe.json rename to litellm-rust/crates/db/.sqlx/query-623a70738ef371dca69ec8c04e63ceebb7cf0688affae913dd9d25bddcd2a4d6.json index c9805198db5..de185711b56 100644 --- a/litellm-rust/crates/db/.sqlx/query-a907e2963ee0cf25b34af9ed477f92fc7030de81c903adb593a5eca47237cfbe.json +++ b/litellm-rust/crates/db/.sqlx/query-623a70738ef371dca69ec8c04e63ceebb7cf0688affae913dd9d25bddcd2a4d6.json @@ -1,6 +1,6 @@ { "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", + "query": "SELECT e.spend AS \"spend?\",\n e.blocked AS \"blocked?\",\n b.max_budget AS \"max_budget?\",\n b.rpm_limit AS \"rpm_limit?\",\n b.tpm_limit AS \"tpm_limit?\"\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 = COALESCE($3, e.budget_id, $2)", "describe": { "columns": [ { @@ -38,13 +38,30 @@ }, { "ordinal": 3, - "name": "default_max_budget?", - "type_info": "Float8", - "origin": "Expression" + "name": "rpm_limit?", + "type_info": "Int8", + "origin": { + "Table": { + "table": "\"LiteLLM_BudgetTable\"", + "name": "rpm_limit" + } + } + }, + { + "ordinal": 4, + "name": "tpm_limit?", + "type_info": "Int8", + "origin": { + "Table": { + "table": "\"LiteLLM_BudgetTable\"", + "name": "tpm_limit" + } + } } ], "parameters": { "Left": [ + "Text", "Text", "Text" ] @@ -53,8 +70,9 @@ true, true, true, - null + true, + true ] }, - "hash": "a907e2963ee0cf25b34af9ed477f92fc7030de81c903adb593a5eca47237cfbe" + "hash": "623a70738ef371dca69ec8c04e63ceebb7cf0688affae913dd9d25bddcd2a4d6" } diff --git a/litellm-rust/crates/db/.sqlx/query-08aa4e832176f266c2200c676ece399bc28aaeab7ee4da8e53f2ce87a52a3da9.json b/litellm-rust/crates/db/.sqlx/query-77f81c78f5498fc5485226ccad7c598589a6afcbd4fb2cf1c38c6f1fa46d7773.json similarity index 64% rename from litellm-rust/crates/db/.sqlx/query-08aa4e832176f266c2200c676ece399bc28aaeab7ee4da8e53f2ce87a52a3da9.json rename to litellm-rust/crates/db/.sqlx/query-77f81c78f5498fc5485226ccad7c598589a6afcbd4fb2cf1c38c6f1fa46d7773.json index 1a9ad366cfa..d8f743ae315 100644 --- a/litellm-rust/crates/db/.sqlx/query-08aa4e832176f266c2200c676ece399bc28aaeab7ee4da8e53f2ce87a52a3da9.json +++ b/litellm-rust/crates/db/.sqlx/query-77f81c78f5498fc5485226ccad7c598589a6afcbd4fb2cf1c38c6f1fa46d7773.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 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", + "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 COALESCE(v.rpm_limit, kb.rpm_limit) AS rpm_limit,\n COALESCE(v.tpm_limit, kb.tpm_limit) AS tpm_limit,\n v.max_parallel_requests,\n (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 AS \"team_rpm_limit?\",\n t.tpm_limit AS \"team_tpm_limit?\",\n (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 AS \"member_rpm_limit?\",\n mb.tpm_limit AS \"member_tpm_limit?\",\n (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 AS \"user_rpm_limit?\",\n u.tpm_limit AS \"user_tpm_limit?\",\n (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": [ { @@ -105,54 +105,77 @@ }, { "ordinal": 10, + "name": "rpm_limit", + "type_info": "Int8", + "origin": "Expression" + }, + { + "ordinal": 11, + "name": "tpm_limit", + "type_info": "Int8", + "origin": "Expression" + }, + { + "ordinal": 12, + "name": "max_parallel_requests", + "type_info": "Int4", + "origin": { + "Table": { + "table": "\"LiteLLM_VerificationToken\"", + "name": "max_parallel_requests" + } + } + }, + { + "ordinal": 13, "name": "key_rate_limits!", "type_info": "Bool", "origin": "Expression" }, { - "ordinal": 11, + "ordinal": 14, "name": "key_policies!", "type_info": "Bool", "origin": "Expression" }, { - "ordinal": 12, + "ordinal": 15, "name": "key_model_aliases!", "type_info": "Bool", "origin": "Expression" }, { - "ordinal": 13, + "ordinal": 16, "name": "key_object_permissions!", "type_info": "Bool", "origin": "Expression" }, { - "ordinal": 14, + "ordinal": 17, "name": "key_cache_controls!", "type_info": "Bool", "origin": "Expression" }, { - "ordinal": 15, + "ordinal": 18, "name": "key_project!", "type_info": "Bool", "origin": "Expression" }, { - "ordinal": 16, + "ordinal": 19, "name": "key_organization!", "type_info": "Bool", "origin": "Expression" }, { - "ordinal": 17, + "ordinal": 20, "name": "key_agent!", "type_info": "Bool", "origin": "Expression" }, { - "ordinal": 18, + "ordinal": 21, "name": "team_id", "type_info": "Text", "origin": { @@ -163,7 +186,7 @@ } }, { - "ordinal": 19, + "ordinal": 22, "name": "team_row_id?", "type_info": "Text", "origin": { @@ -174,7 +197,7 @@ } }, { - "ordinal": 20, + "ordinal": 23, "name": "team_alias?", "type_info": "Text", "origin": { @@ -185,7 +208,7 @@ } }, { - "ordinal": 21, + "ordinal": 24, "name": "team_blocked?", "type_info": "Bool", "origin": { @@ -196,7 +219,7 @@ } }, { - "ordinal": 22, + "ordinal": 25, "name": "team_models?", "type_info": "TextArray", "origin": { @@ -207,7 +230,7 @@ } }, { - "ordinal": 23, + "ordinal": 26, "name": "team_spend?", "type_info": "Float8", "origin": { @@ -218,7 +241,7 @@ } }, { - "ordinal": 24, + "ordinal": 27, "name": "team_max_budget?", "type_info": "Float8", "origin": { @@ -229,7 +252,7 @@ } }, { - "ordinal": 25, + "ordinal": 28, "name": "team_metadata?", "type_info": "Jsonb", "origin": { @@ -240,43 +263,87 @@ } }, { - "ordinal": 26, + "ordinal": 29, + "name": "team_rpm_limit?", + "type_info": "Int8", + "origin": { + "Table": { + "table": "\"LiteLLM_TeamTable\"", + "name": "rpm_limit" + } + } + }, + { + "ordinal": 30, + "name": "team_tpm_limit?", + "type_info": "Int8", + "origin": { + "Table": { + "table": "\"LiteLLM_TeamTable\"", + "name": "tpm_limit" + } + } + }, + { + "ordinal": 31, "name": "team_rate_limits!", "type_info": "Bool", "origin": "Expression" }, { - "ordinal": 27, + "ordinal": 32, "name": "team_policies!", "type_info": "Bool", "origin": "Expression" }, { - "ordinal": 28, + "ordinal": 33, "name": "team_model_aliases!", "type_info": "Bool", "origin": "Expression" }, { - "ordinal": 29, + "ordinal": 34, "name": "team_object_permissions!", "type_info": "Bool", "origin": "Expression" }, { - "ordinal": 30, + "ordinal": 35, "name": "team_organization!", "type_info": "Bool", "origin": "Expression" }, { - "ordinal": 31, + "ordinal": 36, + "name": "member_rpm_limit?", + "type_info": "Int8", + "origin": { + "Table": { + "table": "\"LiteLLM_BudgetTable\"", + "name": "rpm_limit" + } + } + }, + { + "ordinal": 37, + "name": "member_tpm_limit?", + "type_info": "Int8", + "origin": { + "Table": { + "table": "\"LiteLLM_BudgetTable\"", + "name": "tpm_limit" + } + } + }, + { + "ordinal": 38, "name": "member_rate_limits!", "type_info": "Bool", "origin": "Expression" }, { - "ordinal": 32, + "ordinal": 39, "name": "member_allowed_models?", "type_info": "TextArray", "origin": { @@ -287,7 +354,7 @@ } }, { - "ordinal": 33, + "ordinal": 40, "name": "member_spend?", "type_info": "Float8", "origin": { @@ -298,7 +365,7 @@ } }, { - "ordinal": 34, + "ordinal": 41, "name": "member_max_budget?", "type_info": "Float8", "origin": { @@ -309,13 +376,13 @@ } }, { - "ordinal": 35, + "ordinal": 42, "name": "member_budget_increase?", "type_info": "Float8", "origin": "Expression" }, { - "ordinal": 36, + "ordinal": 43, "name": "member_default_budget?", "type_info": "Float8", "origin": { @@ -326,7 +393,7 @@ } }, { - "ordinal": 37, + "ordinal": 44, "name": "user_id", "type_info": "Text", "origin": { @@ -337,7 +404,7 @@ } }, { - "ordinal": 38, + "ordinal": 45, "name": "user_row_id?", "type_info": "Text", "origin": { @@ -348,7 +415,7 @@ } }, { - "ordinal": 39, + "ordinal": 46, "name": "user_models?", "type_info": "TextArray", "origin": { @@ -359,7 +426,7 @@ } }, { - "ordinal": 40, + "ordinal": 47, "name": "user_spend?", "type_info": "Float8", "origin": { @@ -370,7 +437,7 @@ } }, { - "ordinal": 41, + "ordinal": 48, "name": "user_max_budget?", "type_info": "Float8", "origin": { @@ -381,19 +448,41 @@ } }, { - "ordinal": 42, + "ordinal": 49, + "name": "user_rpm_limit?", + "type_info": "Int8", + "origin": { + "Table": { + "table": "\"LiteLLM_UserTable\"", + "name": "rpm_limit" + } + } + }, + { + "ordinal": 50, + "name": "user_tpm_limit?", + "type_info": "Int8", + "origin": { + "Table": { + "table": "\"LiteLLM_UserTable\"", + "name": "tpm_limit" + } + } + }, + { + "ordinal": 51, "name": "user_rate_limits!", "type_info": "Bool", "origin": "Expression" }, { - "ordinal": 43, + "ordinal": 52, "name": "user_policies!", "type_info": "Bool", "origin": "Expression" }, { - "ordinal": 44, + "ordinal": 53, "name": "user_object_permissions!", "type_info": "Bool", "origin": "Expression" @@ -417,6 +506,9 @@ false, null, null, + true, + null, + null, null, null, null, @@ -431,7 +523,8 @@ true, true, true, - null, + true, + true, null, null, null, @@ -439,6 +532,9 @@ null, true, true, + null, + true, + true, true, null, true, @@ -447,10 +543,12 @@ true, true, true, + true, + true, null, null, null ] }, - "hash": "08aa4e832176f266c2200c676ece399bc28aaeab7ee4da8e53f2ce87a52a3da9" + "hash": "77f81c78f5498fc5485226ccad7c598589a6afcbd4fb2cf1c38c6f1fa46d7773" } diff --git a/litellm-rust/crates/db/src/end_user.rs b/litellm-rust/crates/db/src/end_user.rs index a4d98eef2aa..4f7d05a32ec 100644 --- a/litellm-rust/crates/db/src/end_user.rs +++ b/litellm-rust/crates/db/src/end_user.rs @@ -1,6 +1,6 @@ use sqlx::PgPool; -use crate::Error; +use crate::{Error, RateLimits}; /// An end user's spend and budget, as the proxy's end-user budget checks read them. #[derive(Clone, Debug, Default, PartialEq)] @@ -8,33 +8,42 @@ 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. + /// The budget's `max_budget`: the override's, else the row's, else the default's. pub max_budget: Option, + /// That budget's `rpm_limit` and `tpm_limit`. + pub limits: RateLimits, } -/// Reads `end_user_id`, with `default_budget_id` as the budget of an end user without its own. +/// Reads `end_user_id`, with `default_budget_id` as the budget of an end user without its own +/// and `override_budget_id`, when given, as its budget whatever the row says. pub async fn find_end_user( pool: &PgPool, end_user_id: &str, default_budget_id: Option<&str>, + override_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?" + b.rpm_limit AS "rpm_limit?", + b.tpm_limit AS "tpm_limit?" 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"#, + LEFT JOIN "LiteLLM_BudgetTable" AS b ON b.budget_id = COALESCE($3, e.budget_id, $2)"#, end_user_id, default_budget_id, + override_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), + max_budget: row.max_budget, + limits: RateLimits { + rpm: row.rpm_limit, + tpm: row.tpm_limit, + }, }) } diff --git a/litellm-rust/crates/db/src/lib.rs b/litellm-rust/crates/db/src/lib.rs index b41ee5ace78..20d7d7616bb 100644 --- a/litellm-rust/crates/db/src/lib.rs +++ b/litellm-rust/crates/db/src/lib.rs @@ -16,5 +16,6 @@ pub use spend::{ }; pub use sqlx::PgPool; pub use verification_token::{ - Restrictions, Team, TeamMembership, User, VerificationToken, find_verification_token, + RateLimits, Restrictions, Team, TeamMembership, User, VerificationToken, + find_verification_token, }; diff --git a/litellm-rust/crates/db/src/verification_token.rs b/litellm-rust/crates/db/src/verification_token.rs index 3ff008b0cff..5db37ae5c9b 100644 --- a/litellm-rust/crates/db/src/verification_token.rs +++ b/litellm-rust/crates/db/src/verification_token.rs @@ -19,6 +19,9 @@ pub struct VerificationToken { pub key_alias: Option, pub key_name: Option, pub metadata: Value, + /// The key's `rpm_limit` and `tpm_limit`, else its budget table's. + pub limits: RateLimits, + pub max_parallel_requests: Option, pub restrictions: Restrictions, pub team_id: Option, /// `None` when the key has no team or the team row is missing. @@ -38,9 +41,17 @@ pub struct Team { pub spend: f64, pub max_budget: Option, pub metadata: Value, + pub limits: RateLimits, pub restrictions: Restrictions, } +/// Requests and tokens per rate limit window. +#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)] +pub struct RateLimits { + pub rpm: Option, + pub tpm: Option, +} + /// 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 { @@ -50,7 +61,9 @@ pub struct TeamMembership { 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`. + /// The membership budget table's `rpm_limit` and `tpm_limit`. + pub limits: RateLimits, + /// An `rpm_limit` or `tpm_limit` inside the budget table's `model_max_budget`. pub rate_limits: bool, /// The member's own model scope; empty inherits the team's. pub allowed_models: Vec, @@ -61,6 +74,7 @@ pub struct User { pub spend: f64, pub max_budget: Option, pub models: Vec, + pub limits: RateLimits, pub restrictions: Restrictions, } @@ -68,8 +82,8 @@ pub struct User { /// `allowed_routes`. #[derive(Clone, Debug, Default, Eq, PartialEq)] pub struct Restrictions { - /// `rpm_limit`, `tpm_limit`, `tpd_limit` or `max_parallel_requests`, on the row or on a key's - /// budget table, or an `rpm_limit` or `tpm_limit` inside `model_max_budget`. + /// Per-model limits: an `rpm_limit` or `tpm_limit` inside `model_max_budget`, on the row or + /// on a key's budget table. pub rate_limits: bool, pub policies: bool, /// `aliases`, `router_settings`, or a team's `model_id`. @@ -98,11 +112,10 @@ pub async fn find_verification_token( v.key_alias, v.key_name, v.metadata, - (v.rpm_limit IS NOT NULL OR v.tpm_limit IS NOT NULL OR v.tpd_limit IS NOT NULL - OR v.max_parallel_requests IS NOT NULL - OR kb.rpm_limit IS NOT NULL OR kb.tpm_limit IS NOT NULL - OR kb.tpd_limit IS NOT NULL OR kb.max_parallel_requests IS NOT NULL - OR COALESCE(jsonb_path_exists(v.model_max_budget, + COALESCE(v.rpm_limit, kb.rpm_limit) AS rpm_limit, + COALESCE(v.tpm_limit, kb.tpm_limit) AS tpm_limit, + v.max_parallel_requests, + (COALESCE(jsonb_path_exists(v.model_max_budget, '$.* ? (@.rpm_limit != null || @.tpm_limit != null)'), false) OR COALESCE(jsonb_path_exists(kb.model_max_budget, '$.* ? (@.rpm_limit != null || @.tpm_limit != null)'), false)) @@ -123,9 +136,9 @@ pub async fn find_verification_token( t.spend AS "team_spend?", t.max_budget AS "team_max_budget?", t.metadata AS "team_metadata?", - (t.rpm_limit IS NOT NULL OR t.tpm_limit IS NOT NULL OR t.tpd_limit IS NOT NULL - OR t.max_parallel_requests IS NOT NULL - OR COALESCE(jsonb_path_exists(t.model_max_budget, + t.rpm_limit AS "team_rpm_limit?", + t.tpm_limit AS "team_tpm_limit?", + (COALESCE(jsonb_path_exists(t.model_max_budget, '$.* ? (@.rpm_limit != null || @.tpm_limit != null)'), false)) AS "team_rate_limits!", COALESCE(cardinality(t.policies) > 0, false) AS "team_policies!", @@ -133,9 +146,9 @@ pub async fn find_verification_token( OR COALESCE(t.router_settings NOT IN ('{}', 'null'), false)) AS "team_model_aliases!", (t.object_permission_id IS NOT NULL) AS "team_object_permissions!", (t.organization_id IS NOT NULL) AS "team_organization!", - (mb.rpm_limit IS NOT NULL OR mb.tpm_limit IS NOT NULL OR mb.tpd_limit IS NOT NULL - OR mb.max_parallel_requests IS NOT NULL - OR COALESCE(jsonb_path_exists(mb.model_max_budget, + mb.rpm_limit AS "member_rpm_limit?", + mb.tpm_limit AS "member_tpm_limit?", + (COALESCE(jsonb_path_exists(mb.model_max_budget, '$.* ? (@.rpm_limit != null || @.tpm_limit != null)'), false)) AS "member_rate_limits!", mb.allowed_models AS "member_allowed_models?", @@ -149,9 +162,9 @@ pub async fn find_verification_token( 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, + u.rpm_limit AS "user_rpm_limit?", + u.tpm_limit AS "user_tpm_limit?", + (COALESCE(jsonb_path_exists(u.model_max_budget, '$.* ? (@.rpm_limit != null || @.tpm_limit != null)'), false)) AS "user_rate_limits!", COALESCE(cardinality(u.policies) > 0, false) AS "user_policies!", @@ -181,6 +194,11 @@ pub async fn find_verification_token( key_alias: row.key_alias, key_name: row.key_name, metadata: row.metadata, + limits: RateLimits { + rpm: row.rpm_limit, + tpm: row.tpm_limit, + }, + max_parallel_requests: row.max_parallel_requests.map(i64::from), restrictions: Restrictions { rate_limits: row.key_rate_limits, policies: row.key_policies, @@ -199,6 +217,10 @@ pub async fn find_verification_token( spend: row.team_spend.unwrap_or(0.0), max_budget: row.team_max_budget, metadata: row.team_metadata.unwrap_or(Value::Null), + limits: RateLimits { + rpm: row.team_rpm_limit, + tpm: row.team_tpm_limit, + }, restrictions: Restrictions { rate_limits: row.team_rate_limits, policies: row.team_policies, @@ -214,6 +236,10 @@ pub async fn find_verification_token( .member_max_budget .map(|budget| budget + row.member_budget_increase.unwrap_or(0.0)), default_max_budget: row.member_default_budget, + limits: RateLimits { + rpm: row.member_rpm_limit, + tpm: row.member_tpm_limit, + }, rate_limits: row.member_rate_limits, allowed_models: row.member_allowed_models.unwrap_or_default(), }, @@ -222,6 +248,10 @@ pub async fn find_verification_token( spend: row.user_spend.unwrap_or(0.0), max_budget: row.user_max_budget, models: row.user_models.unwrap_or_default(), + limits: RateLimits { + rpm: row.user_rpm_limit, + tpm: row.user_tpm_limit, + }, restrictions: Restrictions { rate_limits: row.user_rate_limits, policies: row.user_policies, diff --git a/litellm-rust/crates/db/tests/verification_token.rs b/litellm-rust/crates/db/tests/verification_token.rs index ffe0b4d66f6..fc4cb83e04f 100644 --- a/litellm-rust/crates/db/tests/verification_token.rs +++ b/litellm-rust/crates/db/tests/verification_token.rs @@ -1,8 +1,8 @@ use std::time::{Duration, SystemTime}; use litellm_db::{ - EndUser, Restrictions, Team, TeamMembership, User, VerificationToken, find_end_user, - find_verification_token, + EndUser, RateLimits, Restrictions, Team, TeamMembership, User, VerificationToken, + find_end_user, find_verification_token, }; use litellm_db_testing::MigratedPostgres; use rstest::rstest; @@ -93,6 +93,8 @@ async fn reads_the_gating_columns_of_a_verification_token_and_its_owners() { key_alias: None, key_name: None, metadata: json!({}), + limits: RateLimits::default(), + max_parallel_requests: None, restrictions: Restrictions::default(), team_id: None, team: None, @@ -115,6 +117,8 @@ async fn reads_the_gating_columns_of_a_verification_token_and_its_owners() { key_alias: Some("prod".into()), key_name: Some("sk-...abcd".into()), metadata: json!({"tags": ["a"]}), + limits: RateLimits::default(), + max_parallel_requests: None, restrictions: Restrictions::default(), team_id: Some("team-1".into()), team: Some(Team { @@ -124,16 +128,21 @@ async fn reads_the_gating_columns_of_a_verification_token_and_its_owners() { spend: 2.5, max_budget: Some(10.0), metadata: json!({"require_trace_id": true}), - restrictions: Restrictions { - rate_limits: true, - ..Restrictions::default() + limits: RateLimits { + rpm: None, + tpm: Some(100), }, + restrictions: Restrictions::default(), }), team_membership: TeamMembership { spend: Some(0.0), max_budget: None, default_max_budget: None, - rate_limits: true, + limits: RateLimits { + rpm: Some(3), + tpm: None, + }, + rate_limits: false, allowed_models: vec![], }, user_id: Some("user-1".into()), @@ -141,8 +150,11 @@ async fn reads_the_gating_columns_of_a_verification_token_and_its_owners() { spend: 0.0, max_budget: Some(1.0), models: vec!["gpt-b".into()], + limits: RateLimits { + rpm: Some(5), + tpm: None, + }, restrictions: Restrictions { - rate_limits: true, policies: true, ..Restrictions::default() }, @@ -152,7 +164,7 @@ async fn reads_the_gating_columns_of_a_verification_token_and_its_owners() { assert_eq!( find("restricted").await.restrictions, Restrictions { - rate_limits: true, + rate_limits: false, policies: true, model_aliases: true, object_permissions: true, @@ -162,16 +174,33 @@ async fn reads_the_gating_columns_of_a_verification_token_and_its_owners() { agent: true, } ); + let limited = find("limited-only").await; assert_eq!( - find("limited-only").await.restrictions, - Restrictions { - rate_limits: true, - ..Restrictions::default() + (limited.max_parallel_requests, limited.restrictions), + (Some(2), Restrictions::default()) + ); + // A key's own limits win; its budget row's fill the gaps. + assert_eq!( + find("restricted").await.limits, + RateLimits { + rpm: Some(10), + tpm: None } ); - // Budgets are not restrictions; limits held in a budget row or per model are. + let budget_row = find("budget-row-limits").await; + assert_eq!( + (budget_row.limits, budget_row.restrictions), + ( + RateLimits { + rpm: Some(3), + tpm: None + }, + Restrictions::default() + ) + ); + // Budgets and plain limits are not restrictions; limits per model are. assert_eq!(find("budgeted").await.restrictions, Restrictions::default()); - for token in ["budget-row-limits", "budget-row-model-limits", "model-rpm"] { + for token in ["budget-row-model-limits", "model-rpm"] { assert_eq!( find(token).await.restrictions, Restrictions { @@ -188,6 +217,7 @@ async fn reads_the_gating_columns_of_a_verification_token_and_its_owners() { spend: Some(0.75), max_budget: Some(5.0), default_max_budget: None, + limits: RateLimits::default(), rate_limits: false, allowed_models: vec!["gpt-a".into()], } @@ -216,15 +246,17 @@ 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(); + let find = async |id: &str, default: Option<&str>| { + find_end_user(pool, id, default, None).await.unwrap() + }; assert_eq!( find("end-1", Some("b-2")).await, EndUser { spend: Some(1.5), blocked: false, - max_budget: Some(5.0) + max_budget: Some(5.0), + limits: RateLimits::default(), } ); assert_eq!( @@ -232,16 +264,37 @@ async fn reads_end_user_budgets_with_their_default() { EndUser { spend: Some(0.25), blocked: true, - max_budget: Some(5.0) + max_budget: Some(5.0), + limits: RateLimits::default(), } ); assert_eq!(find("end-2", None).await.max_budget, None); + // An override replaces the row's own budget. + assert_eq!( + find_end_user(pool, "end-1", None, Some("b-2")) + .await + .unwrap() + .limits, + RateLimits { + rpm: Some(3), + tpm: None + } + ); + // A row without its own budget takes the default's limits too. + assert_eq!( + find("end-2", Some("b-2")).await.limits, + RateLimits { + rpm: Some(3), + tpm: None + } + ); assert_eq!( find("unknown", Some("b-1")).await, EndUser { spend: None, blocked: false, - max_budget: Some(5.0) + max_budget: Some(5.0), + limits: RateLimits::default(), } ); } diff --git a/litellm-rust/crates/gateway-auth/src/error.rs b/litellm-rust/crates/gateway-auth/src/error.rs index 6310e2eda90..30494a7e1a3 100644 --- a/litellm-rust/crates/gateway-auth/src/error.rs +++ b/litellm-rust/crates/gateway-auth/src/error.rs @@ -269,7 +269,7 @@ mod tests { restriction: Restriction::RateLimits, owner: Owner::TeamMembership, }), - "Authentication Error, keys with rate limits set on its team membership are not supported by this gateway yet." + "Authentication Error, keys with per-model, per-tag or dynamic rate limits set on its team membership are not supported by this gateway yet." )] #[case::banned_param( Error::BannedParam("aws_bedrock_runtime_endpoint"), diff --git a/litellm-rust/crates/gateway-auth/src/keys.rs b/litellm-rust/crates/gateway-auth/src/keys.rs index ff99572a1ee..654b8503d3e 100644 --- a/litellm-rust/crates/gateway-auth/src/keys.rs +++ b/litellm-rust/crates/gateway-auth/src/keys.rs @@ -241,7 +241,7 @@ pub enum Refusal { #[derive(Clone, Copy, Debug, Eq, PartialEq, strum::Display, strum::VariantArray)] pub enum Restriction { - #[strum(serialize = "rate limits")] + #[strum(serialize = "per-model, per-tag or dynamic rate limits")] RateLimits, #[strum(serialize = "guardrails or policies")] Guardrails, @@ -298,6 +298,34 @@ pub struct KeyProfile { pub metadata: Arc>, pub team_metadata: Arc>, pub budgets: KeyBudgets, + pub limits: KeyLimits, +} + +/// Requests and tokens per rate limit window. +#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)] +pub struct Limits { + pub rpm: Option, + pub tpm: Option, +} + +impl Limits { + pub fn is_set(&self) -> bool { + self.rpm.is_some() || self.tpm.is_some() + } +} + +/// The rate limits of a key and of who it belongs to, read with the key and refreshed with it. +/// End users' limits are read on their own. +#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)] +pub struct KeyLimits { + /// The key's own limits, else its budget table's. + pub key: Limits, + /// Requests the key may have in flight at once. + pub max_parallel_requests: Option, + pub user: Limits, + pub team: Limits, + /// The key's user within its team, from the membership's budget table. + pub team_member: Limits, } /// Stored spend against an optional budget, as the key store last read it. diff --git a/litellm-rust/crates/gateway-hooks/src/facts.rs b/litellm-rust/crates/gateway-hooks/src/facts.rs index 437768707d0..d696427fd8d 100644 --- a/litellm-rust/crates/gateway-hooks/src/facts.rs +++ b/litellm-rust/crates/gateway-hooks/src/facts.rs @@ -8,7 +8,7 @@ use bytes::Bytes; use http::HeaderMap; use serde_json::{Map, Value}; -use litellm_gateway_auth::keys::KeyBudgets; +use litellm_gateway_auth::keys::{KeyBudgets, KeyLimits}; use crate::usage::Usage; @@ -104,6 +104,7 @@ pub struct KeyFacts { pub metadata: Arc>, pub team_metadata: Arc>, pub budgets: KeyBudgets, + pub limits: KeyLimits, } /// 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 3e57e6b12c8..9246f77d1a2 100644 --- a/litellm-rust/crates/gateway-hooks/src/hook.rs +++ b/litellm-rust/crates/gateway-hooks/src/hook.rs @@ -100,6 +100,8 @@ pub struct Refusal { pub param: Option<&'static str>, pub message: String, pub headers: Vec<(HeaderName, HeaderValue)>, + /// Send Anthropic's error envelope, with `kind` as its error type, instead of the proxy's. + pub anthropic: bool, } impl Refusal { @@ -114,9 +116,17 @@ impl Refusal { param: Some("None"), message: message.into(), headers: Vec::new(), + anthropic: false, } } + /// Sends `{"type": "error", "error": {"type": kind, "message": message}}`, as the proxy's + /// `/v1/messages` route does for some errors. + pub fn anthropic(mut self) -> Self { + self.anthropic = true; + self + } + /// Sends a `null` `param`, as the proxy's budget errors do. pub fn without_param(mut self) -> Self { self.param = None; diff --git a/litellm-rust/crates/gateway-inference/src/error.rs b/litellm-rust/crates/gateway-inference/src/error.rs index 05009604fba..3aa4b44f9aa 100644 --- a/litellm-rust/crates/gateway-inference/src/error.rs +++ b/litellm-rust/crates/gateway-inference/src/error.rs @@ -177,12 +177,18 @@ impl Error { /// A hook's refusal, in the proxy's error shape and with the hook's headers. fn refused(refusal: Refusal) -> Response { let status = StatusCode::from_u16(refusal.status).unwrap_or(StatusCode::INTERNAL_SERVER_ERROR); - let body = json!({"error": { - "message": refusal.message, - "type": refusal.kind, - "param": refusal.param, - "code": status.as_u16().to_string(), - }}); + let body = match refusal.anthropic { + true => json!({"type": "error", "error": { + "type": refusal.kind, + "message": refusal.message, + }}), + false => json!({"error": { + "message": refusal.message, + "type": refusal.kind, + "param": refusal.param, + "code": status.as_u16().to_string(), + }}), + }; let mut response = (status, Json(body)).into_response(); response.headers_mut().extend(refusal.headers); response diff --git a/litellm-rust/crates/gateway-inference/src/hooks.rs b/litellm-rust/crates/gateway-inference/src/hooks.rs index e60bcd140bc..07833d8cc93 100644 --- a/litellm-rust/crates/gateway-inference/src/hooks.rs +++ b/litellm-rust/crates/gateway-inference/src/hooks.rs @@ -121,6 +121,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 index be04ae028be..7196bf132da 100644 --- a/litellm-rust/crates/gateway/src/budgets.rs +++ b/litellm-rust/crates/gateway/src/budgets.rs @@ -17,22 +17,18 @@ use std::{ use futures_util::future::BoxFuture; use litellm_config::Config; -use litellm_db::{EndUser, PgPool}; +use litellm_db::PgPool; use litellm_gateway_hooks::{ Facts, Hook, Prices, Record, Refusal, RequestEdit, RequestFacts, Settle, Targets, Usage, }; use redis::aio::ConnectionManager; -use crate::coordination::Coordination; +use crate::{coordination::Coordination, end_users::EndUsers}; /// 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 @@ -600,45 +596,6 @@ async fn admit_redis( } } -/// 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, @@ -670,10 +627,7 @@ impl Budgets { fail_closed: settings.fail_closed, }), settings, - end_users: database.map(|pool| EndUsers { - pool, - cache: Mutex::default(), - }), + end_users: database.map(EndUsers::new), groups, prices: OnceLock::new(), } @@ -817,7 +771,7 @@ impl Budgets { .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, + Some(end_users) => end_users.get(end_user_id, default_budget_id, None).await, None => None, } .unwrap_or_default(); diff --git a/litellm-rust/crates/gateway/src/end_users.rs b/litellm-rust/crates/gateway/src/end_users.rs new file mode 100644 index 00000000000..5c23e4f6328 --- /dev/null +++ b/litellm-rust/crates/gateway/src/end_users.rs @@ -0,0 +1,70 @@ +//! End users' budgets and limits, read from the proxy's database and reused for a minute. + +use std::{ + collections::HashMap, + sync::Mutex, + time::{Duration, Instant}, +}; + +use litellm_db::{EndUser, PgPool}; + +/// 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; + +/// An end user and the default and override budgets it was read with. +type EndUserKey = (String, Option, Option); + +pub(crate) struct EndUsers { + pool: PgPool, + cache: Mutex>, +} + +impl EndUsers { + pub(crate) fn new(pool: PgPool) -> Self { + Self { + pool, + cache: Mutex::default(), + } + } + + /// `None` when the database could not be read. + pub(crate) async fn get( + &self, + id: &str, + default_budget_id: Option<&str>, + override_budget_id: Option<&str>, + ) -> Option { + let key = ( + id.to_owned(), + default_budget_id.map(str::to_owned), + override_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, override_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 + } + } + } +} diff --git a/litellm-rust/crates/gateway/src/key_status.rs b/litellm-rust/crates/gateway/src/key_status.rs index 50d8bf65567..20e2e795336 100644 --- a/litellm-rust/crates/gateway/src/key_status.rs +++ b/litellm-rust/crates/gateway/src/key_status.rs @@ -1,20 +1,21 @@ -use litellm_db::{Restrictions, VerificationToken}; +use litellm_db::{RateLimits, Restrictions, VerificationToken}; use std::sync::Arc; use litellm_gateway_auth::keys::{ - AllowedRoutes, KeyBudgets, KeyModels, KeyProfile, KeyStatus, MemberSpend, ModelAccess, Owner, - Refusal, Restriction, Spend, + AllowedRoutes, KeyBudgets, KeyLimits, KeyModels, KeyProfile, KeyStatus, Limits, MemberSpend, + ModelAccess, Owner, Refusal, Restriction, Spend, }; use serde_json::Value; /// Key and team `metadata` fields the proxy enforces and this gateway does not. -const METADATA_RESTRICTIONS: [(&str, Restriction); 11] = [ +const METADATA_RESTRICTIONS: [(&str, Restriction); 12] = [ ("guardrails", Restriction::Guardrails), ("policies", Restriction::Guardrails), ("tags", Restriction::Tags), ("model_rpm_limit", Restriction::RateLimits), ("model_tpm_limit", Restriction::RateLimits), ("tag_rpm_limit", Restriction::RateLimits), + ("throttle_on_budget_exceeded", Restriction::RateLimits), ("allowed_tools", Restriction::ToolAllowlists), ("enforced_params", Restriction::EnforcedParams), ("priority", Restriction::Priority), @@ -39,6 +40,7 @@ pub(crate) fn status(row: VerificationToken, personal_keys_use_teams: bool) -> K _ => Arc::default(), }; let budgets = budgets(&row); + let limits = limits(&row); let (team_models, team_alias, team_metadata) = match row.team { Some(team) => ( Some(ModelAccess::new(team.models)), @@ -62,6 +64,7 @@ pub(crate) fn status(row: VerificationToken, personal_keys_use_teams: bool) -> K metadata: object(row.metadata), team_metadata, budgets, + limits, }), } } @@ -96,6 +99,29 @@ fn budgets(row: &VerificationToken) -> KeyBudgets { } } +/// The rate limits the rate limits hook enforces for this key. +fn limits(row: &VerificationToken) -> KeyLimits { + let limits = |limits: RateLimits| Limits { + rpm: limits.rpm, + tpm: limits.tpm, + }; + KeyLimits { + key: limits(row.limits), + max_parallel_requests: row.max_parallel_requests, + user: row + .user + .as_ref() + .map(|user| limits(user.limits)) + .unwrap_or_default(), + team: row + .team + .as_ref() + .map(|team| limits(team.limits)) + .unwrap_or_default(), + team_member: limits(row.team_membership.limits), + } +} + fn refusal(row: &VerificationToken, personal_keys_use_teams: bool) -> Option { let team = match (&row.team_id, &row.team) { (Some(team_id), None) => { @@ -145,7 +171,7 @@ fn unsupported( }) } -/// What the membership's budget table sets: rate limits, or an `allowed_models` list the proxy's +/// What the membership's budget table sets: per-model rate limits, or an `allowed_models` list the proxy's /// `check_team_member_model_access` can deny with (it skips an empty list, and `*` and /// `all-proxy-models` grant every model). fn membership(row: &VerificationToken) -> Option { @@ -195,6 +221,14 @@ fn metadata_restriction(metadata: &Value) -> Option { .iter() .find(|(field, _)| metadata.get(field).is_some_and(is_set)) .map(|(_, restriction)| *restriction) + .or_else(|| { + // Dynamic limits apply only while a deployment is failing, which the gateway does + // not track. + ["rpm_limit_type", "tpm_limit_type"] + .iter() + .any(|field| metadata.get(field).and_then(Value::as_str) == Some("dynamic")) + .then_some(Restriction::RateLimits) + }) } fn is_set(value: &Value) -> bool { @@ -236,6 +270,8 @@ mod tests { key_alias: None, key_name: None, metadata: json!({}), + limits: RateLimits::default(), + max_parallel_requests: None, restrictions: Restrictions::default(), team_id: None, team: None, @@ -253,6 +289,7 @@ mod tests { spend, max_budget, metadata: json!({}), + limits: RateLimits::default(), restrictions: Restrictions::default(), }) } @@ -297,6 +334,46 @@ mod tests { assert_eq!(status(row, false), active(true)); } + #[rstest] + fn carries_the_key_and_owner_rate_limits(row: VerificationToken) { + let limits = |rpm, tpm| RateLimits { rpm, tpm }; + let row = VerificationToken { + limits: limits(Some(10), None), + max_parallel_requests: Some(2), + team_id: Some("team-1".into()), + team: team(false, None, 0.0).map(|team| Team { + limits: limits(None, Some(1_000)), + ..team + }), + team_membership: TeamMembership { + limits: limits(Some(3), Some(30)), + ..TeamMembership::default() + }, + ..on_user(Some("team-1"), &[], NONE, row) + }; + let row = VerificationToken { + user: row.user.clone().map(|user| User { + limits: limits(Some(5), None), + ..user + }), + ..row + }; + let KeyStatus::Active { profile, .. } = status(row, false) else { + panic!("expected an active key"); + }; + let limits = |rpm, tpm| Limits { rpm, tpm }; + assert_eq!( + profile.limits, + KeyLimits { + key: limits(Some(10), None), + max_parallel_requests: Some(2), + user: limits(Some(5), None), + team: limits(None, Some(1_000)), + team_member: limits(Some(3), Some(30)), + } + ); + } + #[rstest] #[case::blocked_key(true, None, None, None, 0.0, KeyStatus::Blocked)] #[case::team_row_missing( @@ -376,6 +453,7 @@ mod tests { spend: 0.0, max_budget: None, models: strings(models), + limits: RateLimits::default(), restrictions, }), ..row @@ -473,6 +551,18 @@ mod tests { |row| key_metadata(json!({"tag_rpm_limit": {"a": 1}}), row), Some((Restriction::RateLimits, Owner::Key)) )] + #[case::metadata_dynamic_rpm( + |row| key_metadata(json!({"rpm_limit_type": "dynamic"}), row), + Some((Restriction::RateLimits, Owner::Key)) + )] + #[case::metadata_guaranteed_tpm( + |row| key_metadata(json!({"tpm_limit_type": "guaranteed_throughput"}), row), + None + )] + #[case::metadata_budget_throttle( + |row| key_metadata(json!({"throttle_on_budget_exceeded": true}), row), + Some((Restriction::RateLimits, Owner::Key)) + )] #[case::metadata_allowed_tools( |row| key_metadata(json!({"allowed_tools": ["search"]}), row), Some((Restriction::ToolAllowlists, Owner::Key)) diff --git a/litellm-rust/crates/gateway/src/lib.rs b/litellm-rust/crates/gateway/src/lib.rs index dcd3aa994a3..57f05ec9322 100644 --- a/litellm-rust/crates/gateway/src/lib.rs +++ b/litellm-rust/crates/gateway/src/lib.rs @@ -2,12 +2,14 @@ mod auth; pub mod budgets; pub mod coordination; pub mod drain; +mod end_users; mod error; pub mod health; pub mod hooks; mod key_status; mod keys; pub mod metrics; +pub mod rate_limits; mod secrets; pub mod spend; pub use error::Error; diff --git a/litellm-rust/crates/gateway/src/main.rs b/litellm-rust/crates/gateway/src/main.rs index 0954fd7634b..3e648f5a758 100644 --- a/litellm-rust/crates/gateway/src/main.rs +++ b/litellm-rust/crates/gateway/src/main.rs @@ -120,6 +120,13 @@ async fn main() -> Result<(), Box> { coordination.clone(), database.clone(), ))); + hooks.push(Arc::new( + litellm_gateway::rate_limits::RateLimiter::for_gateway( + &config, + coordination.clone(), + database.clone(), + ), + )); } let mut observers: Vec> = Vec::new(); let spend = database.zip(litellm_gateway::spend::SpendSettings::from_config(&config)); diff --git a/litellm-rust/crates/gateway/src/rate_limits.rs b/litellm-rust/crates/gateway/src/rate_limits.rs new file mode 100644 index 00000000000..2b9ca57cbbc --- /dev/null +++ b/litellm-rust/crates/gateway/src/rate_limits.rs @@ -0,0 +1,1193 @@ +//! The rate limits hook: refuses a request when its key, user, team, team member or end user is +//! over its requests or tokens per window, or its key has too many requests in flight, with the +//! proxy's counters, order, messages and headers. +//! +//! The counters are the proxy's own Redis keys, `{:}:window|requests|tokens` and +//! `{api_key:}:max_parallel_requests`, changed by scripts with the proxy's semantics, so +//! gateway replicas and Python proxies sharing that Redis enforce one limit. Windows are fixed and +//! start on their first request. A request counts against requests whether or not it is let +//! through; tokens are reserved up front from an estimate and settled to the real count when the +//! response ends. Without Redis, or while it is unreachable, the counters live in this replica's +//! memory. + +use std::{ + collections::HashMap, + sync::{Arc, Mutex}, + time::{Duration, Instant, SystemTime, UNIX_EPOCH}, +}; + +use futures_util::future::BoxFuture; +use http::{HeaderName, HeaderValue}; +use litellm_config::Config; +use litellm_db::PgPool; +use litellm_gateway_auth::keys::Limits; +use litellm_gateway_hooks::{ + Facts, Hook, KeyFacts, Record, Refusal, RequestEdit, RequestFacts, Route, Settle, Targets, +}; +use redis::aio::ConnectionManager; +use serde_json::{Map, Value}; + +use crate::{coordination::Coordination, end_users::EndUsers, spend::civil}; + +const DEFAULT_WINDOW: u64 = 60; +/// How long an in-flight slot counts before it is presumed leaked. +const SLOT_TTL: u64 = 3600; +/// The output tokens assumed at most when a request does not cap them. +const MAX_OUTPUT_FLOOR: i64 = 1024; +/// Windows and slots held in memory before expired ones are dropped. +const MAX_LOCAL: usize = 100_000; +const FAIL_CLOSED: &str = "Rate limit enforcement unavailable: request counters could not be verified against Redis, and fail_closed_rate_limit_enforcement is enabled, so the request was rejected to avoid exceeding the configured rate limit. Retry shortly."; + +/// Counts a request on every `(window, requests)` pair, starting a window that is missing or +/// over, and returns each window's start and count. The caller compares; the count always rises. +const REQUESTS: &str = r#" +local results = {} +local now = tonumber(ARGV[1]) +local window_size = tonumber(ARGV[2]) +for i = 1, #KEYS, 2 do + local window_key = KEYS[i] + local counter_key = KEYS[i + 1] + local window_start = redis.call('GET', window_key) + if not window_start or (now - tonumber(window_start)) >= window_size then + local prefix = string.sub(window_key, 1, -(#':window') - 1) + redis.call('DEL', prefix .. ':requests', prefix .. ':tokens') + redis.call('SET', window_key, tostring(now)) + redis.call('SET', counter_key, 1) + redis.call('EXPIRE', window_key, window_size) + redis.call('EXPIRE', counter_key, window_size) + table.insert(results, tostring(now)) + table.insert(results, 1) + else + local counter = redis.call('INCR', counter_key) + if redis.call('TTL', counter_key) == -1 then + redis.call('EXPIRE', counter_key, window_size) + end + table.insert(results, window_start) + table.insert(results, counter) + end +end +return results +"#; + +/// Adds the increment to every `(window, tokens)` pair, or to none when one would go over its +/// limit: `{0, count, start, ...}` or `{1, index, current, limit}`. +const RESERVE: &str = r#" +local now = tonumber(redis.call('TIME')[1]) +local count = #KEYS / 2 +local state = {} +for i = 1, count do + local window_key = KEYS[(i - 1) * 2 + 1] + local counter_key = KEYS[(i - 1) * 2 + 2] + local base = (i - 1) * 4 + 1 + local limit = tonumber(ARGV[base]) + local increment = tonumber(ARGV[base + 1]) + local window_size = tonumber(ARGV[base + 3]) + local window_start = redis.call('GET', window_key) + local expired = (not window_start) or ((now - tonumber(window_start)) >= window_size) + local current = 0 + if not expired then + current = tonumber(redis.call('GET', counter_key) or 0) + end + local blocked + if increment > 0 then + blocked = current + increment > limit + else + blocked = current >= limit + end + if blocked then + return { 1, i, current, limit } + end + state[i] = { expired, window_start } +end +local results = { 0 } +local reset = {} +for i = 1, count do + local window_key = KEYS[(i - 1) * 2 + 1] + local counter_key = KEYS[(i - 1) * 2 + 2] + local base = (i - 1) * 4 + 1 + local increment = tonumber(ARGV[base + 1]) + local ttl = tonumber(ARGV[base + 2]) + local window_size = tonumber(ARGV[base + 3]) + if state[i][1] then + if not reset[window_key] then + local prefix = string.sub(window_key, 1, -(#':window') - 1) + redis.call('DEL', prefix .. ':requests', prefix .. ':tokens') + reset[window_key] = true + end + redis.call('SET', window_key, tostring(now)) + redis.call('SET', counter_key, increment) + redis.call('EXPIRE', window_key, window_size) + if ttl > 0 then + redis.call('EXPIRE', counter_key, ttl) + end + table.insert(results, increment) + table.insert(results, now) + else + local counter = redis.call('INCRBY', counter_key, increment) + if redis.call('TTL', counter_key) == -1 and ttl > 0 then + redis.call('EXPIRE', counter_key, ttl) + end + table.insert(results, counter) + table.insert(results, tonumber(state[i][2])) + end +end +return results +"#; + +/// Takes an in-flight slot unless the gauge is full: `{0, in_flight}` or `{1, 1, in_flight, limit}`. +const ACQUIRE: &str = r#" +local now = tonumber(redis.call('TIME')[1]) +local limit = tonumber(ARGV[1]) +local slot_ttl = tonumber(ARGV[2]) +redis.call('ZREMRANGEBYSCORE', KEYS[1], '-inf', now - slot_ttl) +local in_flight = redis.call('ZCARD', KEYS[1]) +if in_flight + 1 > limit then + return { 1, 1, in_flight, limit } +end +redis.call('ZADD', KEYS[1], now, ARGV[3]) +redis.call('EXPIRE', KEYS[1], slot_ttl) +return { 0, redis.call('ZCARD', KEYS[1]) } +"#; + +/// Adds each token delta, giving a counter without an expiry the window's. +const CHARGE: &str = r#" +local results = {} +for i = 1, #KEYS do + local value = redis.call('INCRBYFLOAT', KEYS[i], tonumber(ARGV[i * 2 - 1])) + local ttl = tonumber(ARGV[i * 2]) + if ttl and ttl > 0 and redis.call('TTL', KEYS[i]) == -1 then + redis.call('EXPIRE', KEYS[i], ttl) + end + table.insert(results, value) +end +return results +"#; + +#[derive(Clone, Debug, PartialEq)] +pub struct RateLimitSettings { + /// `LITELLM_RATE_LIMIT_WINDOW_SIZE`, in seconds; 60 by default. + pub window: u64, + /// `general_settings.fail_closed_rate_limit_enforcement`: refuse with a 503 rather than fall + /// back to this replica's counters while Redis is unreachable. + pub fail_closed: bool, + /// `litellm_settings.max_end_user_budget_id`: the limits of end users without their own. + pub end_user_budget_id: Option, +} + +impl Default for RateLimitSettings { + fn default() -> Self { + Self { + window: DEFAULT_WINDOW, + fail_closed: false, + end_user_budget_id: None, + } + } +} + +impl RateLimitSettings { + pub fn from_config(config: &Config) -> Self { + let general = &config.general_settings.additional_fields; + Self { + window: std::env::var("LITELLM_RATE_LIMIT_WINDOW_SIZE") + .ok() + .and_then(|window| window.trim().parse().ok()) + .filter(|window| *window > 0) + .unwrap_or(DEFAULT_WINDOW), + fail_closed: match general.get("fail_closed_rate_limit_enforcement") { + Some(litellm_config::Value::Bool(flag)) => *flag, + Some(litellm_config::Value::String(text)) => text.eq_ignore_ascii_case("true"), + _ => false, + }, + end_user_budget_id: config + .litellm_settings + .additional_fields + .get("max_end_user_budget_id") + .and_then(litellm_config::Value::as_str) + .filter(|id| !id.is_empty()) + .map(str::to_owned), + } + } +} + +/// One scope a request counts against: `{key}:{value}` with its limits. +#[derive(Clone, Debug, PartialEq)] +struct Descriptor { + key: &'static str, + value: String, + limits: Limits, + max_parallel_requests: Option, +} + +impl Descriptor { + fn prefix(&self) -> String { + format!("{{{}:{}}}", self.key, self.value) + } +} + +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +enum Kind { + Requests, + Tokens, + Parallel, +} + +impl Kind { + fn name(self) -> &'static str { + match self { + Self::Requests => "requests", + Self::Tokens => "tokens", + Self::Parallel => "max_parallel_requests", + } + } +} + +/// Where a descriptor stands after this request. +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +struct Status { + descriptor: usize, + kind: Kind, + limit: i64, + remaining: i64, +} + +/// The token reservation went over on the descriptor at `index`, which held `current`. +#[derive(Debug, PartialEq)] +enum Reserved { + Counts(Vec), + Over { index: usize, current: i64 }, +} + +/// Redis could not be used and the settings say to refuse rather than fall back. +struct Unavailable; + +#[derive(Debug)] +struct Window { + start: u64, + requests: i64, + tokens: i64, +} + +/// This replica's counters, used without Redis or while it is unreachable. +#[derive(Default)] +struct Local { + windows: HashMap, + slots: HashMap>, +} + +impl Local { + fn window(&mut self, prefix: &str, now: u64, size: u64) -> (&mut Window, bool) { + if self.windows.len() >= MAX_LOCAL { + self.windows + .retain(|_, window| now.saturating_sub(window.start) < size); + } + let window = self.windows.entry(prefix.to_owned()).or_insert(Window { + start: now, + requests: 0, + tokens: 0, + }); + let expired = now.saturating_sub(window.start) >= size; + (window, expired) + } + + fn requests(&mut self, prefixes: &[String], now: u64, size: u64) -> Vec { + prefixes + .iter() + .map(|prefix| { + let (window, expired) = self.window(prefix, now, size); + if expired { + *window = Window { + start: now, + requests: 0, + tokens: 0, + }; + } + window.requests += 1; + window.requests + }) + .collect() + } + + fn reserve( + &mut self, + scopes: &[(String, i64)], + increment: i64, + now: u64, + size: u64, + ) -> Reserved { + for (index, (prefix, limit)) in scopes.iter().enumerate() { + let current = match self.windows.get(prefix) { + Some(window) if now.saturating_sub(window.start) < size => window.tokens, + _ => 0, + }; + if blocked(current, increment, *limit) { + return Reserved::Over { index, current }; + } + } + Reserved::Counts( + scopes + .iter() + .map(|(prefix, _)| { + let (window, expired) = self.window(prefix, now, size); + if expired { + *window = Window { + start: now, + requests: 0, + tokens: 0, + }; + } + window.tokens += increment; + window.tokens + }) + .collect(), + ) + } + + fn charge(&mut self, deltas: &[(String, i64)]) { + for (prefix, delta) in deltas { + if let Some(window) = self.windows.get_mut(prefix) { + window.tokens += delta; + } + } + } + + fn acquire(&mut self, prefix: &str, limit: i64, slot: &str) -> Result { + let ttl = Duration::from_secs(SLOT_TTL); + if self.slots.len() >= MAX_LOCAL { + self.slots.retain(|_, slots| !slots.is_empty()); + } + let slots = self.slots.entry(prefix.to_owned()).or_default(); + slots.retain(|_, taken| taken.elapsed() < ttl); + let in_flight = slots.len() as i64; + if in_flight + 1 > limit { + return Err(in_flight); + } + slots.insert(slot.to_owned(), Instant::now()); + Ok(in_flight + 1) + } + + fn release(&mut self, prefix: &str, slot: &str) { + if let Some(slots) = self.slots.get_mut(prefix) { + slots.remove(slot); + } + } +} + +fn blocked(current: i64, increment: i64, limit: i64) -> bool { + match increment > 0 { + true => current + increment > limit, + false => current >= limit, + } +} + +/// Where the counters live. +struct Store { + redis: Option>, + local: Mutex, + window: u64, + fail_closed: bool, +} + +impl Store { + fn connection(&self) -> Option<(ConnectionManager, &Coordination)> { + let redis = self.redis.as_deref()?; + Some((redis.connection()?, redis)) + } + + fn local(&self) -> std::sync::MutexGuard<'_, Local> { + self.local.lock().unwrap_or_else(|e| e.into_inner()) + } + + /// Logs why Redis was not used; refuses when the settings say to. + fn fall_back(&self, error: Option) -> Result<(), Unavailable> { + match error { + Some(error) => { + tracing::warn!(%error, "rate limit counters failed; using this replica's") + } + None => { + tracing::debug!("Redis not connected; using this replica's rate limit counters") + } + } + match self.fail_closed { + true => Err(Unavailable), + false => Ok(()), + } + } + + /// Counts the request on each prefix's window and returns the counts. + async fn requests(&self, prefixes: &[String]) -> Result, Unavailable> { + let now = epoch_seconds(); + if self.redis.is_some() { + let error = match self.connection() { + Some((mut connection, redis)) => { + let mut command = redis::cmd("EVAL"); + command.arg(REQUESTS).arg(prefixes.len() * 2); + for prefix in prefixes { + command + .arg(redis.namespaced(&format!("{prefix}:window"))) + .arg(redis.namespaced(&format!("{prefix}:requests"))); + } + command.arg(now).arg(self.window); + match integers(command.query_async(&mut connection).await) { + Ok(reply) if reply.len() == prefixes.len() * 2 => { + return Ok(reply.iter().skip(1).step_by(2).copied().collect()); + } + Ok(_) => Some(malformed()), + Err(error) => Some(error), + } + } + None => None, + }; + self.fall_back(error)?; + } + Ok(self.local().requests(prefixes, now, self.window)) + } + + /// Takes an in-flight slot: the count with it, or the full count without. `true` when the + /// slot is in Redis. + async fn acquire( + &self, + prefix: &str, + limit: i64, + slot: &str, + ) -> Result<(Result, bool), Unavailable> { + if self.redis.is_some() { + let error = match self.connection() { + Some((mut connection, redis)) => { + let reply = redis::cmd("EVAL") + .arg(ACQUIRE) + .arg(1) + .arg(redis.namespaced(&format!("{prefix}:max_parallel_requests"))) + .arg(limit) + .arg(SLOT_TTL) + .arg(slot) + .query_async(&mut connection) + .await; + match integers(reply) { + Ok(reply) => match reply.as_slice() { + [0, in_flight] => return Ok((Ok(*in_flight), true)), + [1, _, in_flight, _] => return Ok((Err(*in_flight), true)), + _ => Some(malformed()), + }, + Err(error) => Some(error), + } + } + None => None, + }; + self.fall_back(error)?; + } + Ok((self.local().acquire(prefix, limit, slot), false)) + } + + async fn release(&self, prefix: &str, slot: &str, shared: bool) { + if shared && let Some((mut connection, redis)) = self.connection() { + let key = redis.namespaced(&format!("{prefix}:max_parallel_requests")); + if let Err(error) = redis::cmd("ZREM") + .arg(key) + .arg(slot) + .query_async::(&mut connection) + .await + { + tracing::warn!(%error, "releasing an in-flight slot failed; it expires on its own"); + } + return; + } + self.local().release(prefix, slot); + } + + /// Reserves `increment` tokens on every `(prefix, limit)`, or on none. `true` when the + /// reservation is in Redis. + async fn reserve( + &self, + scopes: &[(String, i64)], + increment: i64, + ) -> Result<(Reserved, bool), Unavailable> { + if self.redis.is_some() { + let error = match self.connection() { + Some((mut connection, redis)) => { + let mut command = redis::cmd("EVAL"); + command.arg(RESERVE).arg(scopes.len() * 2); + for (prefix, _) in scopes { + command + .arg(redis.namespaced(&format!("{prefix}:window"))) + .arg(redis.namespaced(&format!("{prefix}:tokens"))); + } + for (_, limit) in scopes { + command + .arg(*limit) + .arg(increment) + .arg(self.window) + .arg(self.window); + } + match integers(command.query_async(&mut connection).await) { + Ok(reply) if reply.first() == Some(&0) => { + return Ok(( + Reserved::Counts( + reply.iter().skip(1).step_by(2).copied().collect(), + ), + true, + )); + } + Ok(reply) => match reply.as_slice() { + [1, index, current, _] if (1..=scopes.len() as i64).contains(index) => { + return Ok(( + Reserved::Over { + index: (*index - 1) as usize, + current: *current, + }, + true, + )); + } + _ => Some(malformed()), + }, + Err(error) => Some(error), + } + } + None => None, + }; + self.fall_back(error)?; + } + Ok(( + self.local() + .reserve(scopes, increment, epoch_seconds(), self.window), + false, + )) + } + + /// Adds each delta to its prefix's tokens. + async fn charge(&self, deltas: Vec<(String, i64)>, shared: bool) { + if shared && let Some((mut connection, redis)) = self.connection() { + let mut command = redis::cmd("EVAL"); + command.arg(CHARGE).arg(deltas.len()); + for (prefix, _) in &deltas { + command.arg(redis.namespaced(&format!("{prefix}:tokens"))); + } + for (_, delta) in &deltas { + command.arg(*delta).arg(self.window); + } + if let Err(error) = command.query_async::(&mut connection).await { + tracing::warn!(%error, "settling reserved tokens failed"); + } + return; + } + self.local().charge(&deltas); + } +} + +fn malformed() -> redis::RedisError { + redis::RedisError::from(( + redis::ErrorKind::UnexpectedReturnType, + "unexpected rate limit script reply", + )) +} + +/// A script's array reply as integers; counters may come back as strings. +fn integers(reply: redis::RedisResult>) -> redis::RedisResult> { + reply? + .iter() + .map(|value| { + let text = match value { + redis::Value::Int(number) => return Ok(*number), + redis::Value::BulkString(bytes) => std::str::from_utf8(bytes).ok(), + redis::Value::SimpleString(text) => Some(text.as_str()), + _ => None, + }; + text.and_then(|text| text.parse::().ok()) + .map(|number| number as i64) + .ok_or_else(malformed) + }) + .collect() +} + +fn epoch_seconds() -> u64 { + SystemTime::now() + .duration_since(UNIX_EPOCH) + .unwrap_or_default() + .as_secs() +} + +/// `%Y-%m-%d %H:%M:%S UTC`, as the proxy prints when a limit resets. +fn utc(seconds: u64) -> String { + let (year, month, day) = civil((seconds / 86_400) as i64); + let time = seconds % 86_400; + format!( + "{year:04}-{month:02}-{day:02} {:02}:{:02}:{:02} UTC", + time / 3600, + time % 3600 / 60, + time % 60 + ) +} + +/// The proxy's token estimate for a request: a token per four characters of message text, +/// plus the output it may produce for every candidate. +#[derive(Debug, PartialEq)] +struct Estimate { + tokens: i64, + /// The output cap to send upstream when the request has none and the limits are small. + implicit_max_tokens: Option, +} + +fn estimate( + body: &Map, + key: Option<&KeyFacts>, + model: &str, + min_tpm: i64, +) -> Estimate { + let chars: usize = body + .get("messages") + .and_then(Value::as_array) + .into_iter() + .flatten() + .map(|message| match message.get("content") { + Some(Value::String(text)) => text.chars().count(), + Some(Value::Array(blocks)) => blocks + .iter() + .filter_map(|block| block.get("text").and_then(Value::as_str)) + .map(|text| text.chars().count()) + .sum(), + _ => 0, + }) + .sum(); + let chars = i64::try_from(chars).unwrap_or(i64::MAX); + let input = match chars { + 0 => 0, + chars => (chars / 4).max(1), + }; + let integer = |value: &Value| match value { + Value::Number(number) => number.as_f64().map(|number| number as i64), + Value::String(text) => text.trim().parse::().ok().map(|number| number as i64), + _ => None, + }; + let cap = ["max_tokens", "max_completion_tokens"] + .iter() + .filter_map(|field| body.get(*field).and_then(integer)) + .max(); + let configured = key.and_then(|key| { + [&key.metadata, &key.team_metadata] + .into_iter() + .find_map(|metadata| { + metadata + .get("default_estimated_output_tokens_per_model") + .and_then(|models| models.get(model)) + .and_then(integer) + .or_else(|| { + metadata + .get("default_estimated_output_tokens") + .and_then(integer) + }) + }) + }); + let floor = (min_tpm / 4).clamp(1, MAX_OUTPUT_FLOOR); + let output = match (cap, configured) { + (Some(cap), _) => cap, + (None, Some(configured)) => configured, + (None, None) if chars == 0 => 0, + (None, None) => input.max(floor), + }; + let candidates = ["n", "best_of"] + .iter() + .filter_map(|field| body.get(*field).and_then(integer)) + .fold(1, i64::max); + Estimate { + tokens: input + .saturating_add(output.saturating_mul(candidates)) + .max(1), + implicit_max_tokens: (floor < MAX_OUTPUT_FLOOR && cap.is_none()) + .then(|| floor.max(configured.unwrap_or(0))), + } +} + +/// The `rate_limits` hook. +pub struct RateLimiter { + settings: RateLimitSettings, + store: Arc, + end_users: Option, +} + +impl RateLimiter { + /// `redis`, when given, holds the shared counters; `database` gives end users' limits. + pub fn new( + settings: RateLimitSettings, + redis: Option>, + database: Option, + ) -> Self { + Self { + store: Arc::new(Store { + redis, + local: Mutex::default(), + window: settings.window, + fail_closed: settings.fail_closed, + }), + settings, + end_users: database.map(EndUsers::new), + } + } + + pub fn for_gateway( + config: &Config, + redis: Option>, + database: Option, + ) -> Self { + Self::new(RateLimitSettings::from_config(config), redis, database) + } + + /// The scopes this request counts against, in the proxy's order. + async fn descriptors(&self, request: &RequestFacts) -> Vec { + let Some(key) = request.key.as_deref() else { + return Vec::new(); + }; + let limits = &key.limits; + let mut descriptors = Vec::new(); + let mut push = |key: &'static str, value: String, limits: Limits, parallel: Option| { + if limits.is_set() || parallel.is_some() { + descriptors.push(Descriptor { + key, + value, + limits, + max_parallel_requests: parallel, + }); + } + }; + if !key.hash.is_empty() { + push( + "api_key", + key.hash.clone(), + limits.key, + limits.max_parallel_requests, + ); + } + if let Some(user_id) = &key.user_id { + push("user", user_id.clone(), limits.user, None); + } + if let Some(team_id) = &key.team_id { + push("team", team_id.clone(), limits.team, None); + } + if let (Some(team_id), Some(user_id)) = (&key.team_id, &key.user_id) { + push( + "team_member", + format!("{team_id}:{user_id}"), + limits.team_member, + None, + ); + } + if let (Some(end_user_id), Some(end_users)) = (request.end_user.as_deref(), &self.end_users) + && let Some(end_user) = end_users + .get( + end_user_id, + self.settings.end_user_budget_id.as_deref(), + key.budgets.end_user_budget_id.as_deref(), + ) + .await + { + push( + "end_user", + end_user_id.to_owned(), + Limits { + rpm: end_user.limits.rpm, + tpm: end_user.limits.tpm, + }, + None, + ); + } + descriptors + } + + fn refusal(&self, request: &RequestFacts, descriptor: &Descriptor, status: Status) -> Refusal { + let reset_at = utc(epoch_seconds() + self.settings.window); + let message = format!( + "Rate limit exceeded for {}: {}. Limit type: {}. Current limit: {}, Remaining: {}. Limit resets at: {reset_at}", + descriptor.key, + descriptor.value, + status.kind.name(), + status.limit, + status.remaining.max(0), + ); + let mut refusal = match request.route { + Route::Messages => Refusal::new(429, "rate_limit_error", message).anthropic(), + Route::ChatCompletions => { + let mut refusal = Refusal::new(429, "throttling_error", message).without_param(); + for (name, value) in [ + ("retry-after", self.settings.window.to_string()), + ("rate_limit_type", status.kind.name().to_owned()), + ("reset_at", reset_at), + ] { + if let Ok(value) = HeaderValue::from_str(&value) { + refusal = refusal.with_header(HeaderName::from_static(name), value); + } + } + refusal + } + }; + for (name, value) in key_limit_headers(request.key.as_deref()) { + refusal = refusal.with_header(name, value); + } + refusal + } + + fn unavailable(&self, request: &RequestFacts) -> Refusal { + match request.route { + Route::Messages => Refusal::new(503, "api_error", FAIL_CLOSED).anthropic(), + Route::ChatCompletions => { + Refusal::new(503, "internal_server_error", FAIL_CLOSED).without_param() + } + } + } +} + +/// `x-litellm-key-tpm-limit` and `x-litellm-key-rpm-limit`, which the proxy sends whenever the +/// key sets them. +fn key_limit_headers(key: Option<&KeyFacts>) -> Vec<(HeaderName, HeaderValue)> { + let Some(key) = key else { + return Vec::new(); + }; + [ + ("x-litellm-key-tpm-limit", key.limits.key.tpm), + ("x-litellm-key-rpm-limit", key.limits.key.rpm), + ] + .into_iter() + .filter_map(|(name, limit)| Some((HeaderName::from_static(name), HeaderValue::from(limit?)))) + .collect() +} + +impl Hook for RateLimiter { + fn name(&self) -> &'static str { + "rate_limits" + } + + fn reads(&self) -> Facts { + Facts::KEY | Facts::END_USER | Facts::REQUEST_BODY | Facts::USAGE + } + + fn writes(&self) -> Targets { + Targets::ADMISSION | Targets::UPSTREAM_BODY | Targets::RESPONSE_HEADERS + } + + fn run<'a>( + &'a self, + request: &'a Arc, + edit: &'a mut RequestEdit, + ) -> BoxFuture<'a, Result>, Refusal>> { + Box::pin(async move { + let descriptors = self.descriptors(request).await; + if descriptors.is_empty() { + return Ok(None); + } + let unavailable = |Unavailable| self.unavailable(request); + let mut statuses = Vec::new(); + + // Requests per window, counted whether or not the request goes through. + let windowed: Vec = (0..descriptors.len()) + .filter(|index| descriptors[*index].limits.rpm.is_some()) + .collect(); + if !windowed.is_empty() { + let prefixes: Vec = windowed + .iter() + .map(|index| descriptors[*index].prefix()) + .collect(); + let counts = self.store.requests(&prefixes).await.map_err(unavailable)?; + for (index, count) in windowed.iter().zip(counts) { + let limit = descriptors[*index].limits.rpm.unwrap_or_default(); + statuses.push(Status { + descriptor: *index, + kind: Kind::Requests, + limit, + remaining: limit - count, + }); + } + if let Some(over) = statuses.iter().find(|status| status.remaining < 0) { + return Err(self.refusal(request, &descriptors[over.descriptor], *over)); + } + } + + // Requests in flight, for the key only. + let mut slot = None; + if let Some(index) = descriptors + .iter() + .position(|descriptor| descriptor.max_parallel_requests.is_some()) + { + let descriptor = &descriptors[index]; + let limit = descriptor.max_parallel_requests.unwrap_or_default(); + let prefix = descriptor.prefix(); + let id = uuid::Uuid::new_v4().simple().to_string(); + let (taken, shared) = self + .store + .acquire(&prefix, limit, &id) + .await + .map_err(unavailable)?; + let in_flight = match taken { + Ok(in_flight) => in_flight, + Err(in_flight) => { + let status = Status { + descriptor: index, + kind: Kind::Parallel, + limit, + remaining: limit - in_flight, + }; + return Err(self.refusal(request, descriptor, status)); + } + }; + statuses.push(Status { + descriptor: index, + kind: Kind::Parallel, + limit, + remaining: (limit - in_flight).max(0), + }); + slot = Some(Slot { prefix, id, shared }); + } + + // Tokens per window, reserved from an estimate and settled when the response ends. + let limited: Vec<(usize, i64)> = descriptors + .iter() + .enumerate() + .filter_map(|(index, descriptor)| Some((index, descriptor.limits.tpm?))) + .collect(); + let mut tokens = None; + if let Some(min_tpm) = limited.iter().map(|(_, limit)| *limit).min() { + let body: Map = + serde_json::from_slice(request.body()).unwrap_or_default(); + let estimate = estimate( + &body, + request.key.as_deref(), + &request.requested_model, + min_tpm, + ); + if let Some(max_tokens) = estimate.implicit_max_tokens { + edit.set_upstream_body_field("max_tokens", Value::from(max_tokens)); + } + let scopes: Vec<(String, i64)> = limited + .iter() + .map(|(index, limit)| (descriptors[*index].prefix(), *limit)) + .collect(); + let reserved = self.store.reserve(&scopes, estimate.tokens).await; + let release = |slot: Option| async { + if let Some(slot) = slot { + self.store + .release(&slot.prefix, &slot.id, slot.shared) + .await; + } + }; + let (reserved, shared) = match reserved { + Ok(reserved) => reserved, + Err(Unavailable) => { + release(slot).await; + return Err(self.unavailable(request)); + } + }; + match reserved { + Reserved::Counts(counts) => { + for ((index, limit), count) in limited.iter().zip(counts) { + statuses.push(Status { + descriptor: *index, + kind: Kind::Tokens, + limit: *limit, + remaining: (limit - count).max(0), + }); + } + } + Reserved::Over { index, current } => { + release(slot).await; + let (descriptor, limit) = limited[index]; + let status = Status { + descriptor, + kind: Kind::Tokens, + limit, + remaining: limit - current, + }; + return Err(self.refusal(request, &descriptors[descriptor], status)); + } + } + tokens = Some(Tokens { + prefixes: scopes.into_iter().map(|(prefix, _)| prefix).collect(), + reserved: estimate.tokens, + shared, + }); + } + + // The proxy sends where each limit stands on complete responses only. + if !request.stream { + for status in &statuses { + let key = descriptors[status.descriptor].key; + let kind = status.kind.name(); + for (name, value) in [ + ( + format!("x-ratelimit-{key}-remaining-{kind}"), + status.remaining, + ), + (format!("x-ratelimit-{key}-limit-{kind}"), status.limit), + ] { + if let Ok(name) = HeaderName::from_bytes(name.as_bytes()) { + edit.add_response_header(name, HeaderValue::from(value)); + } + } + } + } + for (name, value) in key_limit_headers(request.key.as_deref()) { + edit.add_response_header(name, value); + } + if slot.is_none() && tokens.is_none() { + return Ok(None); + } + Ok(Some(Box::new(Held { + store: self.store.clone(), + slot, + tokens, + }) as Box)) + }) + } +} + +struct Slot { + prefix: String, + id: String, + shared: bool, +} + +struct Tokens { + prefixes: Vec, + reserved: i64, + shared: bool, +} + +/// What a request holds until it ends: an in-flight slot and reserved tokens. +struct Held { + store: Arc, + slot: Option, + tokens: Option, +} + +impl Settle for Held { + fn settle(self: Box, record: &Record) { + let Held { + store, + slot, + tokens, + } = *self; + let charge = tokens.and_then(|tokens| { + // The proxy counts every token but cache reads. Without usage, a success keeps its + // reservation and a failure returns it. + let actual = match (&record.outcome.usage, record.succeeded()) { + (Some(usage), _) => i64::try_from(usage.total_tokens() - usage.cache_read_tokens) + .unwrap_or(i64::MAX), + (None, true) => tokens.reserved, + (None, false) => 0, + }; + let delta = actual - tokens.reserved; + (delta != 0).then(|| { + ( + tokens + .prefixes + .into_iter() + .map(|prefix| (prefix, delta)) + .collect::>(), + tokens.shared, + ) + }) + }); + tokio::spawn(async move { + if let Some(slot) = slot { + store.release(&slot.prefix, &slot.id, slot.shared).await; + } + if let Some((deltas, shared)) = charge { + store.charge(deltas, shared).await; + } + }); + } +} + +#[cfg(test)] +mod tests { + use rstest::rstest; + use serde_json::json; + + use super::*; + + fn body(value: Value) -> Map { + match value { + Value::Object(object) => object, + _ => unreachable!(), + } + } + + #[rstest] + #[case::utc_epoch(0, "1970-01-01 00:00:00 UTC")] + #[case::utc_later(1_792_238_465, "2026-10-17 12:01:05 UTC")] + fn prints_reset_times_like_the_proxy(#[case] seconds: u64, #[case] expected: &str) { + assert_eq!(utc(seconds), expected); + } + + #[rstest] + // 40 characters: 10 input tokens; the small limit floors output at 100_000 / 4 = 1024. + #[case::uncapped(json!({"messages": [{"role": "user", "content": "a".repeat(40)}]}), 100_000, 10 + 1024, None)] + // A small limit floors output at 400 / 4 = 100 and caps the request there. + #[case::small_limit(json!({"messages": [{"role": "user", "content": "a".repeat(40)}]}), 400, 110, Some(100))] + #[case::capped(json!({"max_tokens": 50, "messages": [{"role": "user", "content": [{"type": "text", "text": "a".repeat(8)}]}]}), 400, 52, None)] + #[case::candidates(json!({"max_completion_tokens": 5, "n": 3, "messages": [{"role": "user", "content": "abcd"}]}), 100_000, 16, None)] + #[case::no_text(json!({"messages": []}), 100_000, 1, None)] + fn estimates_tokens_like_the_proxy( + #[case] request: Value, + #[case] min_tpm: i64, + #[case] tokens: i64, + #[case] implicit_max_tokens: Option, + ) { + assert_eq!( + estimate(&body(request), None, "m", min_tpm), + Estimate { + tokens, + implicit_max_tokens + } + ); + } + + #[rstest] + fn configured_output_estimates_come_from_the_key_then_the_team() { + let metadata = |value: Value| Arc::new(body(value)); + let key = KeyFacts { + metadata: metadata(json!({"default_estimated_output_tokens_per_model": {"m": 7}})), + team_metadata: metadata(json!({"default_estimated_output_tokens": 9})), + ..KeyFacts::default() + }; + let request = body(json!({"messages": [{"role": "user", "content": "abcd"}]})); + assert_eq!(estimate(&request, Some(&key), "m", 100_000).tokens, 8); + assert_eq!(estimate(&request, Some(&key), "other", 100_000).tokens, 10); + } + + #[rstest] + fn local_windows_count_reserve_and_reset() { + let mut local = Local::default(); + let prefixes = vec!["{api_key:h}".to_owned()]; + assert_eq!(local.requests(&prefixes, 100, 60), [1]); + assert_eq!(local.requests(&prefixes, 159, 60), [2]); + let scopes = vec![("{api_key:h}".to_owned(), 10)]; + assert_eq!( + local.reserve(&scopes, 6, 159, 60), + Reserved::Counts(vec![6]) + ); + assert_eq!( + local.reserve(&scopes, 6, 159, 60), + Reserved::Over { + index: 0, + current: 6 + } + ); + local.charge(&[("{api_key:h}".to_owned(), -4)]); + assert_eq!( + local.reserve(&scopes, 8, 159, 60), + Reserved::Counts(vec![10]) + ); + // The window started at 100 is over at 160: requests and tokens start again. + assert_eq!(local.requests(&prefixes, 160, 60), [1]); + assert_eq!( + local.reserve(&scopes, 10, 160, 60), + Reserved::Counts(vec![10]) + ); + } + + #[rstest] + fn local_slots_fill_and_release() { + let mut local = Local::default(); + assert_eq!(local.acquire("{api_key:h}", 2, "a"), Ok(1)); + assert_eq!(local.acquire("{api_key:h}", 2, "b"), Ok(2)); + assert_eq!(local.acquire("{api_key:h}", 2, "c"), Err(2)); + local.release("{api_key:h}", "a"); + assert_eq!(local.acquire("{api_key:h}", 2, "c"), Ok(2)); + } +} diff --git a/litellm-rust/crates/gateway/src/spend.rs b/litellm-rust/crates/gateway/src/spend.rs index fc373307bec..a2f0d077819 100644 --- a/litellm-rust/crates/gateway/src/spend.rs +++ b/litellm-rust/crates/gateway/src/spend.rs @@ -648,7 +648,7 @@ fn usage_values(usage: &Usage) -> (Value, Value) { } /// Days since 1970-01-01 as a UTC calendar date. -fn civil(days: i64) -> (i64, u32, u32) { +pub(crate) fn civil(days: i64) -> (i64, u32, u32) { let days = days + 719_468; let era = days.div_euclid(146_097); let day_of_era = days.rem_euclid(146_097); diff --git a/litellm-rust/crates/gateway/tests/rate_limits.rs b/litellm-rust/crates/gateway/tests/rate_limits.rs new file mode 100644 index 00000000000..2e7101bd89c --- /dev/null +++ b/litellm-rust/crates/gateway/tests/rate_limits.rs @@ -0,0 +1,515 @@ +//! The rate limits hook on real requests: refusals and headers as the proxy sends them, tokens +//! reserved and settled to usage, in-flight slots, 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, Mutex, + atomic::{AtomicUsize, Ordering}, + }, + time::Duration, +}; + +use axum::{ + Router, + body::{Body, to_bytes}, + http::{HeaderMap, Request, header::CONTENT_TYPE}, + response::IntoResponse, +}; +use litellm_auth_types::SecretValue; +use litellm_config::Config; +use litellm_gateway::{ + coordination::{Coordination, RedisSettings}, + rate_limits::{RateLimitSettings, RateLimiter}, +}; +use litellm_gateway_auth::keys::{ + AllowedRoutes, KeyHash, KeyLimits, KeyLookup, KeyModels, KeyProfile, KeyStatus, Limits, + LookupFuture, +}; +use rstest::rstest; +use serde_json::{Value, json}; +use tokio::net::TcpListener; +use tower::ServiceExt; + +/// 3 input and 4 output tokens: 7 count against a token limit. +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}}"#; + +#[derive(Clone, Default)] +struct Upstream { + calls: Arc, + bodies: Arc>>, +} + +async fn upstream(delay: Duration) -> (String, Upstream) { + let upstream = Upstream::default(); + let seen = upstream.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: Request| { + let seen = seen.clone(); + async move { + seen.calls.fetch_add(1, Ordering::SeqCst); + let body = to_bytes(request.into_body(), 1 << 20).await.unwrap(); + seen.bodies + .lock() + .unwrap() + .push(serde_json::from_slice(&body).unwrap_or(Value::Null)); + tokio::time::sleep(delay).await; + ([(CONTENT_TYPE, "application/json")], MESSAGE).into_response() + } + }); + tokio::spawn(async move { axum::serve(listener, app).await }); + (format!("http://{address}"), upstream) +} + +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>, limits: KeyLimits) -> 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 { + user_id: Some("user-1".into()), + limits, + ..KeyProfile::default() + }), + } +} + +fn rpm(limit: i64) -> Limits { + Limits { + rpm: Some(limit), + tpm: None, + } +} + +fn tpm(limit: i64) -> Limits { + Limits { + rpm: None, + tpm: Some(limit), + } +} + +struct Gateway { + app: Router, + upstream: Upstream, +} + +struct Setup { + redis: Option>, + database: Option, + delay: Duration, +} + +impl Default for Setup { + fn default() -> Self { + Self { + redis: None, + database: None, + delay: Duration::ZERO, + } + } +} + +async fn gateway(keys: Vec<(&str, KeyStatus)>, setup: Setup) -> Gateway { + let (address, upstream) = upstream(setup.delay).await; + let config = Config::from_yaml(&format!( + "model_list: + - model_name: claude + litellm_params: + model: anthropic/claude-x + api_base: {address} + api_key: sk-upstream +router_settings: + num_retries: 0 +general_settings: + master_key: sk-master +" + )) + .unwrap(); + let inference = litellm_gateway::build_inference(&config).unwrap(); + let limiter = RateLimiter::new(RateLimitSettings::default(), setup.redis, setup.database); + litellm_gateway::hooks::install(&inference, &config, vec![Arc::new(limiter)], 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, upstream } +} + +struct Reply { + status: u16, + headers: HeaderMap, + body: Value, +} + +async fn send(app: &Router, path: &str, body: Value) -> Reply { + let response = app + .clone() + .oneshot( + Request::post(path) + .header("content-type", "application/json") + .header("authorization", "Bearer sk-key") + .body(Body::from(body.to_string())) + .unwrap(), + ) + .await + .unwrap(); + let status = response.status().as_u16(); + let headers = response.headers().clone(); + let body = to_bytes(response.into_body(), 1 << 20).await.unwrap(); + Reply { + status, + headers, + body: serde_json::from_slice(&body).unwrap_or(Value::Null), + } +} + +fn message(extra: Value) -> Value { + let mut body = json!({ + "model": "claude", + "max_tokens": 8, + "messages": [{"role": "user", "content": "hi"}], + }); + for (field, value) in extra.as_object().unwrap() { + body[field] = value.clone(); + } + body +} + +async fn messages(app: &Router) -> Reply { + send(app, "/v1/messages", message(json!({}))).await +} + +async fn chat(app: &Router) -> Reply { + send(app, "/v1/chat/completions", message(json!({}))).await +} + +/// Settling runs after the response, off the request path. +async fn settled() { + tokio::time::sleep(Duration::from_millis(100)).await; +} + +fn header<'a>(reply: &'a Reply, name: &str) -> Option<&'a str> { + reply.headers.get(name).map(|value| value.to_str().unwrap()) +} + +/// The message up to the reset time, which moves with the clock. +fn refusal_message(reply: &Reply, path: &[&str]) -> String { + let mut value = &reply.body; + for part in path { + value = &value[*part]; + } + let message = value.as_str().unwrap(); + let (message, reset) = message.split_once(" Limit resets at: ").unwrap(); + assert!(reset.ends_with(" UTC"), "{reset}"); + message.to_owned() +} + +#[rstest] +#[tokio::test] +async fn requests_over_the_key_limit_are_refused_like_the_proxy() { + let limits = KeyLimits { + key: rpm(2), + ..KeyLimits::default() + }; + let gateway = gateway(vec![("sk-key", key(None, limits))], Setup::default()).await; + let first = messages(&gateway.app).await; + assert_eq!(first.status, 200); + assert_eq!( + header(&first, "x-ratelimit-api_key-remaining-requests"), + Some("1") + ); + assert_eq!( + header(&first, "x-ratelimit-api_key-limit-requests"), + Some("2") + ); + assert_eq!(header(&first, "x-litellm-key-rpm-limit"), Some("2")); + assert_eq!(chat(&gateway.app).await.status, 200); + + let expected = format!( + "Rate limit exceeded for api_key: {}. Limit type: requests. Current limit: 2, Remaining: 0.", + hash("sk-key").as_str() + ); + // Chat Completions: the proxy's error body and its retry headers. + let refused = chat(&gateway.app).await; + assert_eq!(refused.status, 429); + assert_eq!(refusal_message(&refused, &["error", "message"]), expected); + assert_eq!(refused.body["error"]["type"], "throttling_error"); + assert_eq!(refused.body["error"]["param"], Value::Null); + assert_eq!(refused.body["error"]["code"], "429"); + assert_eq!(header(&refused, "retry-after"), Some("60")); + assert_eq!(header(&refused, "rate_limit_type"), Some("requests")); + assert!(header(&refused, "reset_at").unwrap().ends_with(" UTC")); + assert_eq!(header(&refused, "x-litellm-key-rpm-limit"), Some("2")); + // Messages: Anthropic's envelope, without the retry headers. + let refused = messages(&gateway.app).await; + assert_eq!(refused.status, 429); + assert_eq!(refused.body["type"], "error"); + assert_eq!(refused.body["error"]["type"], "rate_limit_error"); + assert_eq!(refusal_message(&refused, &["error", "message"]), expected); + assert_eq!(header(&refused, "retry-after"), None); + assert_eq!(gateway.upstream.calls.load(Ordering::SeqCst), 2); +} + +#[rstest] +#[case::user(KeyLimits { user: rpm(1), ..KeyLimits::default() }, "user: user-1")] +#[case::team(KeyLimits { team: rpm(1), ..KeyLimits::default() }, "team: team-1")] +#[case::team_member(KeyLimits { team_member: rpm(1), ..KeyLimits::default() }, "team_member: team-1:user-1")] +#[tokio::test] +async fn owners_are_limited_under_their_own_names( + #[case] limits: KeyLimits, + #[case] descriptor: &str, +) { + let gateway = gateway( + vec![("sk-key", key(Some("team-1"), limits))], + Setup::default(), + ) + .await; + assert_eq!(messages(&gateway.app).await.status, 200); + let refused = messages(&gateway.app).await; + assert_eq!(refused.status, 429); + assert_eq!( + refusal_message(&refused, &["error", "message"]), + format!( + "Rate limit exceeded for {descriptor}. Limit type: requests. Current limit: 1, Remaining: 0." + ) + ); +} + +#[rstest] +#[tokio::test] +async fn tokens_are_reserved_and_settled_to_usage() { + // Each request reserves 1 input + 8 output tokens and settles to 7. + let limits = KeyLimits { + key: tpm(20), + ..KeyLimits::default() + }; + let gateway = gateway(vec![("sk-key", key(None, limits))], Setup::default()).await; + let first = messages(&gateway.app).await; + assert_eq!(first.status, 200); + assert_eq!( + header(&first, "x-ratelimit-api_key-remaining-tokens"), + Some("11") + ); + assert_eq!(header(&first, "x-litellm-key-tpm-limit"), Some("20")); + settled().await; + assert_eq!(messages(&gateway.app).await.status, 200); + settled().await; + let refused = messages(&gateway.app).await; + assert_eq!(refused.status, 429); + assert_eq!( + refusal_message(&refused, &["error", "message"]), + format!( + "Rate limit exceeded for api_key: {}. Limit type: tokens. Current limit: 20, Remaining: 6.", + hash("sk-key").as_str() + ) + ); +} + +#[rstest] +#[tokio::test] +async fn small_token_limits_cap_uncapped_requests() { + let limits = KeyLimits { + key: tpm(400), + ..KeyLimits::default() + }; + let gateway = gateway(vec![("sk-key", key(None, limits))], Setup::default()).await; + let body = json!({"model": "claude", "messages": [{"role": "user", "content": "hi"}]}); + assert_eq!(send(&gateway.app, "/v1/messages", body).await.status, 200); + assert_eq!( + gateway.upstream.bodies.lock().unwrap()[0]["max_tokens"], + 100 + ); +} + +#[rstest] +#[tokio::test] +async fn in_flight_requests_are_capped_per_key() { + let limits = KeyLimits { + max_parallel_requests: Some(1), + ..KeyLimits::default() + }; + let gateway = gateway( + vec![("sk-key", key(None, limits))], + Setup { + delay: Duration::from_millis(300), + ..Setup::default() + }, + ) + .await; + let (first, second) = tokio::join!(messages(&gateway.app), async { + tokio::time::sleep(Duration::from_millis(100)).await; + messages(&gateway.app).await + }); + assert_eq!(first.status, 200); + assert_eq!( + header( + &first, + "x-ratelimit-api_key-remaining-max_parallel_requests" + ), + Some("0") + ); + assert_eq!(second.status, 429); + assert_eq!( + refusal_message(&second, &["error", "message"]), + format!( + "Rate limit exceeded for api_key: {}. Limit type: max_parallel_requests. Current limit: 1, Remaining: 0.", + hash("sk-key").as_str() + ) + ); + settled().await; + assert_eq!(messages(&gateway.app).await.status, 200); +} + +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_the_proxy_counters() { + let Some(url) = redis_url() else { return }; + let namespace = format!("rate-limits-{}", uuid::Uuid::new_v4()); + let limits = KeyLimits { + key: Limits { + rpm: Some(3), + tpm: Some(1_000), + }, + max_parallel_requests: Some(5), + ..KeyLimits::default() + }; + let mut replicas = Vec::new(); + for _ in 0..2 { + let setup = Setup { + redis: Some(redis(&url, &namespace).await), + ..Setup::default() + }; + replicas.push(gateway(vec![("sk-key", key(None, limits))], setup).await); + } + assert_eq!(messages(&replicas[0].app).await.status, 200); + assert_eq!(messages(&replicas[1].app).await.status, 200); + let third = messages(&replicas[0].app).await; + assert_eq!( + header(&third, "x-ratelimit-api_key-remaining-requests"), + Some("0") + ); + settled().await; + assert_eq!(messages(&replicas[1].app).await.status, 429); + + let mut connection = redis::Client::open(url.as_str()) + .unwrap() + .get_multiplexed_async_connection() + .await + .unwrap(); + let prefix = format!("{namespace}:{{api_key:{}}}", hash("sk-key").as_str()); + let get = async |connection: &mut redis::aio::MultiplexedConnection, suffix: &str| { + redis::cmd("GET") + .arg(format!("{prefix}:{suffix}")) + .query_async::>(connection) + .await + .unwrap() + }; + // The refused request still counted; three settled requests used 7 tokens each. + assert_eq!(get(&mut connection, "requests").await.as_deref(), Some("4")); + assert_eq!(get(&mut connection, "tokens").await.as_deref(), Some("21")); + assert!(get(&mut connection, "window").await.is_some()); + let in_flight: i64 = redis::cmd("ZCARD") + .arg(format!("{prefix}:max_parallel_requests")) + .query_async(&mut connection) + .await + .unwrap(); + assert_eq!(in_flight, 0); +} + +#[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, rpm_limit, created_by, updated_by) + VALUES ('one-a-minute', 1, 'x', 'x'); +INSERT INTO "LiteLLM_EndUserTable" (user_id, spend, budget_id) + VALUES ('limited', 0, 'one-a-minute'), ('free', 0, NULL); +"#, + ) + .execute(pool) + .await + .unwrap(); +} + +#[rstest] +#[tokio::test] +async fn end_users_are_held_to_their_budget_limits() { + let postgres = litellm_db_testing::MigratedPostgres::start().await.unwrap(); + end_users(postgres.pool()).await; + let gateway = gateway( + vec![("sk-key", key(None, KeyLimits::default()))], + Setup { + database: Some(postgres.pool().clone()), + ..Setup::default() + }, + ) + .await; + let as_end_user = |user: &str| message(json!({"metadata": {"user_id": user}})); + for user in ["limited", "free", "free"] { + assert_eq!( + send(&gateway.app, "/v1/messages", as_end_user(user)) + .await + .status, + 200 + ); + } + let refused = send(&gateway.app, "/v1/messages", as_end_user("limited")).await; + assert_eq!(refused.status, 429); + assert_eq!( + refusal_message(&refused, &["error", "message"]), + "Rate limit exceeded for end_user: limited. Limit type: requests. Current limit: 1, Remaining: 0." + ); +} diff --git a/litellm-rust/docs/customer-proxy/rate-limits.svg b/litellm-rust/docs/customer-proxy/rate-limits.svg new file mode 100644 index 00000000000..beff70d8482 --- /dev/null +++ b/litellm-rust/docs/customer-proxy/rate-limits.svg @@ -0,0 +1,55 @@ + + + + + + + Rate limits hook: the proxy's limiter on the proxy's shared counters + + + Scopes, in the proxy's order + api_key: key rpm/tpm (else budget row's), + max_parallel_requests + user, team, team_member (team:user) + end_user: its budget row, or the default + read with the key; end users cached 60 s + no limits set: no Redis call at all + + + Three Redis scripts, in the proxy's phases + 1. requests: +1 on every {scope:id}:requests in one call, fixed window from the first hit; + over when count > limit; the count is never given back + 2. in flight (key only): ZSET of slots on Redis TIME, pruned at 3600 s, full at limit + 3. tokens: reserve chars/4 + output x candidates on every scope, all or none; + output = max_tokens, else the configured estimate, else max(input, min(1024, tpm/4)); + a limit under 4096 caps an uncapped request at that floor upstream + + + + Refusal: 429, first scope over + "Rate limit exceeded for api_key: <hash>. Limit type: requests. + Current limit: 10, Remaining: 0. Limit resets at: ... UTC" + chat: throttling_error + retry-after, rate_limit_type, reset_at + messages: Anthropic rate_limit_error envelope, no retry headers + a refused token reservation gives its slot back + + + + Admitted: request runs + complete responses carry where each limit stands: + x-ratelimit-{scope}-remaining-{type}, -limit-{type} + every response: x-litellm-key-rpm-limit / -tpm-limit + streams get no x-ratelimit-* headers, as in the proxy + + + + Settle (off the request path) + tokens: INCRBYFLOAT (actual - reserved) on each reserved scope, actual = all tokens but cache reads; a failure returns the reservation + in flight: ZREM the request's slot when the response ends, fails, or the client goes away + + + + Same keys, scripts and windows as the Python proxy, under its namespace + ns:{api_key:<hash>}:window | requests | tokens | max_parallel_requests: gateway replicas and Python proxies enforce one limit. + No Redis, or Redis unreachable: the same logic on this replica's counters (fail_closed_rate_limit_enforcement: 503 instead) +